mirror of
https://github.com/0xWheatyz/handler.git
synced 2026-08-31 06:06:24 +00:00
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).
This commit is contained in:
@@ -220,3 +220,6 @@ __marimo__/
|
||||
# Streamlit
|
||||
.streamlit/secrets.toml
|
||||
.omc/
|
||||
|
||||
# Handler runtime artifacts
|
||||
/handler.db
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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()]
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
+168
-25
@@ -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"
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
+21
-3
@@ -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(
|
||||
|
||||
@@ -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") == ()
|
||||
Reference in New Issue
Block a user