Spaces:
Paused
Paused
| """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-<n>) | |
| 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 | |
| def state_path(self) -> Path: | |
| return self.sessions_dir / f"{self.name}.state.json" | |
| 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 | |
| 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, | |
| ) | |