Files
LLeMbas/tests/test_agent_jobs.py
T
Jaroslav Beneš 3fc3449726 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>
2026-08-03 14:15:49 +02:00

284 lines
10 KiB
Python

"""Background jobs, exercised against a real local shell.
The sshd fixture elsewhere is a *fake* shell (it echoes the command), which
cannot run `setsid`, `base64`, a wait loop, or an exit-file. So these run the
wrappers through the machine's own `/bin/sh` with the same contract
`SshExecutor.run` has — cd-prefix, interleaved output, a timeout that leaves the
detached child alive — because the shell logic is the whole of the risk here.
"""
from __future__ import annotations
import asyncio
import shutil
import time
import pytest
from lembas.services import tools as tools_service
from lembas.services.agent import jobs, policy
from lembas.services.agent import ssh as ssh_service
from lembas.services.agent.base import ExecRequest, ExecResult, clean_output
from lembas.services.agent.session import AgentContext
from lembas.services.agent.tools import _run_shell
pytestmark = pytest.mark.skipif(
shutil.which("setsid") is None or shutil.which("base64") is None,
reason="needs setsid and base64 (Linux)",
)
class LocalExecutor:
"""`SshExecutor.run`'s contract, run against the local shell.
On timeout it kills only the outer shell -- exactly what an SSH channel
teardown does -- so a `setsid`-detached child survives, which is the whole
behaviour a background job depends on.
"""
def __init__(self, project_dir: str) -> None:
self.project_dir = project_dir
async def run(self, request: ExecRequest) -> ExecResult:
directory = request.cwd or self.project_dir
command = request.command
if directory:
command = f"cd {ssh_service._quote(directory)} && {command}"
started = time.monotonic()
proc = await asyncio.create_subprocess_shell(
command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.STDOUT,
stdin=asyncio.subprocess.DEVNULL,
)
try:
out, _ = await asyncio.wait_for(proc.communicate(), timeout=request.timeout)
except TimeoutError:
proc.kill()
await proc.wait()
return ExecResult(
exit_status=-1,
output=f"The command was still running after {request.timeout:g}s.",
timed_out=True,
duration_ms=int((time.monotonic() - started) * 1000),
)
output, truncated = clean_output(out.decode("utf-8", "replace"), limit=request.max_bytes)
return ExecResult(
exit_status=proc.returncode if proc.returncode is not None else -1,
output=output,
truncated=truncated,
duration_ms=int((time.monotonic() - started) * 1000),
)
@pytest.fixture(autouse=True)
def _clear_jobs():
jobs.clear()
yield
jobs.clear()
def _agent(tmp_path, **over) -> AgentContext:
fields = {
"chat_id": "chat0001",
"label": "Box",
"project_dir": str(tmp_path),
"mode": policy.MODE_AUTO,
"max_output": 65536,
"timeout": 60.0,
"max_timeout": 600.0,
"background": True,
"background_on_timeout": True,
"background_notify": True,
}
fields.update(over)
ctx = AgentContext(**fields)
ctx.executor = lambda: LocalExecutor(fields["project_dir"]) # type: ignore[method-assign]
return ctx
def _context(agent):
return tools_service.ToolContext(owner_id="u", agent=agent)
# --- The wrappers, as strings --------------------------------------------------
def test_the_command_is_never_in_a_quoted_context():
"""git commit -m 'fix' would shatter sh -c '<cmd>'. It is base64'd instead."""
import base64
cmd = "git commit -m 'fix: it'"
wrapper = jobs.launch_and_wait_command("chatx", "abc123abc123", cmd, 4096)
assert cmd not in wrapper, "the command leaked into the wrapper as shell text"
blob = wrapper.split("printf %s '", 1)[1].split("'", 1)[0]
inner = base64.b64decode(blob).decode()
assert cmd in inner, "the command was lost in the base64 round-trip"
def test_parse_reads_the_last_sentinel():
out = ("line one\n__LEMBAS_jobjobjob01__:0\nmore\n__LEMBAS_jobjobjob01__:0")
done = jobs.parse_completed(out, "jobjobjob01")
assert done.exit_status == 0
def test_valid_id_refuses_a_path():
assert jobs.valid_id("deadbeef0000")
assert not jobs.valid_id("../../etc/passwd")
assert not jobs.valid_id("")
# --- Real shell: launch-and-wait -----------------------------------------------
async def test_a_fast_command_completes_like_a_foreground_one(tmp_path):
agent = _agent(tmp_path)
outcome = await _run_shell(_context(agent), {"command": "echo hello"})
assert "hello" in outcome.content
assert outcome.event["status"] == "ok"
assert not jobs.for_chat(agent.chat_id), "a finished command is not a job"
async def test_a_nonzero_exit_is_read_from_the_exit_file_not_the_wrapper(tmp_path):
"""The wrapper's own status is ~0 from its trailing rm; the command's real
status is in the exit-file."""
agent = _agent(tmp_path)
outcome = await _run_shell(_context(agent), {"command": "sh -c 'exit 3'"})
assert "exited 3" in outcome.content
assert outcome.event["status"] == "error"
async def test_a_command_with_single_quotes_runs(tmp_path):
agent = _agent(tmp_path)
outcome = await _run_shell(_context(agent), {"command": "echo 'a'\\''b'"})
assert "a'b" in outcome.content
async def test_output_is_captured(tmp_path):
agent = _agent(tmp_path)
outcome = await _run_shell(_context(agent), {"command": "printf 'one\\ntwo\\n'"})
assert "one" in outcome.content and "two" in outcome.content
# --- Real shell: auto-background on timeout -------------------------------------
async def test_a_slow_command_is_kept_running_when_it_times_out(tmp_path):
"""The apt-install case. It is not killed; it becomes a job that finishes
on its own, and its result is readable afterwards."""
agent = _agent(tmp_path)
outcome = await _run_shell(
_context(agent), {"command": "sleep 2; echo done-late", "timeout": 1}
)
assert "background" in outcome.content.lower()
running = jobs.for_chat(agent.chat_id)
assert len(running) == 1
job_id = running[0].id
# The detached child survived the wait being cut off; give it time to finish.
for _ in range(40):
reading = await jobs.read(agent, job_id)
if reading.status == "done":
break
await asyncio.sleep(0.2)
assert reading.status == "done"
assert "done-late" in reading.body
async def test_off_timeout_keeps_the_hard_stop(tmp_path):
"""With the auto-convert off, a timed-out command is killed as before and no
job is left running."""
agent = _agent(tmp_path, background_on_timeout=False)
outcome = await _run_shell(
_context(agent), {"command": "sleep 2; echo late", "timeout": 1}
)
assert "stopped after" in outcome.content
assert not jobs.for_chat(agent.chat_id)
async def test_feature_off_is_the_plain_path(tmp_path):
"""No wrapper, no job, byte-for-byte the old wording."""
agent = _agent(tmp_path, background=False)
outcome = await _run_shell(_context(agent), {"command": "echo hi"})
assert "hi" in outcome.content
assert not jobs.for_chat(agent.chat_id)
# --- Explicit background=true --------------------------------------------------
async def test_background_true_returns_at_once_with_a_job(tmp_path):
agent = _agent(tmp_path)
started = time.monotonic()
outcome = await _run_shell(
_context(agent), {"command": "sleep 5", "background": True}
)
assert time.monotonic() - started < 3, "it should not wait for the command"
running = jobs.for_chat(agent.chat_id)
assert len(running) == 1
assert "job " in outcome.content
await jobs.stop(agent, running[0].id)
# --- The job tools -------------------------------------------------------------
async def test_job_output_reads_a_running_job(tmp_path):
from lembas.services.agent.tools import _run_job_output
agent = _agent(tmp_path)
job = await jobs.launch(agent, "sleep 5")
outcome = await _run_job_output(_context(agent), {"id": job.id})
assert "running" in outcome.content.lower()
await jobs.stop(agent, job.id)
async def test_job_stop_ends_it(tmp_path):
from lembas.services.agent.tools import _run_job_output, _run_job_stop
agent = _agent(tmp_path)
job = await jobs.launch(agent, "sleep 30")
stopped = await _run_job_stop(_context(agent), {"id": job.id})
assert "Stopped" in stopped.content
reading = await _run_job_output(_context(agent), {"id": job.id})
assert "running" not in reading.content.lower()
async def test_job_output_refuses_a_bad_id(tmp_path):
from lembas.services.agent.tools import _run_job_output
agent = _agent(tmp_path)
outcome = await _run_job_output(_context(agent), {"id": "../../etc/passwd"})
assert outcome.event["status"] == "error"
async def test_a_job_is_only_visible_to_its_own_chat(tmp_path):
"""for_chat is keyed on chat id; the file paths are too, so a model in
another chat cannot even name it."""
agent_a = _agent(tmp_path, chat_id="chataaaa")
await jobs.launch(agent_a, "sleep 5")
assert jobs.for_chat("chatbbbb") == []
for job in jobs.for_chat("chataaaa"):
await jobs.stop(agent_a, job.id)
# --- Gating --------------------------------------------------------------------
def test_job_tools_appear_only_when_enabled(tmp_path):
from lembas.services.agent.tools import tool_defs
on = {t.name for t in tool_defs(_agent(tmp_path, background=True))}
off = {t.name for t in tool_defs(_agent(tmp_path, background=False))}
assert {"job_output", "job_list", "job_stop"} <= on
assert not ({"job_output", "job_list", "job_stop"} & off)
def test_background_param_appears_only_when_enabled(tmp_path):
from lembas.services.agent.tools import tool_defs
def shell(agent):
return next(t for t in tool_defs(agent) if t.name == "shell_run")
on = shell(_agent(tmp_path, background=True)).parameters["properties"]
off = shell(_agent(tmp_path, background=False)).parameters["properties"]
assert "background" in on
assert "background" not in off