"""Pi coding agent interface — v1's only coding agent. Runs `pi -p --mode json` and tails its JSONL stdout line by line, forwarding each event to a callback WHILE the agent works (the streaming crack, solved by construction). `--session-id` creates-or-continues, so running and continuing an agent are the same call: same session id = same context window. """ from __future__ import annotations import json import os import queue import signal import subprocess import threading import time from enum import Enum from functools import lru_cache from pathlib import Path from typing import Any, Callable, Optional from .data_types import PiRequest, PiResult from .utils import now_iso, operator_env PI_PATH = os.environ.get("PI_PATH", "pi") MODELS_JSON = os.environ.get("PI_MODELS_PATH", str(Path.home() / ".pi" / "agent" / "models.json")) RESULT_SNIPPET_CHARS = 20_000 # tool output rides along whole; clip only guards pathological cases ARG_VALUE_CHARS = 20_000 # args too — the UI scrolls, it must not be handed cut-off data LABEL_CHARS = 80 # "bash: " shown as the event name DEFAULT_AGENT_TIMEOUT_SECONDS = 1800 DEFAULT_AGENT_IDLE_TIMEOUT_SECONDS = 600 TERMINATE_GRACE_SECONDS = 2.0 # The arg that identifies a call at a glance, in the order tools tend to use. PRIMARY_ARGS = ("command", "path", "file_path", "pattern", "query", "url") class TimeoutLifecycle(str, Enum): STARTING = "starting" STREAMING = "streaming" TERMINATING = "terminating" TERMINAL = "terminal" def _count(value: str) -> int: """Parse pi's compact model-list counts (`272K`, `1.0M`).""" suffixes = {"K": 1_000, "M": 1_000_000} suffix = value[-1:].upper() if suffix in suffixes: return int(float(value[:-1]) * suffixes[suffix]) return int(value) @lru_cache(maxsize=1) def _pi_catalog() -> list[tuple[str, str, int]]: """Read pi's merged catalog, including built-in providers and custom models.""" try: result = subprocess.run( [PI_PATH, "--list-models"], capture_output=True, text=True, timeout=30, env=operator_env(), check=False, ) except (OSError, subprocess.TimeoutExpired): return [] if result.returncode != 0: return [] rows = [] for line in result.stdout.splitlines()[1:]: columns = line.split() if len(columns) < 3: continue try: rows.append((columns[0], columns[1], _count(columns[2]))) except ValueError: continue return rows def resolve_model(pattern: str) -> tuple[str, str]: """Resolve a model pattern to an explicit ``(provider, model_id)`` pair. Pi's catalog merges built-in models with ``~/.pi/agent/models.json``. Using that same merged view lets SSSF target direct providers such as ``openai/gpt-5.6-terra`` without re-registering built-in models locally. """ catalog = [(provider, model_id) for provider, model_id, _ in _pi_catalog()] if "/" in pattern: provider, model_id = pattern.split("/", 1) if (provider, model_id) in catalog: return provider, model_id matches = [(provider, model_id) for provider, model_id in catalog if pattern == model_id or pattern in model_id] exact = [match for match in matches if match[1] == pattern or match[1].endswith("/" + pattern)] if len(exact) == 1: return exact[0] if len(matches) == 1: return matches[0] if not matches: raise ValueError(f"model pattern {pattern!r} not found in pi --list-models — " "authenticate/register it or fix the config") raise ValueError(f"model pattern {pattern!r} is ambiguous: {matches}") def _context_tokens(usage: dict) -> int: """Tokens occupying the window after a turn. Mirrors pi's own `calculateContextTokens` (coding-agent `core/compaction/compaction.ts`), which is what pi compacts against and shows in its footer: prefer the provider's `totalTokens`, else sum the parts. Cache reads count — cached prompt is still prompt. """ total = usage.get("totalTokens") or 0 if total: return int(total) return int(sum(usage.get(part) or 0 for part in ("input", "output", "cacheRead", "cacheWrite"))) def context_window(provider: str, model_id: str) -> int: """The model's context ceiling from pi's merged model catalog.""" registry = json.loads(Path(MODELS_JSON).read_text()) for model in registry.get("providers", {}).get(provider, {}).get("models", []): if model.get("id") == model_id: return int(model.get("contextWindow") or 0) for listed_provider, listed_model, window in _pi_catalog(): if listed_provider == provider and listed_model == model_id: return window return 0 def _text_of(container: dict) -> str: """Join the text blocks of anything pi shapes as {content: [...]} — a message or a tool result.""" return "".join(part.get("text", "") for part in container.get("content", []) or [] if isinstance(part, dict) and part.get("type") == "text") def _clip(text: str, limit: int) -> str: return text if len(text) <= limit else text[:limit].rstrip() + "…" def _label(tool: str, args: dict) -> str: """One-line human name for a tool call: `bash: ls -la src`.""" value = next((args[key] for key in PRIMARY_ARGS if isinstance(args.get(key), str) and args[key].strip()), "") if not value: value = next((v for v in args.values() if isinstance(v, str) and v.strip()), "") value = " ".join(str(value).split()) return f"{tool}: {_clip(value, LABEL_CHARS)}" if value else tool class ToolCallTracker: """Folds pi's tool stream into ONE normalized record per completed call. pi announces a call as a `toolCall` content block, then emits tool_execution_start / _update / _end for it. Only the end carries the result, so that is where a record is emitted — one trace event per real tool call, the moment it returns, instead of three shapeless ones. The record carries the call's real span (`started_at`/`ended_at`), which the tracer writes to columns so the UI can lay tool calls on a time axis without parsing every payload. """ def __init__(self) -> None: self._open: dict[str, dict] = {} def observe(self, event: dict) -> Optional[dict]: """Returns the record for a finished tool call, else None.""" etype = event.get("type", "") if etype == "message_end": for block in event.get("message", {}).get("content", []) or []: if isinstance(block, dict) and block.get("type") == "toolCall": self._announce(block.get("id"), block.get("name"), block.get("arguments")) return None if etype == "tool_execution_start": self._announce(event.get("toolCallId"), event.get("toolName"), event.get("args")) return None if etype != "tool_execution_end": return None call_id = str(event.get("toolCallId") or "") opened = self._open.pop(call_id, {}) tool = str(event.get("toolName") or opened.get("tool") or "tool") args = event.get("args") or opened.get("args") or {} record = { "tool": tool, "tool_call_id": call_id, "args": {key: _clip(value, ARG_VALUE_CHARS) if isinstance(value, str) else value for key, value in args.items()}, "ok": not event.get("isError", False), "label": _label(tool, args), } result_text = _text_of(event.get("result") or {}) if result_text: record["result_snippet"] = _clip(result_text, RESULT_SNIPPET_CHARS) record["ended_at"] = now_iso() if opened.get("clock"): record["duration_ms"] = int((time.monotonic() - opened["clock"]) * 1000) if opened.get("started_at"): record["started_at"] = opened["started_at"] return record def _announce(self, call_id, tool, args) -> None: """First sighting starts the clock; a later sighting only fills gaps.""" if not call_id: return known = self._open.get(str(call_id), {}) self._open[str(call_id)] = { "tool": tool or known.get("tool", ""), "args": args or known.get("args", {}), "started_at": known.get("started_at") or now_iso(), # wall clock, for the row "clock": known.get("clock") or time.monotonic(), # monotonic, for duration } def run(request: PiRequest, on_event: Optional[Callable[[dict], None]] = None, on_spawn: Optional[Callable[[int], None]] = None, on_exit: Optional[Callable[[int], None]] = None, timeout_seconds: float | None = None, idle_timeout_seconds: float | None = None, *, clock: Callable[[], float] | None = None, sleeper: Callable[[float], None] | None = None, popen_factory: Callable[..., Any] | None = None, kill_process_group: Callable[[int, int], None] | None = None, thread_factory: Callable[..., Any] | None = None) -> PiResult: """Run one non-interactive pi turn.""" clock = clock or time.monotonic sleeper = sleeper or time.sleep popen_factory = popen_factory or subprocess.Popen kill_process_group = kill_process_group or os.killpg thread_factory = thread_factory or threading.Thread provider, model_id = resolve_model(request.model) cmd = [ PI_PATH, "-p", "--mode", "json", "--provider", provider, "--model", model_id, "--thinking", request.thinking, "--session-id", request.session_id, "--session-dir", request.session_dir, "--system-prompt", request.system_prompt, ] if request.tools: cmd += ["--tools", ",".join(request.tools)] for extension in request.extensions: cmd += ["-e", extension] cmd.append(request.prompt) raw_path = Path(request.raw_output_path) raw_path.parent.mkdir(parents=True, exist_ok=True) stderr_path = raw_path.with_name("stderr.log") timeout = (timeout_seconds if timeout_seconds is not None else float( os.environ.get("FACTORY_AGENT_TIMEOUT_SECONDS", DEFAULT_AGENT_TIMEOUT_SECONDS))) idle_timeout = (idle_timeout_seconds if idle_timeout_seconds is not None else float( os.environ.get("FACTORY_AGENT_IDLE_TIMEOUT_SECONDS", DEFAULT_AGENT_IDLE_TIMEOUT_SECONDS))) result = PiResult(session_id=request.session_id, context_window=context_window(provider, model_id)) process: Any = None timed_out = False timeout_kind: str | None = None error_text: str | None = None lifecycle = TimeoutLifecycle.STARTING terminal_receipt_emitted = False try: process = popen_factory(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, cwd=request.cwd, env=operator_env(), start_new_session=True) if on_spawn: on_spawn(process.pid) stream_records: queue.Queue[dict] = queue.Queue() last_output_clock: float | None = None def consume_stdout() -> None: nonlocal last_output_clock, lifecycle with raw_path.open("a") as raw: assert process is not None and process.stdout is not None for line in process.stdout: last_output_clock = clock() lifecycle = TimeoutLifecycle.STREAMING raw.write(line) raw.flush() try: event = json.loads(line.strip()) except json.JSONDecodeError: continue if event.get("type") == "message_end": message = event.get("message", {}) if message.get("role") == "assistant": text = _text_of(message) if text: result.text = text usage = message.get("usage", {}) or {} turn = _context_tokens(usage) result.tokens += turn result.usage.add_turn(usage, turn) if turn and message.get("stopReason") not in ("aborted", "error"): result.context_tokens = turn result.cost += (usage.get("cost", {}) or {}).get("total", 0.0) or 0.0 stream_records.put(event) stderr_chunks: list[str] = [] def consume_stderr() -> None: assert process is not None and process.stderr is not None for line in process.stderr: stderr_chunks.append(line) reader = thread_factory(target=consume_stdout, daemon=True) errors = thread_factory(target=consume_stderr, daemon=True) reader.start() errors.start() started_clock = clock() def forward_stream_records() -> None: while True: try: event = stream_records.get_nowait() except queue.Empty: return if on_event: on_event(event) while process.poll() is None: forward_stream_records() now = clock() if now - started_clock >= timeout: timeout_kind = "wall" break if (lifecycle is TimeoutLifecycle.STREAMING and last_output_clock is not None and now - last_output_clock >= idle_timeout): timeout_kind = "idle" break sleep_seconds = timeout - (now - started_clock) if lifecycle is TimeoutLifecycle.STREAMING and last_output_clock is not None: sleep_seconds = min(sleep_seconds, idle_timeout - (now - last_output_clock)) sleeper(min(0.05, max(0.0, sleep_seconds))) if timeout_kind is not None: timed_out = True lifecycle = TimeoutLifecycle.TERMINATING grace_started_clock = clock() kill_process_group(process.pid, signal.SIGTERM) try: process.wait(timeout=TERMINATE_GRACE_SECONDS) except subprocess.TimeoutExpired: pass grace_remaining = TERMINATE_GRACE_SECONDS - (clock() - grace_started_clock) if grace_remaining > 0: sleeper(grace_remaining) try: kill_process_group(process.pid, signal.SIGKILL) except ProcessLookupError: pass try: result.returncode = process.wait(timeout=TERMINATE_GRACE_SECONDS) except subprocess.TimeoutExpired as exc: raise TimeoutError("pi process group did not exit after SIGKILL") from exc else: result.returncode = process.wait() lifecycle = TimeoutLifecycle.TERMINAL reader.join(timeout=TERMINATE_GRACE_SECONDS) errors.join(timeout=TERMINATE_GRACE_SECONDS) if reader.is_alive() or errors.is_alive(): raise TimeoutError("pi stream drain timed out") forward_stream_records() stderr = "".join(stderr_chunks) stderr_path.write_text(stderr) if timed_out: limit = timeout if timeout_kind == "wall" else idle_timeout error_text = f"pi {timeout_kind} timeout: timed out after {limit:g}s" raise TimeoutError(error_text) if result.returncode != 0 and not result.text: error_text = f"pi exited {result.returncode}: {stderr.strip()[-800:]}" raise RuntimeError(error_text) return result except BaseException as exc: error_text = error_text or str(exc) raise finally: lifecycle = TimeoutLifecycle.TERMINAL if process is not None and on_exit: on_exit(process.pid) if on_event and not terminal_receipt_emitted: terminal_receipt_emitted = True returncode = result.returncode if process is not None else None sig = -returncode if returncode is not None and returncode < 0 else None usage = result.usage.model_dump() if result.tokens else None on_event({"type": "agent_attempt_end", "returncode": returncode, "signal": sig, "timedOut": timed_out, "timeoutKind": timeout_kind, "stderrPath": str(stderr_path), "tokens": result.tokens or None, "usage": usage, "error": error_text, "lifecycle": lifecycle.value})