"""Long-lived Pi RPC client for gateway services. Extracted and generalised from the memo-inbox gateway (``pi-workspaces/memo-inbox/telegram-gateway/gateway.py``, class ``PiRPC``), which has run this pattern in production since 2026-07. Compared with the original it adds the four gaps that both existing scenarios shared: * an explicit minimal ``env`` so provider and backend API keys never reach the node process (``docs/gateway-patterns.md`` pattern 8); * ``start_new_session`` plus process-group termination so a timeout cannot leave orphaned node children burning provider quota (pattern 7); * the loading-isolation flags from ``docs/isolation-baseline.md`` are part of the launch contract rather than something each caller remembers; * a per-turn deadline enforced with RPC ``abort`` instead of killing the process. It deliberately keeps the original's proven mechanics: strict ``\\n``-only JSONL framing, request/response correlation by ``id``, ``agent_settled`` as the terminal event, and deterministic receipts harvested from ``tool_execution_end`` rather than from model prose (pattern 3). Stdlib only. Requires Python 3.11+ (``asyncio.timeout``). Protocol reference: ``docs/pi-runtime-notes.md`` §15. """ from __future__ import annotations import asyncio import contextlib import json import logging import os import signal import time from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, Awaitable, Callable, Iterable, Sequence LOG = logging.getLogger("pi_rpc") __all__ = [ "PiLaunchConfig", "PiRpcClient", "PiRpcError", "PiRpcRejected", "PiTurnAborted", "TurnResult", "Usage", ] # Environment variables a pi subprocess legitimately needs. Everything else -- # in particular every ``*_API_KEY`` and ``*_TOKEN`` belonging to the host # service -- is withheld. The provider credential is read by pi itself from # ~/.pi/agent/models.json and must not be passed through the environment. DEFAULT_ENV_ALLOWLIST: tuple[str, ...] = ( "PATH", "HOME", "LANG", "LC_ALL", "LC_CTYPE", "TZ", "TERM", "NODE_OPTIONS", "SSL_CERT_FILE", "SSL_CERT_DIR", "NO_PROXY", "no_proxy", ) class PiRpcError(RuntimeError): """Transport or lifecycle failure.""" class PiRpcRejected(PiRpcError): """Pi returned ``success: false`` for a command.""" class PiTurnAborted(PiRpcError): """The turn exceeded its deadline and was aborted.""" @dataclass(frozen=True) class Usage: """Token and cost accounting for one turn. Populated from ``message_update.usage`` deltas. See ``docs/pi-runtime-notes.md`` §16 -- observability needs no extra instrumentation because the stream already carries this. """ input: int = 0 output: int = 0 cache_read: int = 0 cache_write: int = 0 total_tokens: int = 0 cost_total: float = 0.0 def merged(self, raw: dict[str, Any]) -> "Usage": cost = raw.get("cost") or {} return Usage( input=self.input + int(raw.get("input") or 0), output=self.output + int(raw.get("output") or 0), cache_read=self.cache_read + int(raw.get("cacheRead") or 0), cache_write=self.cache_write + int(raw.get("cacheWrite") or 0), total_tokens=self.total_tokens + int(raw.get("totalTokens") or 0), cost_total=self.cost_total + float(cost.get("total") or 0.0), ) @dataclass class ToolCallRecord: """One completed tool execution.""" tool_name: str args: dict[str, Any] text: str details: Any is_error: bool @dataclass class TurnResult: """Everything one ``prompt`` produced. ``receipts`` holds the text of tool results whose tool is named in ``PiLaunchConfig.receipt_tools``. Callers must render user-visible outcomes of state changes from ``receipts`` (or from ``tool_calls``), never from ``replies`` -- see ``docs/gateway-patterns.md`` pattern 3 for the production failure this prevents. """ replies: list[str] = field(default_factory=list) receipts: list[str] = field(default_factory=list) tool_calls: list[ToolCallRecord] = field(default_factory=list) usage: Usage = field(default_factory=Usage) latency_seconds: float = 0.0 model: str = "" thinking: str = "" aborted: bool = False extension_errors: list[dict[str, Any]] = field(default_factory=list) @property def text(self) -> str: """The final assistant message, or an empty string.""" return self.replies[-1] if self.replies else "" @dataclass(frozen=True) class PiLaunchConfig: """Launch contract for a dedicated Pi agent. The isolation defaults implement ``docs/isolation-baseline.md``. Overriding them widens what the agent loads or can call, so each override should be justified in the scenario's ``profile.toml``. """ pi_bin: str workspace: Path session_dir: Path provider: str model: str thinking: str = "medium" display_name: str = "Pi Agent" # --- session identity ------------------------------------------------- # A stable prefix; the client appends a rotation counter so that history # stays auditable instead of being summarised away. Empty means --no-session. session_id_prefix: str = "" # --- layer 1: loading isolation -------------------------------------- extensions: tuple[Path, ...] = () skills: tuple[Path, ...] = () no_extensions: bool = True no_skills: bool = True no_prompt_templates: bool = True no_themes: bool = True no_context_files: bool = False approve: bool = True # --- layer 3: capability --------------------------------------------- # Prefer no_builtin_tools over an explicit ``tools`` allowlist: --tools is a # registry-level filter, which prevents an extension from registering tools # at runtime. See docs/pi-runtime-notes.md section 4. no_builtin_tools: bool = True tools: tuple[str, ...] = () exclude_tools: tuple[str, ...] = () # --- budgets ---------------------------------------------------------- turn_deadline_seconds: float = 180.0 startup_timeout_seconds: float = 60.0 stop_timeout_seconds: float = 10.0 # --- rotation --------------------------------------------------------- rotate_after_prompts: int = 24 rotate_after_messages: int = 60 # --- receipts --------------------------------------------------------- receipt_tools: frozenset[str] = frozenset() # --- environment ------------------------------------------------------ env_allowlist: tuple[str, ...] = DEFAULT_ENV_ALLOWLIST extra_env: tuple[tuple[str, str], ...] = () def __post_init__(self) -> None: if not self.no_builtin_tools and not self.tools: LOG.warning( "PiLaunchConfig for %r enables built-in tools with no allowlist; " "bash/edit/write will be active", self.display_name, ) if self.skills and not self._read_reachable(): LOG.warning( "PiLaunchConfig for %r loads skills but no 'read' tool is reachable; " "the skills section will be omitted from the system prompt and the " "skill bodies will be unloadable " "(see docs/pi-runtime-notes.md section 1)", self.display_name, ) def _read_reachable(self) -> bool: """Whether an active tool named ``read`` can plausibly exist. Pi only emits the skills section when ``read`` is active. With ``no_builtin_tools`` the extension is expected to register a restricted ``read`` override; with an explicit allowlist ``read`` must be named. """ if self.tools: return "read" in self.tools return self.no_builtin_tools or not self.no_extensions or bool(self.extensions) def build_env(self) -> dict[str, str]: env = {k: os.environ[k] for k in self.env_allowlist if k in os.environ} env.setdefault("HOME", str(Path.home())) env.update(dict(self.extra_env)) return env def build_args(self, session_id: str | None) -> list[str]: args: list[str] = [self.pi_bin, "--mode", "rpc"] if session_id: args += ["--session-id", session_id, "--session-dir", str(self.session_dir)] else: args += ["--no-session"] if self.no_builtin_tools: args += ["--no-builtin-tools"] if self.tools: args += ["--tools", ",".join(self.tools)] if self.exclude_tools: args += ["--exclude-tools", ",".join(self.exclude_tools)] if self.no_extensions: args += ["--no-extensions"] for path in self.extensions: args += ["-e", str(path)] if self.no_skills: args += ["--no-skills"] for path in self.skills: args += ["--skill", str(path)] if self.no_prompt_templates: args += ["--no-prompt-templates"] if self.no_themes: args += ["--no-themes"] if self.no_context_files: args += ["--no-context-files"] if self.approve: args += ["--approve"] args += [ "--provider", self.provider, "--model", self.model, "--thinking", self.thinking, "--name", self.display_name, ] return args class PiRpcClient: """A single long-lived ``pi --mode rpc`` process. One client owns one conversation. Instantiate one per Telegram chat, per ticket, or per whatever your unit of continuity is; do not multiplex unrelated conversations through one client, because they would share a session and contaminate each other. Not safe for concurrent ``prompt`` calls -- an internal lock serialises commands, so callers queue rather than interleave. """ def __init__( self, config: PiLaunchConfig, *, on_event: Callable[[dict[str, Any]], Awaitable[None] | None] | None = None, ) -> None: self.config = config self._on_event = on_event self._process: asyncio.subprocess.Process | None = None self._events: asyncio.Queue[dict[str, Any]] = asyncio.Queue() self._lock = asyncio.Lock() self._stdout_task: asyncio.Task[None] | None = None self._stderr_task: asyncio.Task[None] | None = None self._stderr_tail: list[str] = [] self._prompt_count = 0 self._rotation = 0 # -- lifecycle --------------------------------------------------------- @property def running(self) -> bool: return self._process is not None and self._process.returncode is None @property def pid(self) -> int | None: return self._process.pid if self._process else None def _session_id(self) -> str | None: if not self.config.session_id_prefix: return None return f"{self.config.session_id_prefix}-{self._rotation:04d}" async def start(self) -> None: if self.running: return cfg = self.config cfg.session_dir.mkdir(parents=True, exist_ok=True) args = cfg.build_args(self._session_id()) LOG.info("Starting pi: %s", " ".join(args)) self._process = await asyncio.create_subprocess_exec( *args, cwd=str(cfg.workspace), env=cfg.build_env(), stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, # Own process group so a stuck node tree can be killed wholesale. start_new_session=True, ) self._stdout_task = asyncio.create_task(self._read_stdout(), name="pi-rpc-stdout") self._stderr_task = asyncio.create_task(self._read_stderr(), name="pi-rpc-stderr") try: async with asyncio.timeout(cfg.startup_timeout_seconds): await self._command({"type": "get_state"}, wait_for_settled=False, use_lock=False) except TimeoutError as exc: tail = " | ".join(self._stderr_tail[-5:]) await self.stop() raise PiRpcError(f"pi did not answer get_state within " f"{cfg.startup_timeout_seconds}s: {tail}") from exc LOG.info("pi RPC ready pid=%s session=%s", self.pid, self._session_id()) async def stop(self) -> None: process, self._process = self._process, None for task in (self._stdout_task, self._stderr_task): if task: task.cancel() for task in (self._stdout_task, self._stderr_task): if task: with contextlib.suppress(asyncio.CancelledError, Exception): await task self._stdout_task = self._stderr_task = None if process is None or process.returncode is not None: return # Terminate the whole group: pi is a node CLI and may have children. with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(os.getpgid(process.pid), signal.SIGTERM) try: async with asyncio.timeout(self.config.stop_timeout_seconds): await process.wait() except TimeoutError: LOG.warning("pi pid=%s ignored SIGTERM, sending SIGKILL to group", process.pid) with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(os.getpgid(process.pid), signal.SIGKILL) with contextlib.suppress(Exception): await process.wait() async def __aenter__(self) -> "PiRpcClient": await self.start() return self async def __aexit__(self, *_exc: object) -> None: await self.stop() # -- stream plumbing --------------------------------------------------- async def _read_stdout(self) -> None: assert self._process and self._process.stdout buffer = b"" # Strict JSONL: split on \n only. A generic line reader is not # protocol-compliant because U+2028/U+2029 are legal inside JSON # strings. See docs/pi-runtime-notes.md section 15. while chunk := await self._process.stdout.read(65536): buffer += chunk while b"\n" in buffer: line, buffer = buffer.split(b"\n", 1) line = line.rstrip(b"\r") if not line: continue try: event = json.loads(line) except json.JSONDecodeError: LOG.warning("Discarding malformed pi RPC record (%d bytes)", len(line)) continue await self._events.put(event) if self._on_event is not None: try: result = self._on_event(event) if asyncio.iscoroutine(result): await result except Exception: LOG.exception("on_event callback failed") LOG.warning("pi RPC stdout closed") async def _read_stderr(self) -> None: assert self._process and self._process.stderr while line := await self._process.stderr.readline(): text = line.decode(errors="replace").rstrip() if not text: continue # Kept even on success: the memo/curator scenarios both discarded # stderr and then could not explain model failures afterwards. self._stderr_tail.append(text) del self._stderr_tail[:-50] LOG.info("pi stderr: %s", text) async def _send(self, payload: dict[str, Any]) -> None: await self.start() assert self._process and self._process.stdin blob = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode() self._process.stdin.write(blob + b"\n") await self._process.stdin.drain() # -- commands ---------------------------------------------------------- async def command(self, payload: dict[str, Any], *, wait_for_settled: bool = False, deadline: float | None = None) -> tuple[dict[str, Any], TurnResult]: async with self._lock: return await self._command(payload, wait_for_settled=wait_for_settled, use_lock=False, deadline=deadline) async def _command( self, payload: dict[str, Any], *, wait_for_settled: bool, use_lock: bool = True, deadline: float | None = None, ) -> tuple[dict[str, Any], TurnResult]: if use_lock: async with self._lock: return await self._command(payload, wait_for_settled=wait_for_settled, use_lock=False, deadline=deadline) request_id = payload.setdefault("id", f"req-{time.monotonic_ns():x}") # Drop anything left over from an aborted turn so correlation is clean. while not self._events.empty(): self._events.get_nowait() started = time.monotonic() result = TurnResult(model=self.config.model, thinking=self.config.thinking) await self._send(payload) budget = deadline if deadline is not None else self.config.turn_deadline_seconds response: dict[str, Any] | None = None try: async with asyncio.timeout(budget): response = await self._consume(request_id, result, wait_for_settled) except TimeoutError: result.aborted = True result.latency_seconds = time.monotonic() - started LOG.warning("pi turn exceeded %.0fs; aborting", budget) with contextlib.suppress(Exception): await self._send({"id": f"{request_id}-abort", "type": "abort"}) raise PiTurnAborted(f"pi turn exceeded {budget:.0f}s") from None result.latency_seconds = time.monotonic() - started assert response is not None return response, result async def _consume(self, request_id: str, result: TurnResult, wait_for_settled: bool) -> dict[str, Any]: """Collect events until the command is complete. ``success: true`` only means *accepted*; post-acceptance failures arrive as events. So for a prompt we wait for ``agent_settled`` as well. """ response: dict[str, Any] | None = None pending_args: dict[str, dict[str, Any]] = {} while True: event = await self._events.get() kind = event.get("type") if kind == "response" and event.get("id") == request_id: response = event if not wait_for_settled or not event.get("success"): return response elif kind == "message_update": raw = event.get("usage") if isinstance(raw, dict): result.usage = result.usage.merged(raw) elif kind == "message_end": text = _assistant_text(event.get("message") or {}) if text: result.replies.append(text) elif kind == "tool_execution_start": call_id = str(event.get("toolCallId") or "") if call_id: pending_args[call_id] = event.get("args") or {} elif kind == "tool_execution_end": name = str(event.get("toolName") or "") payload = event.get("result") or {} text = _result_text(payload) record = ToolCallRecord( tool_name=name, args=pending_args.pop(str(event.get("toolCallId") or ""), {}), text=text, details=payload.get("details"), is_error=bool(event.get("isError")), ) result.tool_calls.append(record) # Deterministic receipts, not model prose. Pattern 3. if name in self.config.receipt_tools and text and not record.is_error: result.receipts.append(text) elif kind == "extension_error": LOG.error("pi extension error: %s", event) result.extension_errors.append(event) elif kind == "agent_settled" and response is not None: return response # -- high level -------------------------------------------------------- async def prompt(self, message: str, *, deadline: float | None = None, images: Sequence[dict[str, Any]] | None = None) -> TurnResult: """Send one user turn and return everything it produced. Rotates the session first when either rotation bound is reached: pi has no session TTL and auto-compaction does not fire on large-context models (``docs/pi-runtime-notes.md`` sections 17-18). """ await self._maybe_rotate() payload: dict[str, Any] = {"type": "prompt", "message": message} if images: payload["images"] = list(images) response, result = await self.command(payload, wait_for_settled=True, deadline=deadline) if not response.get("success"): raise PiRpcRejected(str(response.get("error") or "pi rejected the prompt")) self._prompt_count += 1 return result async def _maybe_rotate(self) -> None: cfg = self.config if not cfg.session_id_prefix: return if cfg.rotate_after_prompts and self._prompt_count >= cfg.rotate_after_prompts: LOG.info("Rotating pi session: prompt_count=%d", self._prompt_count) await self.rotate_session() return if cfg.rotate_after_messages: state = await self.get_state() count = int(state.get("messageCount") or 0) if count >= cfg.rotate_after_messages: LOG.info("Rotating pi session: message_count=%d", count) await self.rotate_session() async def rotate_session(self) -> None: """Start a fresh session, preserving the previous file for audit. Restarting the process with the next ``--session-id`` is preferred over the ``new_session`` command because the session id then encodes the rotation, so history remains greppable on disk. """ self._rotation += 1 self._prompt_count = 0 await self.stop() await self.start() async def get_state(self) -> dict[str, Any]: response, _ = await self.command({"type": "get_state"}) return response.get("data") or {} async def get_session_stats(self) -> dict[str, Any]: response, _ = await self.command({"type": "get_session_stats"}) return response.get("data") or {} async def set_model(self, model: str) -> None: response, _ = await self.command({"type": "set_model", "model": model}) if not response.get("success"): raise PiRpcRejected(str(response.get("error") or f"cannot select model {model}")) self.config = replace(self.config, model=model) async def set_thinking_level(self, level: str) -> None: response, _ = await self.command({"type": "set_thinking_level", "level": level}) if not response.get("success"): raise PiRpcRejected(str(response.get("error") or f"cannot set thinking {level}")) self.config = replace(self.config, thinking=level) async def abort(self) -> None: await self.command({"type": "abort"}) async def compact(self, instructions: str | None = None) -> None: payload: dict[str, Any] = {"type": "compact"} if instructions: payload["customInstructions"] = instructions await self.command(payload, wait_for_settled=False) @property def stderr_tail(self) -> list[str]: """Recent stderr lines, retained even on success for diagnostics.""" return list(self._stderr_tail) # -- helpers --------------------------------------------------------------- def _assistant_text(message: dict[str, Any]) -> str: """Concatenate the text blocks of an assistant message. Thinking blocks are excluded: they are reasoning, not output. """ content = message.get("content") if isinstance(content, str): return content.strip() if not isinstance(content, Iterable): return "" parts = [ block["text"] for block in content if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str) ] return "\n".join(parts).strip() def _result_text(result: dict[str, Any]) -> str: content = result.get("content") or [] if not isinstance(content, Iterable): return "" parts = [ block["text"] for block in content if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str) ] return "\n".join(parts).strip()