pidriver 0.0.1__py3-none-any.whl

pidriver/__init__.py

@@ -0,0 +1,132 @@
+"""pidriver — an async Python driver for the pi coding agent.
+
+Drives ``pi --mode rpc`` over a JSONL stdin/stdout protocol, with first-class
+support for *isolated* agents (private config dir, scrubbed environment) so a pi
+instance developing a project can't touch the operator's global pi setup.
+
+Layers:
+
+* :mod:`pidriver.config` — :class:`PiConfig` (isolation knobs, ``to_argv``/``to_env``).
+* :mod:`pidriver._transport` — :class:`PiTransport` protocol + :class:`SubprocessTransport`
+  (the swap boundary).
+* :mod:`pidriver.events` — the typed :class:`Event` hierarchy + :func:`parse_event`.
+* :mod:`pidriver.interaction` — answering the agent's questions (handlers + policies).
+* :mod:`pidriver.session` — :class:`PiSession` (the live RPC session).
+* :mod:`pidriver.client` — :class:`PiClient` (start sessions).
+* :mod:`pidriver.manager` — :class:`PiSessionManager` (registry, idle reaper, ``stop_all``).
+* :mod:`pidriver.usage` — :class:`UsageTotals`.
+* :mod:`pidriver.errors` — exception hierarchy.
+
+Quick start::
+
+    from pidriver import PiClient, PiConfig, AutoApprove, MessageDelta, AgentEnd
+
+    client = PiClient(PiConfig(binary="omp", api_key=KEY, model="anthropic/claude-..."))
+    session = await client.start("Add a healthcheck endpoint", cwd="/srv/proj",
+                                 interaction_handler=AutoApprove())
+    async for event in session:
+        match event:
+            case MessageDelta(text=text):
+                print(text, end="")
+            case AgentEnd():
+                break
+"""
+
+from __future__ import annotations
+
+from ._transport import PiTransport, SubprocessTransport
+from .client import PiClient, TransportFactory
+from .config import PROVIDER_API_KEY_ENV, PiConfig
+from .errors import (
+    PiCommandError,
+    PiDriverError,
+    PiNotFoundError,
+    PiProtocolError,
+    PiStartError,
+    PiTimeoutError,
+    SessionNotFoundError,
+    TransportClosedError,
+)
+from .events import (
+    AgentEnd,
+    AgentStart,
+    Channel,
+    Decision,
+    Error,
+    Event,
+    InteractionKind,
+    InteractionRequest,
+    MessageComplete,
+    MessageDelta,
+    ToolEnd,
+    ToolStart,
+    ToolUpdate,
+    Usage,
+    parse_event,
+)
+from .interaction import (
+    AskHost,
+    AutoApprove,
+    DenyAll,
+    InteractionHandler,
+    InteractionResponse,
+    chain,
+    response_to_wire,
+)
+from .manager import ManagedSession, PiSessionManager
+from .session import PiSession
+from .usage import UsageTotals
+
+__version__ = "0.0.1"
+
+__all__ = [
+    "__version__",
+    # config
+    "PiConfig",
+    "PROVIDER_API_KEY_ENV",
+    # transport
+    "PiTransport",
+    "SubprocessTransport",
+    # client / session
+    "PiClient",
+    "TransportFactory",
+    "PiSession",
+    # manager
+    "PiSessionManager",
+    "ManagedSession",
+    # usage
+    "UsageTotals",
+    # events
+    "Event",
+    "AgentStart",
+    "MessageDelta",
+    "MessageComplete",
+    "ToolStart",
+    "ToolUpdate",
+    "ToolEnd",
+    "Usage",
+    "AgentEnd",
+    "Error",
+    "Channel",
+    "parse_event",
+    # interaction
+    "InteractionRequest",
+    "InteractionKind",
+    "InteractionResponse",
+    "InteractionHandler",
+    "Decision",
+    "AskHost",
+    "AutoApprove",
+    "DenyAll",
+    "chain",
+    "response_to_wire",
+    # errors
+    "PiDriverError",
+    "PiNotFoundError",
+    "PiStartError",
+    "PiProtocolError",
+    "PiCommandError",
+    "PiTimeoutError",
+    "TransportClosedError",
+    "SessionNotFoundError",
+]

pidriver/_transport.py

@@ -0,0 +1,272 @@
+"""The transport boundary: process lifecycle + JSONL framing for ``pi --mode rpc``.
+
+This module is intentionally **dumb**. It spawns the pi RPC subprocess, writes
+JSON commands as newline-delimited lines, and yields parsed JSON objects (events
+*and* responses, undifferentiated) from stdout. It does **not** correlate
+request/response ids, dispatch events, or know any command semantics — those live
+in the higher session/client layer.
+
+Keeping this layer thin makes it the **swap boundary**: the public surface is the
+:class:`PiTransport` protocol, and :class:`SubprocessTransport` is one concrete
+implementation. A future omp-rpc-backed transport can satisfy the same protocol
+without touching the session/manager code above it.
+
+Framing follows pi's RPC contract strictly: records are split on ``\\n`` only, and
+a trailing ``\\r`` is stripped. We never split on U+2028/U+2029 (which are valid
+inside JSON strings), so we do manual buffering rather than a generic line reader.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import shutil
+from collections.abc import AsyncIterator, Callable
+from typing import Any, Protocol, runtime_checkable
+
+from .config import PiConfig
+from .errors import (
+    PiNotFoundError,
+    PiProtocolError,
+    PiStartError,
+    TransportClosedError,
+)
+
+__all__ = ["PiTransport", "SubprocessTransport"]
+
+# Generous per-line cap. pi can emit large tool-result lines; we still bound memory
+# to avoid an unbounded buffer if the stream is malformed and never sees a newline.
+_MAX_LINE_BYTES = 64 * 1024 * 1024
+_READ_CHUNK = 1 << 16
+
+
+@runtime_checkable
+class PiTransport(Protocol):
+    """The minimal surface every transport must provide.
+
+    Contract:
+
+    * exactly **one** consumer may call :meth:`receive` / iterate at a time;
+    * :meth:`receive` returns ``None`` once the process's stdout reaches EOF;
+    * :meth:`send` and :meth:`receive` raise :class:`TransportClosedError` if
+      called before :meth:`start` or after the process has gone.
+    """
+
+    async def start(self) -> None: ...
+    async def send(self, obj: dict[str, Any]) -> None: ...
+    async def receive(self) -> dict[str, Any] | None: ...
+    def __aiter__(self) -> AsyncIterator[dict[str, Any]]: ...
+    async def aclose(self) -> None: ...
+
+    @property
+    def pid(self) -> int | None: ...
+    @property
+    def returncode(self) -> int | None: ...
+
+
+class SubprocessTransport:
+    """:class:`PiTransport` backed by a ``pi --mode rpc`` child process.
+
+    Args:
+        config: how to spawn pi (argv + env + cwd).
+        on_stderr: optional callback invoked with each decoded stderr line; pi
+            writes diagnostics there. stderr is never interleaved into the dict
+            stream returned by :meth:`receive`.
+        spawn: injection seam for tests — defaults to
+            :func:`asyncio.create_subprocess_exec`.
+        terminate_timeout: seconds to wait after SIGTERM before SIGKILL in
+            :meth:`aclose`.
+    """
+
+    def __init__(
+        self,
+        config: PiConfig,
+        *,
+        on_stderr: Callable[[str], None] | None = None,
+        spawn: Callable[..., Any] = asyncio.create_subprocess_exec,
+        terminate_timeout: float = 5.0,
+    ) -> None:
+        self._config = config
+        self._on_stderr = on_stderr
+        self._spawn = spawn
+        self._terminate_timeout = terminate_timeout
+
+        self._proc: Any | None = None
+        self._buf = bytearray()
+        self._eof = False
+        self._stderr_task: asyncio.Task[None] | None = None
+        # Serialize stdin writes so concurrent sends (e.g. a steer during a prompt)
+        # can't interleave bytes of two JSONL records. Mirrors omp's own client lock.
+        self._write_lock = asyncio.Lock()
+
+    # ---------------------------------------------------------------- props
+    @property
+    def pid(self) -> int | None:
+        return self._proc.pid if self._proc is not None else None
+
+    @property
+    def returncode(self) -> int | None:
+        return self._proc.returncode if self._proc is not None else None
+
+    # -------------------------------------------------------------- lifecycle
+    async def start(self) -> None:
+        """Resolve the pi binary and spawn the RPC subprocess."""
+        if self._proc is not None:
+            raise PiStartError("transport already started")
+
+        argv = self._config.to_argv()
+        if shutil.which(argv[0]) is None and not os.path.exists(argv[0]):
+            raise PiNotFoundError(argv[0])
+
+        env = self._config.to_env()
+        cwd = str(self._config.workspace) if self._config.workspace is not None else None
+
+        try:
+            self._proc = await self._spawn(
+                *argv,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                cwd=cwd,
+                env=env,
+            )
+        except (OSError, ValueError) as exc:  # pragma: no cover - spawn failure
+            raise PiStartError(f"failed to spawn pi: {exc}") from exc
+
+        if self._proc.stdin is None or self._proc.stdout is None:
+            raise PiStartError("pi subprocess started without stdin/stdout pipes")
+
+        if self._proc.stderr is not None:
+            self._stderr_task = asyncio.ensure_future(self._drain_stderr(self._proc.stderr))
+
+    async def _drain_stderr(self, stream: asyncio.StreamReader) -> None:
+        try:
+            while True:
+                line = await stream.readline()
+                if not line:
+                    break
+                if self._on_stderr is not None:
+                    self._on_stderr(line.decode("utf-8", errors="replace").rstrip("\r\n"))
+        except asyncio.CancelledError:  # pragma: no cover - shutdown path
+            raise
+        except Exception:  # pragma: no cover - never let stderr drain crash us
+            pass
+
+    # ------------------------------------------------------------------- io
+    async def send(self, obj: dict[str, Any]) -> None:
+        """Serialize ``obj`` to one JSONL line and write it to pi's stdin."""
+        proc = self._proc
+        if proc is None or proc.stdin is None:
+            raise TransportClosedError("transport not started")
+        if proc.returncode is not None:
+            raise TransportClosedError(f"pi process has exited (code {proc.returncode})")
+
+        line = json.dumps(obj, ensure_ascii=False, separators=(",", ":")) + "\n"
+        async with self._write_lock:
+            try:
+                proc.stdin.write(line.encode("utf-8"))
+                await proc.stdin.drain()
+            except (ConnectionResetError, BrokenPipeError) as exc:
+                raise TransportClosedError(f"pi stdin closed: {exc}") from exc
+
+    async def receive(self) -> dict[str, Any] | None:
+        """Return the next parsed inbound JSON object, or ``None`` at EOF.
+
+        Only one coroutine may call this at a time (single-reader contract).
+        """
+        proc = self._proc
+        if proc is None or proc.stdout is None:
+            raise TransportClosedError("transport not started")
+
+        while True:
+            line = self._take_buffered_line()
+            if line is not None:
+                if not line.strip():
+                    continue  # tolerate blank keepalive lines
+                return self._parse(line)
+
+            if self._eof:
+                # Flush any trailing partial line at EOF.
+                if self._buf:
+                    rest = bytes(self._buf)
+                    self._buf.clear()
+                    text = rest.decode("utf-8", errors="replace").rstrip("\r")
+                    if text.strip():
+                        return self._parse(text)
+                return None
+
+            chunk = await proc.stdout.read(_READ_CHUNK)
+            if not chunk:
+                self._eof = True
+                continue
+            self._buf.extend(chunk)
+            if len(self._buf) > _MAX_LINE_BYTES and b"\n" not in self._buf:
+                raise PiProtocolError(
+                    f"line exceeded {_MAX_LINE_BYTES} bytes without a newline"
+                )
+
+    def _take_buffered_line(self) -> str | None:
+        idx = self._buf.find(b"\n")
+        if idx == -1:
+            return None
+        raw = bytes(self._buf[:idx])
+        del self._buf[: idx + 1]
+        return raw.decode("utf-8", errors="replace").rstrip("\r")
+
+    @staticmethod
+    def _parse(line: str) -> dict[str, Any]:
+        try:
+            obj = json.loads(line)
+        except json.JSONDecodeError as exc:
+            raise PiProtocolError(f"invalid JSON from pi: {exc}", raw=line) from exc
+        if not isinstance(obj, dict):
+            raise PiProtocolError(
+                f"expected JSON object from pi, got {type(obj).__name__}", raw=line
+            )
+        return obj
+
+    def __aiter__(self) -> AsyncIterator[dict[str, Any]]:
+        return self._iter()
+
+    async def _iter(self) -> AsyncIterator[dict[str, Any]]:
+        while True:
+            obj = await self.receive()
+            if obj is None:
+                return
+            yield obj
+
+    # --------------------------------------------------------------- shutdown
+    async def aclose(self) -> None:
+        """Close stdin, then terminate and reap the process (SIGKILL fallback)."""
+        proc = self._proc
+        if proc is None:
+            return
+
+        if proc.stdin is not None:
+            try:
+                proc.stdin.close()
+            except Exception:  # pragma: no cover
+                pass
+
+        if proc.returncode is None:
+            try:
+                proc.terminate()
+            except ProcessLookupError:  # pragma: no cover - already gone
+                pass
+            try:
+                await asyncio.wait_for(proc.wait(), timeout=self._terminate_timeout)
+            except TimeoutError:
+                try:
+                    proc.kill()
+                except ProcessLookupError:  # pragma: no cover
+                    pass
+                await proc.wait()
+
+        if self._stderr_task is not None:
+            self._stderr_task.cancel()
+            try:
+                await self._stderr_task
+            except (asyncio.CancelledError, Exception):  # pragma: no cover
+                pass
+            self._stderr_task = None

pidriver/client.py

@@ -0,0 +1,92 @@
+"""Starting sessions.
+
+:class:`PiClient` is the entry point: configure it once, then call
+:meth:`PiClient.start` per task to launch an isolated ``pi --mode rpc`` process
+and get back a started :class:`~pidriver.session.PiSession`.
+
+``PiClient.start`` has the signature a :class:`~pidriver.manager.PiSessionManager`
+``factory`` expects, so you can wire the two together directly::
+
+    client = PiClient(PiConfig(...))
+    manager = PiSessionManager(factory=client.start)
+    session = await manager.create("Add a healthcheck endpoint", cwd="/srv/proj")
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import replace
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from ._transport import PiTransport, SubprocessTransport
+from .config import PiConfig
+from .interaction import InteractionHandler
+from .session import PiSession
+
+if TYPE_CHECKING:
+    pass
+
+__all__ = ["PiClient", "TransportFactory"]
+
+# Builds a (not-yet-started) transport for a resolved config. Override in tests.
+TransportFactory = Callable[[PiConfig], PiTransport]
+
+
+class PiClient:
+    """Launches isolated pi/omp sessions from a base :class:`PiConfig`."""
+
+    def __init__(
+        self,
+        config: PiConfig,
+        *,
+        transport_factory: TransportFactory | None = None,
+        interaction_handler: InteractionHandler | None = None,
+    ) -> None:
+        self._config = config
+        self._make_transport: TransportFactory = transport_factory or SubprocessTransport
+        self._default_handler = interaction_handler
+
+    async def start(
+        self,
+        task: str,
+        *,
+        cwd: str | Path | None = None,
+        config: PiConfig | None = None,
+        interaction_handler: InteractionHandler | None = None,
+        resume: str | None = None,
+        session_id: str | None = None,
+        send_initial: bool = True,
+    ) -> PiSession:
+        """Launch a session and (by default) send ``task`` as the first prompt.
+
+        :param task: the initial instruction for the agent.
+        :param cwd: working directory the agent operates in (the target project);
+            overrides ``config.workspace``.
+        :param config: per-session override of the client's base config.
+        :param interaction_handler: auto-answer policy; defaults to the client's
+            handler, else :class:`~pidriver.interaction.AskHost`.
+        :param resume: a prior pi session path/id to ``--continue`` from.
+        :param session_id: stable local id for the manager registry key
+            (auto-generated when omitted).
+        :param send_initial: set ``False`` to open the session without prompting.
+        """
+        resolved = config or self._config
+        if resolved is not None:
+            if cwd is not None:
+                resolved = replace(resolved, workspace=Path(cwd))
+            if resume:
+                resolved = replace(resolved, session_path=resume, continue_session=True)
+
+        transport = self._make_transport(resolved)
+        await transport.start()
+
+        session = PiSession(
+            transport,
+            session_id=session_id,
+            config=resolved,
+            interaction_handler=interaction_handler or self._default_handler,
+        )
+        if task and send_initial:
+            await session.prompt(task)
+        return session