"""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 ''. 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_every_wrapper_is_valid_shell(): """`sh -n` parses without executing. The one that would have caught it: `{log}` for `{logf}` formatted the module logger into the launch-and-wait wrapper, and `` is shell syntax. The error landed after the sentinel, where nothing reads it, so the command still worked and the cleanup silently never ran. """ import subprocess wrappers = [ jobs.launch_command("chatx", "abc123abc123", "echo hi"), jobs.launch_and_wait_command("chatx", "abc123abc123", "echo hi", 4096), jobs.read_command("chatx", "abc123abc123", 4096), jobs.stop_command("chatx", "abc123abc123"), jobs.cleanup_command("chatx", "abc123abc123"), ] for wrapper in wrappers: done = subprocess.run( ["sh", "-n"], input=wrapper, capture_output=True, text=True, check=False ) assert done.returncode == 0, f"not valid shell:\n{wrapper}\n{done.stderr}" def test_launch_and_wait_removes_every_file_it_made(): """Its last line is the only cleanup on the fast path -- nothing calls `_cleanup` when a command finishes in time.""" wrapper = jobs.launch_and_wait_command("chatx", "abc123abc123", "echo hi", 4096) removal = next(line for line in wrapper.splitlines() if line.startswith("rm -f")) for extension in ("sh", "pid", "log", "exit"): assert jobs._file("chatx", "abc123abc123", extension) in removal, extension assert " str: from lembas.db.models import Chat chat = Chat(user_id=user_id, model_id="m") db.add(chat) db.commit() return chat.id def test_the_completion_names_itself_a_machine_event(): text = jobs._completion_text("abc123abc123", "apt-get install -y x", "done", 0, "ok\n") assert "machine event" in text assert "[job abc123abc123]" in text assert "apt-get install" in text assert "finished successfully" in text.lower() def test_a_nonzero_completion_reports_the_code(): text = jobs._completion_text("j", "build", "done", 2, "") assert "exited 2" in text def test_backticks_in_output_cannot_close_the_fence(): text = jobs._completion_text("j", "c", "done", 0, "see ```code``` here") assert "```code```" not in text async def test_an_idle_chat_gets_a_fresh_reply(db, user_id, registered, monkeypatch): from lembas.services import generation as generation_service chat_id = _chat(db, user_id) monkeypatch.setattr(generation_service, "running_for", lambda _c: None) started: list[str] = [] monkeypatch.setattr(generation_service, "ensure", lambda c, m: started.append(m)) await jobs.wake(chat_id, "abc123abc123", "sleep 1", "done", 0, "hi") from lembas.db.models import Message db.expire_all() msgs = db.query(Message).filter(Message.chat_id == chat_id).all() users = [m for m in msgs if m.role == "user"] assert len(users) == 1 assert users[0].queued is False, "an idle chat's completion is delivered, not queued" assert any(m.role == "assistant" and not m.complete for m in msgs) assert len(started) == 1, "a reply was started" async def test_a_busy_chat_gets_a_queued_turn_and_no_new_reply( db, user_id, registered, monkeypatch ): from lembas.services import generation as generation_service chat_id = _chat(db, user_id) monkeypatch.setattr(generation_service, "running_for", lambda _c: object()) started: list[str] = [] monkeypatch.setattr(generation_service, "ensure", lambda c, m: started.append(m)) await jobs.wake(chat_id, "abc123abc123", "sleep 1", "done", 0, "hi") from lembas.db.models import Message db.expire_all() users = [ m for m in db.query(Message).filter(Message.chat_id == chat_id).all() if m.role == "user" ] assert len(users) == 1 assert users[0].queued is True, "a running reply's _inject/_drain will deliver it" assert started == [], "no second reply for a chat already writing one" async def test_two_jobs_finishing_at_once_start_one_reply(db, user_id, registered, monkeypatch): """The per-chat lock. Without it, both wakes would each start a generation.""" from lembas.services import generation as generation_service chat_id = _chat(db, user_id) live = {"running": False} monkeypatch.setattr( generation_service, "running_for", lambda _c: object() if live["running"] else None ) started: list[str] = [] def _ensure(c, m): live["running"] = True # what running_for will now see started.append(m) monkeypatch.setattr(generation_service, "ensure", _ensure) await asyncio.gather( jobs.wake(chat_id, "aaaaaaaaaaaa", "a", "done", 0, "x"), jobs.wake(chat_id, "bbbbbbbbbbbb", "b", "done", 0, "y"), ) assert len(started) == 1, "the lock made the second wake see the first's reply" # --- The watcher, end to end --------------------------------------------------- async def test_the_watcher_end_to_end(db, user_id, registered, monkeypatch, tmp_path): from lembas.services.agent import ssh as ssh_service chat_id = _chat(db, user_id) class _Local(LocalExecutor): def __init__(self, spec, project_dir): super().__init__(str(tmp_path)) monkeypatch.setattr(ssh_service, "SshExecutor", _Local) monkeypatch.setattr(jobs, "_poll_interval", lambda _e: 0.1) woken: list = [] async def _wake(chat, job_id, command, status, code, body): woken.append((status, code, body)) monkeypatch.setattr(jobs, "wake", _wake) agent = _agent(tmp_path, chat_id=chat_id) agent.executor = lambda: _Local(None, str(tmp_path)) job = await jobs.launch(agent, "sleep 0.5; echo finished") jobs.start_watch(agent, job) for _ in range(50): if woken: break await asyncio.sleep(0.1) assert woken, "the watcher never woke the model" status, code, body = woken[0] assert status == "done" assert "finished" in body