Let a command run in the background instead of being killed
An agent command is one blocking conn.run over a per-call connection, killed the
moment it hits its timeout -- so a ten-minute apt install is impossible, which is
exactly what a user hit. This is the substrate for running it detached instead:
the model can ask for background=true, or a command that outlasts its timeout is
kept running rather than killed, and either way the model gets tools to read and
stop it. Opt-in, off by default, under Admin -> Agents; off is byte-for-byte the
old behaviour.
The mechanism has to survive the connection closing (that is the whole premise
of the per-call model), so a job is a setsid-detached process on the far side,
redirected to a remote logfile and an exit-file; LLeMbas reconnects, as always,
to read it later. services/agent/jobs.py holds the wrappers.
Three things in those wrappers are load-bearing and each was got wrong in the
first sketch:
- The command never touches a quoted shell context. sh -c '<cmd>' shatters the
instant the command contains a quote -- git commit -m 'fix', awk '{…}', sed
's/…/…/' are the common case, and it is an injection hole besides. So the
command is base64-encoded in Python and decoded on the far side into a script
file; it is bytes, never shell syntax.
- The child records its own pid via $$ as its first act, under setsid where it
is the session leader, so job_stop can kill the whole process group. echo $!
from the launcher captures the wrong pid.
- The command's exit status comes from the exit-file, never the wrapper's own
status -- which is ~0 from its trailing rm. Reading the wrapper's status would
mark every job a success.
A command that finishes in time is indistinguishable from a foreground one --
same output, same wording; the difference shows only when it does not, where
instead of "stopped after Ns" it becomes a job id. Auto-convert is its own
sub-switch: with it off, a timeout stays a hard stop and nothing is left
running, because routing the plain case through the detached wrapper would leave
an orphan running past a stop an administrator asked for.
New agent tools job_output/job_list/job_stop, offered only when the feature is
on (the plan_submit gating pattern); job_stop is RISK_EXECUTE since it kills a
process. A job's files are namespaced by the calling chat's id and the wrappers
are always built from it, so a model in one chat cannot even name another's job.
Tested against a real local /bin/sh rather than the fake echo-the-command sshd
fixture, because the shell logic -- setsid, base64, the wait loop, the child
surviving the wait being cut off -- is the whole of the risk. The auto-wake that
prompts the model back when a job finishes is the next commit.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
@@ -0,0 +1,361 @@
|
||||
"""Commands that outlive the reply that started them.
|
||||
|
||||
An ordinary `shell_run` is one blocking `conn.run` over a per-call connection
|
||||
(`ssh.py`): when it hits its timeout the command is killed, so a ten-minute
|
||||
`apt install` is impossible. A background job is the same command launched
|
||||
*detached* on the far side -- `setsid`, redirected to a remote logfile and an
|
||||
exit-file -- so it survives the connection closing. LLeMbas reconnects (a fresh
|
||||
connection, as always) to read the log and the exit code later.
|
||||
|
||||
This is the opposite of `terminal.py`, which survives by *holding* a connection
|
||||
open. Here we hold nothing: the whole point of `ssh.py`/`base.py` is that no live
|
||||
connection is kept, and a job that needed one would be a job that broke that.
|
||||
|
||||
**The command never touches a quoted shell context.** `sh -c '<cmd>'` shatters
|
||||
the instant the command contains a `'` -- `git commit -m 'fix'`, `awk '{…}'`,
|
||||
`sed 's/…/…/'` are the common case, not an edge one, and would also be an
|
||||
injection hole. So the command is base64-encoded here in Python and decoded on
|
||||
the far side into a script file; it is bytes, never shell syntax. Only
|
||||
server-generated hex ids and a fixed root ever reach a path.
|
||||
|
||||
Three things make the wrappers correct, and each was got wrong in an earlier
|
||||
sketch:
|
||||
|
||||
* **The child records its own pid via `$$`**, as its first act, under `setsid`
|
||||
where it is the session/group leader -- so `job_stop` can `kill -<pid>` the
|
||||
whole process group. `echo $!` from the launcher captures the wrong pid.
|
||||
* **The exit-file is the primary signal.** An empty pid-file means "still
|
||||
starting", not "dead"; reading liveness first would race the launch and report
|
||||
a job lost the instant it began.
|
||||
* **The command's exit status comes from the exit-file, never from the wrapper's
|
||||
own status** -- which is ~0 from the trailing `rm`. Reading the wrapper's
|
||||
status would mark every job a success.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import contextlib
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from lembas.services.agent.base import ExecError, ExecRequest, clean_output
|
||||
|
||||
# Where a job's files live on the far side. `${TMPDIR:-/tmp}` so a host that
|
||||
# puts scratch space elsewhere is honoured, and it clears on reboot -- a job
|
||||
# does not survive a reboot of its own host either. The chat id namespaces it,
|
||||
# which is also what makes cross-chat access structurally impossible: a path is
|
||||
# only ever built from the *calling* chat's id, so a model in one chat cannot
|
||||
# name another chat's files.
|
||||
JOB_ROOT = "${TMPDIR:-/tmp}/lembas-jobs"
|
||||
|
||||
# A job id is our own short hex; anything else is refused before it reaches a
|
||||
# path, so `job_output("../../etc/passwd")` cannot walk out of the job root.
|
||||
_ID = re.compile(r"^[a-f0-9]{12}$")
|
||||
|
||||
# How long the fire-and-return launcher waits for the shell to accept the
|
||||
# command. Not the command's own timeout -- it returns the moment the process is
|
||||
# detached, which is immediate.
|
||||
LAUNCH_GRACE = 10.0
|
||||
|
||||
# In-process, keyed by job id, lost on restart -- the durable record is the `Job`
|
||||
# row (added with the watcher). This holds the metadata `job_list` shows within
|
||||
# a session and, later, the watcher task.
|
||||
_JOBS: dict[str, JobState] = {}
|
||||
|
||||
|
||||
def new_id() -> str:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobState:
|
||||
"""What LLeMbas remembers about one background job, in this process."""
|
||||
|
||||
id: str
|
||||
chat_id: str
|
||||
command: str
|
||||
status: str = "running" # running | done | killed | lost
|
||||
exit_status: int | None = None
|
||||
started_at: float = field(default_factory=time.monotonic)
|
||||
finished_at: float = 0.0
|
||||
|
||||
|
||||
# --- Paths and the wrappers ----------------------------------------------------
|
||||
def _dir(chat_id: str) -> str:
|
||||
return f'"{JOB_ROOT}/{chat_id}"'
|
||||
|
||||
|
||||
def _file(chat_id: str, job_id: str, ext: str) -> str:
|
||||
# Double-quoted so `${TMPDIR:-/tmp}` still expands while the whole path stays
|
||||
# one word. The chat id and job id are hex, so nothing here needs escaping.
|
||||
return f'"{JOB_ROOT}/{chat_id}/{job_id}.{ext}"'
|
||||
|
||||
|
||||
def _sentinel(job_id: str) -> str:
|
||||
return f"__LEMBAS_{job_id}__"
|
||||
|
||||
|
||||
def _inner_script(chat_id: str, job_id: str, command: str) -> str:
|
||||
"""The detached program: record the pid, run the command, record the status.
|
||||
|
||||
base64-encoded before it leaves, so `command` is bytes and never shell
|
||||
syntax. `$$` first, because it is the session leader's pid under setsid and
|
||||
`job_stop` kills the group by it. `$?` last, capturing the command's status;
|
||||
it is the file `run`'s own exit status must never be read in place of.
|
||||
"""
|
||||
return (
|
||||
f"echo $$ > {_file(chat_id, job_id, 'pid')}\n"
|
||||
f"{command}\n"
|
||||
f"echo $? > {_file(chat_id, job_id, 'exit')}\n"
|
||||
)
|
||||
|
||||
|
||||
def _blob(chat_id: str, job_id: str, command: str) -> str:
|
||||
raw = _inner_script(chat_id, job_id, command).encode("utf-8")
|
||||
return base64.b64encode(raw).decode("ascii")
|
||||
|
||||
|
||||
def _launch_lines(chat_id: str, job_id: str, command: str) -> str:
|
||||
"""Create the job dir, drop the script, and detach it. No wait."""
|
||||
blob = _blob(chat_id, job_id, command)
|
||||
return (
|
||||
f"mkdir -p {_dir(chat_id)} 2>/dev/null\n"
|
||||
f"printf %s '{blob}' | base64 -d > {_file(chat_id, job_id, 'sh')}\n"
|
||||
f"setsid sh {_file(chat_id, job_id, 'sh')} "
|
||||
f"> {_file(chat_id, job_id, 'log')} 2>&1 < /dev/null &\n"
|
||||
)
|
||||
|
||||
|
||||
def launch_command(chat_id: str, job_id: str, command: str) -> str:
|
||||
"""Fire-and-return: detach the command and stop. Run with a short timeout."""
|
||||
return _launch_lines(chat_id, job_id, command) + "printf started\n"
|
||||
|
||||
|
||||
def launch_and_wait_command(chat_id: str, job_id: str, command: str, max_bytes: int) -> str:
|
||||
"""Detach the command AND wait up to the (asyncssh) timeout for it.
|
||||
|
||||
If it finishes, stdout is the log tail plus a sentinel line carrying the exit
|
||||
code, and the files are removed. If asyncssh times out first the channel is
|
||||
torn down before the `rm`, so the files survive for a later read and the
|
||||
detached process -- new session, redirected, stdin from /dev/null -- keeps
|
||||
running. That torn-down-mid-wait case is exactly "it became a background
|
||||
job".
|
||||
"""
|
||||
s = _sentinel(job_id)
|
||||
pid = _file(chat_id, job_id, "pid")
|
||||
exit_ = _file(chat_id, job_id, "exit")
|
||||
log = _file(chat_id, job_id, "log")
|
||||
return (
|
||||
_launch_lines(chat_id, job_id, command)
|
||||
+ "while :; do\n"
|
||||
f" [ -f {exit_} ] && break\n"
|
||||
f" __p=$(cat {pid} 2>/dev/null)\n"
|
||||
' [ -n "$__p" ] && ! kill -0 "$__p" 2>/dev/null && break\n'
|
||||
# 0.2s: with the feature on, every ordinary command waits one poll for
|
||||
# the exit-file, so this is added latency on the hot path. Short enough
|
||||
# not to be felt, long enough not to spin.
|
||||
" sleep 0.2\n"
|
||||
"done\n"
|
||||
f"tail -c {max_bytes} {log} 2>/dev/null\n"
|
||||
f"printf '\\n{s}:'\n"
|
||||
f"cat {exit_} 2>/dev/null || printf LOST\n"
|
||||
f"rm -f {_file(chat_id, job_id, 'sh')} {pid} {log} {exit_}\n"
|
||||
)
|
||||
|
||||
|
||||
def read_command(chat_id: str, job_id: str, max_bytes: int) -> str:
|
||||
"""The log so far, and whether the job is still running."""
|
||||
s = _sentinel(job_id)
|
||||
pid = _file(chat_id, job_id, "pid")
|
||||
exit_ = _file(chat_id, job_id, "exit")
|
||||
return (
|
||||
f"tail -c {max_bytes} {_file(chat_id, job_id, 'log')} 2>/dev/null\n"
|
||||
f"printf '\\n{s}:'\n"
|
||||
f"if [ -f {exit_} ]; then printf 'done '; cat {exit_};\n"
|
||||
f'elif __p=$(cat {pid} 2>/dev/null); [ -n "$__p" ] && kill -0 "$__p" 2>/dev/null;'
|
||||
" then printf running;\n"
|
||||
"else printf lost; fi\n"
|
||||
)
|
||||
|
||||
|
||||
def stop_command(chat_id: str, job_id: str) -> str:
|
||||
"""Kill the whole process group, then record an exit so a reader is not told
|
||||
the job is merely lost. A killed process never writes its own exit file."""
|
||||
pid = _file(chat_id, job_id, "pid")
|
||||
exit_ = _file(chat_id, job_id, "exit")
|
||||
return (
|
||||
f'__p=$(cat {pid} 2>/dev/null); [ -n "$__p" ] && kill -TERM -"$__p" 2>/dev/null\n'
|
||||
"sleep 0.3\n"
|
||||
f'[ -n "$__p" ] && kill -KILL -"$__p" 2>/dev/null\n'
|
||||
f"[ -f {exit_} ] || echo 143 > {exit_}\n"
|
||||
"printf stopped\n"
|
||||
)
|
||||
|
||||
|
||||
def cleanup_command(chat_id: str, job_id: str) -> str:
|
||||
return (
|
||||
f"rm -f {_file(chat_id, job_id, 'sh')} {_file(chat_id, job_id, 'pid')} "
|
||||
f"{_file(chat_id, job_id, 'log')} {_file(chat_id, job_id, 'exit')}\n"
|
||||
)
|
||||
|
||||
|
||||
# --- Parsing what a wrapper printed --------------------------------------------
|
||||
@dataclass(frozen=True)
|
||||
class Completed:
|
||||
body: str
|
||||
exit_status: int | None # None ⇒ the job was lost (killed without an exit)
|
||||
|
||||
|
||||
def parse_completed(output: str, job_id: str) -> Completed:
|
||||
"""Split a launch-and-wait result into the command's output and its status.
|
||||
|
||||
On the *last* sentinel, because the command's own output could contain a
|
||||
line that looks like one; everything before it is the body, everything after
|
||||
is the exit code the file held.
|
||||
"""
|
||||
marker = f"\n{_sentinel(job_id)}:"
|
||||
at = output.rfind(marker)
|
||||
if at == -1:
|
||||
return Completed(body=output.strip(), exit_status=None)
|
||||
body = output[:at].strip()
|
||||
tail = output[at + len(marker) :].strip()
|
||||
if tail.upper() == "LOST" or not tail:
|
||||
return Completed(body=body, exit_status=None)
|
||||
try:
|
||||
return Completed(body=body, exit_status=int(tail.split()[0]))
|
||||
except (ValueError, IndexError):
|
||||
return Completed(body=body, exit_status=None)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Reading:
|
||||
body: str
|
||||
status: str # running | done | lost
|
||||
exit_status: int | None
|
||||
|
||||
|
||||
def parse_reading(output: str, job_id: str) -> Reading:
|
||||
marker = f"\n{_sentinel(job_id)}:"
|
||||
at = output.rfind(marker)
|
||||
if at == -1:
|
||||
return Reading(body=output.strip(), status="lost", exit_status=None)
|
||||
body = output[:at].strip()
|
||||
tail = output[at + len(marker) :].strip()
|
||||
if tail.startswith("done"):
|
||||
parts = tail.split()
|
||||
code = int(parts[1]) if len(parts) > 1 and parts[1].lstrip("-").isdigit() else None
|
||||
return Reading(body=body, status="done", exit_status=code)
|
||||
if tail == "running":
|
||||
return Reading(body=body, status="running", exit_status=None)
|
||||
return Reading(body=body, status="lost", exit_status=None)
|
||||
|
||||
|
||||
# --- Operations against the machine --------------------------------------------
|
||||
async def launch(agent, command: str, cwd: str = "") -> JobState:
|
||||
"""Detach a command and return immediately. Raises ExecError if it will not
|
||||
even start."""
|
||||
job_id = new_id()
|
||||
result = await agent.executor().run(
|
||||
ExecRequest(
|
||||
command=launch_command(agent.chat_id, job_id, command),
|
||||
cwd=cwd,
|
||||
timeout=LAUNCH_GRACE,
|
||||
max_bytes=agent.max_output,
|
||||
)
|
||||
)
|
||||
if result.timed_out:
|
||||
raise ExecError("The machine did not accept the command in time.")
|
||||
job = JobState(id=job_id, chat_id=agent.chat_id, command=command)
|
||||
_JOBS[job_id] = job
|
||||
return job
|
||||
|
||||
|
||||
async def read(agent, job_id: str) -> Reading:
|
||||
output, _ = _clean(
|
||||
await agent.executor().run(
|
||||
ExecRequest(
|
||||
command=read_command(agent.chat_id, job_id, agent.max_output),
|
||||
timeout=agent.timeout,
|
||||
max_bytes=agent.max_output,
|
||||
)
|
||||
),
|
||||
agent.max_output,
|
||||
)
|
||||
reading = parse_reading(output, job_id)
|
||||
_record(job_id, reading.status, reading.exit_status)
|
||||
if reading.status in ("done", "lost"):
|
||||
await _cleanup(agent, job_id)
|
||||
return reading
|
||||
|
||||
|
||||
async def stop(agent, job_id: str) -> None:
|
||||
await agent.executor().run(
|
||||
ExecRequest(command=stop_command(agent.chat_id, job_id), timeout=agent.timeout)
|
||||
)
|
||||
_record(job_id, "killed", 143)
|
||||
|
||||
|
||||
async def _cleanup(agent, job_id: str) -> None:
|
||||
with contextlib.suppress(ExecError):
|
||||
await agent.executor().run(
|
||||
ExecRequest(command=cleanup_command(agent.chat_id, job_id), timeout=agent.timeout)
|
||||
)
|
||||
|
||||
|
||||
def _clean(result, limit: int) -> tuple[str, bool]:
|
||||
if result.timed_out:
|
||||
return result.output, False
|
||||
return clean_output(result.output or "", limit=limit)
|
||||
|
||||
|
||||
# --- The registry --------------------------------------------------------------
|
||||
def register(job: JobState) -> None:
|
||||
_JOBS[job.id] = job
|
||||
|
||||
|
||||
def get(job_id: str) -> JobState | None:
|
||||
return _JOBS.get(job_id)
|
||||
|
||||
|
||||
def for_chat(chat_id: str) -> list[JobState]:
|
||||
return [j for j in _JOBS.values() if j.chat_id == chat_id]
|
||||
|
||||
|
||||
def valid_id(job_id: str) -> bool:
|
||||
return bool(_ID.match(job_id or ""))
|
||||
|
||||
|
||||
def _record(job_id: str, status: str, exit_status: int | None) -> None:
|
||||
job = _JOBS.get(job_id)
|
||||
if job is None:
|
||||
return
|
||||
if status in ("done", "lost", "killed") and job.status == "running":
|
||||
job.status = status
|
||||
job.exit_status = exit_status
|
||||
job.finished_at = time.monotonic()
|
||||
|
||||
|
||||
def clear() -> None:
|
||||
_JOBS.clear()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"JOB_ROOT",
|
||||
"Completed",
|
||||
"JobState",
|
||||
"Reading",
|
||||
"clear",
|
||||
"for_chat",
|
||||
"get",
|
||||
"launch",
|
||||
"launch_and_wait_command",
|
||||
"new_id",
|
||||
"parse_completed",
|
||||
"read",
|
||||
"register",
|
||||
"stop",
|
||||
"valid_id",
|
||||
]
|
||||
Reference in New Issue
Block a user