Jaroslav Beneš
89d2d6ebfd
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>
284 lines
10 KiB
Python
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
|