diff --git a/CLAUDE.md b/CLAUDE.md index a23cf50..7ee36ec 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -20,7 +20,7 @@ lembas info # paths + counts, useful when confused lembas secret-key # generate LEMBAS_SECRET_KEY lembas create-admin # create or promote an admin -pytest # 1206 tests, ~70s +pytest # 1231 tests, ~75s # PLAN.md tracks what is and is not built ruff check . # lint (line length 100) python scripts/build_artwork.py # regenerate artwork (SVG + PWA icons; @@ -354,6 +354,39 @@ cause of "the agent seems stupid", and the harness says it out loud. So does the other one: on a Debian-derived host `apt-get install` reports the package missing until `apt-get update` has run. +**A command can outlive the reply, and that is the one place the fresh-shell +model is fought rather than obeyed.** `services/agent/jobs.py`: a background job +is a `setsid`-detached process on the far side, redirected to a remote logfile +and an exit-file, so it survives the connection closing; LLeMbas reconnects (a +fresh connection, as always) to read it. Opt-in, off by default. When on, the +same wrapper runs *every* command: it launches detached and waits, and a command +that outlasts its timeout is kept running as a job rather than killed. Three +things in the wrappers are load-bearing and were each got wrong first: the +command is **base64'd into a script file**, never put in a quoted `sh -c '…'` +(which shatters on `git commit -m 'fix'` and is an injection hole); the child +records its **own pid via `$$`** under `setsid` as the group leader, so +`job_stop` kills the whole group; and the exit status is read from the +**exit-file, not the wrapper's own status**, which is ~0 from its trailing `rm`. +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. + +**"Prompt the model back when a job finishes" reuses the queue.** A per-job +poller (`jobs._watch`, a fresh connection per tick — never a held one, that +being the thing the whole subsystem forbids) notices completion and calls +`jobs.wake`. Wake writes the completion as a **user-role turn whose content names +itself a machine event** — `_inject` sends a queued turn verbatim, so the framing +lives in the words, the way `execute_plan` quotes the plan, and `tool.background` +tells the model these arrive. If a reply is running the completion is left +`queued` for its `_inject`/`_drain`; if the chat is idle a fresh reply is started +(the `send_queued_now` move). All of it is under a **per-chat `asyncio.Lock` with +no `await` between the running-check and `ensure`**, so two jobs finishing at +once cannot each spin up a generation — the second sees the first's reply live +and leaves its completion for it. The `Job` table exists for one reason the +terminal/generation "lost on restart" precedent does *not* cover: a job runs for +hours with nobody watching, so a restart rehydrates its watcher from the row +(`jobs.rehydrate`, in the lifespan) rather than forgetting the one thing the +feature promises. Cancelling a watcher never stops the detached remote job. + **Files never go through a shell.** The SSH exec protocol carries one command *string* that the far side parses, with no argv form at all, so a model-supplied path in a command line is unavoidably a quoting problem. `file_read`/`file_write` diff --git a/src/lembas/db/models/__init__.py b/src/lembas/db/models/__init__.py index 8f7a313..91601bf 100644 --- a/src/lembas/db/models/__init__.py +++ b/src/lembas/db/models/__init__.py @@ -9,6 +9,7 @@ from lembas.db.models.agent import ( AUTH_KEY, AUTH_METHODS, AUTH_PASSWORD, + Job, SshProfile, ) from lembas.db.models.attachment import ( @@ -111,6 +112,7 @@ __all__ = [ "SOURCE_LINK", "SOURCE_UPLOAD", "Chat", + "Job", "Connection", "CustomTool", "Document", diff --git a/src/lembas/db/models/agent.py b/src/lembas/db/models/agent.py index 47d4463..c391f88 100644 --- a/src/lembas/db/models/agent.py +++ b/src/lembas/db/models/agent.py @@ -100,4 +100,34 @@ class SshProfile(UUIDPrimaryKey, Timestamps, Base): return f"" -__all__ = ["AUTH_KEY", "AUTH_METHODS", "AUTH_PASSWORD", "SshProfile"] +class Job(Timestamps, Base): + """A command left running on the far side after the reply that started it. + + The durable record behind `services/agent/jobs.py`, which otherwise keeps + only an in-process registry lost on restart. A background job runs for + minutes to hours with nobody watching -- exactly the case a restart must not + forget -- so the row lets a startup hook re-poll the job's deterministic + exit-file and wake the model as if nothing had happened. + + The id is `jobs`'s own short hex, not a UUIDPrimaryKey, because the same id + names the files on the machine and is quoted back by the model. + """ + + __tablename__ = "agent_jobs" + + id: Mapped[str] = mapped_column(String(32), primary_key=True) + chat_id: Mapped[str] = mapped_column( + String(32), ForeignKey("chats.id", ondelete="CASCADE"), index=True, nullable=False + ) + command: Mapped[str] = mapped_column(Text, default="") + # running | done | killed | lost. `lost` means it stopped without an exit + # code being recorded -- killed out of band, or the host rebooted under it. + status: Mapped[str] = mapped_column(String(16), default="running", nullable=False) + exit_status: Mapped[int | None] = mapped_column(Integer) + finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + + def __repr__(self) -> str: + return f"" + + +__all__ = ["AUTH_KEY", "AUTH_METHODS", "AUTH_PASSWORD", "Job", "SshProfile"] diff --git a/src/lembas/main.py b/src/lembas/main.py index 07c31a1..fc86502 100644 --- a/src/lembas/main.py +++ b/src/lembas/main.py @@ -85,12 +85,24 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: except Exception: # noqa: BLE001 - housekeeping must never block startup log.exception("orphaned upload sweep failed") + # Background jobs that were still running when we last stopped keep running + # on their own hosts; pick their watchers back up so the model is still + # woken when they finish. Best-effort, and inside the loop so its tasks land + # in this event loop. + try: + from lembas.services.agent.jobs import rehydrate as rehydrate_jobs + + rehydrate_jobs() + except Exception: # noqa: BLE001 - a job that cannot be rehydrated is not fatal + log.exception("could not rehydrate background jobs") + log.info("LLeMbas %s starting on http://%s:%s", __version__, settings.host, settings.port) log.info("data directory: %s", settings.data_dir.resolve()) yield # Replies still being written are cancelled and persisted with whatever # they have, rather than left as permanently unfinished rows. + from lembas.services.agent.jobs import shutdown as stop_jobs from lembas.services.agent.terminal import shutdown as stop_terminals from lembas.services.generation import shutdown as stop_generations @@ -99,6 +111,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # is cut off mid-command. Every deploy does this, and the panel is told why # rather than left to guess -- see deploy/README.md. await stop_terminals() + # Background jobs are the exception: cancelling a watcher does NOT stop the + # detached remote job, which keeps running and is rehydrated on the next + # start. Only the watching stops here. + await stop_jobs() log.info("LLeMbas stopped") diff --git a/src/lembas/services/agent/jobs.py b/src/lembas/services/agent/jobs.py index 0919f44..02b5fdc 100644 --- a/src/lembas/services/agent/jobs.py +++ b/src/lembas/services/agent/jobs.py @@ -34,8 +34,10 @@ sketch: from __future__ import annotations +import asyncio import base64 import contextlib +import logging import re import time import uuid @@ -43,6 +45,8 @@ from dataclasses import dataclass, field from lembas.services.agent.base import ExecError, ExecRequest, clean_output +log = logging.getLogger(__name__) + # 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, @@ -60,10 +64,23 @@ _ID = re.compile(r"^[a-f0-9]{12}$") # 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. +# The working set, keyed by job id: what `job_list` shows this session. Mirrored +# to a `Job` row for jobs that are watched, so a restart can rehydrate them. _JOBS: dict[str, JobState] = {} +# One watcher task per job being polled to completion. +_WATCHERS: dict[str, asyncio.Task] = {} +# One lock per chat, so two jobs finishing at once cannot each start a reply -- +# see `wake`. +_WAKE_LOCKS: dict[str, asyncio.Lock] = {} + +# Stop watching a job after this. The remote process may keep running; we simply +# stop holding a watcher for it and mark it lost. A job that runs longer than +# this is beyond what auto-wake promises. +MAX_WATCH_SECONDS = 6 * 3600 + +# How much of a finished job's output is put in front of the model when it is +# woken. Capped so a job that printed a gigabyte does not blow the window. +MAX_COMPLETION_CHARS = 4000 def new_id() -> str: @@ -147,7 +164,7 @@ def launch_and_wait_command(chat_id: str, job_id: str, command: str, max_bytes: 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") + logf = _file(chat_id, job_id, "log") return ( _launch_lines(chat_id, job_id, command) + "while :; do\n" @@ -159,7 +176,7 @@ def launch_and_wait_command(chat_id: str, job_id: str, command: str, max_bytes: # 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"tail -c {max_bytes} {logf} 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" @@ -330,18 +347,256 @@ def valid_id(job_id: str) -> bool: def _record(job_id: str, status: str, exit_status: int | None) -> None: job = _JOBS.get(job_id) - if job is None: + if job is None or job.status != "running": return - if status in ("done", "lost", "killed") and job.status == "running": + if status in ("done", "lost", "killed"): job.status = status job.exit_status = exit_status job.finished_at = time.monotonic() + _persist_row(job) def clear() -> None: _JOBS.clear() +# --- Durable record ------------------------------------------------------------ +# Best-effort throughout: a job whose row cannot be written (a test with no real +# chat, a transient database hiccup) still runs and is still tracked in-process; +# it just will not survive a restart, which is the row's only purpose. +def _persist_row(job: JobState) -> None: + from datetime import UTC, datetime + + from lembas.db.models import Job + from lembas.db.session import session_scope + + try: + with session_scope() as db: + row = db.get(Job, job.id) + if row is None: + row = Job(id=job.id, chat_id=job.chat_id) + db.add(row) + row.command = job.command[:4000] + row.status = job.status + row.exit_status = job.exit_status + row.finished_at = None if job.status == "running" else datetime.now(UTC) + except Exception: # noqa: BLE001 - the row is a convenience, not the job + log.debug("could not persist job %s", job.id, exc_info=True) + + +# --- The watcher --------------------------------------------------------------- +def _poll_interval(elapsed: float) -> float: + if elapsed < 30: + return 3.0 + if elapsed < 300: + return 10.0 + return 25.0 + + +def start_watch(agent, job: JobState) -> None: + """Poll a job to completion and, when it finishes, wake the model. + + Only when notify is on -- the watcher's whole job is the wake and the status + update, and without notify the model reads `job_output` itself, which + updates the status anyway. Capped by `background_max_jobs`: past it a job + still runs and can be read, it simply is not watched. + + The credential is copied, not referenced: `generation` clears the agent's + `spec` when the reply ends, and the watcher outlives the reply. Holding the + copy for the job's life is the same trade the terminal makes for a held + shell. + """ + _persist_row(job) + if not agent.background_notify or len(_WATCHERS) >= agent.background_max_jobs: + return + task = asyncio.create_task( + _watch( + dict(agent.spec), + agent.project_dir, + job.chat_id, + job.id, + job.command, + agent.max_output, + ) + ) + _WATCHERS[job.id] = task + + +async def _watch( + spec: dict, project_dir: str, chat_id: str, job_id: str, command: str, max_output: int +) -> None: + from lembas.services.agent.ssh import SshExecutor + + started = time.monotonic() + try: + while True: + await asyncio.sleep(_poll_interval(time.monotonic() - started)) + if time.monotonic() - started > MAX_WATCH_SECONDS: + _record(job_id, "lost", None) + return + try: + result = await SshExecutor(spec, project_dir).run( + ExecRequest( + command=read_command(chat_id, job_id, max_output), + timeout=30, + max_bytes=max_output, + ) + ) + except ExecError: + continue # transient -- the host is briefly unreachable; retry + if result.timed_out: + continue + output, _ = clean_output(result.output or "", limit=max_output) + reading = parse_reading(output, job_id) + if reading.status in ("done", "lost"): + _record(job_id, reading.status, reading.exit_status) + with contextlib.suppress(ExecError): + await SshExecutor(spec, project_dir).run( + ExecRequest(command=cleanup_command(chat_id, job_id), timeout=30) + ) + await wake(chat_id, job_id, command, reading.status, reading.exit_status, + reading.body) + return + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 - a watcher that dies must not take others + log.exception("job watcher for %s raised", job_id) + finally: + _WATCHERS.pop(job_id, None) + + +# --- Waking the model ---------------------------------------------------------- +def _lock(chat_id: str) -> asyncio.Lock: + lock = _WAKE_LOCKS.get(chat_id) + if lock is None: + lock = _WAKE_LOCKS[chat_id] = asyncio.Lock() + return lock + + +def _completion_text( + job_id: str, command: str, status: str, exit_status: int | None, output: str +) -> str: + if status == "done" and exit_status == 0: + line = "It finished successfully." + elif status == "done": + line = f"It exited {exit_status}." + else: + line = "It stopped without an exit status (it may have been killed)." + body = (output or "").strip()[:MAX_COMPLETION_CHARS] + # A fence for the model's benefit; backticks in the output are neutralised so + # they cannot close it, the same move `instructions.clean` makes. + fenced = f"\n\n```\n{body.replace('```', chr(39) * 3)}\n```" if body else "" + return ( + f"A background job you started has finished — this is a machine event, " + f"not the person speaking.\n\n" + f"[job {job_id}] `{command}`\n{line}{fenced}" + ) + + +async def wake( + chat_id: str, job_id: str, command: str, status: str, exit_status: int | None, output: str +) -> None: + """Tell the model a job finished, as a new turn. + + Reuses the queue: if a reply is being written, the completion is left + `queued` for that reply's `_inject`/`_drain` to deliver; if the chat is idle, + a fresh reply is started to answer it, the `send_queued_now` move. + + The whole thing is under a per-chat lock, and there is no `await` between the + running-check and starting the reply, so two jobs finishing at once cannot + each spin up a generation -- the second sees the first's reply already live + and leaves its completion for it. That is the invariant the queue exists to + hold, reached here from outside a request. + + The completion is a user-role turn whose *content* names itself a machine + event -- `_inject` sends a queued turn verbatim, so the framing cannot live + there; it lives in the words, the way `execute_plan` quotes the plan. + """ + from lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Chat + from lembas.db.session import session_scope + from lembas.services import chat as chat_service + from lembas.services import generation as generation_service + + content = _completion_text(job_id, command, status, exit_status, output) + async with _lock(chat_id): + running = generation_service.running_for(chat_id) is not None + assistant_id = "" + try: + with session_scope() as db: + chat = db.get(Chat, chat_id) + if chat is None: + return + chat_service.create_message(db, chat, ROLE_USER, content, queued=running) + if not running: + assistant = chat_service.create_message( + db, chat, ROLE_ASSISTANT, "", complete_=False, model_id=chat.model_id + ) + assistant_id = assistant.id + except Exception: # noqa: BLE001 - a failed wake must not crash the watcher + log.exception("could not wake chat %s for job %s", chat_id, job_id) + return + if assistant_id: + generation_service.ensure(chat_id, assistant_id) + + +# --- Rehydration and shutdown -------------------------------------------------- +def rehydrate() -> None: + """After a restart, watch again the jobs that were still running. + + Their remote files are keyed deterministically on chat and id, so a fresh + watcher re-polls them and wakes the model as if nothing happened -- which is + the whole reason the row exists. Best-effort per job: a host that is down, a + profile that is gone, a chat that was deleted each just drop that one. + """ + from lembas.db.models import Chat, SshProfile + from lembas.db.session import session_scope + from lembas.services import settings_store + from lembas.services.agent import ssh as ssh_service + + with session_scope() as db: + values = settings_store.agents(db) + if not values.get("enabled") or not values.get("background_notify"): + return + max_output = int(values.get("max_output_bytes") or 64 * 1024) + running = list(db.scalars(_running_rows())) + for row in running: + chat = db.get(Chat, row.chat_id) + if chat is None or not chat.ssh_profile_id: + continue + profile = db.get(SshProfile, chat.ssh_profile_id) + if profile is None or not profile.enabled: + continue + spec = ssh_service.spec_from(profile) + project_dir = chat.project_dir or profile.default_dir or "" + job = JobState(id=row.id, chat_id=row.chat_id, command=row.command) + _JOBS[job.id] = job + if len(_WATCHERS) >= int(values.get("background_max_jobs") or 5): + break + _WATCHERS[job.id] = asyncio.create_task( + _watch(spec, project_dir, row.chat_id, row.id, row.command, max_output) + ) + + +def _running_rows(): + from sqlalchemy import select + + from lembas.db.models import Job + + return select(Job).where(Job.status == "running") + + +async def shutdown() -> None: + """Cancel every watcher. The detached remote jobs are unaffected -- they run + on, and a later start rehydrates them from their rows.""" + tasks = list(_WATCHERS.values()) + _WATCHERS.clear() + for task in tasks: + task.cancel() + for task in tasks: + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + + __all__ = [ "JOB_ROOT", "Completed", diff --git a/src/lembas/services/agent/session.py b/src/lembas/services/agent/session.py index d6cb9e4..97ea919 100644 --- a/src/lembas/services/agent/session.py +++ b/src/lembas/services/agent/session.py @@ -96,6 +96,9 @@ class AgentContext: # Whether a finished job wakes the model on its own, rather than only being # seen when it next runs. Read by the wording here and by the watcher. background_notify: bool = True + # Most jobs watched at once. A watcher is a periodic reconnect, so this is a + # real resource; past it a job still runs but is not watched or woken for. + background_max_jobs: int = 5 def executor(self) -> Executor: return ssh_service.SshExecutor(self.spec, self.project_dir) @@ -192,6 +195,7 @@ def resolve(db: DBSession, chat: Chat, user: User | None) -> AgentContext | None background=bool(values.get("background_enabled")), background_on_timeout=bool(values.get("background_on_timeout", True)), background_notify=bool(values.get("background_notify", True)), + background_max_jobs=int(values.get("background_max_jobs") or 5), spec=ssh_service.spec_from(profile), ) diff --git a/src/lembas/services/agent/tools.py b/src/lembas/services/agent/tools.py index e1799ac..490fcb3 100644 --- a/src/lembas/services/agent/tools.py +++ b/src/lembas/services/agent/tools.py @@ -190,6 +190,7 @@ async def _run_convertible( if result.timed_out: job = jobs.JobState(id=job_id, chat_id=agent.chat_id, command=command) jobs.register(job) + jobs.start_watch(agent, job) return _backgrounded(agent, command, job, converted=True, timeout=timeout) # It finished. The wrapper already tailed the log remotely; strip ANSI here. @@ -208,6 +209,7 @@ async def _run_background(agent: AgentContext, command: str, cwd: str) -> ToolOu return ToolOutcome( exc.message, _event("shell_run", agent, command, status="error", error=exc.message) ) + jobs.start_watch(agent, job) return _backgrounded(agent, command, job, converted=False, timeout=0) diff --git a/src/lembas/services/harness.py b/src/lembas/services/harness.py index 02f7eb8..e0b8f82 100644 --- a/src/lembas/services/harness.py +++ b/src/lembas/services/harness.py @@ -148,6 +148,7 @@ def context_variables( "agent_dir": "", "agent_mode": "", "agent_rewound": "", + "background": "", "project_files": "", "agent_instructions": "", "agent_instructions_file": "", @@ -197,6 +198,9 @@ def _agent_values(db: DBSession, chat, user) -> dict[str, str]: "agent_dir": context.project_dir or "the login directory", "agent_mode": policy.MODE_GUIDANCE.get(context.mode, ""), "agent_rewound": rewound, + # Non-empty only when commands may run in the background, which is what + # gates the fragment telling the model so. + "background": "on" if context.background else "", "max_rounds": str(context.limits.steps), # Blanked, which is what makes `core.rounds` vanish here: `steps` is a # runaway backstop and telling a model it has a budget of two hundred diff --git a/src/lembas/services/prompts.py b/src/lembas/services/prompts.py index 90171fb..67b3d90 100644 --- a/src/lembas/services/prompts.py +++ b/src/lembas/services/prompts.py @@ -170,6 +170,12 @@ VARIABLES: tuple[Variable, ...] = ( "built, when the feature is off, or when the directory could not be " "read -- and the section it lives in disappears with it.", ), + Variable( + "background", + "Background commands allowed", + "Non-empty when a command may run detached. Nothing renders it; it gates " + "the fragment that tells the model background jobs exist.", + ), Variable( "plan", "The current plan", @@ -956,6 +962,28 @@ BUILTIN: tuple[Fragment, ...] = ( "not look for another way round it." ), ), + Fragment( + key="tool.background", + label="Long commands", + group=GROUP_TOOLS, + order=251, + families=("agent",), + requires=("background",), + hint="Appears only when background commands are enabled. Tells the model " + "the long-command escape hatch exists and that a completion arrives as " + "a new turn -- and that that turn is a machine event, not the person, " + "the same distinction core.interjection draws for a typed message.", + default=( + "- A command that would take a while — an install, a build, a download — " + "can run in the background: pass `background: true`, or just let it run and " + "it is kept going rather than killed when it reaches its timeout. It keeps " + "running after this reply. Read it with job_output, stop it with job_stop.\n" + "- When a background job finishes you are told in a new turn that begins " + "\"A background job you started has finished\". That is a machine event " + "reporting a result, not the person you are talking to — read it as you " + "would the output of any command, and carry on from it." + ), + ), Fragment( key="tool.project_files", label="What is in the project directory", diff --git a/tests/test_agent_jobs.py b/tests/test_agent_jobs.py index 5695c8f..c4f37dc 100644 --- a/tests/test_agent_jobs.py +++ b/tests/test_agent_jobs.py @@ -281,3 +281,134 @@ def test_background_param_appears_only_when_enabled(tmp_path): off = shell(_agent(tmp_path, background=False)).parameters["properties"] assert "background" in on assert "background" not in off + + +# --- Waking the model when a job finishes -------------------------------------- +def _chat(db, user_id) -> 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 diff --git a/tests/test_harness.py b/tests/test_harness.py index 57ce0ea..6a4ccc1 100644 --- a/tests/test_harness.py +++ b/tests/test_harness.py @@ -397,3 +397,38 @@ def test_an_agent_chat_is_told_to_keep_going_instead(db, owner): def test_the_fetch_guidance_appears_only_with_the_tool(db, owner): assert "read one web page at a time" in harness.compose(db, owner, _tools("fetch")) assert "read one web page at a time" not in harness.compose(db, owner, _tools("web_search")) + + +def test_the_background_guidance_appears_only_when_enabled(db, owner): + """Gated on the feature, so an agent chat without background commands is not + told about a tool it does not have.""" + from lembas.db.models import Chat, Connection, Model, SshProfile + + owner.role = "admin" # agent tools need tools.agent, which admins pass + db.commit() + profile = SshProfile( + owner_id=owner.id, name="Box", host="127.0.0.1", port=22, username="t", + host_key="k", host_fingerprint="f", default_dir="/work", + ) + connection = Connection(name="cbg", base_url="http://127.0.0.1:1", api_key_encrypted="") + db.add_all([profile, connection]) + db.commit() + db.add(Model(connection_id=connection.id, model_id="mbg", capabilities_json={"tools": True})) + db.commit() + settings_store.update(db, {"enabled": True}, key=settings_store.AGENTS) + chat = Chat( + user_id=owner.id, model_id="mbg", connection_id=connection.id, kind="agent", + ssh_profile_id=profile.id, project_dir="/work", + ) + db.add(chat) + db.commit() + + def _text(): + offered = tools_service.resolve_tools(db, chat, owner).schemas + return harness.compose(db, owner, offered, chat) + + settings_store.update(db, {"background_enabled": False}, key=settings_store.AGENTS) + assert "run in the background" not in _text() + + settings_store.update(db, {"background_enabled": True}, key=settings_store.AGENTS) + assert "run in the background" in _text()