Jaroslav Beneš
a9aa89b2c1
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>
617 lines
23 KiB
Python
617 lines
23 KiB
Python
"""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",
|
|
]
|