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

pi_agent/__init__.py

@@ -3,6 +3,7 @@
 Importing this package never launches Pi or changes its configuration.
 """
 
+from ._observation import ObservationStatus, ProcessOutput
 from .client import AsyncPiClient
 from .errors import (
     PiBusyError,
@@ -27,6 +28,8 @@ __all__ = [
     "Event",
     "ImageContent",
     "Limits",
+    "ObservationStatus",
+    "ProcessOutput",
     "PiBusyError",
     "PiCommandError",
     "PiClient",

pi_agent/_collections.py

@@ -0,0 +1,59 @@
+"""Bounded observation through the next settlement, without run ownership."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+
+from ._events import EventSubscription
+from .errors import PiProcessError, PiResultOverflow, PiTimeoutError
+from .types import Event, Limits, _validate_timeout
+
+
+def deadline(timeout: float | None) -> float | None:
+    if timeout is None:
+        return None
+    _validate_timeout(timeout, "Collection timeout")
+    return asyncio.get_running_loop().time() + timeout
+
+
+def start_collection(
+    subscription: EventSubscription,
+    limits: Limits,
+    *,
+    timeout: float | None,
+    retain: bool,
+) -> asyncio.Task[list[Event] | None]:
+    expires = deadline(timeout)
+    # Register now, not when the task first runs: callers may submit immediately.
+    subscription._enter()
+
+    async def consume() -> list[Event] | None:
+        events: list[Event] = []
+        size = 0
+        timer = asyncio.timeout_at(expires)
+        try:
+            async with timer:
+                async for event in subscription:
+                    if retain:
+                        size += len(json.dumps(event.raw).encode("utf-8"))
+                        if (
+                            len(events) >= limits.collection_event_count
+                            or size > limits.collection_event_bytes
+                        ):
+                            raise PiResultOverflow("Event collection exceeded its count/byte limit")
+                        events.append(event)
+                    if event.type == "agent_settled":
+                        return events if retain else None
+                raise PiProcessError("Event collection closed before settlement")
+        except TimeoutError as exc:
+            if not timer.expired():
+                raise
+            raise PiTimeoutError("Event collection deadline elapsed", uncertain=False) from exc
+        finally:
+            await subscription.aclose()
+
+    task = asyncio.create_task(consume(), name="pi-event-collection")
+    # Cancellation before the task's first instruction cannot run its finally.
+    task.add_done_callback(lambda _: subscription._discard())
+    return task

pi_agent/_events.py

@@ -6,13 +6,15 @@ import asyncio
 from collections import deque
 from collections.abc import Callable
 from types import TracebackType
-from typing import Self
+from typing import Generic, Self, TypeVar
 
 from .errors import PiSubscriptionOverflow
 from .types import Event, Limits
 
+T = TypeVar("T")
 
-class EventSubscription:
+
+class _Subscription(Generic[T]):
     """Enter before submitting work to observe its events without a history buffer.
 
     A slow consumer receives PiSubscriptionOverflow rather than silently losing
@@ -22,13 +24,13 @@ class EventSubscription:
     def __init__(
         self,
         limits: Limits,
-        register: Callable[[EventSubscription], None],
-        unregister: Callable[[EventSubscription], None],
+        register: Callable[[], None],
+        unregister: Callable[[], None],
     ) -> None:
         self._limits = limits
         self._register = register
         self._unregister = unregister
-        self._records: deque[tuple[Event, int]] = deque()
+        self._records: deque[tuple[T, int]] = deque()
         self._bytes = 0
         self._ready = asyncio.Event()
         self._entered = False
@@ -37,9 +39,12 @@ class EventSubscription:
         self._reading = False
 
     async def __aenter__(self) -> Self:
+        return self._enter()
+
+    def _enter(self) -> Self:
         if self._entered or self._closed:
             raise RuntimeError("Event subscriptions are single-use")
-        self._register(self)
+        self._register()
         self._entered = True
         return self
 
@@ -53,22 +58,29 @@ class EventSubscription:
 
     async def aclose(self) -> None:
         """Unsubscribe and release queued payloads; safe to call repeatedly."""
+        self._discard()
+
+    def _discard(self, *, buffered: bool = False) -> None:
+        """Also used by the sync facade after closing the subscription."""
         self._finish()
         self._records.clear()
         self._bytes = 0
+        self._error = None
 
     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._unregister()
+        if isinstance(error, PiSubscriptionOverflow):
+            # Overflow is loss, not a complete prefix ending at a process failure.
+            # Keep it immediate; a healthy terminal queue can still be drained.
             self._records.clear()
             self._bytes = 0
         self._ready.set()
 
-    def _put(self, event: Event, size: int) -> Exception | None:
+    def _put(self, event: T, size: int) -> Exception | None:
         if self._closed:
             return self._error
         if (
@@ -86,7 +98,7 @@ class EventSubscription:
     def __aiter__(self) -> Self:
         return self
 
-    async def __anext__(self) -> Event:
+    async def __anext__(self) -> T:
         if not self._entered:
             raise RuntimeError("Enter the event subscription context before iterating")
         if self._reading:
@@ -94,15 +106,44 @@ class EventSubscription:
         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
+                event = self._next_nowait()
+                if event is not None:
                     return event
-                if self._closed:
-                    raise StopAsyncIteration
                 self._ready.clear()
                 await self._ready.wait()
         finally:
             self._reading = False
+
+    async def _next_batch(self) -> list[T]:
+        records = [await self.__anext__()]
+        # Transfer only records already buffered; never wait to fill a batch.
+        # This retains at most one queue's byte budget in the blocking facade.
+        while self._records and len(records) < 64:
+            record = self._next_nowait()
+            assert record is not None
+            records.append(record)
+        return records
+
+    def _next_nowait(self) -> T | None:
+        """Read on the owning loop, or after that loop has completely stopped."""
+        if self._records:
+            event, size = self._records.popleft()
+            self._bytes -= size
+            return event
+        if self._error is not None:
+            raise self._error
+        if self._closed:
+            raise StopAsyncIteration
+        return None
+
+
+class EventSubscription(_Subscription[Event]):
+    """Bounded future events; enter before startup to include startup events."""
+
+    def __init__(
+        self,
+        limits: Limits,
+        register: Callable[[EventSubscription], None],
+        unregister: Callable[[EventSubscription], None],
+    ) -> None:
+        super().__init__(limits, lambda: register(self), lambda: unregister(self))

pi_agent/_observation.py

@@ -0,0 +1,73 @@
+"""Opt-in process output delivery; storage and forwarding belong to callers."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field, replace
+from typing import Any, Literal
+
+from ._events import _Subscription
+from .errors import PiSubscriptionOverflow
+from .types import Limits
+
+OutputSource = Literal["stderr", "stdout", "rpc"]
+
+
+@dataclass(frozen=True)
+class ProcessOutput:
+    """Original bytes or a parsed object from the RPC child, excluding the version probe."""
+
+    source: OutputSource
+    time_ns: int
+    data: bytes | dict[str, Any] = field(repr=False)
+
+
+@dataclass(frozen=True)
+class ObservationStatus:
+    """Delivery coverage, independent of RPC success and consumer storage."""
+
+    started_at_ns: int | None = None
+    ended_at_ns: int | None = None
+    from_start: bool = False
+    complete: bool = False
+    lost: bool = False
+    stderr_eof: bool = False
+    stdout_eof: bool = False
+    rpc_complete: bool = True
+    error: Exception | None = field(default=None, repr=False)
+
+
+class ProcessObservation(_Subscription[ProcessOutput]):
+    """Bounded output iterator. Check status after draining to detect incomplete output."""
+
+    def __init__(
+        self,
+        limits: Limits,
+        register: Callable[[ProcessObservation], None],
+        unregister: Callable[[ProcessObservation], None],
+        sources: frozenset[OutputSource],
+    ) -> None:
+        super().__init__(limits, lambda: register(self), lambda: unregister(self))
+        self._status = ObservationStatus()
+        self._sources = sources
+
+    @property
+    def status(self) -> ObservationStatus:
+        return self._status
+
+    def _finish(self, error: Exception | None = None) -> None:
+        if not self._closed and self._status.ended_at_ns is None:
+            self._status = ObservationStatus(
+                started_at_ns=self._status.started_at_ns,
+                ended_at_ns=time.time_ns(),
+                from_start=self._status.from_start,
+                lost=self._status.lost or isinstance(error, PiSubscriptionOverflow),
+                error=error,
+            )
+        super()._finish(error)
+
+    def _discard(self, *, buffered: bool = False) -> None:
+        if self._records or buffered:
+            self._status = replace(self._status, complete=False, lost=True)
+        super()._discard()

pi_agent/_runs.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-import math
 import time
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Self
@@ -20,14 +19,14 @@ from .errors import (
     PiRunStartTimeout,
     PiTimeoutError,
 )
-from .types import Event, ImageContent, RunResult
+from .types import Event, ImageContent, RunResult, _validate_timeout
 
 if TYPE_CHECKING:
     from .client import AsyncPiClient, Timeout
 
 
 class RunStream:
-    """A single-use owned run. Enter, then iterate or await result() to drain.
+    """A single-use owned run. Enter on acceptance or start, then consume events.
 
     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
@@ -43,10 +42,8 @@ class RunStream:
         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")
+        if timeout is not None:
+            _validate_timeout(timeout, "run timeout")
         self._client = client
         self._message = message
         self._images = images
@@ -55,10 +52,12 @@ class RunStream:
         self._events = EventSubscription(
             client.limits, lambda subscription: None, lambda subscription: None
         )
+        self._retain_events = True
         self._started = asyncio.Event()
         self._settled = asyncio.Event()
         self._task: asyncio.Task[RunResult] | None = None
-        self._accepted: asyncio.Future[None] | None = None
+        self._ready: asyncio.Future[None] | None = None
+        self._accepted = False
         self._error: Exception | None = None
         self._messages: list[dict[str, Any]] = []
         self._message_bytes = 0
@@ -86,12 +85,14 @@ class RunStream:
                 "run()/stream() because delayed session events cannot be attributed safely"
             )
         self._client._owner = self
-        self._accepted = asyncio.get_running_loop().create_future()
+        self._ready = 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)
+            await self._ready
+            if self._error is not None:
+                raise self._error
         except BaseException:
             await self.aclose()
             raise
@@ -120,8 +121,8 @@ class RunStream:
                 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()
+        if self._ready is not None and self._ready.done() and not self._ready.cancelled():
+            self._ready.exception()
         await self._events.aclose()
 
     def _fail(self, error: Exception) -> None:
@@ -137,6 +138,10 @@ class RunStream:
             return
         if event.type == "agent_start":
             self._started.set()
+            # Extensions may wait for the entire run before acknowledging the
+            # prompt. Make consumption possible without declaring acceptance.
+            if self._ready is not None and not self._ready.done():
+                self._ready.set_result(None)
         elif event.type == "agent_settled" and self._started.is_set():
             self._end = time.monotonic()
             self._settled.set()
@@ -162,12 +167,13 @@ class RunStream:
             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)
+        if self._retain_events:
+            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
+        assert self._ready is not None
         failure: BaseException | None = None
         try:
             async with asyncio.timeout(self._timeout):
@@ -182,7 +188,9 @@ class RunStream:
                 await self._client._request(
                     "prompt", fields, timeout=self._command_timeout, owner=self
                 )
-                self._accepted.set_result(None)
+                self._accepted = True
+                if not self._ready.done():
+                    self._ready.set_result(None)
                 try:
                     async with asyncio.timeout(self._client.limits.run_start_timeout):
                         await self._started.wait()
@@ -211,27 +219,34 @@ class RunStream:
                 failure = PiTimeoutError(
                     "Pi run deadline elapsed", command="prompt", uncertain=self._submitted
                 )
-            accepted = self._accepted.done() and not self._accepted.cancelled()
+            if self._client._ui_error is not None and failure is self._client._ui_error:
+                # The owner now reports this failure. Do not report it again as
+                # a failure of an otherwise successful cleanup command.
+                self._client._ui_error = None
             if (
                 self._submitted
-                and not isinstance(failure, PiCommandError)
-                and (not self._settled.is_set() or not accepted)
+                and (not isinstance(failure, PiCommandError) or self._started.is_set())
+                and (not self._settled.is_set() or not self._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():
+                if not self._accepted or not self._started.is_set():
                     await self._client.aclose()
                 else:
                     await self._cleanup()
             if failure is exc:
                 raise
+            if self._error is not None and failure is self._error:
+                # Cancellation only woke the driver; retain the recorded
+                # error's original cause (for example the UI callback error).
+                raise failure from failure.__cause__
             raise failure from exc
         finally:
-            if not self._accepted.done():
+            if not self._ready.done():
                 if isinstance(failure, asyncio.CancelledError):
-                    self._accepted.cancel()
+                    self._ready.cancel()
                 elif failure is not None:
-                    self._accepted.set_exception(failure)
+                    self._ready.set_exception(failure)
             # Final model errors are delivered by result(), after their events are consumed.
             event_error = (
                 failure
@@ -300,6 +315,16 @@ class RunStream:
             self._iteration_done = True
             raise
 
+    async def _next_batch(self) -> list[Event]:
+        if self._draining:
+            raise PiBusyError("Cannot iterate while result() drains the stream")
+        self._iterating = True
+        try:
+            return await self._events._next_batch()
+        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: