From 650f376934eece8ecdc11f122976ace09c3f1444 Mon Sep 17 00:00:00 2001 From: Wyatt Date: Tue, 21 Jul 2026 22:55:28 -0400 Subject: [PATCH] feat(control): flag-gated headless runner with cross-worker resume (phase 2) Wires the phase-1 headless machinery behind runner=headless (default stays tmux; legacy agents, session_id null, keep the tmux paths): - spawn: branches tmux vs headless.launch; extracts _agent_env (shared with resume - a headless resume is a new process needing identity/ credential env); headless spawns require a task (no idle-REPL mode), enforced at spawn and as a 400 in the API - resume: headless path materializes the session archive from the DB onto whichever worker claimed the command, then claude -p --resume; falls back to a fresh session with DB-re-injected context (visible worker event) when no transcript survives anywhere; refuses while a run is live. Undeliverable resumes now raise -> command FAILED, fixing silent input loss on both runners - kill: headless path flags cancel_requested; the owning supervisor SIGTERMs its own child (cross-worker safe) - worker: stable per-container ids, DB-driven run slots (full workers skip claiming spawn/resume/mise_init, leaving them for less-loaded workers), credsync refresh in the main loop - settings_gen: permissions block (defaultMode + allowlist) for headless runs - -p auto-denies anything that would prompt; hooks remain the hard gate - credsync + migration 0009 (runtime_secrets): login publishes the Fernet-encrypted claude credential bundle; every worker materializes it (merge-safe for local trust state); login_submit pinned to the login_start worker via commands.target_worker Suite 270 -> 290 green, including the cross-worker resume linchpin (clean-HOME materialize + --resume against the fake binary). --- .gitignore | 3 + src/handler/api/routes/agents.py | 11 +- src/handler/api/routes/login.py | 10 +- src/handler/config.py | 4 + src/handler/control/credsync.py | 174 +++++++++++ src/handler/control/settings_gen.py | 21 +- src/handler/control/spawn.py | 193 +++++++++++-- src/handler/control/worker.py | 73 ++++- src/handler/db/repository.py | 32 +++ src/handler/db/tables.py | 12 + .../versions/0009_runtime_secrets.py | 39 +++ tests/test_credsync.py | 110 +++++++ tests/test_headless_run.py | 272 ++++++++++++++++++ tests/test_worker.py | 24 +- tests/test_worker_concurrency.py | 81 ++++++ 15 files changed, 1016 insertions(+), 43 deletions(-) create mode 100644 src/handler/control/credsync.py create mode 100644 src/handler/migrations/versions/0009_runtime_secrets.py create mode 100644 tests/test_credsync.py create mode 100644 tests/test_headless_run.py create mode 100644 tests/test_worker_concurrency.py diff --git a/.gitignore b/.gitignore index e42ddb6..ff72444 100644 --- a/.gitignore +++ b/.gitignore @@ -220,3 +220,6 @@ __marimo__/ # Streamlit .streamlit/secrets.toml .omc/ + +# Handler runtime artifacts +/handler.db diff --git a/src/handler/api/routes/agents.py b/src/handler/api/routes/agents.py index 7de33e3..056eae8 100644 --- a/src/handler/api/routes/agents.py +++ b/src/handler/api/routes/agents.py @@ -12,6 +12,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query, status from sqlalchemy import Connection from sqlalchemy.exc import IntegrityError +from ...config import get_settings from ...db import repository as repo from ..deps import db_conn, require_admin, require_auth from ..schemas import AgentIn, AgentOut, CheckmarkOut, CommandOut, LogEntryOut, SpawnIn @@ -63,13 +64,21 @@ def create_agent(project: str, body: AgentIn, conn: Connection = Depends(db_conn dependencies=[Depends(require_admin)], ) def enqueue_spawn(project: str, body: SpawnIn, conn: Connection = Depends(db_conn)) -> dict: - """Enqueue a spawn; the worker creates the agent row + tmux session and reports back.""" + """Enqueue a spawn; the worker creates the agent row + claude process and reports back.""" _require_project(conn, project) if repo.get_agent_by_name(conn, project, body.name) is not None: raise HTTPException( status.HTTP_409_CONFLICT, detail=f"agent '{body.name}' already exists in project '{project}'", ) + if get_settings().runner == "headless" and not body.task: + # A tmux agent can idle at the REPL awaiting input; a headless `claude -p` run + # with no prompt exits immediately. Reject here (400) instead of letting the + # command fail asynchronously in the worker. + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + detail="a task is required: the headless runner has no idle-REPL mode", + ) payload = body.model_dump(exclude={"name"}, exclude_none=True) return repo.enqueue_command( conn, diff --git a/src/handler/api/routes/login.py b/src/handler/api/routes/login.py index a682d02..fbec472 100644 --- a/src/handler/api/routes/login.py +++ b/src/handler/api/routes/login.py @@ -43,10 +43,18 @@ def enqueue_login_start(conn: Connection = Depends(db_conn)) -> dict: dependencies=[Depends(require_admin)], ) def enqueue_login_submit(body: LoginSubmitIn, conn: Connection = Depends(db_conn)) -> dict: - """Feed the pasted authorization code back into the waiting login session.""" + """Feed the pasted authorization code back into the waiting login session. + + Pinned to the worker that ran ``login_start`` — the live tmux login session exists + only in that container, so with multiple workers any other claimant would find + nothing to paste into. No prior login_start leaves the pin empty (single-worker + deployments behave exactly as before). + """ + started = repo.get_latest_finished_command(conn, "login_start") return repo.enqueue_command( conn, "login_submit", payload={"code": body.code}, requested_by="operator:web", + target_worker=started["claimed_by"] if started else None, ) diff --git a/src/handler/config.py b/src/handler/config.py index eaebd14..bb618a7 100644 --- a/src/handler/config.py +++ b/src/handler/config.py @@ -95,6 +95,10 @@ class Settings(BaseSettings): def protected_branch_set(self) -> set[str]: return {b.strip() for b in self.protected_branches.split(",") if b.strip()} + @property + def headless_allowed_tools_list(self) -> list[str]: + return [t.strip() for t in self.headless_allowed_tools.split(",") if t.strip()] + @property def cors_origin_list(self) -> list[str]: return [o.strip() for o in self.cors_origins.split(",") if o.strip()] diff --git a/src/handler/control/credsync.py b/src/handler/control/credsync.py new file mode 100644 index 0000000..f253b41 --- /dev/null +++ b/src/handler/control/credsync.py @@ -0,0 +1,174 @@ +"""Distribute claude's OAuth credentials to every worker through the database. + +The interactive ``/login`` flow (see :mod:`~handler.control.login`) completes on exactly +one worker container and writes credentials under that container's ``$HOME``. With +multiple headless workers, the others need those files too — and the deployment rule is +"no shared filesystems, the DB is the single source of truth". So: + +- ``upload()`` bundles the credential files, encrypts the bundle with the existing + Fernet secret store (``HANDLER_SECRET_KEY``), and upserts it into ``runtime_secrets``. + Called after a confirmed login, and after a run if claude refreshed the token on disk. +- ``refresh()`` (workers, at startup and periodically) materializes the bundle locally + when the DB copy is newer than what this worker last saw. ``~/.claude.json`` is merged + (only credential-ish top-level keys are taken) so a worker's own onboarding/trust + state written by ``claude_config.ensure_onboarded`` survives; the pure credential file + is written verbatim. + +Everything is best-effort and silently off when no ``HANDLER_SECRET_KEY`` is configured +(single-worker deployments work exactly as before). +""" + +from __future__ import annotations + +import json +import os + +from .. import secretstore +from ..db import repository as repo +from ..db.engine import connection + +SECRET_KEY = "claude_credentials" + +# ~/.claude.json keys that belong to the *account*, not this machine's UI/trust state. +_ACCOUNT_KEYS = ("oauthAccount", "userID", "organization", "account") + + +def _credential_files() -> dict[str, str]: + """Relative path -> absolute path of every claude credential file we bundle.""" + home = os.path.expanduser("~") + return { + ".claude.json": os.path.join(home, ".claude.json"), + ".claude/.credentials.json": os.path.join(home, ".claude", ".credentials.json"), + } + + +def fingerprint() -> tuple: + """(path, mtime_ns, size) of the on-disk credential files — cheap change detection.""" + fp = [] + for path in sorted(_credential_files().values()): + try: + st = os.stat(path) + fp.append((path, st.st_mtime_ns, st.st_size)) + except OSError: + continue + return tuple(fp) + + +def upload() -> bool: + """Encrypt the local credential files into ``runtime_secrets``. Returns whether a + bundle was stored (False when disabled or there is nothing to store).""" + if not secretstore.enabled(): + return False + bundle: dict[str, str] = {} + for rel, path in _credential_files().items(): + try: + with open(path) as fh: + bundle[rel] = fh.read() + except OSError: + continue + if not bundle: + return False + value_enc = secretstore.encrypt(json.dumps(bundle)) + with connection() as conn: + repo.upsert_runtime_secret(conn, SECRET_KEY, value_enc) + return True + + +def _merge_claude_json(path: str, incoming: str) -> None: + """Take the account keys from ``incoming`` into ``path``, preserving local state. + + A worker's ``~/.claude.json`` also carries per-directory trust + onboarding flags for + *its* clones; clobbering those would re-wedge spawns on the trust dialog. + """ + try: + new_data = json.loads(incoming) + except (ValueError, TypeError): + return + if not isinstance(new_data, dict): + return + current: dict = {} + if os.path.exists(path): + try: + with open(path) as fh: + current = json.load(fh) + except (OSError, ValueError): + current = {} + if not isinstance(current, dict): + current = {} + if not current: + merged = new_data + else: + merged = current + for key in _ACCOUNT_KEYS: + if key in new_data: + merged[key] = new_data[key] + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + tmp = f"{path}.credsync.tmp" + with open(tmp, "w") as fh: + json.dump(merged, fh, indent=2) + os.replace(tmp, path) + + +class _State: + """Per-process sync cursor: the DB updated_at we last materialized, and the local + file fingerprint after the last upload/materialize (to detect claude refreshing + the token mid-run so the next refresh() re-uploads).""" + + def __init__(self) -> None: + self.seen_updated_at = None + self.last_fingerprint: tuple | None = None + + +_state = _State() + + +def refresh() -> str | None: + """One sync pass; returns "materialized", "uploaded", or None (no-op). + + - A DB bundle newer than what we've seen wins: materialize it locally. + - Otherwise, local credential files that changed since our last pass (a login on + this worker, or claude refreshing a token during a run) get uploaded. + """ + if not secretstore.enabled(): + return None + with connection() as conn: + row = repo.get_runtime_secret(conn, SECRET_KEY) + + if row is not None and row["updated_at"] != _state.seen_updated_at: + try: + bundle = json.loads(secretstore.decrypt(row["value_enc"])) + except (secretstore.SecretStoreError, ValueError): + return None + files = _credential_files() + for rel, content in bundle.items(): + path = files.get(rel) + if path is None: + continue # never write outside the known credential set + if rel == ".claude.json": + _merge_claude_json(path, content) + else: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + tmp = f"{path}.credsync.tmp" + with open(tmp, "w") as fh: + fh.write(content) + os.replace(tmp, path) + _state.seen_updated_at = row["updated_at"] + _state.last_fingerprint = fingerprint() + return "materialized" + + current = fingerprint() + if current and current != _state.last_fingerprint: + if upload(): + _state.last_fingerprint = current + with connection() as conn: + stored = repo.get_runtime_secret(conn, SECRET_KEY) + _state.seen_updated_at = stored["updated_at"] if stored else None + return "uploaded" + return None + + +def note_local_write() -> None: + """Record that this process just changed local credentials deliberately (e.g. the + claude_config onboarding merge at spawn), so refresh() doesn't misread the mtime + bump as a new login and ping-pong uploads between workers.""" + _state.last_fingerprint = fingerprint() diff --git a/src/handler/control/settings_gen.py b/src/handler/control/settings_gen.py index c268e4a..f8c1ea1 100644 --- a/src/handler/control/settings_gen.py +++ b/src/handler/control/settings_gen.py @@ -12,6 +12,8 @@ import json import os import sys +from ..config import get_settings + def _hook_command(event: str) -> str: # Use the exact interpreter the control layer runs under, so the hook resolves the @@ -19,8 +21,8 @@ def _hook_command(event: str) -> str: return f"{sys.executable} -m handler.hooks {event}" -def build_settings() -> dict: - return { +def build_settings(headless: bool = False) -> dict: + settings = { "hooks": { "Stop": [{"hooks": [{"type": "command", "command": _hook_command("stop")}]}], "SessionEnd": [ @@ -37,13 +39,24 @@ def build_settings() -> dict: ], } } + if headless: + # ``claude -p`` never prompts — anything that would ask for permission is + # auto-denied. The allowlist is therefore what lets normal work (git, mise, the + # project's own tooling) proceed; the PreToolUse/Stop hooks above remain the hard + # gate either way, since a hook deny overrides any allow. + s = get_settings() + settings["permissions"] = { + "defaultMode": s.headless_permission_mode, + "allow": s.headless_allowed_tools_list, + } + return settings -def write_settings(working_dir: str) -> str: +def write_settings(working_dir: str, headless: bool = False) -> str: """Write ``.claude/settings.json`` under the agent's working dir; return its path.""" claude_dir = os.path.join(working_dir, ".claude") os.makedirs(claude_dir, exist_ok=True) path = os.path.join(claude_dir, "settings.json") with open(path, "w") as fh: - json.dump(build_settings(), fh, indent=2) + json.dump(build_settings(headless=headless), fh, indent=2) return path diff --git a/src/handler/control/spawn.py b/src/handler/control/spawn.py index 21f1573..008fd7c 100644 --- a/src/handler/control/spawn.py +++ b/src/handler/control/spawn.py @@ -17,8 +17,10 @@ from ..db.engine import connection from . import ( claude_config, credentials, + credsync, forge, gitops, + headless, mise, reposync, settings_gen, @@ -84,6 +86,7 @@ def spawn( role: str | None = None, require_tests: bool = True, mise_init: bool = False, + worker_id: str | None = None, ) -> dict: """Create and launch an agent. Returns the agent row. @@ -91,8 +94,13 @@ def spawn( it off, because a project with no ``.mise.toml`` yet is exactly what it exists to fix. ``mise_init`` marks the launched agent (via ``HANDLER_MISE_INIT``) so its hooks enforce the bootstrap contract — create the test task, commit, and push — instead of the normal - test gate. + test gate. ``worker_id`` identifies the calling worker container (headless runs record + it on the run row; the CLI defaults to a pid-scoped id). """ + if get_settings().runner == "headless" and not task: + # A tmux agent without a task idles at the REPL waiting for input; ``claude -p`` + # has no such mode — an empty prompt would exit immediately having done nothing. + raise SpawnError("a headless agent requires a task (the runner is 'headless')") sync_note = None with connection() as conn: project = repo.get_project(conn, project_id) @@ -146,14 +154,56 @@ def spawn( role=role, ) - settings_path = settings_gen.write_settings(working_dir) + headless_run = get_settings().runner == "headless" + settings_path = settings_gen.write_settings(working_dir, headless=headless_run) + env = _agent_env(project, agent, token, role=role, mise_init=mise_init) + # Verify the pinned forge version, if one is configured. Non-fatal: a version drift + # is recorded as a warning rather than blocking the spawn, since not every agent + # touches forge and the base image is the real pin (README 3.6, Phase 2). + forge_note = _check_forge_version(working_dir) + + # Mark Claude Code onboarding complete + trust the working dir before launching. The + # tmux path needs both (no human at the TTY to answer the theme/trust screens); + # ``-p`` skips the trust dialog but still reads onboarding state, so keep it for both. + claude_config.ensure_onboarded(working_dir) + credsync.note_local_write() + + if headless_run: + headless.launch( + agent, + kind="spawn", + prompt=task, + settings_path=settings_path, + env=env, + worker_id=worker_id or f"cli-{os.getpid()}", + ) + else: + session = tmux.session_name(project_id, name) + command = _claude_command(task, settings_path) + tmux.new_session(session, cwd=working_dir, command=command, env=env) + agent = {**agent, "forge_note": forge_note, "sync_note": sync_note} + return agent + + +def _agent_env( + project: dict, + agent: dict, + token: str | None, + *, + role: str | None = None, + mise_init: bool = False, +) -> dict[str, str]: + """The environment an agent process (and therefore its hooks) runs with: identity, + ``DATABASE_URL``, and resolved forge/git credentials. Shared by spawn and the + headless resume path (a resume is a brand-new process needing the same env).""" env = { - "HANDLER_PROJECT_ID": project_id, - "HANDLER_AGENT_NAME": name, + "HANDLER_PROJECT_ID": project["id"], + "HANDLER_AGENT_NAME": agent["name"], "HANDLER_AGENT_ID": str(agent["id"]), "DATABASE_URL": get_settings().database_url, } + role = role or agent.get("role") if role: env["HANDLER_AGENT_ROLE"] = role if mise_init: @@ -164,28 +214,13 @@ def spawn( with connection() as conn: env.update(credentials.credential_env(token, project.get("git_remote"), conn)) if token: - _install_git_credentials(working_dir, project.get("git_remote"), conn) + _install_git_credentials(agent["working_dir"], project.get("git_remote"), conn) # SSH remotes: pin the agent's git to the server's deploy key, when one is stored. try: env.update(reposync.ssh_env(project.get("git_remote"), conn)) except reposync.SyncError as exc: raise SpawnError(str(exc)) from exc - - # Verify the pinned forge version, if one is configured. Non-fatal: a version drift - # is recorded as a warning rather than blocking the spawn, since not every agent - # touches forge and the base image is the real pin (README 3.6, Phase 2). - forge_note = _check_forge_version(working_dir) - - # Mark Claude Code onboarding complete + trust the working dir before launching, so the - # detached agent boots straight to the REPL instead of wedging on the first-run theme - # picker / trust prompt with no human at the tmux TTY to answer it. - claude_config.ensure_onboarded(working_dir) - - session = tmux.session_name(project_id, name) - command = _claude_command(task, settings_path) - tmux.new_session(session, cwd=working_dir, command=command, env=env) - agent = {**agent, "forge_note": forge_note, "sync_note": sync_note} - return agent + return env def _check_forge_version(working_dir: str) -> str | None: @@ -203,20 +238,128 @@ def kill(project_id: str, name: str) -> None: agent = repo.get_agent_by_name(conn, project_id, name) if agent is None: raise SpawnError(f"agent '{name}' not found in project '{project_id}'") + if agent.get("session_id"): + # Headless agent: flag the running run for cancel; the owning worker's + # supervisor polls the flag and SIGTERMs its own child (cross-worker safe — + # nobody signals a process they don't own). No running run = already dead. + run = repo.get_latest_run(conn, agent["id"]) + if run is not None and run["status"] == "running": + repo.request_run_cancel(conn, run["id"]) + repo.set_agent_status(conn, agent["id"], "done") + return session = tmux.session_name(project_id, name) if tmux.has_session(session): tmux.kill_session(session) repo.set_agent_status(conn, agent["id"], "done") -def resume(agent: dict, answer: str) -> tuple[bool, str]: - """Feed an operator's answer back to a live agent. +def resume(agent: dict, answer: str, worker_id: str | None = None) -> tuple[bool, str]: + """Feed an operator's answer back to an agent. - The seam the API's ``/resume`` route calls (and the one tests mock). Sends the - answer into the agent's tmux session so the waiting ``claude`` process receives it. + The seam the API's ``/resume`` route calls (and the one tests mock). Legacy tmux + agents (``session_id`` null) get the answer typed into their live session; headless + agents get a brand-new ``claude -p --resume`` run on this worker, with the session + transcript materialized from the DB archive first so any worker can serve the resume. """ + if agent.get("session_id"): + return _resume_headless(agent, answer, worker_id or f"cli-{os.getpid()}") session = tmux.session_name(agent["project_id"], agent["name"]) if not tmux.has_session(session): return False, f"no live session '{session}' to resume" tmux.send_keys(session, answer) return True, f"answer delivered to session '{session}'" + + +def _resume_headless(agent: dict, answer: str, worker_id: str) -> tuple[bool, str]: + """Launch a ``--resume`` run for a headless agent, materializing its session first. + + Refuses while a run is still live (two concurrent processes on one session would + corrupt it). When neither this worker nor the DB has the transcript — the owning + worker died before its first archive — falls back to a *fresh* session whose prompt + re-injects context from the DB (checkmark + open question + answer), recorded as a + ``worker`` event so the UI shows the degraded continuity. + """ + with connection() as conn: + run = repo.get_latest_run(conn, agent["id"]) + if run is not None and run["status"] == "running": + return False, "agent already has a live run; wait for it to finish or kill it" + project = repo.get_project(conn, agent["project_id"]) + archive = repo.get_session_archive(conn, agent["id"]) + if project is None: + return False, f"project '{agent['project_id']}' not registered" + + working_dir = agent["working_dir"] + settings_path = settings_gen.write_settings(working_dir, headless=True) + try: + token = None + with connection() as conn: + token = credentials.resolve_for_project(project, conn) + except credentials.CredentialError as exc: + return False, str(exc) + env = _agent_env(project, agent, token) + + transcript = headless.session_dir(working_dir) / f"{agent['session_id']}.jsonl" + if archive is not None: + try: + headless.materialize_session(working_dir, bytes(archive["archive"])) + except (OSError, ValueError) as exc: + return False, f"could not materialize session archive: {exc}" + elif not transcript.exists(): + return _resume_reinjected(agent, answer, settings_path, env, worker_id) + + run = headless.launch( + agent, + kind="resume", + prompt=answer, + settings_path=settings_path, + env=env, + worker_id=worker_id, + ) + return True, f"headless resume run {run['id']} started for session {agent['session_id']}" + + +def _resume_reinjected( + agent: dict, answer: str, settings_path: str, env: dict, worker_id: str +) -> tuple[bool, str]: + """Degraded resume: no transcript anywhere, so start a fresh session with the + context rebuilt from the DB. Continuity is approximate — say so in the event log.""" + with connection() as conn: + checkmark = repo.get_checkmark(conn, agent["id"]) + recent = repo.get_log(conn, agent["id"], limit=5) + parts = [ + "You are resuming work you started in an earlier session whose transcript is " + "unavailable. Reconstruct context from your checkpoint below, then continue.", + ] + if checkmark: + if checkmark.get("where_it_stopped"): + parts.append(f"Where you stopped: {checkmark['where_it_stopped']}") + if checkmark.get("next_steps"): + parts.append(f"Planned next steps: {checkmark['next_steps']}") + if checkmark.get("open_question"): + parts.append(f"You had asked: {checkmark['open_question']}") + for entry in reversed(recent): + if entry.get("summary"): + parts.append(f"Earlier log: {entry['summary']}") + parts.append(f"The operator's answer/instruction: {answer}") + run = headless.launch( + agent, + kind="spawn", # a genuinely new session (new UUID) — --resume has nothing to load + prompt="\n\n".join(parts), + settings_path=settings_path, + env=env, + worker_id=worker_id, + ) + with connection() as conn: + repo.insert_agent_event( + conn, + agent["id"], + run["id"], + seq=0, + type="worker", + payload={ + "notice": "resume without transcript — context re-injected from DB", + "previous_session_id": agent["session_id"], + }, + session_id=run["session_id"], + ) + return True, f"transcript unavailable; started fresh run {run['id']} with re-injected context" diff --git a/src/handler/control/worker.py b/src/handler/control/worker.py index 2adc1bd..aeff4e2 100644 --- a/src/handler/control/worker.py +++ b/src/handler/control/worker.py @@ -15,12 +15,24 @@ returns a JSON-safe result or raises); ``drain``/``run`` are the claim+finish pl from __future__ import annotations import os +import secrets +import socket import time from datetime import UTC, datetime, timedelta +from ..config import get_settings from ..db import repository as repo from ..db.engine import connection -from . import gitops, login, poller, reposync, skills_gen, spawn, tmux +from . import credsync, gitops, login, poller, reposync, skills_gen, spawn, tmux + +# Command types that launch a claude run and therefore need a free slot on this worker. +# A worker with all slots busy leaves these queued for a less-loaded worker to claim. +_RUN_COMMANDS = ("spawn", "resume", "mise_init") + + +def make_worker_id() -> str: + """A stable-for-this-process, unique-across-containers worker id.""" + return f"worker-{socket.gethostname()}-{os.getpid()}-{secrets.token_hex(2)}" class CommandError(Exception): @@ -43,6 +55,7 @@ def _cmd_spawn(command: dict) -> dict: worktree_branch=p.get("worktree"), task=p.get("task"), role=p.get("role"), + worker_id=command.get("claimed_by"), ) result = { "agent_id": agent["id"], @@ -73,11 +86,15 @@ def _cmd_resume(command: dict) -> dict: agent = repo.get_agent_by_name(conn, command["project_id"], name) if agent is None: raise CommandError(f"agent '{name}' not found in project '{command['project_id']}'") - ok, detail = spawn.resume(agent, answer) - if ok: - with connection() as conn: - repo.set_agent_status(conn, agent["id"], "working") - return {"resumed": ok, "detail": detail} + ok, detail = spawn.resume(agent, answer, worker_id=command.get("claimed_by")) + if not ok: + # A resume that couldn't be delivered is a FAILED command, not a quiet no-op — + # the operator's input must never vanish silently (the old tmux path typed into + # dead panes and reported success). + raise CommandError(detail) + with connection() as conn: + repo.set_agent_status(conn, agent["id"], "working") + return {"resumed": True, "detail": detail} def _record_verdict(command: dict, status: str) -> dict: @@ -219,6 +236,13 @@ def _cmd_login_submit(command: dict) -> dict: # Surface the pane tail so the operator can see why claude rejected the code. detail = result.get("output") or "claude did not confirm a successful login" raise CommandError(f"login not confirmed — {detail}") + # Publish the fresh credentials so every other worker container can materialize + # them (multi-worker: the login ran here, but any worker may run the next agent). + try: + result["credentials_published"] = credsync.upload() + except Exception as exc: # noqa: BLE001 - login succeeded; publishing is best-effort + result["credentials_published"] = False + result["credsync_note"] = str(exc) return result @@ -360,7 +384,9 @@ def drain(worker_id: str, limit: int | None = None) -> int: processed = 0 while limit is None or processed < limit: with connection() as conn: - command = repo.claim_next_command(conn, worker_id) + command = repo.claim_next_command( + conn, worker_id, types_excluded=_full_slot_exclusions(worker_id) + ) if command is None: break _run_one(command) @@ -368,22 +394,43 @@ def drain(worker_id: str, limit: int | None = None) -> int: return processed +def _full_slot_exclusions(worker_id: str) -> tuple[str, ...]: + """Command types this worker must not claim right now. + + Headless runs are supervised in-process, so a worker at ``max_concurrent_runs`` skips + claiming run-starting commands — they stay queued for a worker with a free slot. Slot + accounting is DB-driven (this worker's ``running`` runs), so it needs no in-memory + registry and survives restarts (a fresh process gets a fresh worker id; stale rows + belong to the old id and are the reaper's problem). Tmux runs are fire-and-forget and + never consume a slot. + """ + if get_settings().runner != "headless": + return () + with connection() as conn: + active = len(repo.list_running_runs(conn, worker_id=worker_id)) + if active >= get_settings().max_concurrent_runs: + return _RUN_COMMANDS + return () + + def run( worker_id: str | None = None, poll_interval: float = 2.0, ci_interval: float = 30.0, capture_interval: float = 2.0, + credsync_interval: float = 30.0, iterations: int | None = None, ) -> None: """The control-container main loop: drain the command queue, snapshot live agent - output, and sweep CI periodically. + output, sync claude credentials, and sweep CI periodically. ``iterations`` bounds the loop for tests; production runs unbounded. Sleeps ``poll_interval`` only when a pass found no commands, so bursts drain promptly. """ - worker_id = worker_id or f"worker-{os.getpid()}" + worker_id = worker_id or make_worker_id() last_ci = 0.0 last_capture = 0.0 + last_credsync = 0.0 count = 0 while iterations is None or count < iterations: try: @@ -398,6 +445,14 @@ def run( except Exception: # noqa: BLE001 - a capture hiccup must not kill the worker pass last_capture = now + if credsync_interval > 0 and (last_credsync == 0.0 or now - last_credsync >= credsync_interval): + # First pass runs immediately: a fresh worker container must materialize the + # claude credentials before it claims its first spawn. + try: + credsync.refresh() + except Exception: # noqa: BLE001 - cred sync must not kill the worker + pass + last_credsync = now if ci_interval > 0 and now - last_ci >= ci_interval: try: poller.sweep() diff --git a/src/handler/db/repository.py b/src/handler/db/repository.py index ecf4dc0..3684c27 100644 --- a/src/handler/db/repository.py +++ b/src/handler/db/repository.py @@ -31,6 +31,7 @@ from .tables import ( forge_hosts, log_entries, projects, + runtime_secrets, schedules, session_archives, shared_context, @@ -874,3 +875,34 @@ def get_session_archive(conn: Connection, agent_id: int) -> dict | None: select(session_archives).where(session_archives.c.agent_id == agent_id) ).first() return _row_to_dict(row) + + +def upsert_runtime_secret(conn: Connection, key: str, value_enc: str) -> None: + """Store an encrypted control-plane secret (the caller encrypts via secretstore).""" + now = _now() + result = conn.execute( + runtime_secrets.update() + .where(runtime_secrets.c.key == key) + .values(value_enc=value_enc, updated_at=now) + ) + if result.rowcount == 0: + conn.execute( + runtime_secrets.insert().values(key=key, value_enc=value_enc, updated_at=now) + ) + + +def get_runtime_secret(conn: Connection, key: str) -> dict | None: + row = conn.execute(select(runtime_secrets).where(runtime_secrets.c.key == key)).first() + return _row_to_dict(row) + + +def get_latest_finished_command(conn: Connection, type: str) -> dict | None: + """The most recent ``done`` command of a type (login pinning reads login_start's + ``claimed_by`` to route login_submit to the same worker container).""" + row = conn.execute( + select(commands) + .where(commands.c.type == type, commands.c.status == "done") + .order_by(commands.c.id.desc()) + .limit(1) + ).first() + return _row_to_dict(row) diff --git a/src/handler/db/tables.py b/src/handler/db/tables.py index 843839c..de180a3 100644 --- a/src/handler/db/tables.py +++ b/src/handler/db/tables.py @@ -243,6 +243,18 @@ forge_hosts = Table( CheckConstraint(_in("forge_type", FORGE_TYPES), name="ck_forge_hosts_type"), ) +# Control-plane key/value secrets, Fernet-encrypted like forge_hosts tokens (see +# ``handler.secretstore``). Holds the claude OAuth credential bundle so every worker +# container can materialize it locally — the login flow runs on ONE worker, but all of +# them need to run ``claude``. Never exposed by any API route. +runtime_secrets = Table( + "runtime_secrets", + metadata, + Column("key", String, primary_key=True), + Column("value_enc", String, nullable=False), + Column("updated_at", PortableTimestamp, nullable=False, server_default=func.now()), +) + # Worker registry + heartbeat (headless runner). Each worker container upserts its row # every loop pass; the reaper marks a worker's running runs (and their agents) ``crashed`` # when ``heartbeat_at`` goes stale — the positive-liveness replacement for tmux scraping. diff --git a/src/handler/migrations/versions/0009_runtime_secrets.py b/src/handler/migrations/versions/0009_runtime_secrets.py new file mode 100644 index 0000000..8d26401 --- /dev/null +++ b/src/handler/migrations/versions/0009_runtime_secrets.py @@ -0,0 +1,39 @@ +"""runtime_secrets: encrypted control-plane key/value store + +Revision ID: 0009_runtime_secrets +Revises: 0008_headless_runs +Create Date: 2026-07-21 + +A small Fernet-encrypted key/value table (values encrypted with ``HANDLER_SECRET_KEY`` +before insert, same policy as ``forge_hosts.token_enc``). First use: the claude OAuth +credential bundle — the interactive ``/login`` flow completes on one worker container, +and every other worker materializes the bundle from here at startup/refresh so any of +them can run headless ``claude -p`` without shared files. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +from handler.db.types import PortableTimestamp + +revision: str = "0009_runtime_secrets" +down_revision: str | None = "0008_headless_runs" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "runtime_secrets", + sa.Column("key", sa.String(), primary_key=True), + sa.Column("value_enc", sa.String(), nullable=False), + sa.Column("updated_at", PortableTimestamp, nullable=False, server_default=sa.func.now()), + ) + + +def downgrade() -> None: + op.drop_table("runtime_secrets") diff --git a/tests/test_credsync.py b/tests/test_credsync.py new file mode 100644 index 0000000..cb748c4 --- /dev/null +++ b/tests/test_credsync.py @@ -0,0 +1,110 @@ +"""Credential distribution through runtime_secrets: the login completes on one worker, +every other worker materializes the encrypted bundle from the DB — no shared files.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +from cryptography.fernet import Fernet + +from handler.control import credsync +from handler.db import repository as repo +from handler.db.engine import get_engine + + +@pytest.fixture +def secret_env(env, monkeypatch): + from handler import config + + monkeypatch.setenv("HANDLER_SECRET_KEY", Fernet.generate_key().decode()) + config.get_settings.cache_clear() + credsync._state.__init__() # fresh sync cursor per test + yield env + config.get_settings.cache_clear() + + +def _write_local_credentials(home: Path) -> None: + (home / ".claude").mkdir(parents=True, exist_ok=True) + (home / ".claude.json").write_text( + json.dumps({"oauthAccount": {"email": "op@example.com"}, "theme": "light"}) + ) + (home / ".claude" / ".credentials.json").write_text('{"token": "secret-oauth-token"}') + + +def test_upload_stores_encrypted_bundle(secret_env, tmp_path): + _write_local_credentials(tmp_path) + assert credsync.upload() is True + with get_engine().begin() as conn: + row = repo.get_runtime_secret(conn, credsync.SECRET_KEY) + assert row is not None + # Ciphertext at rest — the raw token must not appear in the DB value. + assert "secret-oauth-token" not in row["value_enc"] + + +def test_refresh_materializes_on_fresh_worker(secret_env, tmp_path, monkeypatch): + _write_local_credentials(tmp_path) + assert credsync.upload() is True + + other_home = tmp_path / "worker-b" + other_home.mkdir() + monkeypatch.setenv("HOME", str(other_home)) + credsync._state.__init__() # worker B's process state + + assert credsync.refresh() == "materialized" + creds = json.loads((other_home / ".claude" / ".credentials.json").read_text()) + assert creds["token"] == "secret-oauth-token" + data = json.loads((other_home / ".claude.json").read_text()) + assert data["oauthAccount"]["email"] == "op@example.com" + # A second pass is a no-op — nothing changed anywhere. + assert credsync.refresh() is None + + +def test_materialize_merges_claude_json_preserving_local_state(secret_env, tmp_path, monkeypatch): + _write_local_credentials(tmp_path) + credsync.upload() + + other_home = tmp_path / "worker-c" + (other_home / ".claude").mkdir(parents=True) + (other_home / ".claude.json").write_text( + json.dumps( + { + "hasCompletedOnboarding": True, + "theme": "dark", + "projects": {"/projects/p/a": {"hasTrustDialogAccepted": True}}, + } + ) + ) + monkeypatch.setenv("HOME", str(other_home)) + credsync._state.__init__() + + assert credsync.refresh() == "materialized" + data = json.loads((other_home / ".claude.json").read_text()) + # Account arrived... + assert data["oauthAccount"]["email"] == "op@example.com" + # ...but this worker's own onboarding/trust state (claude_config's writes) survived. + assert data["theme"] == "dark" + assert data["projects"]["/projects/p/a"]["hasTrustDialogAccepted"] is True + + +def test_refresh_uploads_local_change(secret_env, tmp_path): + _write_local_credentials(tmp_path) + assert credsync.refresh() == "uploaded" # bootstrap: local creds, empty DB + # A token refresh on disk (mtime/size change) re-publishes. + os.utime(tmp_path / ".claude" / ".credentials.json", ns=(1, 1)) + assert credsync.refresh() == "uploaded" + + +def test_disabled_without_secret_key(env, tmp_path): + _write_local_credentials(tmp_path) + assert credsync.upload() is False + assert credsync.refresh() is None + + +def test_note_local_write_suppresses_upload(secret_env, tmp_path): + _write_local_credentials(tmp_path) + credsync.note_local_write() + # The deliberate local write (e.g. ensure_onboarded at spawn) is not re-published. + assert credsync.refresh() is None diff --git a/tests/test_headless_run.py b/tests/test_headless_run.py new file mode 100644 index 0000000..ad9600c --- /dev/null +++ b/tests/test_headless_run.py @@ -0,0 +1,272 @@ +"""End-to-end headless runner tests against the fake ``claude`` binary. + +Real subprocesses, real threads, real SQLite: ``headless.launch`` starts +``tests/fixtures/fake_claude.py`` (selected via the ``claude_bin`` setting, the same +seam the tmux fakes used), the supervisor streams its stdout into ``agent_events``, and +the tests assert on what landed in the DB — exactly what the API/UI will read.""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import pytest + +from handler.control import headless, settings_gen, spawn +from handler.db import repository as repo +from handler.db.engine import get_engine + +REPO_ROOT = Path(__file__).resolve().parents[1] +FAKE_CLAUDE = str(REPO_ROOT / "tests" / "fixtures" / "fake_claude.py") + + +@pytest.fixture +def headless_env(env, monkeypatch): + from handler import config + + monkeypatch.setenv("CLAUDE_BIN", FAKE_CLAUDE) + monkeypatch.setenv("RUNNER", "headless") + config.get_settings.cache_clear() + yield env + config.get_settings.cache_clear() + + +def _make_agent(tmp_path, name="h1"): + working_dir = tmp_path / "projects" / "p" / name + working_dir.mkdir(parents=True) + with get_engine().begin() as conn: + repo.create_project(conn, "p", str(tmp_path / "projects" / "p")) + agent = repo.create_agent(conn, "p", name, str(working_dir)) + return agent + + +def _wait_for(predicate, timeout=20.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + result = predicate() + if result: + return result + time.sleep(0.1) + return None + + +def _finished_run(run_id): + def check(): + with get_engine().begin() as conn: + run = repo.get_run(conn, run_id) + return run if run["status"] != "running" else None + + return check + + +def test_spawn_run_streams_events_and_completes(headless_env, tmp_path): + agent = _make_agent(tmp_path) + run = headless.launch( + agent, kind="spawn", prompt="build the thing", + settings_path=str(tmp_path / "settings.json"), env={}, worker_id="w1", + ) + finished = _wait_for(_finished_run(run["id"])) + assert finished is not None, "run never finished" + assert finished["status"] == "completed" + assert finished["exit_code"] == 0 + assert finished["result"]["is_error"] is False + + with get_engine().begin() as conn: + events = repo.list_agent_events(conn, agent["id"]) + updated = repo.get_agent_by_id(conn, agent["id"]) + archive = repo.get_session_archive(conn, agent["id"]) + + types = [e["type"] for e in events] + assert types == ["system", "assistant", "result"] + assert [e["seq"] for e in events] == [1, 2, 3] + # last_output is now derived from assistant text — the log the UI shows is real. + assert updated["last_output"] == "working on: build the thing" + assert updated["status"] == "done" + assert updated["session_id"] == run["session_id"] + assert updated["worker_id"] == "w1" + # The session archive was uploaded at exit for cross-worker resume. + assert archive is not None + assert archive["session_id"] == run["session_id"] + + +def test_failed_run_marks_blocked_with_worker_event(headless_env, tmp_path, monkeypatch): + monkeypatch.setenv("FAKE_CLAUDE_MODE", "error") + agent = _make_agent(tmp_path, "h-err") + run = headless.launch( + agent, kind="spawn", prompt="boom", + settings_path=str(tmp_path / "s.json"), env={}, worker_id="w1", + ) + finished = _wait_for(_finished_run(run["id"])) + assert finished["status"] == "failed" + assert finished["exit_code"] == 2 + + with get_engine().begin() as conn: + events = repo.list_agent_events(conn, agent["id"]) + updated = repo.get_agent_by_id(conn, agent["id"]) + types = [e["type"] for e in events] + # The unparseable stdout line is preserved verbatim as a raw event, and the runner + # records why the run failed as a worker event. + assert "raw" in types + raw = next(e for e in events if e["type"] == "raw") + assert "this is not json" in raw["payload"]["line"] + worker_ev = next(e for e in events if e["type"] == "worker") + assert worker_ev["payload"]["exit_code"] == 2 + assert updated["status"] == "blocked" + + +def test_hook_written_status_survives_reconciliation(headless_env, tmp_path, monkeypatch): + """Hooks are the status authority: if one set paused_for_input during the run, the + supervisor's exit pass must not overwrite it with done.""" + monkeypatch.setenv("FAKE_CLAUDE_MODE", "slow") + monkeypatch.setenv("FAKE_CLAUDE_SLOW_SECONDS", "1.5") + agent = _make_agent(tmp_path, "h-hook") + run = headless.launch( + agent, kind="spawn", prompt="ask me something", + settings_path=str(tmp_path / "s.json"), env={}, worker_id="w1", + ) + # Simulate a hook (inside the run) recording an open question. + with get_engine().begin() as conn: + repo.set_agent_status(conn, agent["id"], "paused_for_input") + finished = _wait_for(_finished_run(run["id"])) + assert finished["status"] == "completed" + with get_engine().begin() as conn: + assert repo.get_agent_by_id(conn, agent["id"])["status"] == "paused_for_input" + + +def test_cancel_terminates_hanging_run(headless_env, tmp_path, monkeypatch): + monkeypatch.setenv("FAKE_CLAUDE_MODE", "hang") + agent = _make_agent(tmp_path, "h-hang") + run = headless.launch( + agent, kind="spawn", prompt="hang forever", + settings_path=str(tmp_path / "s.json"), env={}, worker_id="w1", + ) + # Give the supervisor a moment to start the process, then flag the cancel the same + # way a cross-worker kill would. + _wait_for(lambda: _events_count(agent["id"]) >= 1) + with get_engine().begin() as conn: + assert repo.request_run_cancel(conn, run["id"]) is True + finished = _wait_for(_finished_run(run["id"]), timeout=30.0) + assert finished is not None, "cancel never terminated the run" + assert finished["status"] == "canceled" + + +def _events_count(agent_id): + with get_engine().begin() as conn: + return len(repo.list_agent_events(conn, agent_id)) + + +def test_kill_headless_agent_requests_cancel(headless_env, tmp_path, monkeypatch): + monkeypatch.setenv("FAKE_CLAUDE_MODE", "hang") + agent = _make_agent(tmp_path, "h-kill") + run = headless.launch( + agent, kind="spawn", prompt="hang", + settings_path=str(tmp_path / "s.json"), env={}, worker_id="w1", + ) + _wait_for(lambda: _events_count(agent["id"]) >= 1) + spawn.kill("p", "h-kill") + finished = _wait_for(_finished_run(run["id"]), timeout=30.0) + assert finished["status"] == "canceled" + with get_engine().begin() as conn: + assert repo.get_agent_by_id(conn, agent["id"])["status"] == "done" + + +def test_cross_worker_resume_materializes_archive(headless_env, tmp_path, monkeypatch): + """The linchpin: worker B resumes a session it never ran, from the DB archive alone.""" + agent = _make_agent(tmp_path, "h-resume") + run = headless.launch( + agent, kind="spawn", prompt="first pass", + settings_path=str(tmp_path / "s.json"), env={}, worker_id="worker-a", + ) + assert _wait_for(_finished_run(run["id"]))["status"] == "completed" + + # "Worker B": a clean HOME with no local claude state at all. + other_home = tmp_path / "worker-b-home" + other_home.mkdir() + monkeypatch.setenv("HOME", str(other_home)) + + with get_engine().begin() as conn: + agent = repo.get_agent_by_id(conn, agent["id"]) # refetch: has session_id now + ok, detail = spawn.resume(agent, "the operator's answer", worker_id="worker-b") + assert ok, detail + + with get_engine().begin() as conn: + resumed = repo.get_latest_run(conn, agent["id"]) + assert resumed["kind"] == "resume" + assert resumed["worker_id"] == "worker-b" + finished = _wait_for(_finished_run(resumed["id"])) + # fake_claude exits 3 if the transcript was NOT materialized where claude looks. + assert finished["status"] == "completed", f"exit={finished['exit_code']}" + assert finished["session_id"] == run["session_id"] # same session, continued + + +def test_resume_without_transcript_reinjects_context(headless_env, tmp_path, monkeypatch): + """Owning worker died before its first archive: resume degrades to a fresh session + with DB-rebuilt context, visibly marked as such.""" + agent = _make_agent(tmp_path, "h-fallback") + with get_engine().begin() as conn: + repo.set_agent_session(conn, agent["id"], "lost-session-uuid", "worker-dead") + repo.upsert_checkmark_row( + conn, agent["id"], status="paused_for_input", + where_it_stopped="mid-refactor", open_question="which db?", + ) + agent = repo.get_agent_by_id(conn, agent["id"]) + + ok, detail = spawn.resume(agent, "use postgres", worker_id="worker-b") + assert ok + assert "re-injected" in detail + + with get_engine().begin() as conn: + new_run = repo.get_latest_run(conn, agent["id"]) + events = repo.list_agent_events(conn, agent["id"]) + assert new_run["kind"] == "spawn" # genuinely new session + assert new_run["session_id"] != "lost-session-uuid" + notice = next(e for e in events if e["type"] == "worker") + assert notice["payload"]["previous_session_id"] == "lost-session-uuid" + finished = _wait_for(_finished_run(new_run["id"])) + assert finished["status"] == "completed" + with get_engine().begin() as conn: + assistant = [ + e for e in repo.list_agent_events(conn, agent["id"]) if e["type"] == "assistant" + ] + # The re-injected prompt (with the operator's answer) reached the fresh claude. + assert any("use postgres" in json.dumps(e["payload"]) for e in assistant) + + +def test_resume_refused_while_run_live(headless_env, tmp_path, monkeypatch): + monkeypatch.setenv("FAKE_CLAUDE_MODE", "hang") + agent = _make_agent(tmp_path, "h-busy") + run = headless.launch( + agent, kind="spawn", prompt="hang", + settings_path=str(tmp_path / "s.json"), env={}, worker_id="w1", + ) + with get_engine().begin() as conn: + agent = repo.get_agent_by_id(conn, agent["id"]) + ok, detail = spawn.resume(agent, "answer", worker_id="w1") + assert not ok + assert "live run" in detail + with get_engine().begin() as conn: + repo.request_run_cancel(conn, run["id"]) + _wait_for(_finished_run(run["id"]), timeout=30.0) + + +def test_headless_settings_include_permissions(headless_env, tmp_path): + path = settings_gen.write_settings(str(tmp_path / "wd"), headless=True) + data = json.loads(Path(path).read_text()) + assert data["permissions"]["defaultMode"] == "acceptEdits" + assert "Bash(git *)" in data["permissions"]["allow"] + assert "hooks" in data # the hard gate is untouched + + tmux_path = settings_gen.write_settings(str(tmp_path / "wd2")) + assert "permissions" not in json.loads(Path(tmux_path).read_text()) + + +def test_api_rejects_empty_task_headless_spawn(headless_env, client, auth): + client.post( + "/projects", json={"id": "p2", "root_dir": "/tmp/p2"}, headers=auth + ) + resp = client.post( + "/projects/p2/agents/spawn", json={"name": "idle"}, headers=auth + ) + assert resp.status_code == 400 + assert "task is required" in resp.json()["detail"] diff --git a/tests/test_worker.py b/tests/test_worker.py index bdbd7e6..6ac847f 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -65,8 +65,8 @@ def test_resume_command_feeds_answer_and_sets_working(env, monkeypatch): _seed_project("api") seen = {} - def fake_resume(agent, ans): - seen.update(name=agent["name"], ans=ans) + def fake_resume(agent, ans, worker_id=None): + seen.update(name=agent["name"], ans=ans, worker_id=worker_id) return True, "ok" monkeypatch.setattr(spawn, "resume", fake_resume) @@ -74,11 +74,29 @@ def test_resume_command_feeds_answer_and_sets_working(env, monkeypatch): worker.drain("w") assert _get(cmd["id"])["status"] == "done" - assert seen == {"name": "api", "ans": "Postgres"} + assert seen == {"name": "api", "ans": "Postgres", "worker_id": "w"} with get_engine().begin() as conn: assert repo.get_agent_by_name(conn, "p", "api")["status"] == "working" +def test_resume_command_fails_loudly_when_undeliverable(env, monkeypatch): + """An undeliverable answer must surface as a FAILED command — never a silent 'done' + (the original bug: send-keys into a dead tmux pane reported success).""" + _seed_project("api") + monkeypatch.setattr( + spawn, "resume", lambda agent, ans, worker_id=None: (False, "no live session") + ) + cmd = _enqueue(type="resume", project_id="p", agent_name="api", payload={"answer": "x"}) + + worker.drain("w") + failed = _get(cmd["id"]) + assert failed["status"] == "failed" + assert "no live session" in failed["error"] + with get_engine().begin() as conn: + # The agent must NOT be flipped to working when nothing was delivered. + assert repo.get_agent_by_name(conn, "p", "api")["status"] != "working" + + def test_approve_command_records_operator_verdict_with_head_sha(env, fake_gitops): _seed_project("senior") cmd = _enqueue( diff --git a/tests/test_worker_concurrency.py b/tests/test_worker_concurrency.py new file mode 100644 index 0000000..b96fe9d --- /dev/null +++ b/tests/test_worker_concurrency.py @@ -0,0 +1,81 @@ +"""Slot-aware command claiming: a worker at max_concurrent_runs must leave run-starting +commands queued (for a less-loaded worker) while still processing everything else. Slot +accounting is DB-driven — this worker's ``running`` agent_runs rows — so it needs no +in-memory registry and is exercised here without real subprocesses.""" + +from __future__ import annotations + +import pytest + +from handler.control import spawn, worker +from handler.db import repository as repo +from handler.db.engine import get_engine + + +@pytest.fixture +def headless_env(env, monkeypatch): + from handler import config + + monkeypatch.setenv("RUNNER", "headless") + monkeypatch.setenv("MAX_CONCURRENT_RUNS", "2") + config.get_settings.cache_clear() + yield env + config.get_settings.cache_clear() + + +def _seed(conn_count_running_for=None): + with get_engine().begin() as conn: + repo.create_project(conn, "p", "/tmp/p") + agent = repo.create_agent(conn, "p", "a", "/tmp/p/a") + return agent + + +def _running_run(agent_id, worker_id): + with get_engine().begin() as conn: + return repo.create_run(conn, agent_id, f"sid-{worker_id}", worker_id, "spawn") + + +def test_full_worker_skips_run_commands_but_processes_others(headless_env, monkeypatch): + agent = _seed() + _running_run(agent["id"], "w-full") + _running_run(agent["id"], "w-full") # 2 running == MAX_CONCURRENT_RUNS + + spawned = {} + monkeypatch.setattr( + spawn, "spawn", + lambda project_id, name, **kw: spawned.update(name=name, **kw) + or {"id": 1, "name": name, "working_dir": "/tmp/p/x", "forge_note": None}, + ) + with get_engine().begin() as conn: + spawn_cmd = repo.enqueue_command(conn, "spawn", project_id="p", agent_name="x") + kill_cmd = repo.enqueue_command(conn, "kill", project_id="p", agent_name="a") + monkeypatch.setattr(spawn, "kill", lambda p, n: None) + + # The full worker processes the kill but leaves the spawn queued. + assert worker.drain("w-full") == 1 + with get_engine().begin() as conn: + assert repo.get_command(conn, kill_cmd["id"])["status"] == "done" + assert repo.get_command(conn, spawn_cmd["id"])["status"] == "queued" + assert spawned == {} + + # A worker with free slots picks the spawn up. + assert worker.drain("w-free") == 1 + with get_engine().begin() as conn: + assert repo.get_command(conn, spawn_cmd["id"])["status"] == "done" + assert spawned["name"] == "x" + assert spawned["worker_id"] == "w-free" + + +def test_slot_frees_when_run_finishes(headless_env, monkeypatch): + agent = _seed() + run1 = _running_run(agent["id"], "w1") + _running_run(agent["id"], "w1") + assert worker._full_slot_exclusions("w1") == worker._RUN_COMMANDS + + with get_engine().begin() as conn: + repo.finish_run(conn, run1["id"], "completed", exit_code=0) + assert worker._full_slot_exclusions("w1") == () + + +def test_tmux_runner_never_excludes(env): + assert worker._full_slot_exclusions("w") == ()