Compare commits
2 Commits
a5fa982ae3
...
6cffcb357d
| Author | SHA1 | Date | |
|---|---|---|---|
| 6cffcb357d | |||
| 3fc3449726 |
@@ -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`
|
||||
|
||||
@@ -79,6 +79,10 @@ async def save_agents(
|
||||
instructions_enabled: bool = Form(False),
|
||||
instructions_chars: int = Form(4000),
|
||||
nudge_unfinished: bool = Form(False),
|
||||
background_enabled: bool = Form(False),
|
||||
background_on_timeout: bool = Form(False),
|
||||
background_notify: bool = Form(False),
|
||||
background_max_jobs: int = Form(5),
|
||||
) -> Response:
|
||||
settings_store.update(
|
||||
db,
|
||||
@@ -112,6 +116,10 @@ async def save_agents(
|
||||
"instructions_enabled": instructions_enabled,
|
||||
"instructions_chars": min(max(instructions_chars, 0), 20_000),
|
||||
"nudge_unfinished": nudge_unfinished,
|
||||
"background_enabled": background_enabled,
|
||||
"background_on_timeout": background_on_timeout,
|
||||
"background_notify": background_notify,
|
||||
"background_max_jobs": min(max(background_max_jobs, 1), 100),
|
||||
},
|
||||
key=settings_store.AGENTS,
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -100,4 +100,34 @@ class SshProfile(UUIDPrimaryKey, Timestamps, Base):
|
||||
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
|
||||
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")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,616 @@
|
||||
"""Commands that outlive the reply that started them.
|
||||
|
||||
An ordinary `shell_run` is one blocking `conn.run` over a per-call connection
|
||||
(`ssh.py`): when it hits its timeout the command is killed, so a ten-minute
|
||||
`apt install` is impossible. A background job is the same command launched
|
||||
*detached* on the far side -- `setsid`, redirected to a remote logfile and an
|
||||
exit-file -- so it survives the connection closing. LLeMbas reconnects (a fresh
|
||||
connection, as always) to read the log and the exit code later.
|
||||
|
||||
This is the opposite of `terminal.py`, which survives by *holding* a connection
|
||||
open. Here we hold nothing: the whole point of `ssh.py`/`base.py` is that no live
|
||||
connection is kept, and a job that needed one would be a job that broke that.
|
||||
|
||||
**The command never touches a quoted shell context.** `sh -c '<cmd>'` shatters
|
||||
the instant the command contains a `'` -- `git commit -m 'fix'`, `awk '{…}'`,
|
||||
`sed 's/…/…/'` are the common case, not an edge one, and would also be an
|
||||
injection hole. So the command is base64-encoded here in Python and decoded on
|
||||
the far side into a script file; it is bytes, never shell syntax. Only
|
||||
server-generated hex ids and a fixed root ever reach a path.
|
||||
|
||||
Three things make the wrappers correct, and each was got wrong in an earlier
|
||||
sketch:
|
||||
|
||||
* **The child records its own pid via `$$`**, as its first act, under `setsid`
|
||||
where it is the session/group leader -- so `job_stop` can `kill -<pid>` the
|
||||
whole process group. `echo $!` from the launcher captures the wrong pid.
|
||||
* **The exit-file is the primary signal.** An empty pid-file means "still
|
||||
starting", not "dead"; reading liveness first would race the launch and report
|
||||
a job lost the instant it began.
|
||||
* **The command's exit status comes from the exit-file, never from the wrapper's
|
||||
own status** -- which is ~0 from the trailing `rm`. Reading the wrapper's
|
||||
status would mark every job a success.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
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,
|
||||
# which is also what makes cross-chat access structurally impossible: a path is
|
||||
# only ever built from the *calling* chat's id, so a model in one chat cannot
|
||||
# name another chat's files.
|
||||
JOB_ROOT = "${TMPDIR:-/tmp}/lembas-jobs"
|
||||
|
||||
# A job id is our own short hex; anything else is refused before it reaches a
|
||||
# path, so `job_output("../../etc/passwd")` cannot walk out of the job root.
|
||||
_ID = re.compile(r"^[a-f0-9]{12}$")
|
||||
|
||||
# How long the fire-and-return launcher waits for the shell to accept the
|
||||
# command. Not the command's own timeout -- it returns the moment the process is
|
||||
# detached, which is immediate.
|
||||
LAUNCH_GRACE = 10.0
|
||||
|
||||
# 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:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobState:
|
||||
"""What LLeMbas remembers about one background job, in this process."""
|
||||
|
||||
id: str
|
||||
chat_id: str
|
||||
command: str
|
||||
status: str = "running" # running | done | killed | lost
|
||||
exit_status: int | None = None
|
||||
started_at: float = field(default_factory=time.monotonic)
|
||||
finished_at: float = 0.0
|
||||
|
||||
|
||||
# --- Paths and the wrappers ----------------------------------------------------
|
||||
def _dir(chat_id: str) -> str:
|
||||
return f'"{JOB_ROOT}/{chat_id}"'
|
||||
|
||||
|
||||
def _file(chat_id: str, job_id: str, ext: str) -> str:
|
||||
# Double-quoted so `${TMPDIR:-/tmp}` still expands while the whole path stays
|
||||
# one word. The chat id and job id are hex, so nothing here needs escaping.
|
||||
return f'"{JOB_ROOT}/{chat_id}/{job_id}.{ext}"'
|
||||
|
||||
|
||||
def _sentinel(job_id: str) -> str:
|
||||
return f"__LEMBAS_{job_id}__"
|
||||
|
||||
|
||||
def _inner_script(chat_id: str, job_id: str, command: str) -> str:
|
||||
"""The detached program: record the pid, run the command, record the status.
|
||||
|
||||
base64-encoded before it leaves, so `command` is bytes and never shell
|
||||
syntax. `$$` first, because it is the session leader's pid under setsid and
|
||||
`job_stop` kills the group by it. `$?` last, capturing the command's status;
|
||||
it is the file `run`'s own exit status must never be read in place of.
|
||||
"""
|
||||
return (
|
||||
f"echo $$ > {_file(chat_id, job_id, 'pid')}\n"
|
||||
f"{command}\n"
|
||||
f"echo $? > {_file(chat_id, job_id, 'exit')}\n"
|
||||
)
|
||||
|
||||
|
||||
def _blob(chat_id: str, job_id: str, command: str) -> str:
|
||||
raw = _inner_script(chat_id, job_id, command).encode("utf-8")
|
||||
return base64.b64encode(raw).decode("ascii")
|
||||
|
||||
|
||||
def _launch_lines(chat_id: str, job_id: str, command: str) -> str:
|
||||
"""Create the job dir, drop the script, and detach it. No wait."""
|
||||
blob = _blob(chat_id, job_id, command)
|
||||
return (
|
||||
f"mkdir -p {_dir(chat_id)} 2>/dev/null\n"
|
||||
f"printf %s '{blob}' | base64 -d > {_file(chat_id, job_id, 'sh')}\n"
|
||||
f"setsid sh {_file(chat_id, job_id, 'sh')} "
|
||||
f"> {_file(chat_id, job_id, 'log')} 2>&1 < /dev/null &\n"
|
||||
)
|
||||
|
||||
|
||||
def launch_command(chat_id: str, job_id: str, command: str) -> str:
|
||||
"""Fire-and-return: detach the command and stop. Run with a short timeout."""
|
||||
return _launch_lines(chat_id, job_id, command) + "printf started\n"
|
||||
|
||||
|
||||
def launch_and_wait_command(chat_id: str, job_id: str, command: str, max_bytes: int) -> str:
|
||||
"""Detach the command AND wait up to the (asyncssh) timeout for it.
|
||||
|
||||
If it finishes, stdout is the log tail plus a sentinel line carrying the exit
|
||||
code, and the files are removed. If asyncssh times out first the channel is
|
||||
torn down before the `rm`, so the files survive for a later read and the
|
||||
detached process -- new session, redirected, stdin from /dev/null -- keeps
|
||||
running. That torn-down-mid-wait case is exactly "it became a background
|
||||
job".
|
||||
"""
|
||||
s = _sentinel(job_id)
|
||||
pid = _file(chat_id, job_id, "pid")
|
||||
exit_ = _file(chat_id, job_id, "exit")
|
||||
logf = _file(chat_id, job_id, "log")
|
||||
return (
|
||||
_launch_lines(chat_id, job_id, command)
|
||||
+ "while :; do\n"
|
||||
f" [ -f {exit_} ] && break\n"
|
||||
f" __p=$(cat {pid} 2>/dev/null)\n"
|
||||
' [ -n "$__p" ] && ! kill -0 "$__p" 2>/dev/null && break\n'
|
||||
# 0.2s: with the feature on, every ordinary command waits one poll for
|
||||
# the exit-file, so this is added latency on the hot path. Short enough
|
||||
# not to be felt, long enough not to spin.
|
||||
" sleep 0.2\n"
|
||||
"done\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"
|
||||
)
|
||||
|
||||
|
||||
def read_command(chat_id: str, job_id: str, max_bytes: int) -> str:
|
||||
"""The log so far, and whether the job is still running."""
|
||||
s = _sentinel(job_id)
|
||||
pid = _file(chat_id, job_id, "pid")
|
||||
exit_ = _file(chat_id, job_id, "exit")
|
||||
return (
|
||||
f"tail -c {max_bytes} {_file(chat_id, job_id, 'log')} 2>/dev/null\n"
|
||||
f"printf '\\n{s}:'\n"
|
||||
f"if [ -f {exit_} ]; then printf 'done '; cat {exit_};\n"
|
||||
f'elif __p=$(cat {pid} 2>/dev/null); [ -n "$__p" ] && kill -0 "$__p" 2>/dev/null;'
|
||||
" then printf running;\n"
|
||||
"else printf lost; fi\n"
|
||||
)
|
||||
|
||||
|
||||
def stop_command(chat_id: str, job_id: str) -> str:
|
||||
"""Kill the whole process group, then record an exit so a reader is not told
|
||||
the job is merely lost. A killed process never writes its own exit file."""
|
||||
pid = _file(chat_id, job_id, "pid")
|
||||
exit_ = _file(chat_id, job_id, "exit")
|
||||
return (
|
||||
f'__p=$(cat {pid} 2>/dev/null); [ -n "$__p" ] && kill -TERM -"$__p" 2>/dev/null\n'
|
||||
"sleep 0.3\n"
|
||||
f'[ -n "$__p" ] && kill -KILL -"$__p" 2>/dev/null\n'
|
||||
f"[ -f {exit_} ] || echo 143 > {exit_}\n"
|
||||
"printf stopped\n"
|
||||
)
|
||||
|
||||
|
||||
def cleanup_command(chat_id: str, job_id: str) -> str:
|
||||
return (
|
||||
f"rm -f {_file(chat_id, job_id, 'sh')} {_file(chat_id, job_id, 'pid')} "
|
||||
f"{_file(chat_id, job_id, 'log')} {_file(chat_id, job_id, 'exit')}\n"
|
||||
)
|
||||
|
||||
|
||||
# --- Parsing what a wrapper printed --------------------------------------------
|
||||
@dataclass(frozen=True)
|
||||
class Completed:
|
||||
body: str
|
||||
exit_status: int | None # None ⇒ the job was lost (killed without an exit)
|
||||
|
||||
|
||||
def parse_completed(output: str, job_id: str) -> Completed:
|
||||
"""Split a launch-and-wait result into the command's output and its status.
|
||||
|
||||
On the *last* sentinel, because the command's own output could contain a
|
||||
line that looks like one; everything before it is the body, everything after
|
||||
is the exit code the file held.
|
||||
"""
|
||||
marker = f"\n{_sentinel(job_id)}:"
|
||||
at = output.rfind(marker)
|
||||
if at == -1:
|
||||
return Completed(body=output.strip(), exit_status=None)
|
||||
body = output[:at].strip()
|
||||
tail = output[at + len(marker) :].strip()
|
||||
if tail.upper() == "LOST" or not tail:
|
||||
return Completed(body=body, exit_status=None)
|
||||
try:
|
||||
return Completed(body=body, exit_status=int(tail.split()[0]))
|
||||
except (ValueError, IndexError):
|
||||
return Completed(body=body, exit_status=None)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Reading:
|
||||
body: str
|
||||
status: str # running | done | lost
|
||||
exit_status: int | None
|
||||
|
||||
|
||||
def parse_reading(output: str, job_id: str) -> Reading:
|
||||
marker = f"\n{_sentinel(job_id)}:"
|
||||
at = output.rfind(marker)
|
||||
if at == -1:
|
||||
return Reading(body=output.strip(), status="lost", exit_status=None)
|
||||
body = output[:at].strip()
|
||||
tail = output[at + len(marker) :].strip()
|
||||
if tail.startswith("done"):
|
||||
parts = tail.split()
|
||||
code = int(parts[1]) if len(parts) > 1 and parts[1].lstrip("-").isdigit() else None
|
||||
return Reading(body=body, status="done", exit_status=code)
|
||||
if tail == "running":
|
||||
return Reading(body=body, status="running", exit_status=None)
|
||||
return Reading(body=body, status="lost", exit_status=None)
|
||||
|
||||
|
||||
# --- Operations against the machine --------------------------------------------
|
||||
async def launch(agent, command: str, cwd: str = "") -> JobState:
|
||||
"""Detach a command and return immediately. Raises ExecError if it will not
|
||||
even start."""
|
||||
job_id = new_id()
|
||||
result = await agent.executor().run(
|
||||
ExecRequest(
|
||||
command=launch_command(agent.chat_id, job_id, command),
|
||||
cwd=cwd,
|
||||
timeout=LAUNCH_GRACE,
|
||||
max_bytes=agent.max_output,
|
||||
)
|
||||
)
|
||||
if result.timed_out:
|
||||
raise ExecError("The machine did not accept the command in time.")
|
||||
job = JobState(id=job_id, chat_id=agent.chat_id, command=command)
|
||||
_JOBS[job_id] = job
|
||||
return job
|
||||
|
||||
|
||||
async def read(agent, job_id: str) -> Reading:
|
||||
output, _ = _clean(
|
||||
await agent.executor().run(
|
||||
ExecRequest(
|
||||
command=read_command(agent.chat_id, job_id, agent.max_output),
|
||||
timeout=agent.timeout,
|
||||
max_bytes=agent.max_output,
|
||||
)
|
||||
),
|
||||
agent.max_output,
|
||||
)
|
||||
reading = parse_reading(output, job_id)
|
||||
_record(job_id, reading.status, reading.exit_status)
|
||||
if reading.status in ("done", "lost"):
|
||||
await _cleanup(agent, job_id)
|
||||
return reading
|
||||
|
||||
|
||||
async def stop(agent, job_id: str) -> None:
|
||||
await agent.executor().run(
|
||||
ExecRequest(command=stop_command(agent.chat_id, job_id), timeout=agent.timeout)
|
||||
)
|
||||
_record(job_id, "killed", 143)
|
||||
|
||||
|
||||
async def _cleanup(agent, job_id: str) -> None:
|
||||
with contextlib.suppress(ExecError):
|
||||
await agent.executor().run(
|
||||
ExecRequest(command=cleanup_command(agent.chat_id, job_id), timeout=agent.timeout)
|
||||
)
|
||||
|
||||
|
||||
def _clean(result, limit: int) -> tuple[str, bool]:
|
||||
if result.timed_out:
|
||||
return result.output, False
|
||||
return clean_output(result.output or "", limit=limit)
|
||||
|
||||
|
||||
# --- The registry --------------------------------------------------------------
|
||||
def register(job: JobState) -> None:
|
||||
_JOBS[job.id] = job
|
||||
|
||||
|
||||
def get(job_id: str) -> JobState | None:
|
||||
return _JOBS.get(job_id)
|
||||
|
||||
|
||||
def for_chat(chat_id: str) -> list[JobState]:
|
||||
return [j for j in _JOBS.values() if j.chat_id == chat_id]
|
||||
|
||||
|
||||
def valid_id(job_id: str) -> bool:
|
||||
return bool(_ID.match(job_id or ""))
|
||||
|
||||
|
||||
def _record(job_id: str, status: str, exit_status: int | None) -> None:
|
||||
job = _JOBS.get(job_id)
|
||||
if job is None or job.status != "running":
|
||||
return
|
||||
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",
|
||||
"JobState",
|
||||
"Reading",
|
||||
"clear",
|
||||
"for_chat",
|
||||
"get",
|
||||
"launch",
|
||||
"launch_and_wait_command",
|
||||
"new_id",
|
||||
"parse_completed",
|
||||
"read",
|
||||
"register",
|
||||
"stop",
|
||||
"valid_id",
|
||||
]
|
||||
@@ -86,6 +86,19 @@ class AgentContext:
|
||||
# read the same stale plan from the database and the second would lose the
|
||||
# first. This snapshot is what they actually merge into.
|
||||
plan: dict[str, Any] = field(default_factory=dict)
|
||||
# Whether commands may run detached. When off, `shell_run` is byte-for-byte
|
||||
# what it always was and the `job_*` tools are not offered -- a command that
|
||||
# times out is killed, as before. When on, a command can be launched in the
|
||||
# background (or converted to one when it times out) and the model gets the
|
||||
# tools to check on it. `on_timeout` is the sub-switch for the auto-convert.
|
||||
background: bool = False
|
||||
background_on_timeout: bool = True
|
||||
# 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)
|
||||
@@ -179,6 +192,10 @@ def resolve(db: DBSession, chat: Chat, user: User | None) -> AgentContext | None
|
||||
timeout=float(values.get("default_timeout") or 60),
|
||||
max_timeout=float(values.get("max_timeout") or 600),
|
||||
max_output=int(values.get("max_output_bytes") or 64 * 1024),
|
||||
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),
|
||||
)
|
||||
|
||||
|
||||
@@ -23,8 +23,8 @@ import posixpath
|
||||
from typing import Any
|
||||
|
||||
from lembas.services import plans
|
||||
from lembas.services.agent import index, instructions, patch, policy
|
||||
from lembas.services.agent.base import ExecError, ExecRequest
|
||||
from lembas.services.agent import index, instructions, jobs, patch, policy
|
||||
from lembas.services.agent.base import ExecError, ExecRequest, clean_output
|
||||
from lembas.services.agent.session import AgentContext
|
||||
from lembas.services.tools import (
|
||||
RISK_EXECUTE,
|
||||
@@ -126,29 +126,137 @@ async def _run_shell(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
if reason := _permitted(agent, "shell_run", RISK_EXECUTE, command):
|
||||
return _refused("shell_run", agent, command, reason)
|
||||
|
||||
cwd = str(args.get("cwd") or "").strip()
|
||||
timeout = _timeout(args.get("timeout"), agent)
|
||||
background = bool(args.get("background"))
|
||||
|
||||
# The explicit choice: launch detached and return at once.
|
||||
if background and agent.background:
|
||||
return await _run_background(agent, command, cwd)
|
||||
|
||||
# A timed-out command becomes a job only when the feature AND the auto-convert
|
||||
# are both on. Otherwise -- feature off, or auto-convert off -- the plain path
|
||||
# runs, which is byte-for-byte what shell_run always did: killed on timeout,
|
||||
# nothing left running. Routing the plain case through the detached wrapper
|
||||
# would leave an orphan running past a timeout an administrator said to kill.
|
||||
if agent.background and agent.background_on_timeout:
|
||||
return await _run_convertible(agent, command, cwd, timeout)
|
||||
|
||||
return await _run_foreground(agent, command, cwd, timeout)
|
||||
|
||||
|
||||
async def _run_foreground(
|
||||
agent: AgentContext, command: str, cwd: str, timeout: float
|
||||
) -> ToolOutcome:
|
||||
"""One command, run to completion or killed at the timeout. The original."""
|
||||
try:
|
||||
result = await agent.executor().run(
|
||||
ExecRequest(
|
||||
command=command,
|
||||
cwd=str(args.get("cwd") or "").strip(),
|
||||
timeout=timeout,
|
||||
max_bytes=agent.max_output,
|
||||
ExecRequest(command=command, cwd=cwd, timeout=timeout, max_bytes=agent.max_output)
|
||||
)
|
||||
except ExecError as exc:
|
||||
return ToolOutcome(
|
||||
exc.message, _event("shell_run", agent, command, status="error", error=exc.message)
|
||||
)
|
||||
return _shell_outcome(
|
||||
agent,
|
||||
command,
|
||||
result.output.strip(),
|
||||
exit_status=result.exit_status,
|
||||
timed_out=result.timed_out,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
async def _run_convertible(
|
||||
agent: AgentContext, command: str, cwd: str, timeout: float
|
||||
) -> ToolOutcome:
|
||||
"""Launch detached and wait; if it outlasts the timeout, keep it as a job.
|
||||
|
||||
While it finishes in time this is indistinguishable from `_run_foreground` --
|
||||
same output, same wording. The difference is only visible when it does not:
|
||||
instead of being killed, it is left running and handed back as a job id.
|
||||
"""
|
||||
job_id = jobs.new_id()
|
||||
wrapper = jobs.launch_and_wait_command(agent.chat_id, job_id, command, agent.max_output)
|
||||
try:
|
||||
result = await agent.executor().run(
|
||||
ExecRequest(command=wrapper, cwd=cwd, timeout=timeout, max_bytes=agent.max_output)
|
||||
)
|
||||
except ExecError as exc:
|
||||
return ToolOutcome(
|
||||
exc.message, _event("shell_run", agent, command, status="error", error=exc.message)
|
||||
)
|
||||
|
||||
body = result.output.strip()
|
||||
if result.timed_out:
|
||||
head = f"The command was stopped after {timeout:g}s."
|
||||
elif result.exit_status == 0:
|
||||
head = "" if body else "It ran, and printed nothing."
|
||||
else:
|
||||
head = f"It exited {result.exit_status}."
|
||||
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.
|
||||
output, _ = clean_output(result.output or "", limit=agent.max_output)
|
||||
done = jobs.parse_completed(output, job_id)
|
||||
return _shell_outcome(
|
||||
agent, command, done.body, exit_status=done.exit_status, timed_out=False, timeout=timeout
|
||||
)
|
||||
|
||||
|
||||
async def _run_background(agent: AgentContext, command: str, cwd: str) -> ToolOutcome:
|
||||
"""The model asked to background it up front."""
|
||||
try:
|
||||
job = await jobs.launch(agent, command, cwd)
|
||||
except ExecError as exc:
|
||||
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)
|
||||
|
||||
|
||||
def _backgrounded(
|
||||
agent: AgentContext, command: str, job: jobs.JobState, *, converted: bool, timeout: float
|
||||
) -> ToolOutcome:
|
||||
lead = (
|
||||
f"Still running after {timeout:g}s, so it was kept running in the background"
|
||||
if converted
|
||||
else "Started in the background"
|
||||
)
|
||||
tail = (
|
||||
" You will be told when it finishes."
|
||||
if agent.background_notify
|
||||
else f' Check on it with job_output("{job.id}").'
|
||||
)
|
||||
return ToolOutcome(
|
||||
f'{lead} as job {job.id}. It keeps running after this reply.{tail}',
|
||||
_event(
|
||||
"shell_run",
|
||||
agent,
|
||||
command,
|
||||
status="ok",
|
||||
text=f"job {job.id} — running in the background",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _shell_outcome(
|
||||
agent: AgentContext,
|
||||
command: str,
|
||||
body: str,
|
||||
*,
|
||||
exit_status: int | None,
|
||||
timed_out: bool,
|
||||
timeout: float,
|
||||
) -> ToolOutcome:
|
||||
if timed_out:
|
||||
head = f"The command was stopped after {timeout:g}s."
|
||||
elif exit_status == 0:
|
||||
head = "" if body else "It ran, and printed nothing."
|
||||
elif exit_status is None:
|
||||
head = "It stopped before its exit status could be read."
|
||||
else:
|
||||
head = f"It exited {exit_status}."
|
||||
|
||||
ok = exit_status == 0 and not timed_out
|
||||
content = f"{head}\n\n{body}".strip() if head else body
|
||||
return ToolOutcome(
|
||||
content or "It ran, and printed nothing.",
|
||||
@@ -156,8 +264,8 @@ async def _run_shell(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
"shell_run",
|
||||
agent,
|
||||
command,
|
||||
status="ok" if result.ok else "error",
|
||||
error="" if result.ok else head,
|
||||
status="ok" if ok else "error",
|
||||
error="" if ok else head,
|
||||
text=body[:MAX_EVENT_CHARS],
|
||||
),
|
||||
)
|
||||
@@ -501,7 +609,131 @@ async def _run_plan_update(context: ToolContext, args: dict[str, Any]) -> ToolOu
|
||||
)
|
||||
|
||||
|
||||
# --- Background jobs -----------------------------------------------------------
|
||||
# A job's files are namespaced by the *calling* chat's id (see jobs.py), and the
|
||||
# read/stop wrappers are always built from `agent.chat_id`, so a model in one
|
||||
# chat cannot name another chat's job -- the path simply would not exist. The id
|
||||
# is still validated as our own hex first, so a crafted id cannot walk out of the
|
||||
# job directory.
|
||||
async def _run_job_output(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
agent = _agent(context)
|
||||
if agent is None:
|
||||
return _no_machine("job_output")
|
||||
job_id = str(args.get("id") or "").strip()
|
||||
if not jobs.valid_id(job_id):
|
||||
return ToolOutcome(
|
||||
"There is no job with that id.",
|
||||
_event("job_output", agent, job_id, status="error", error="Unknown job."),
|
||||
)
|
||||
if reason := _permitted(agent, "job_output", RISK_READ):
|
||||
return _refused("job_output", agent, job_id, reason)
|
||||
|
||||
try:
|
||||
reading = await jobs.read(agent, job_id)
|
||||
except ExecError as exc:
|
||||
return ToolOutcome(
|
||||
exc.message, _event("job_output", agent, job_id, status="error", error=exc.message)
|
||||
)
|
||||
|
||||
if reading.status == "running":
|
||||
head = "Still running."
|
||||
elif reading.status == "done":
|
||||
head = "Finished, and printed nothing." if not reading.body else (
|
||||
"Finished." if reading.exit_status == 0 else f"Finished, exit {reading.exit_status}."
|
||||
)
|
||||
else:
|
||||
head = "No longer running — no exit status was recorded (it may have been killed)."
|
||||
|
||||
content = f"{head}\n\n{reading.body}".strip() if reading.body else head
|
||||
return ToolOutcome(
|
||||
content,
|
||||
_event(
|
||||
"job_output",
|
||||
agent,
|
||||
job_id,
|
||||
status="ok" if reading.status != "lost" else "error",
|
||||
text=reading.body[:MAX_EVENT_CHARS],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _run_job_list(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
agent = _agent(context)
|
||||
if agent is None:
|
||||
return _no_machine("job_list")
|
||||
running = jobs.for_chat(agent.chat_id)
|
||||
if not running:
|
||||
return ToolOutcome(
|
||||
"No background jobs in this conversation.",
|
||||
_event("job_list", agent, "", status="ok", text="none"),
|
||||
)
|
||||
lines = [
|
||||
f"{job.id} [{job.status}] {job.command}"
|
||||
+ (f" (exit {job.exit_status})" if job.exit_status is not None else "")
|
||||
for job in running
|
||||
]
|
||||
body = "\n".join(lines)
|
||||
return ToolOutcome(
|
||||
f"Background jobs:\n{body}",
|
||||
_event("job_list", agent, f"{len(running)} job(s)", status="ok", text=body),
|
||||
)
|
||||
|
||||
|
||||
async def _run_job_stop(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
agent = _agent(context)
|
||||
if agent is None:
|
||||
return _no_machine("job_stop")
|
||||
job_id = str(args.get("id") or "").strip()
|
||||
if not jobs.valid_id(job_id):
|
||||
return ToolOutcome(
|
||||
"There is no job with that id.",
|
||||
_event("job_stop", agent, job_id, status="error", error="Unknown job."),
|
||||
)
|
||||
if reason := _permitted(agent, "job_stop", RISK_EXECUTE):
|
||||
return _refused("job_stop", agent, job_id, reason)
|
||||
|
||||
try:
|
||||
await jobs.stop(agent, job_id)
|
||||
except ExecError as exc:
|
||||
return ToolOutcome(
|
||||
exc.message, _event("job_stop", agent, job_id, status="error", error=exc.message)
|
||||
)
|
||||
return ToolOutcome(
|
||||
f"Stopped job {job_id}.",
|
||||
_event("job_stop", agent, job_id, status="ok", text="stopped"),
|
||||
)
|
||||
|
||||
|
||||
def _no_machine(name: str) -> ToolOutcome:
|
||||
return ToolOutcome(
|
||||
"This conversation is not connected to a machine.",
|
||||
{"name": name, "status": "error", "error": "No connection.", "results": []},
|
||||
)
|
||||
|
||||
|
||||
# --- The definitions -----------------------------------------------------------
|
||||
def _shell_parameters(background_on: bool) -> dict[str, Any]:
|
||||
properties: dict[str, Any] = {
|
||||
"command": {**_STRING, "description": "The command line to run."},
|
||||
"cwd": {**_STRING, "description": "Where to run it. Defaults to the project directory."},
|
||||
"timeout": {
|
||||
"type": "number",
|
||||
"description": "Seconds to allow. Bounded by the instance settings.",
|
||||
},
|
||||
}
|
||||
if background_on:
|
||||
properties["background"] = {
|
||||
"type": "boolean",
|
||||
"description": (
|
||||
"Run it detached instead of waiting. It keeps running after this "
|
||||
"reply; check on it with job_output. Use it for a long install, "
|
||||
"build or download. A command left waiting is also kept running "
|
||||
"as a job rather than killed when it hits its timeout."
|
||||
),
|
||||
}
|
||||
return {"type": "object", "properties": properties, "required": ["command"]}
|
||||
|
||||
|
||||
def tool_defs(context: AgentContext | None = None) -> list[ToolDef]:
|
||||
"""The agent tools, bound to one chat's machine.
|
||||
|
||||
@@ -525,21 +757,7 @@ def tool_defs(context: AgentContext | None = None) -> list[ToolDef]:
|
||||
"answer a prompt, so pass the flags that make a command "
|
||||
"non-interactive rather than waiting for it to ask."
|
||||
),
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {**_STRING, "description": "The command line to run."},
|
||||
"cwd": {
|
||||
**_STRING,
|
||||
"description": "Where to run it. Defaults to the project directory.",
|
||||
},
|
||||
"timeout": {
|
||||
"type": "number",
|
||||
"description": "Seconds to allow. Bounded by the instance settings.",
|
||||
},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
parameters=_shell_parameters(bool(context and context.background)),
|
||||
run=_run_shell,
|
||||
risk=RISK_EXECUTE,
|
||||
),
|
||||
@@ -779,6 +997,44 @@ def tool_defs(context: AgentContext | None = None) -> list[ToolDef]:
|
||||
# about it would mean an approval card per ticked-off task.
|
||||
risk=RISK_READ,
|
||||
),
|
||||
ToolDef(
|
||||
name="job_output",
|
||||
family=FAMILY_AGENT,
|
||||
description=(
|
||||
"Read what a background job has printed so far, and whether it is "
|
||||
"still running. Give the id shell_run returned when it was "
|
||||
"backgrounded."
|
||||
),
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"id": {**_STRING, "description": "The job id."}},
|
||||
"required": ["id"],
|
||||
},
|
||||
run=_run_job_output,
|
||||
risk=RISK_READ,
|
||||
),
|
||||
ToolDef(
|
||||
name="job_list",
|
||||
family=FAMILY_AGENT,
|
||||
description="List the background jobs in this conversation and their state.",
|
||||
parameters={"type": "object", "properties": {}, "required": []},
|
||||
run=_run_job_list,
|
||||
risk=RISK_READ,
|
||||
),
|
||||
ToolDef(
|
||||
name="job_stop",
|
||||
family=FAMILY_AGENT,
|
||||
description="Stop a background job, killing it and everything it started.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"id": {**_STRING, "description": "The job id."}},
|
||||
"required": ["id"],
|
||||
},
|
||||
run=_run_job_stop,
|
||||
# It terminates a process on the machine, so it goes through the mode
|
||||
# table exactly as shell_run does.
|
||||
risk=RISK_EXECUTE,
|
||||
),
|
||||
]
|
||||
if context is None:
|
||||
return defs
|
||||
@@ -790,6 +1046,10 @@ def tool_defs(context: AgentContext | None = None) -> list[ToolDef]:
|
||||
drop = {"plan_submit"} if context.mode != policy.MODE_PLAN else {"plan_update"}
|
||||
if not context.plan:
|
||||
drop.add("plan_update")
|
||||
# The job tools exist only when commands may run in the background. Offering
|
||||
# them otherwise is a tool for checking on something that can never exist.
|
||||
if not context.background:
|
||||
drop.update({"job_output", "job_list", "job_stop"})
|
||||
return [tool for tool in defs if tool.name not in drop]
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -137,6 +137,21 @@ def _agents_defaults() -> dict[str, Any]:
|
||||
# one thing there is to be objectively wrong about -- a model with no
|
||||
# plan that says it has finished is believed.
|
||||
"nudge_unfinished": True,
|
||||
# Whether a command may run detached, keep running after the reply ends,
|
||||
# and be checked on later. Off by default, and off means byte-for-byte
|
||||
# the old behaviour: a command that times out is killed. See
|
||||
# services/agent/jobs.py.
|
||||
"background_enabled": False,
|
||||
# The sub-switch: a timed-out command is left running as a job instead
|
||||
# of killed. Off leaves the timeout a hard stop and offers only the
|
||||
# model's explicit `background=true`.
|
||||
"background_on_timeout": True,
|
||||
# Whether the model is woken with the result when a job finishes, rather
|
||||
# than only seeing it when it next runs of its own accord.
|
||||
"background_notify": True,
|
||||
# Most background jobs watched at once. Each is a periodic reconnect to
|
||||
# the far side, so it is a real cost, not a scruple.
|
||||
"background_max_jobs": 5,
|
||||
}
|
||||
|
||||
|
||||
@@ -323,4 +338,5 @@ def agents(db: DBSession) -> dict[str, Any]:
|
||||
values["max_completion_tokens"] = min(
|
||||
max(int(values.get("max_completion_tokens") or 0), 0), 5_000_000
|
||||
)
|
||||
values["background_max_jobs"] = min(max(int(values.get("background_max_jobs") or 0), 1), 100)
|
||||
return values
|
||||
|
||||
@@ -47,6 +47,9 @@ LABELS: dict[str, str] = {
|
||||
"file_list": "List",
|
||||
"plan_submit": "Plan",
|
||||
"plan_update": "Plan updated",
|
||||
"job_output": "Job output",
|
||||
"job_list": "Jobs",
|
||||
"job_stop": "Job stopped",
|
||||
# The web.
|
||||
"web_search": "Web search",
|
||||
"fetch": "Fetch",
|
||||
@@ -77,6 +80,9 @@ ICONS: dict[str, str] = {
|
||||
"file_list": "folder",
|
||||
"plan_submit": "check",
|
||||
"plan_update": "check",
|
||||
"job_output": "clock",
|
||||
"job_list": "dots",
|
||||
"job_stop": "stop-circle",
|
||||
"web_search": "globe",
|
||||
"fetch": "link",
|
||||
"knowledge_search": "archive",
|
||||
@@ -126,6 +132,7 @@ ACTIONS: dict[str, str] = {
|
||||
"skill_get": "Read a skill",
|
||||
"skill_create": "Write a skill",
|
||||
"skill_edit": "Change a skill",
|
||||
"job_stop": "Stop a background job",
|
||||
}
|
||||
|
||||
# Which argument is the thing being agreed to. Shown verbatim and escaped on the
|
||||
@@ -141,6 +148,7 @@ DETAIL_KEYS: dict[str, str] = {
|
||||
"web_search": "query",
|
||||
"knowledge_search": "query",
|
||||
"notes_search": "query",
|
||||
"job_stop": "id",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -124,6 +124,61 @@
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<section class="card">
|
||||
<h2 class="card__title">Background commands</h2>
|
||||
<p class="card__lede">
|
||||
A command that would outlast its timeout can be left running instead of
|
||||
killed — detached on the far side, checked on later. It is how a long
|
||||
install, build or download becomes possible at all.
|
||||
</p>
|
||||
<div class="field">
|
||||
<label class="checkbox">
|
||||
<input type="checkbox" name="background_enabled"
|
||||
{{ 'checked' if values.background_enabled }}>
|
||||
<span>Allow commands to run in the background</span>
|
||||
</label>
|
||||
<p class="field__hint">
|
||||
Off means byte-for-byte the old behaviour: a command that hits its
|
||||
timeout is killed. On, a command can be launched detached (or kept
|
||||
running when it times out), writing to a file under
|
||||
<code>/tmp</code> on the machine, and the model gets tools to read and
|
||||
stop it. A detached command's log can grow without bound on the host —
|
||||
that is the host's to contain, as with everything an agent runs there.
|
||||
</p>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label class="checkbox">
|
||||
<input type="checkbox" name="background_on_timeout"
|
||||
{{ 'checked' if values.background_on_timeout }}>
|
||||
<span>Keep a timed-out command running instead of killing it</span>
|
||||
</label>
|
||||
<p class="field__hint">
|
||||
Off leaves the timeout a hard stop; the model can still choose to
|
||||
background a command up front.
|
||||
</p>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label class="checkbox">
|
||||
<input type="checkbox" name="background_notify"
|
||||
{{ 'checked' if values.background_notify }}>
|
||||
<span>Wake the model when a background job finishes</span>
|
||||
</label>
|
||||
<p class="field__hint">
|
||||
On, a finished job starts (or joins) a reply carrying its result. Off,
|
||||
the model only sees it the next time it runs of its own accord.
|
||||
</p>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label class="field__label" for="background_max_jobs">Most jobs watched at once</label>
|
||||
<input class="input" id="background_max_jobs" name="background_max_jobs"
|
||||
value="{{ values.background_max_jobs }}" inputmode="numeric">
|
||||
<p class="field__hint">
|
||||
Each is a periodic reconnect to the machine. Jobs past this still run;
|
||||
they are simply not watched, and the model is not woken for them.
|
||||
</p>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<section class="card">
|
||||
<h2 class="card__title">What one reply may spend</h2>
|
||||
<p class="field__hint">
|
||||
|
||||
@@ -0,0 +1,414 @@
|
||||
"""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
|
||||
|
||||
|
||||
# --- 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):
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user