pi-agent-python-sdk 0.1.0__py3-none-any.whl

pi_agent/__init__.py

@@ -0,0 +1,47 @@
+"""Python access to an existing Pi agent runtime.
+
+Importing this package never launches Pi or changes its configuration.
+"""
+
+from .client import AsyncPiClient
+from .errors import (
+    PiBusyError,
+    PiCommandError,
+    PiError,
+    PiProcessError,
+    PiProtocolError,
+    PiResultOverflow,
+    PiRunError,
+    PiRunOwnershipError,
+    PiRunStartTimeout,
+    PiSubscriptionOverflow,
+    PiTimeoutError,
+    PiUIHandlerError,
+    PiVersionError,
+)
+from .sync import PiClient
+from .types import Event, ImageContent, Limits, RunResult, SessionInfo, UsageSummary
+
+__all__ = [
+    "AsyncPiClient",
+    "Event",
+    "ImageContent",
+    "Limits",
+    "PiBusyError",
+    "PiCommandError",
+    "PiClient",
+    "PiError",
+    "PiProcessError",
+    "PiProtocolError",
+    "PiResultOverflow",
+    "PiRunError",
+    "PiRunOwnershipError",
+    "PiRunStartTimeout",
+    "PiSubscriptionOverflow",
+    "PiTimeoutError",
+    "PiUIHandlerError",
+    "PiVersionError",
+    "RunResult",
+    "SessionInfo",
+    "UsageSummary",
+]

pi_agent/_events.py

@@ -0,0 +1,108 @@
+"""Bounded future-only event subscriptions; response routing never waits here."""
+
+from __future__ import annotations
+
+import asyncio
+from collections import deque
+from collections.abc import Callable
+from types import TracebackType
+from typing import Self
+
+from .errors import PiSubscriptionOverflow
+from .types import Event, Limits
+
+
+class EventSubscription:
+    """Enter before submitting work to observe its events without a history buffer.
+
+    A slow consumer receives PiSubscriptionOverflow rather than silently losing
+    events. Closing the context unregisters immediately and releases its buffer.
+    """
+
+    def __init__(
+        self,
+        limits: Limits,
+        register: Callable[[EventSubscription], None],
+        unregister: Callable[[EventSubscription], None],
+    ) -> None:
+        self._limits = limits
+        self._register = register
+        self._unregister = unregister
+        self._records: deque[tuple[Event, int]] = deque()
+        self._bytes = 0
+        self._ready = asyncio.Event()
+        self._entered = False
+        self._closed = False
+        self._error: Exception | None = None
+        self._reading = False
+
+    async def __aenter__(self) -> Self:
+        if self._entered or self._closed:
+            raise RuntimeError("Event subscriptions are single-use")
+        self._register(self)
+        self._entered = True
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+    async def aclose(self) -> None:
+        """Unsubscribe and release queued payloads; safe to call repeatedly."""
+        self._finish()
+        self._records.clear()
+        self._bytes = 0
+
+    def _finish(self, error: Exception | None = None) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._error = error
+        self._unregister(self)
+        if error is not None:
+            self._records.clear()
+            self._bytes = 0
+        self._ready.set()
+
+    def _put(self, event: Event, size: int) -> Exception | None:
+        if self._closed:
+            return self._error
+        if (
+            len(self._records) >= self._limits.event_queue_size
+            or self._bytes + size > self._limits.event_queue_bytes
+        ):
+            error = PiSubscriptionOverflow("Event consumer exceeded its configured buffer limit")
+            self._finish(error)
+            return error
+        self._records.append((event, size))
+        self._bytes += size
+        self._ready.set()
+        return None
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> Event:
+        if not self._entered:
+            raise RuntimeError("Enter the event subscription context before iterating")
+        if self._reading:
+            raise RuntimeError("Only one reader may iterate an event subscription")
+        self._reading = True
+        try:
+            while True:
+                if self._error is not None:
+                    raise self._error
+                if self._records:
+                    event, size = self._records.popleft()
+                    self._bytes -= size
+                    return event
+                if self._closed:
+                    raise StopAsyncIteration
+                self._ready.clear()
+                await self._ready.wait()
+        finally:
+            self._reading = False

pi_agent/_launch.py

@@ -0,0 +1,203 @@
+"""CLI construction and offline compatibility checks for the owned process."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import re
+import shutil
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+
+from ._transport import _close_process_pipes, _wait_for_exit
+from .errors import PiProcessError, PiVersionError
+
+MINIMUM_PI_VERSION = "0.85.1"
+TESTED_PI_VERSION = "0.85.1"
+
+# The client owns these options. In particular --print consumes an initial prompt.
+_RESERVED = {
+    "--mode",
+    "--print",
+    "-p",
+    "--help",
+    "-h",
+    "--version",
+    "-v",
+    "--export",
+    "--list-models",
+    "--resume",
+    "-r",
+    "--continue",
+    "-c",
+    "--session",
+    "--session-dir",
+    "--session-id",
+    "--fork",
+    "--no-session",
+    "--provider",
+    "--model",
+}
+_VALUES = {
+    "--api-key",
+    "--system-prompt",
+    "--append-system-prompt",
+    "--name",
+    "-n",
+    "--models",
+    "--tools",
+    "-t",
+    "--exclude-tools",
+    "-xt",
+    "--thinking",
+    "--extension",
+    "-e",
+    "--skill",
+    "--prompt-template",
+    "--theme",
+    "--use-theme",
+    "--tui-mode",
+}
+_FLAGS = {
+    "--no-tools",
+    "-nt",
+    "--no-builtin-tools",
+    "-nbt",
+    "--no-extensions",
+    "-ne",
+    "--no-skills",
+    "-ns",
+    "--no-prompt-templates",
+    "-np",
+    "--no-themes",
+    "--no-context-files",
+    "-nc",
+    "--verbose",
+    "--approve",
+    "-a",
+    "--no-approve",
+    "-na",
+    "--offline",
+}
+
+
+def validate_extra_args(args: Sequence[str]) -> None:
+    """Reject positional input and conflicting flags, retaining extension flags."""
+    index = 0
+    while index < len(args):
+        arg = args[index]
+        name = arg.split("=", 1)[0]
+        if name in _RESERVED or arg == "--":
+            raise ValueError(f"Pi option {name} is managed by the client or incompatible with RPC")
+        if name in _VALUES:
+            if "=" in arg or index + 1 == len(args):
+                raise ValueError(f"Pi option {name} requires a separate value")
+            if name in {"--use-theme", "--tui-mode"} and args[index + 1].startswith("-"):
+                raise ValueError(f"Pi option {name} requires a value before another option")
+            index += 2
+        elif name in _FLAGS:
+            if "=" in arg:
+                raise ValueError(f"Pi option {name} takes no value")
+            index += 1
+        elif arg.startswith("--"):
+            # Pi accepts extension flags as --name=value or --name [value].
+            index += 1
+            if "=" not in arg and index < len(args) and not args[index].startswith(("-", "@")):
+                index += 1
+        else:
+            raise ValueError(
+                "extra_args cannot contain startup prompts, attachments, or unknown short flags"
+            )
+
+
+def executable_argv(
+    executable: str | os.PathLike[str] | Sequence[str], env: Mapping[str, str]
+) -> list[str]:
+    """Resolve a binary or explicit argv; npm's Windows shim runs through Node."""
+    if isinstance(executable, (str, os.PathLike)):
+        argv = [os.fspath(executable)]
+    else:
+        argv = list(executable)
+    if not argv or not all(isinstance(item, str) and item for item in argv):
+        raise ValueError("executable must be a path or a nonempty sequence of arguments")
+    resolved = shutil.which(argv[0], path=env.get("PATH"))
+    if resolved is None:
+        raise PiProcessError("Pi executable was not found; install Pi or specify executable")
+    if os.name == "nt" and Path(resolved).suffix.lower() in {".cmd", ".bat"}:
+        # Avoid shell=True and cmd quoting. Only the known npm Pi shim layout is resolved.
+        root = Path(resolved).parent
+        cli = root / "node_modules/@earendil-works/pi-coding-agent/dist/bundle/cli.js"
+        node = root / "node.exe"
+        node_path = str(node) if node.is_file() else shutil.which("node", path=env.get("PATH"))
+        if not cli.is_file() or node_path is None:
+            raise PiProcessError(
+                "Cannot resolve Pi npm shim; pass executable=[node_path, pi_cli_path]"
+            )
+        return [node_path, str(cli), *argv[1:]]
+    return [resolved, *argv[1:]]
+
+
+async def check_version(
+    argv: Sequence[str],
+    *,
+    cwd: str | None,
+    env: Mapping[str, str],
+    timeout: float,
+    allow_unknown: bool,
+    strict: bool,
+) -> tuple[str | None, str]:
+    """Read --version without a network request or loading user configuration."""
+    process: asyncio.subprocess.Process | None = None
+    try:
+        async with asyncio.timeout(timeout):
+            process = await asyncio.create_subprocess_exec(
+                *argv,
+                "--version",
+                cwd=cwd,
+                env=dict(env),
+                stdin=asyncio.subprocess.DEVNULL,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.DEVNULL,
+            )
+            assert process.stdout is not None
+            output = bytearray()
+            while chunk := await process.stdout.read(4097 - len(output)):
+                output.extend(chunk)
+                if len(output) > 4096:
+                    raise PiVersionError("Pi --version returned too much output")
+            code = await process.wait()
+            match = re.fullmatch(rb"\s*(\d+)\.(\d+)\.(\d+)\s*", output)
+            if code != 0 or match is None:
+                if allow_unknown:
+                    return None, "unknown"
+                raise PiVersionError(
+                    "Cannot determine Pi version; custom wrappers may use "
+                    "allow_unknown_version=True"
+                )
+            version = tuple(int(part) for part in match.groups())
+            version_text = ".".join(str(part) for part in version)
+            if version < tuple(map(int, MINIMUM_PI_VERSION.split("."))):
+                raise PiVersionError(
+                    f"Pi {version_text} is older than minimum {MINIMUM_PI_VERSION}"
+                )
+            if version_text != TESTED_PI_VERSION and strict:
+                raise PiVersionError(
+                    f"Pi {version_text} is untested; recorded tested version is {TESTED_PI_VERSION}"
+                )
+            return version_text, "tested" if version_text == TESTED_PI_VERSION else "untested"
+    except TimeoutError as exc:
+        raise PiVersionError("Timed out checking Pi version") from exc
+    except OSError as exc:
+        raise PiProcessError("Could not launch Pi for version checking") from exc
+    finally:
+        if process is not None:
+            if process.returncode is None:
+                try:
+                    process.kill()
+                except ProcessLookupError:
+                    pass
+                # Process.wait() may depend on full or inherited pipes closing.
+                # Wait for the owned child's exit first, then release those pipes.
+                await _wait_for_exit(process)
+            _close_process_pipes(process)
+            await process.wait()

pi_agent/_runs.py

@@ -0,0 +1,317 @@
+"""One owned conversation run, including settlement and cancellation cleanup."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import time
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Self
+
+from ._events import EventSubscription
+from ._usage import UsageAccumulator
+from .errors import (
+    PiBusyError,
+    PiCommandError,
+    PiProtocolError,
+    PiResultOverflow,
+    PiRunError,
+    PiRunOwnershipError,
+    PiRunStartTimeout,
+    PiTimeoutError,
+)
+from .types import Event, ImageContent, RunResult
+
+if TYPE_CHECKING:
+    from .client import AsyncPiClient, Timeout
+
+
+class RunStream:
+    """A single-use owned run. Enter, then iterate or await result() to drain.
+
+    Breaking iteration must be followed by leaving the context so queued work
+    is cleared and Pi is aborted. An accepted command without a start event
+    has an uncertain disposition; its start deadline closes Pi.
+    """
+
+    def __init__(
+        self,
+        client: AsyncPiClient,
+        message: str,
+        *,
+        images: list[ImageContent] | None,
+        timeout: float | None,
+        command_timeout: Timeout,
+    ) -> None:
+        if timeout is not None and (
+            isinstance(timeout, bool) or not math.isfinite(timeout) or timeout <= 0
+        ):
+            raise ValueError("run timeout must be a positive finite number or None")
+        self._client = client
+        self._message = message
+        self._images = images
+        self._timeout = timeout
+        self._command_timeout = command_timeout
+        self._events = EventSubscription(
+            client.limits, lambda subscription: None, lambda subscription: None
+        )
+        self._started = asyncio.Event()
+        self._settled = asyncio.Event()
+        self._task: asyncio.Task[RunResult] | None = None
+        self._accepted: asyncio.Future[None] | None = None
+        self._error: Exception | None = None
+        self._messages: list[dict[str, Any]] = []
+        self._message_bytes = 0
+        self._usage = UsageAccumulator()
+        self._submitted = False
+        self._finishing = False
+        self._closed = False
+        self._iterating = False
+        self._draining = False
+        self._iteration_done = False
+        self._begin = 0.0
+        self._end = 0.0
+
+    async def __aenter__(self) -> Self:
+        self._client._check_loop()
+        if self._task is not None or self._closed:
+            raise RuntimeError("Run streams are single-use")
+        if self._client.busy:
+            raise PiBusyError("Another run owns this Pi conversation")
+        if not self._client.running:
+            raise RuntimeError("Start Pi before opening a run stream")
+        if self._client._unowned_submission:
+            raise PiRunOwnershipError(
+                "This client submitted low-level conversation work; use a fresh client for "
+                "run()/stream() because delayed session events cannot be attributed safely"
+            )
+        self._client._owner = self
+        self._accepted = asyncio.get_running_loop().create_future()
+        await self._events.__aenter__()
+        self._task = asyncio.create_task(self._drive(), name="pi-run")
+        self._task.add_done_callback(self._consume_task_exception)
+        try:
+            await asyncio.shield(self._accepted)
+        except BaseException:
+            await self.aclose()
+            raise
+        return self
+
+    @staticmethod
+    def _consume_task_exception(task: asyncio.Task[RunResult]) -> None:
+        if not task.cancelled():
+            task.exception()
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+    async def aclose(self) -> None:
+        """Finish cancellation cleanup before releasing the owned conversation."""
+        self._closed = True
+        if self._task is not None and not self._task.done():
+            self._task.cancel()
+        if self._task is not None:
+            try:
+                await asyncio.shield(self._task)
+            except (Exception, asyncio.CancelledError):
+                pass
+        if self._accepted is not None and self._accepted.done() and not self._accepted.cancelled():
+            self._accepted.exception()
+        await self._events.aclose()
+
+    def _fail(self, error: Exception) -> None:
+        if self._error is not None or self._finishing:
+            return
+        self._error = error
+        self._events._finish(error)
+        if self._task is not None and not self._task.done():
+            self._task.cancel()
+
+    def _on_event(self, event: Event, size: int) -> None:
+        if not self._submitted or self._settled.is_set() or self._error is not None:
+            return
+        if event.type == "agent_start":
+            self._started.set()
+        elif event.type == "agent_settled" and self._started.is_set():
+            self._end = time.monotonic()
+            self._settled.set()
+        elif event.type == "message_end" and self._started.is_set():
+            message = event.raw.get("message")
+            if not isinstance(message, dict) or not isinstance(message.get("role"), str):
+                self._fail(PiProtocolError("message_end requires a message with a role"))
+                return
+            if message.get("role") == "assistant":
+                if not isinstance(message.get("content"), list) or not isinstance(
+                    message.get("stopReason"), str
+                ):
+                    self._fail(
+                        PiProtocolError("Final assistant message requires content and stopReason")
+                    )
+                    return
+            if (
+                len(self._messages) >= self._client.limits.result_message_count
+                or self._message_bytes + size > self._client.limits.result_message_bytes
+            ):
+                self._fail(PiResultOverflow("Run exceeded its retained message count/byte limit"))
+                return
+            self._message_bytes += size
+            self._messages.append(message)
+            self._usage.add(message)
+        error = self._events._put(event, size)
+        if error is not None:
+            self._fail(error)
+
+    async def _drive(self) -> RunResult:
+        assert self._accepted is not None
+        failure: BaseException | None = None
+        try:
+            async with asyncio.timeout(self._timeout):
+                state = await self._client.get_state()
+                if state["isStreaming"] or state["isCompacting"] or state["pendingMessageCount"]:
+                    raise PiBusyError("Pi already has active or queued low-level work")
+                fields: dict[str, Any] = {"message": self._message}
+                if self._images is not None:
+                    fields["images"] = self._images
+                self._submitted = True
+                self._begin = time.monotonic()
+                await self._client._request(
+                    "prompt", fields, timeout=self._command_timeout, owner=self
+                )
+                self._accepted.set_result(None)
+                try:
+                    async with asyncio.timeout(self._client.limits.run_start_timeout):
+                        await self._started.wait()
+                except TimeoutError as exc:
+                    raise PiRunStartTimeout(
+                        "Prompt accepted but no run was observed; "
+                        "Pi is closed to prevent delayed work",
+                        command="prompt",
+                        uncertain=True,
+                    ) from exc
+                await self._settled.wait()
+                if self._error is not None:
+                    raise self._error
+                # Extensions can switch sessions internally; do not reuse startup identity.
+                await self._client.get_state()
+                result = self._result()
+                if result.stop_reason in {"error", "aborted"}:
+                    raise PiRunError(f"Pi run ended with {result.stop_reason}", result)
+                return result
+        except BaseException as exc:
+            self._finishing = True
+            failure = exc
+            if isinstance(exc, asyncio.CancelledError) and self._error is not None:
+                failure = self._error
+            elif isinstance(exc, TimeoutError) and not isinstance(exc, PiTimeoutError):
+                failure = PiTimeoutError(
+                    "Pi run deadline elapsed", command="prompt", uncertain=self._submitted
+                )
+            accepted = self._accepted.done() and not self._accepted.cancelled()
+            if (
+                self._submitted
+                and not isinstance(failure, PiCommandError)
+                and (not self._settled.is_set() or not accepted)
+            ):
+                # abort() cannot cancel an extension's pending input/UI preflight.
+                # Keep delayed work from escaping a failed owned operation.
+                if not accepted or not self._started.is_set():
+                    await self._client.aclose()
+                else:
+                    await self._cleanup()
+            if failure is exc:
+                raise
+            raise failure from exc
+        finally:
+            if not self._accepted.done():
+                if isinstance(failure, asyncio.CancelledError):
+                    self._accepted.cancel()
+                elif failure is not None:
+                    self._accepted.set_exception(failure)
+            # Final model errors are delivered by result(), after their events are consumed.
+            event_error = (
+                failure
+                if isinstance(failure, Exception) and not isinstance(failure, PiRunError)
+                else None
+            )
+            self._events._finish(event_error)
+            if self._client._owner is self:
+                self._client._owner = None
+
+    async def _cleanup(self) -> None:
+        async def stop() -> None:
+            if not self._client.running:
+                await self._client.aclose()
+                return
+            try:
+                async with asyncio.timeout(self._client.limits.cleanup_timeout):
+                    await self._client._request("clear_queue", owner=self)
+                    await self._client._request("abort", owner=self)
+            except Exception:
+                await self._client.aclose()
+
+        # Shield cleanup from the caller's cancellation, retaining ownership until it completes.
+        task = asyncio.create_task(stop(), name="pi-run-cleanup")
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError:
+            await task
+
+    def _result(self) -> RunResult:
+        assistants = [message for message in self._messages if message["role"] == "assistant"]
+        last = assistants[-1] if assistants else None
+        content = last["content"] if last is not None else []
+        text_parts = []
+        for block in content:
+            if not isinstance(block, dict):
+                raise PiProtocolError("Assistant content blocks must be objects")
+            if block.get("type") == "text":
+                if not isinstance(block.get("text"), str):
+                    raise PiProtocolError("Assistant text blocks require string text")
+                text_parts.append(block["text"])
+        return RunResult(
+            text="".join(text_parts),
+            messages=list(self._messages),
+            stop_reason=last["stopReason"] if last is not None else None,
+            session=self._client.session,
+            elapsed_seconds=max(0.0, self._end - self._begin),
+            usage=self._usage.snapshot(),
+        )
+
+    def __aiter__(self) -> Self:
+        if self._draining:
+            raise PiBusyError("Cannot iterate while result() drains the stream")
+        if self._iterating:
+            raise PiBusyError("A stream permits one event iterator")
+        self._iterating = True
+        return self
+
+    async def __anext__(self) -> Event:
+        if self._draining:
+            raise PiBusyError("Cannot iterate while result() drains the stream")
+        self._iterating = True
+        try:
+            return await self._events.__anext__()
+        except StopAsyncIteration:
+            self._iteration_done = True
+            raise
+
+    async def result(self) -> RunResult:
+        """Drain without retaining events, or return the result after iteration ends."""
+        if self._task is None:
+            raise RuntimeError("Enter the stream context before awaiting its result")
+        if self._draining or (self._iterating and not self._iteration_done):
+            raise PiBusyError("Finish iteration before calling result()")
+        self._draining = True
+        try:
+            if not self._iteration_done:
+                async for _ in self._events:
+                    pass
+                self._iteration_done = True
+            return await asyncio.shield(self._task)
+        finally:
+            self._draining = False