"""AgentRuntime: one per assigned agent; process-per-burst (design.md section 4). Phases of the harness process: none -> starting -> running -> (stopping) -> none. - A prompt with no process starts a burst; prompts arriving while it starts are buffered and forwarded in order once it is up; later prompts go straight to the running process. - The burst ends when the harness has been idle for `idle_grace`, when a turn exceeds its timeout (process killed, agent `stuck` until the next turn), on stop/unassign, or on drain once the in-flight turn ends. The worker never schedules or queues beyond the startup buffer. """ from __future__ import annotations import asyncio import logging import os from collections.abc import Callable from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from flock.models import AgentHeartbeat, Event, TokenUsage, TypeSpec, parse_duration from flock.worker.adapters.base import AdapterEvent, HarnessAdapter, ResumeError, usage_sub from flock.worker.events import Outbox, SessionLog, utcnow if TYPE_CHECKING: from flock.worker.identity import Identities from flock.worker.sync import BucketSync log = logging.getLogger(__name__) INSTRUCTIONS_FILE = {"claude": "CLAUDE.md", "codex": "AGENTS.md"} class AgentState(BaseModel): """What must survive a worker change; synced next to the session log.""" session_id: str | None = None turn_seq: int = 0 # last turn number handed out (turn ids are t-) turns: int = 0 # turns ended (completed, failed or interrupted) tokens: dict[str, TokenUsage] = Field(default_factory=dict) # cumulative, by model harness_cum: dict[str, TokenUsage] = Field(default_factory=dict) # last cumulative snapshot harness_estimate_usd: float | None = None def restart_prompt(old_session: str | None, transcript: Path | None, session_log: Path, why: str) -> str: where = f"{transcript} (harness transcript) and {session_log}" if transcript else str(session_log) return ( f"Your previous session ({old_session}) could not be resumed: {why}. " f"Its transcript is at {where}. Read it to recover your context, then continue your work." ) class AgentRuntime: def __init__( self, name: str, type_name: str, spec: TypeSpec, *, swarm: str, work_root: Path, adapter: HarnessAdapter, outbox: Outbox, sync: BucketSync | None = None, identities: Identities | None = None, working_count: Callable[[], int] | None = None, ) -> None: self.name, self.type_name, self.spec, self.swarm = name, type_name, spec, swarm self.work_root = Path(work_root) self.workdir = self.work_root / name self.sessions_dir = self.work_root / "_sessions" self.adapter = adapter self.outbox = outbox self.sync = sync self.identities = identities self._working_count = working_count self.log = SessionLog(self.sessions_dir / f"{name}.jsonl", swarm=swarm, agent=name, harness=spec.harness.value, cwd=str(self.workdir)) self.state = AgentState() self.phase = "none" self.buffer: list[str] = [] self.turn_id: str | None = None self.awaiting = False # a prompt was forwarded and its turn has not started yet self.stuck = False self.resumed: bool | None = None self.idle_since: datetime = utcnow() self.last_sync: datetime | None = None self.dirty = False self.allow_starts = True self._stop_req: str | None = None self._burst: asyncio.Task | None = None self._ready = asyncio.Event() self._turn_started_at = 0.0 self._last_activity = 0.0 self._turn_timeout = parse_duration(spec.turn_timeout).total_seconds() self._next_timeout: float | None = None self._turn_tokens: dict[str, TokenUsage] = {} self._instructions_written = False self.parent: str | None = None # the spawning agent (AssignCommand.parent) self.flock_env: dict[str, str] = {} # FLOCK_AGENT, FLOCK_AGENT_TOKEN, ... (worker sets) # ------------------------------------------------------------------ paths and state @property def state_path(self) -> Path: return self.sessions_dir / f"{self.name}.state.json" @property def idle_grace(self) -> float: return parse_duration(self.spec.idle_grace).total_seconds() def load_state(self) -> None: if self.state_path.is_file(): try: self.state = AgentState.model_validate_json(self.state_path.read_text()) except ValueError as e: log.warning("%s: unreadable state file: %s", self.name, e) self.log.reload() def save_state(self) -> None: self.sessions_dir.mkdir(parents=True, exist_ok=True) tmp = self.state_path.with_suffix(".tmp") tmp.write_text(self.state.model_dump_json()) os.replace(tmp, self.state_path) def harness_files(self) -> list[Path]: return self.adapter.session_files(self.state.session_id) def _now(self) -> float: return asyncio.get_running_loop().time() # ------------------------------------------------------------------ events def _emit(self, ev: Event) -> None: self.dirty = True self.outbox.add(ev) def emit(self, kind: str, data: dict[str, Any] | None = None) -> None: self._emit(self.log.custom(kind, data, turn_id=self.turn_id)) # ------------------------------------------------------------------ lifecycle async def setup(self, restore: bool = True, resume_prompt: str | None = None) -> None: """Run once on assign: restore from the bucket, load state, then allow starts.""" try: if restore and self.sync is not None: try: await self.sync.restore(self) except Exception as e: # noqa: BLE001 - a failed restore means a fresh start log.warning("%s: restore failed: %s", self.name, e) self.emit("error", {"error": f"restore failed: {e}"}) self.load_state() finally: self._ready.set() if resume_prompt: await self.prompt(resume_prompt) async def prompt(self, text: str, *, prompt_id: str | None = None, timeout: str | None = None) -> None: if timeout: self._next_timeout = parse_duration(timeout).total_seconds() if self.phase == "running" and not self._stop_req: await self._forward(text) return if self.phase == "none" and (not self.allow_starts or self._stop_req): self.emit("prompt_deferred", {"prompt_id": prompt_id, "text": text, "reason": "draining" if not self.allow_starts else self._stop_req}) return self.buffer.append(text) if self.phase == "none": self.phase = "starting" self._burst = asyncio.create_task(self._run_burst(), name=f"burst-{self.name}") async def _forward(self, text: str) -> None: self._last_activity = self._now() if self.turn_id is None: self.awaiting = True try: await self.adapter.send(text) except Exception as e: # noqa: BLE001 - process died under us: keep the prompt for the next burst log.warning("%s: send failed (%s); buffering for the next start", self.name, e) self.buffer.append(text) def _defer_buffer(self, reason: str) -> None: for text in self.buffer: self.emit("prompt_deferred", {"text": text, "reason": reason}) self.buffer.clear() def _prepare_workdir(self) -> None: self.workdir.mkdir(parents=True, exist_ok=True) if self._instructions_written: return target = self.workdir / INSTRUCTIONS_FILE[self.spec.harness.value] if self.spec.system_prompt is not None: target.write_text(self.spec.system_prompt) self._instructions_written = True return if not self.spec.system_prompt_file: return src = Path(self.spec.system_prompt_file) candidates = [src] if src.is_absolute() else [ Path(os.environ.get("FLOCK_PROMPTS_DIR", ".")) / src, Path.cwd() / src] for c in candidates: if c.is_file(): target.write_text(c.read_text()) self._instructions_written = True return log.warning("%s: system_prompt_file %s not found", self.name, src) async def _run_burst(self) -> None: pump: asyncio.Task | None = None try: await self._ready.wait() if not self.allow_starts or self._stop_req: self._defer_buffer("draining" if not self.allow_starts else str(self._stop_req)) return if self._working_count and self._working_count() > self.spec.max_working_per_worker: log.warning("%s: %d live harness processes exceed max_working_per_worker=%d", self.name, self._working_count(), self.spec.max_working_per_worker) self._prepare_workdir() env = dict(self.spec.env) if self.identities is not None: env.update(await self.identities.env_for(self.name, self.spec)) old = self.state.session_id env.update(self.flock_env) first: list[str] = [] data: dict[str, Any] = {"harness": self.spec.harness.value, "model": self.spec.model} try: sid = await self.adapter.start(self.workdir, old, self.spec, env=env) kind, self.resumed = ("session_resumed", True) if old else ("session_started", False) except ResumeError as e: transcript = next((p for p in self.adapter.session_files(old) if old and old in p.name), None) sid = await self.adapter.start(self.workdir, None, self.spec, env=env) kind, self.resumed = "session_restarted", False self.state.harness_cum = {} first.append(restart_prompt(old, transcript, self.log.path, str(e))) data.update(previous_session_id=old, reason=str(e)) self.state.session_id = sid self.save_state() self.emit(kind, {**data, "session_id": sid, "resumed": self.resumed, "parent": self.parent}) self.phase = "running" self._last_activity = self._now() pump = asyncio.create_task(self._pump()) pending, self.buffer = first + self.buffer, [] for text in pending: if self._stop_req and self._stop_req != "drain": self.buffer.append(text) continue await self._forward(text) if self._stop_req and self._stop_req != "drain": self._defer_buffer(self._stop_req) await self._watch(pump) except asyncio.CancelledError: raise except Exception as e: log.exception("%s: burst failed", self.name) self.emit("error", {"error": f"{type(e).__name__}: {e}", "dropped_prompts": len(self.buffer)}) self.buffer.clear() finally: self.phase = "stopping" await self.adapter.stop() if pump is not None: try: await asyncio.wait_for(pump, 10) except (TimeoutError, asyncio.CancelledError): pump.cancel() if self.turn_id is not None: self._interrupt("process exited") self.awaiting = False self.phase = "none" self._burst = None self.dirty = True self.save_state() if self.buffer and self.allow_starts and not self._stop_req: self.phase = "starting" self._burst = asyncio.create_task(self._run_burst(), name=f"burst-{self.name}") async def _watch(self, pump: asyncio.Task) -> None: tick = max(0.01, min(1.0, self.idle_grace / 10, self._turn_timeout / 10)) while True: await asyncio.wait({pump}, timeout=tick) if pump.done(): return now = self._now() busy = self.turn_id is not None or self.awaiting if self._stop_req and (self._stop_req != "drain" or not busy): if busy: self._interrupt(self._stop_req) return timeout = self._turn_timeout started = self._turn_started_at if self.turn_id else self._last_activity if busy and now - started > timeout: log.warning("%s: turn exceeded %.0fs; killing the harness", self.name, timeout) self._interrupt("timeout") self.stuck = True return if not busy and now - self._last_activity >= self.idle_grace: return async def _pump(self) -> None: try: async for ev in self.adapter.events(): self.on_adapter_event(ev) except Exception: log.exception("%s: adapter event stream failed", self.name) # ------------------------------------------------------------------ turns and usage def on_adapter_event(self, ev: AdapterEvent) -> None: k = ev.kind if k == "turn_started": if self.turn_id is None: self._begin_turn(ev.data) elif k == "message" and ev.message is not None: self._emit(self.log.message(ev.message, turn_id=self.turn_id)) elif k == "usage" and ev.usage is not None: self._account(ev) elif k == "turn_ended": if self.turn_id is not None: tokens = {m: u.model_dump() for m, u in self._turn_tokens.items()} self.emit("turn_ended", {"status": ev.status, "error": ev.error, "tokens": tokens, **{k2: v for k2, v in ev.data.items() if k2 != "tokens"}}) self._end_turn() elif k == "error": self.emit("error", {"error": ev.error, **ev.data}) self.dirty = True def _begin_turn(self, data: dict[str, Any]) -> None: self.state.turn_seq += 1 self.turn_id = f"t-{self.state.turn_seq}" self.stuck = False self.awaiting = False self._turn_started_at = self._now() self._turn_timeout = self._next_timeout or parse_duration(self.spec.turn_timeout).total_seconds() self._next_timeout = None self._turn_tokens = {} self.emit("turn_started", dict(data)) def _end_turn(self) -> None: self.state.turns += 1 self.turn_id = None self.awaiting = False self.idle_since = utcnow() self._last_activity = self._now() self.save_state() def _interrupt(self, reason: str) -> None: if self.turn_id is not None: self.emit("turn_interrupted", {"reason": reason}) self._end_turn() self.awaiting = False def _account(self, ev: AdapterEvent) -> None: model = ev.model or self.spec.model usage = ev.usage assert usage is not None if ev.cost_usd is not None: self.state.harness_estimate_usd = ev.cost_usd if ev.cumulative: prev = self.state.harness_cum.get(model) self.state.harness_cum[model] = usage if ev.baseline: return usage = usage_sub(usage, prev) self.state.tokens[model] = self.state.tokens.get(model, TokenUsage()) + usage self._turn_tokens[model] = self._turn_tokens.get(model, TokenUsage()) + usage # ------------------------------------------------------------------ control @property def live(self) -> bool: return self.phase != "none" async def stop(self, reason: str = "stop") -> None: """Kill the harness process if any (interrupting a turn); the agent stays assigned.""" self._stop_req = reason try: if self._burst is not None: await self._burst finally: self._stop_req = None self._defer_buffer(reason) async def drain(self) -> None: """No new starts; let the in-flight turn finish (bounded by its timeout); final sync.""" self.allow_starts = False if self._burst is not None: self._stop_req = "drain" try: await self._burst finally: self._stop_req = None self._defer_buffer("draining") await self.sync_now("drain") async def unassign(self) -> None: self.allow_starts = False await self.stop("unassign") await self.sync_now("unassign") async def sync_now(self, reason: str) -> None: if self.sync is None: return self.emit("sync", {"reason": reason}) self.save_state() try: await self.sync.push(self) self.last_sync = utcnow() self.dirty = False except Exception as e: # noqa: BLE001 log.warning("%s: sync failed: %s", self.name, e) def heartbeat(self) -> AgentHeartbeat: working = self.turn_id is not None or self.awaiting or (self.phase == "starting" and bool(self.buffer)) state = "working" if working else ("stuck" if self.stuck else "idle") return AgentHeartbeat( name=self.name, state=state, turn_id=self.turn_id, idle_since=None if working else self.idle_since, session_id=self.state.session_id, resumed=self.resumed, tokens=dict(self.state.tokens), harness_estimate_usd=self.state.harness_estimate_usd, turns=self.state.turns, last_sync=self.last_sync, msgboard_principal=self.identities.principal(self.name) if self.identities else None, )