Wake the model when a background job finishes
The other half of background execution: a job that finishes while nobody is looking prompts the model back with its result, rather than sitting unread until the model happens to run again. The vehicle is the queue, because it is the only wiring that already delivers a turn into or after a reply. A per-job poller notices completion and calls jobs.wake. If a reply is being written the completion is left queued for that reply's _inject/_drain; if the chat is idle a fresh reply is started to answer it -- the send_queued_now move. All of it under a per-chat 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 already live and leaves its completion for it. That is the invariant the queue exists to hold, reached from outside a request for the first time. The completion is a user-role turn whose content names itself a machine event -- "A background job you started has finished" -- not a bare person turn. _inject sends a queued turn verbatim, so the framing cannot live there; it lives in the words, the way execute_plan quotes the plan, and a tool.background fragment tells the model these arrive and are a machine event rather than the person speaking. The poller reconnects a fresh connection each tick rather than holding one open -- holding one is the exact live-connection state the whole ssh.py/base.py design forbids, and poll is self-healing besides. Bounded by background_max_jobs and a six-hour ceiling, after which the remote job may keep running but we stop watching it. A Job table, and here the terminal/generation "lost on restart" precedent does NOT transfer: those are seconds long with a human watching, a background job is hours long with nobody watching -- the one case a restart forgetting it would silently break the feature's whole promise. So the row lets a lifespan startup hook rehydrate the watcher and wake as if nothing happened. Cancelling a watcher never stops the detached remote job; it runs on and is picked back up. Tested end to end against a real local shell: launch a detached command, poll it to completion through a watcher, and assert the model was woken with the exit code and output -- plus the lock proving two simultaneous completions start one reply, not two. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -20,7 +20,7 @@ lembas info # paths + counts, useful when confused
|
|||||||
lembas secret-key # generate LEMBAS_SECRET_KEY
|
lembas secret-key # generate LEMBAS_SECRET_KEY
|
||||||
lembas create-admin # create or promote an admin
|
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
|
# PLAN.md tracks what is and is not built
|
||||||
ruff check . # lint (line length 100)
|
ruff check . # lint (line length 100)
|
||||||
python scripts/build_artwork.py # regenerate artwork (SVG + PWA icons;
|
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
|
other one: on a Debian-derived host `apt-get install` reports the package missing
|
||||||
until `apt-get update` has run.
|
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
|
**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
|
*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`
|
path in a command line is unavoidably a quoting problem. `file_read`/`file_write`
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ from lembas.db.models.agent import (
|
|||||||
AUTH_KEY,
|
AUTH_KEY,
|
||||||
AUTH_METHODS,
|
AUTH_METHODS,
|
||||||
AUTH_PASSWORD,
|
AUTH_PASSWORD,
|
||||||
|
Job,
|
||||||
SshProfile,
|
SshProfile,
|
||||||
)
|
)
|
||||||
from lembas.db.models.attachment import (
|
from lembas.db.models.attachment import (
|
||||||
@@ -111,6 +112,7 @@ __all__ = [
|
|||||||
"SOURCE_LINK",
|
"SOURCE_LINK",
|
||||||
"SOURCE_UPLOAD",
|
"SOURCE_UPLOAD",
|
||||||
"Chat",
|
"Chat",
|
||||||
|
"Job",
|
||||||
"Connection",
|
"Connection",
|
||||||
"CustomTool",
|
"CustomTool",
|
||||||
"Document",
|
"Document",
|
||||||
|
|||||||
@@ -100,4 +100,34 @@ class SshProfile(UUIDPrimaryKey, Timestamps, Base):
|
|||||||
return f"<SshProfile {self.name} {self.address}>"
|
return f"<SshProfile {self.name} {self.address}>"
|
||||||
|
|
||||||
|
|
||||||
__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"<Job {self.id} {self.status}>"
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["AUTH_KEY", "AUTH_METHODS", "AUTH_PASSWORD", "Job", "SshProfile"]
|
||||||
|
|||||||
@@ -85,12 +85,24 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|||||||
except Exception: # noqa: BLE001 - housekeeping must never block startup
|
except Exception: # noqa: BLE001 - housekeeping must never block startup
|
||||||
log.exception("orphaned upload sweep failed")
|
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("LLeMbas %s starting on http://%s:%s", __version__, settings.host, settings.port)
|
||||||
log.info("data directory: %s", settings.data_dir.resolve())
|
log.info("data directory: %s", settings.data_dir.resolve())
|
||||||
yield
|
yield
|
||||||
|
|
||||||
# Replies still being written are cancelled and persisted with whatever
|
# Replies still being written are cancelled and persisted with whatever
|
||||||
# they have, rather than left as permanently unfinished rows.
|
# 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.agent.terminal import shutdown as stop_terminals
|
||||||
from lembas.services.generation import shutdown as stop_generations
|
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
|
# is cut off mid-command. Every deploy does this, and the panel is told why
|
||||||
# rather than left to guess -- see deploy/README.md.
|
# rather than left to guess -- see deploy/README.md.
|
||||||
await stop_terminals()
|
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")
|
log.info("LLeMbas stopped")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -34,8 +34,10 @@ sketch:
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import base64
|
import base64
|
||||||
import contextlib
|
import contextlib
|
||||||
|
import logging
|
||||||
import re
|
import re
|
||||||
import time
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
@@ -43,6 +45,8 @@ from dataclasses import dataclass, field
|
|||||||
|
|
||||||
from lembas.services.agent.base import ExecError, ExecRequest, clean_output
|
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
|
# 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
|
# 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,
|
# 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.
|
# detached, which is immediate.
|
||||||
LAUNCH_GRACE = 10.0
|
LAUNCH_GRACE = 10.0
|
||||||
|
|
||||||
# In-process, keyed by job id, lost on restart -- the durable record is the `Job`
|
# The working set, keyed by job id: what `job_list` shows this session. Mirrored
|
||||||
# row (added with the watcher). This holds the metadata `job_list` shows within
|
# to a `Job` row for jobs that are watched, so a restart can rehydrate them.
|
||||||
# a session and, later, the watcher task.
|
|
||||||
_JOBS: dict[str, JobState] = {}
|
_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:
|
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)
|
s = _sentinel(job_id)
|
||||||
pid = _file(chat_id, job_id, "pid")
|
pid = _file(chat_id, job_id, "pid")
|
||||||
exit_ = _file(chat_id, job_id, "exit")
|
exit_ = _file(chat_id, job_id, "exit")
|
||||||
log = _file(chat_id, job_id, "log")
|
logf = _file(chat_id, job_id, "log")
|
||||||
return (
|
return (
|
||||||
_launch_lines(chat_id, job_id, command)
|
_launch_lines(chat_id, job_id, command)
|
||||||
+ "while :; do\n"
|
+ "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.
|
# not to be felt, long enough not to spin.
|
||||||
" sleep 0.2\n"
|
" sleep 0.2\n"
|
||||||
"done\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"printf '\\n{s}:'\n"
|
||||||
f"cat {exit_} 2>/dev/null || printf LOST\n"
|
f"cat {exit_} 2>/dev/null || printf LOST\n"
|
||||||
f"rm -f {_file(chat_id, job_id, 'sh')} {pid} {log} {exit_}\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:
|
def _record(job_id: str, status: str, exit_status: int | None) -> None:
|
||||||
job = _JOBS.get(job_id)
|
job = _JOBS.get(job_id)
|
||||||
if job is None:
|
if job is None or job.status != "running":
|
||||||
return
|
return
|
||||||
if status in ("done", "lost", "killed") and job.status == "running":
|
if status in ("done", "lost", "killed"):
|
||||||
job.status = status
|
job.status = status
|
||||||
job.exit_status = exit_status
|
job.exit_status = exit_status
|
||||||
job.finished_at = time.monotonic()
|
job.finished_at = time.monotonic()
|
||||||
|
_persist_row(job)
|
||||||
|
|
||||||
|
|
||||||
def clear() -> None:
|
def clear() -> None:
|
||||||
_JOBS.clear()
|
_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__ = [
|
__all__ = [
|
||||||
"JOB_ROOT",
|
"JOB_ROOT",
|
||||||
"Completed",
|
"Completed",
|
||||||
|
|||||||
@@ -96,6 +96,9 @@ class AgentContext:
|
|||||||
# Whether a finished job wakes the model on its own, rather than only being
|
# 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.
|
# seen when it next runs. Read by the wording here and by the watcher.
|
||||||
background_notify: bool = True
|
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:
|
def executor(self) -> Executor:
|
||||||
return ssh_service.SshExecutor(self.spec, self.project_dir)
|
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=bool(values.get("background_enabled")),
|
||||||
background_on_timeout=bool(values.get("background_on_timeout", True)),
|
background_on_timeout=bool(values.get("background_on_timeout", True)),
|
||||||
background_notify=bool(values.get("background_notify", 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),
|
spec=ssh_service.spec_from(profile),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -190,6 +190,7 @@ async def _run_convertible(
|
|||||||
if result.timed_out:
|
if result.timed_out:
|
||||||
job = jobs.JobState(id=job_id, chat_id=agent.chat_id, command=command)
|
job = jobs.JobState(id=job_id, chat_id=agent.chat_id, command=command)
|
||||||
jobs.register(job)
|
jobs.register(job)
|
||||||
|
jobs.start_watch(agent, job)
|
||||||
return _backgrounded(agent, command, job, converted=True, timeout=timeout)
|
return _backgrounded(agent, command, job, converted=True, timeout=timeout)
|
||||||
|
|
||||||
# It finished. The wrapper already tailed the log remotely; strip ANSI here.
|
# 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(
|
return ToolOutcome(
|
||||||
exc.message, _event("shell_run", agent, command, status="error", error=exc.message)
|
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)
|
return _backgrounded(agent, command, job, converted=False, timeout=0)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -148,6 +148,7 @@ def context_variables(
|
|||||||
"agent_dir": "",
|
"agent_dir": "",
|
||||||
"agent_mode": "",
|
"agent_mode": "",
|
||||||
"agent_rewound": "",
|
"agent_rewound": "",
|
||||||
|
"background": "",
|
||||||
"project_files": "",
|
"project_files": "",
|
||||||
"agent_instructions": "",
|
"agent_instructions": "",
|
||||||
"agent_instructions_file": "",
|
"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_dir": context.project_dir or "the login directory",
|
||||||
"agent_mode": policy.MODE_GUIDANCE.get(context.mode, ""),
|
"agent_mode": policy.MODE_GUIDANCE.get(context.mode, ""),
|
||||||
"agent_rewound": rewound,
|
"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),
|
"max_rounds": str(context.limits.steps),
|
||||||
# Blanked, which is what makes `core.rounds` vanish here: `steps` is a
|
# 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
|
# runaway backstop and telling a model it has a budget of two hundred
|
||||||
|
|||||||
@@ -170,6 +170,12 @@ VARIABLES: tuple[Variable, ...] = (
|
|||||||
"built, when the feature is off, or when the directory could not be "
|
"built, when the feature is off, or when the directory could not be "
|
||||||
"read -- and the section it lives in disappears with it.",
|
"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(
|
Variable(
|
||||||
"plan",
|
"plan",
|
||||||
"The current plan",
|
"The current plan",
|
||||||
@@ -956,6 +962,28 @@ BUILTIN: tuple[Fragment, ...] = (
|
|||||||
"not look for another way round it."
|
"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(
|
Fragment(
|
||||||
key="tool.project_files",
|
key="tool.project_files",
|
||||||
label="What is in the project directory",
|
label="What is in the project directory",
|
||||||
|
|||||||
@@ -281,3 +281,134 @@ def test_background_param_appears_only_when_enabled(tmp_path):
|
|||||||
off = shell(_agent(tmp_path, background=False)).parameters["properties"]
|
off = shell(_agent(tmp_path, background=False)).parameters["properties"]
|
||||||
assert "background" in on
|
assert "background" in on
|
||||||
assert "background" not in off
|
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
|
||||||
|
|||||||
@@ -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):
|
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" in harness.compose(db, owner, _tools("fetch"))
|
||||||
assert "read one web page at a time" not in harness.compose(db, owner, _tools("web_search"))
|
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()
|
||||||
|
|||||||
Reference in New Issue
Block a user