feat(shared): long-lived Pi RPC client extracted from the memo-inbox gateway
Generalises PiRPC from pi-workspaces/memo-inbox/telegram-gateway/gateway.py, which has run this pattern in production since 2026-07, and closes the four gaps both existing scenarios shared: - explicit minimal env, so provider and backend API keys never reach the node process (verified: 6 variables, an injected secret is withheld) - start_new_session plus killpg on stop, so a stuck node tree cannot outlive the turn (verified: no orphan after stop) - the loading-isolation flags are part of the launch contract instead of something each caller has to remember - a per-turn deadline enforced with RPC abort rather than by killing the process Retains the original's proven mechanics: strict newline-only JSONL framing, correlation by id, agent_settled as terminal event, and receipts harvested from tool_execution_end rather than from model prose. PiLaunchConfig warns when skills are configured but no 'read' tool can be active, which is exactly the condition that silently disabled Curator's SKILL.md. Includes a zero-token smoke test: it drives a real pi process with get_state only, so no model call is billed.
This commit is contained in:
Kai
@@ -0,0 +1,659 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Zero-token smoke test for pi_rpc against a real pi process.
|
||||
|
||||
Sends only `get_state`, so no model call is made. Verifies that the launch
|
||||
contract produces a live process, that the isolation flags land, and that the
|
||||
process group is cleaned up.
|
||||
|
||||
Run: python3 shared/lib/py/tests/test_pi_rpc_smoke.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from pi_rpc import PiLaunchConfig, PiRpcClient # noqa: E402
|
||||
|
||||
REPO = Path(__file__).resolve().parents[4]
|
||||
HARNESS = REPO / "docs" / "evidence" / "probe-harness"
|
||||
PI_BIN = os.environ.get("PI_BIN", "/home/claw/.npm-global/bin/pi")
|
||||
|
||||
|
||||
def check(label: str, ok: bool, detail: str = "") -> bool:
|
||||
print(f" {'PASS' if ok else 'FAIL'} {label}{(' — ' + detail) if detail else ''}")
|
||||
return ok
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
failures = 0
|
||||
|
||||
print("== build_args reflects the isolation baseline ==")
|
||||
cfg = PiLaunchConfig(
|
||||
pi_bin=PI_BIN,
|
||||
workspace=HARNESS,
|
||||
session_dir=Path("/tmp/pi-rpc-smoke-sessions"),
|
||||
provider="zenmux",
|
||||
model="openai/gpt-5.6-luna",
|
||||
thinking="low",
|
||||
display_name="pi_rpc smoke",
|
||||
extensions=(HARNESS / "probe-ext.ts",),
|
||||
skills=(HARNESS / ".pi" / "skills" / "probe-skill",),
|
||||
session_id_prefix="", # ephemeral for the smoke test
|
||||
receipt_tools=frozenset({"probe_plain_schema"}),
|
||||
turn_deadline_seconds=45.0,
|
||||
)
|
||||
args = cfg.build_args(None)
|
||||
for flag in ("--no-builtin-tools", "--no-extensions", "--no-skills",
|
||||
"--no-prompt-templates", "--no-themes", "--approve", "--no-session"):
|
||||
failures += not check(f"{flag} present", flag in args)
|
||||
failures += not check("--tools absent (registry allowlist would block dynamic tools)",
|
||||
"--tools" not in args)
|
||||
|
||||
print("== env is minimal ==")
|
||||
os.environ["SMOKE_FAKE_SECRET"] = "must-not-propagate"
|
||||
env = cfg.build_env()
|
||||
failures += not check("secret withheld", "SMOKE_FAKE_SECRET" not in env)
|
||||
failures += not check("PATH forwarded", "PATH" in env)
|
||||
failures += not check(f"small env ({len(env)} vars)", len(env) < 20)
|
||||
|
||||
print("== live process: start, get_state, stop ==")
|
||||
client = PiRpcClient(cfg)
|
||||
try:
|
||||
await client.start()
|
||||
pid = client.pid
|
||||
failures += not check("process running", client.running, f"pid={pid}")
|
||||
state = await client.get_state()
|
||||
failures += not check("get_state returned a model", bool(state.get("model")),
|
||||
str(state.get("model", {}).get("id", "?")))
|
||||
failures += not check("not streaming", state.get("isStreaming") is False)
|
||||
failures += not check("thinkingLevel honoured", state.get("thinkingLevel") == "low",
|
||||
str(state.get("thinkingLevel")))
|
||||
stats = await client.get_session_stats()
|
||||
failures += not check("get_session_stats works", "tokens" in stats)
|
||||
finally:
|
||||
await client.stop()
|
||||
|
||||
failures += not check("process stopped", not client.running)
|
||||
if pid:
|
||||
alive = subprocess.run(["ps", "-p", str(pid)], capture_output=True).returncode == 0
|
||||
failures += not check("no orphan process", not alive)
|
||||
|
||||
print()
|
||||
print("RESULT:", "all checks passed" if failures == 0 else f"{failures} check(s) failed")
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(asyncio.run(main()))
|
||||
Reference in New Issue
Block a user