diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c6b7d0..754a664 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,15 @@ Releases are published automatically on every merge to `main`. Release notes for each version are on the [GitHub releases page](https://github.com/shrijayan/opencode-python-sdk/releases). +## Unreleased + +- `client.events`: a shared event stream with one connection for all subscribers, bounded per-subscriber queues, + `server.connected` replay, `of_type` / `for_session` / `where` filters, `on_activity` keepalive hook, cancellation, + and opt-in reconnection with backoff (`EventHub(client, reconnect=True)`). +- `client.rpc(Definition)`: typed plugin RPC calls and event subscriptions declared with `RpcDefinition`, + `rpc_method` and `rpc_event`; RPC failures raise `RpcCallError` / `RpcCallInternalError`. +- `Stream.on_activity` and `Stream.interrupt()`. + ## 0.1.0 First release, generated from OpenCode 2.1.2: sync and async clients for all 142 operations, typed models and diff --git a/README.md b/README.md index 9fc75ba..c22ee30 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,82 @@ Events are typed (`SessionCreatedEvent`, `SessionTextDeltaEvent`, and so on). Ev not know arrive as `UnknownEvent` with all their fields kept. Streams apply only the connect timeout by default, since they stay open between events. +### Shared event stream + +`client.events` shares one connection among every subscriber, like `client.event` in the TypeScript client. The +connection opens with the first subscriber and closes when the last one leaves. Each subscriber has its own queue of +up to 4096 events; one that falls further behind fails with `EventOverflowError` without affecting the others. The +latest `server.connected` event is replayed to subscribers that join later. + +```python +from opencode.types import SessionIdleEvent + +with client.events.of_type(SessionIdleEvent).for_session(session.id) as idle: + client.session.prompt(session.id, text="Hello") + event = idle.get(timeout=300) # typed as SessionIdleEvent +``` + +- Filters chain: `of_type(*classes)` narrows the yielded type, `for_session(id)` keeps one session's events, and + `where(predicate)` takes any test. Entering the `with` block joins the stream, so nothing published after that is + missed. +- `subscribe(on_activity=callback)` calls `callback` for every chunk received, keepalive comments included, as a + liveness signal. +- `close()` (or leaving the `with` block) ends the subscription, also from another thread. In async code, cancel a wait + with `anyio.fail_after`, `asyncio.timeout` or a trio cancel scope; the subscription stays usable. +- Reconnection is opt-in: `EventHub(client, reconnect=True)` (or `AsyncEventHub`, or a `Reconnect(...)` policy) reopens + the connection with exponential backoff after connection errors, 5xx and 429. Subscribers see `server.connected` + once. **The stream is volatile**: events published while disconnected are lost, so re-read state through the API + after a reconnect when it matters. + +The async hub works the same way under asyncio and trio: + +```python +async with client.events.for_session(session.id) as events: + async for event in events: + print(event.type) +``` + +`client.event.subscribe()` still opens a dedicated connection. + +### Plugin RPC + +Declare a plugin RPC once as a class, then call it with typed inputs and outputs: + +```python +from pydantic import BaseModel +from opencode import RpcCallError, RpcDefinition, rpc_event, rpc_method + + +class AddInput(BaseModel): + amount: int + + +class Total(BaseModel): + total: int + + +class Counter(RpcDefinition, id="counter"): + add = rpc_method(AddInput, Total) + reset = rpc_method(None, Total) + changed = rpc_event(Total) + + +counter = client.rpc(Counter) +try: + print(counter.call(Counter.add, AddInput(amount=2)).total) +except RpcCallError as error: + print(error.type, error.message, error.data) + +with counter.subscribe(Counter.changed) as changes: # only rpc.counter.changed, with validated data + for change in changes: + print(change.data.total) +``` + +Any type pydantic validates works for inputs, outputs and event data. `rpc_method(..., errors={"not_found": Model})` +validates the `data` of declared errors. Failures raise `RpcCallError`; crashes and outputs that do not match the +declared type raise `RpcCallInternalError`. `counter.events()` yields every raw event of the RPC. The untyped +`client.rpc.call(rpc_id, method, input=...)` is still available. + ## Terminals `client.pty.attach()` opens an interactive connection to a PTY over a WebSocket. It needs the `pty` extra: diff --git a/docs/guides/events.md b/docs/guides/events.md index 0aad2ec..1cd224d 100644 --- a/docs/guides/events.md +++ b/docs/guides/events.md @@ -47,3 +47,42 @@ async with await client.event.subscribe() as events: async for event in events: print(event.type) ``` + +## Shared event stream + +`client.events` shares one connection among every subscriber, like `client.event` in the TypeScript client. The +connection opens with the first subscriber and closes when the last one leaves. Each subscriber has its own queue of +up to 4096 events; one that falls further behind fails with `EventOverflowError` without affecting the others. The +latest `server.connected` event is replayed to subscribers that join later. + +```python +from opencode.types import SessionIdleEvent + +with client.events.of_type(SessionIdleEvent).for_session(session.id) as idle: + client.session.prompt(session.id, text="Hello") + event = idle.get(timeout=300) # typed as SessionIdleEvent +``` + +- Filters chain: `of_type(*classes)` narrows the yielded type, `for_session(id)` keeps one session's events, and + `where(predicate)` takes any test. Entering the `with` block joins the stream, so nothing published after that is + missed. +- `subscribe(on_activity=callback)` calls `callback` for every chunk received, keepalive comments included, as a + liveness signal. +- `close()` (or leaving the `with` block) ends the subscription, also from another thread. In async code, cancel a wait + with `anyio.fail_after`, `asyncio.timeout` or a trio cancel scope; the subscription stays usable. +- Reconnection is opt-in: `EventHub(client, reconnect=True)` (or `AsyncEventHub`, or a `Reconnect(...)` policy) reopens + the connection with exponential backoff after connection errors, 5xx and 429. Subscribers see `server.connected` + once. **The stream is volatile** (see [Delivery](#delivery)): events published while disconnected are lost, so re-read state through the API + after a reconnect when it matters. + +The async hub works the same way under asyncio and trio: + +```python +async with client.events.for_session(session.id) as events: + async for event in events: + print(event.type) +``` + +`client.event.subscribe()` still opens a dedicated connection. + +Typed plugin RPC events are covered in the [plugin RPC guide](plugin-rpc.md). diff --git a/docs/guides/index.md b/docs/guides/index.md index 678196a..31cc8ff 100644 --- a/docs/guides/index.md +++ b/docs/guides/index.md @@ -28,7 +28,7 @@ entry; the common ones also have a guide. CI fails if a group on the client is m | `client.provider` | Model providers | - | [`ProviderResource`](../reference/resources.md#opencode.resources.provider.ProviderResource) | | `client.pty` | Terminals | [Interactive terminals](terminals.md) | [`PtyResource`](../reference/resources.md#opencode.resources.pty.PtyResource) | | `client.reference` | References | - | [`ReferenceResource`](../reference/resources.md#opencode.resources.reference.ReferenceResource) | -| `client.rpc` | Plugin RPC | - | [`RpcResource`](../reference/resources.md#opencode.resources.rpc.RpcResource) | +| `client.rpc` | Plugin RPC | [Plugin RPC](plugin-rpc.md) | [`RpcResource`](../reference/resources.md#opencode.resources.rpc.RpcResource) | | `client.server` | Server info and health | - | [`ServerResource`](../reference/resources.md#opencode.resources.server.ServerResource) | | `client.session` | Sessions, prompts and their sub-resources | [Guide](sessions.md) | [`SessionResource`](../reference/resources.md#opencode.resources.session.SessionResource) | | `client.shell` | Shells | - | [`ShellResource`](../reference/resources.md#opencode.resources.shell.ShellResource) | diff --git a/docs/guides/plugin-rpc.md b/docs/guides/plugin-rpc.md new file mode 100644 index 0000000..040ba26 --- /dev/null +++ b/docs/guides/plugin-rpc.md @@ -0,0 +1,53 @@ +# Plugin RPC + +Declare a plugin RPC once as a class, then call it with typed inputs and outputs: + +```python +from pydantic import BaseModel +from opencode import RpcCallError, RpcDefinition, rpc_event, rpc_method + + +class AddInput(BaseModel): + amount: int + + +class Total(BaseModel): + total: int + + +class Counter(RpcDefinition, id="counter"): + add = rpc_method(AddInput, Total) + reset = rpc_method(None, Total) + changed = rpc_event(Total) + + +counter = client.rpc(Counter) +try: + print(counter.call(Counter.add, AddInput(amount=2)).total) +except RpcCallError as error: + print(error.type, error.message, error.data) + +with counter.subscribe(Counter.changed) as changes: # only rpc.counter.changed, with validated data + for change in changes: + print(change.data.total) +``` + +Any type pydantic validates works for inputs, outputs and event data. `rpc_method(..., errors={"not_found": Model})` +validates the `data` of declared errors. Failures raise `RpcCallError`; crashes and outputs that do not match the +declared type raise `RpcCallInternalError`. `counter.events()` yields every raw event of the RPC. The untyped +`client.rpc.call(rpc_id, method, input=...)` is still available. + + +## Async + +The async client returns `AsyncRpcClient`, with the same methods: + +```python +counter = client.rpc(Counter) +total = await counter.call(Counter.add, AddInput(amount=2)) +async with counter.subscribe(Counter.changed) as changes: + async for change in changes: + print(change.data.total) +``` + +RPC events travel on the [shared event stream](events.md#shared-event-stream), so they share its delivery rules. diff --git a/docs/reference/events.md b/docs/reference/events.md new file mode 100644 index 0000000..a5923f2 --- /dev/null +++ b/docs/reference/events.md @@ -0,0 +1,50 @@ +# Shared events and plugin RPC + +The shared event stream behind `client.events`, and the typed plugin RPC helpers behind `client.rpc(Definition)`. +See the [events guide](../guides/events.md) and the [plugin RPC guide](../guides/plugin-rpc.md). + +## Event hubs + +::: opencode.EventHub + +::: opencode.AsyncEventHub + +::: opencode.Subscription + +::: opencode.AsyncSubscription + +::: opencode.Reconnect + +::: opencode.EventOverflowError + +::: opencode.SUBSCRIBER_CAPACITY + +::: opencode.session_id_of + +## Plugin RPC + +::: opencode._rpc.TypedRpcResource + +::: opencode.RpcDefinition + +::: opencode.rpc_method + +::: opencode.rpc_event + +::: opencode.RpcMethod + +::: opencode.RpcEventDefinition + +::: opencode.RpcClient + +::: opencode.AsyncRpcClient + +::: opencode.RpcEvents + +::: opencode.AsyncRpcEvents + +::: opencode.RpcEventPayload + +::: opencode.RpcCallError + +::: opencode.RpcCallInternalError diff --git a/docs/reference/index.md b/docs/reference/index.md index f54eeb5..75627eb 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -8,6 +8,7 @@ Generated from the docstrings in `src/opencode`. | [Server launcher](server.md) | `OpencodeServer`, `AsyncOpencodeServer` | | [Resources](resources.md) | Every `client.` and its methods | | [Pagination and streams](pagination.md) | `SyncCursorPage`, `AsyncPaginator`, `Stream`, `AsyncStream` | +| [Shared events and RPC](events.md) | `EventHub`, `Subscription`, `RpcDefinition`, `RpcClient` | | [Exceptions](exceptions.md) | Status errors and every declared server error | | [Types](types/index.md) | The models in `opencode.types`, grouped by domain | diff --git a/mkdocs.yml b/mkdocs.yml index 8fd157f..35423f5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -139,6 +139,7 @@ nav: - Models and unions: guides/models.md - Pagination: guides/pagination.md - Events and streaming: guides/events.md + - Plugin RPC: guides/plugin-rpc.md - Interactive terminals: guides/terminals.md - Errors: guides/errors.md - Retries and timeouts: guides/retries.md @@ -151,6 +152,7 @@ nav: - Server launcher: reference/server.md - Resources: reference/resources.md - Pagination and streams: reference/pagination.md + - Shared events and RPC: reference/events.md - Terminals: reference/terminals.md - Exceptions: reference/exceptions.md - Types: reference/types/index.md diff --git a/src/opencode/__init__.py b/src/opencode/__init__.py index 6dd71e8..c1f35ef 100644 --- a/src/opencode/__init__.py +++ b/src/opencode/__init__.py @@ -81,6 +81,33 @@ from ._types import NOT_GIVEN, NotGiven from ._version import __version__ +# isort: split +# Hand-written helpers, kept apart from the generated error-import block above. +from ._events import ( + SUBSCRIBER_CAPACITY, + AsyncEventHub, + AsyncSubscription, + EventHub, + EventOverflowError, + Reconnect, + Subscription, + session_id_of, +) +from ._rpc import ( + AsyncRpcClient, + AsyncRpcEvents, + RpcCallError, + RpcCallInternalError, + RpcClient, + RpcDefinition, + RpcEventDefinition, + RpcEventPayload, + RpcEvents, + RpcMethod, + rpc_event, + rpc_method, +) + # API_VERSION = "2.1.2" """The OpenCode server version this client was generated from.""" @@ -123,6 +150,27 @@ "UnprocessableEntityError", "__version__", "types", + # Shared event stream and typed plugin RPC + "SUBSCRIBER_CAPACITY", + "AsyncEventHub", + "AsyncRpcClient", + "AsyncRpcEvents", + "AsyncSubscription", + "EventHub", + "EventOverflowError", + "Reconnect", + "RpcCallError", + "RpcCallInternalError", + "RpcClient", + "RpcDefinition", + "RpcEventDefinition", + "RpcEventPayload", + "RpcEvents", + "RpcMethod", + "Subscription", + "rpc_event", + "rpc_method", + "session_id_of", # "AgentNotFoundError", "CommandExecutionError", diff --git a/src/opencode/_client.py b/src/opencode/_client.py index ac70fdf..fdd5d7a 100644 --- a/src/opencode/_client.py +++ b/src/opencode/_client.py @@ -1,11 +1,14 @@ from __future__ import annotations +from functools import cached_property from typing import TYPE_CHECKING from typing_extensions import Self from ._base_client import AsyncAPIClient, SyncAPIClient from ._constants import DEFAULT_MAX_RETRIES, DEFAULT_TIMEOUT, DEFAULT_USERNAME +from ._events import AsyncEventHub, EventHub +from ._rpc import AsyncTypedRpcResource, TypedRpcResource from ._types import NOT_GIVEN, NotGiven, Timeout from .resources import AsyncResources, SyncResources @@ -70,6 +73,16 @@ def __init__( http_client=http_client, ) + @cached_property + def events(self) -> EventHub: + """The shared event stream: one connection for every subscriber. See :class:`opencode.EventHub`.""" + return EventHub(self) + + @cached_property + def rpc(self) -> TypedRpcResource: + """Plugin RPC calls: ``client.rpc.call(...)``, or ``client.rpc(Definition)`` for typed access.""" + return TypedRpcResource(self) + def with_options( self, *, @@ -138,6 +151,16 @@ def __init__( http_client=http_client, ) + @cached_property + def events(self) -> AsyncEventHub: + """The shared event stream: one connection for every subscriber. See :class:`opencode.AsyncEventHub`.""" + return AsyncEventHub(self) + + @cached_property + def rpc(self) -> AsyncTypedRpcResource: + """Plugin RPC calls: ``client.rpc.call(...)``, or ``client.rpc(Definition)`` for typed access.""" + return AsyncTypedRpcResource(self) + def with_options( self, *, diff --git a/src/opencode/_events.py b/src/opencode/_events.py new file mode 100644 index 0000000..deb2c10 --- /dev/null +++ b/src/opencode/_events.py @@ -0,0 +1,729 @@ +"""A shared, multiplexed event stream. + +``GET /api/event`` delivers every event of the server to every connection. Opening one connection per consumer +wastes server work and sockets, so :class:`EventHub` and :class:`AsyncEventHub` share a single connection among all +of their subscribers, the way ``client.event`` does in the official TypeScript client: + +- The connection opens with the first subscriber and closes when the last one leaves. +- Each subscriber has its own bounded queue (:data:`SUBSCRIBER_CAPACITY` events). A subscriber that falls that far + behind fails with :class:`EventOverflowError`; the others keep receiving. +- The most recent ``server.connected`` event is replayed to subscribers that join after it arrived. + +Beyond the TypeScript client, a hub can reconnect with backoff (:class:`Reconnect`), and subscriptions can be +narrowed with :meth:`Subscription.of_type`, :meth:`Subscription.for_session` and :meth:`Subscription.where`. + +The event stream is volatile by contract: events published while no connection is open, including while a hub +waits to reconnect, are never delivered. Use it to react to changes, and read state through the API when exact +history matters. +""" + +from __future__ import annotations + +import asyncio +import collections +import logging +import random +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast + +import anyio +from typing_extensions import Self + +from ._exceptions import APIConnectionError, APIStatusError, OpencodeError +from .types import ServerConnectedEvent + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Awaitable, Callable, Iterator + from types import TracebackType + + from ._client import AsyncOpencode, Opencode + from ._streaming import AsyncStream, Stream + from .types import Event + +__all__ = [ + "SUBSCRIBER_CAPACITY", + "AsyncEventHub", + "AsyncSubscription", + "EventHub", + "EventOverflowError", + "Reconnect", + "Subscription", +] + +log = logging.getLogger("opencode") + +_E = TypeVar("_E") +_T = TypeVar("_T") + +SUBSCRIBER_CAPACITY = 4096 +"""Events a subscriber may have queued before it fails with :class:`EventOverflowError`. Matches the TS client.""" + + +class EventOverflowError(OpencodeError): + """A subscriber fell more than :data:`SUBSCRIBER_CAPACITY` events behind the shared stream. + + Only that subscriber fails. Its queued events are discarded, since they no longer form a gap-free sequence. + """ + + +@dataclass(frozen=True) +class Reconnect: + """Opt-in reconnection policy for an event hub. + + After the connection ends or fails with a connection error, a 5xx or a 429, the hub waits and connects again. + The delay starts at ``initial_delay`` seconds, doubles per failed attempt up to ``max_delay``, gets up to 25% + jitter, and resets once a connection delivers an event. Other errors (authentication, malformed events) still + end every subscription. + + Events published while disconnected are lost. The server sends ``server.connected`` again on every connection; + subscribers see it only once, and late subscribers get the latest one. + + Args: + initial_delay: Seconds before the first reconnection attempt. + max_delay: Upper bound for the delay between attempts. + max_attempts: Consecutive failed attempts before giving up, or ``None`` to retry forever. + """ + + initial_delay: float = 0.5 + max_delay: float = 30.0 + max_attempts: int | None = None + + def delay(self, attempt: int) -> float: + """Seconds to wait before reconnection attempt ``attempt`` (counting from zero).""" + delay = min(self.initial_delay * 2.0**attempt, self.max_delay) + return delay * (1 - 0.25 * random.random()) + + def allows(self, attempt: int) -> bool: + """Whether reconnection attempt ``attempt`` (counting from zero) may run.""" + return self.max_attempts is None or attempt < self.max_attempts + + +def _policy(reconnect: bool | Reconnect) -> Reconnect | None: + if isinstance(reconnect, Reconnect): + return reconnect + return Reconnect() if reconnect else None + + +def _retryable(error: BaseException | None) -> bool: + """Whether a connection that ended with ``error`` (``None`` for a clean end) may be reopened.""" + if error is None or isinstance(error, APIConnectionError): + return True + if isinstance(error, APIStatusError): + return error.status_code >= 500 or error.status_code == 429 + return False + + +def session_id_of(event: object) -> str | None: + """Returns the session an event belongs to, or ``None`` for events outside any session. + + Looks for ``data.session_id`` on typed events and ``data["sessionID"]`` on raw ones. + """ + data: Any = getattr(event, "data", None) + if isinstance(data, dict): + value: object = cast("dict[str, Any]", data).get("sessionID") + else: + value = getattr(data, "session_id", None) + return value if isinstance(value, str) else None + + +class _Filters(Generic[_E]): + """Narrowing helpers shared by sync and async subscriptions.""" + + _predicates: tuple[Callable[[Any], bool], ...] + + def _accepts(self, event: object) -> bool: + return all(predicate(event) for predicate in self._predicates) + + +_End = anyio.EndOfStream +"""Raised by :meth:`AsyncSubscription.get` once the stream is over.""" + +# Completion states of one subscriber. +_OPEN: object = object() +_DONE: object = object() + + +class _SyncSubscriber: + def __init__(self, accepts: Callable[[object], bool], on_activity: Callable[[], object] | None) -> None: + self.accepts = accepts + self.on_activity = on_activity + self.queue: collections.deque[object] = collections.deque() + self.completion: object = _OPEN + self.cond = threading.Condition() + + def push(self, event: object) -> bool: + """Queues an event. Returns ``False`` when the subscriber overflowed and must be dropped.""" + if not self.accepts(event): + return True + with self.cond: + if self.completion is not _OPEN: + return True + if len(self.queue) >= SUBSCRIBER_CAPACITY: + self._finish(EventOverflowError(f"Event subscriber exceeded its {SUBSCRIBER_CAPACITY}-event capacity")) + return False + self.queue.append(event) + self.cond.notify_all() + return True + + def finish(self, completion: object, *, discard: bool) -> None: + with self.cond: + if self.completion is _OPEN: + self._finish(completion, discard=discard) + + def _finish(self, completion: object, *, discard: bool = True) -> None: + self.completion = completion + if discard: + self.queue.clear() + self.cond.notify_all() + + +class _SyncConnection: + def __init__(self) -> None: + self.subscribers: set[_SyncSubscriber] = set() + self.connected: ServerConnectedEvent | None = None + self.stopped = threading.Event() + self.stream: Stream[Event] | None = None + self.delivered = False + + def activity(self) -> None: + for subscriber in list(self.subscribers): + if subscriber.on_activity is not None: + subscriber.on_activity() + + +class EventHub: + """One shared event connection for any number of synchronous subscribers. + + ``client.events`` is the hub of a client. Create another hub with different options, such as reconnection, with + ``EventHub(client, reconnect=True)``. A background thread reads the connection while it is open. + + Example:: + + with client.events.of_type(SessionIdleEvent).for_session(session.id) as idle: + client.session.prompt(session.id, text="Hello") + next(idle) + + Args: + client: The client that opens the connection. + reconnect: ``True`` or a :class:`Reconnect` policy to reopen the connection after it drops. Off by default. + """ + + def __init__(self, client: Opencode, *, reconnect: bool | Reconnect = False) -> None: + self._client = client + self._reconnect = _policy(reconnect) + self._lock = threading.Lock() + self._current: _SyncConnection | None = None + + def subscribe(self, *, on_activity: Callable[[], object] | None = None) -> Subscription[Event]: + """Returns a subscription to every event. It joins the shared connection when first iterated. + + Args: + on_activity: Called on the reader thread for every chunk received, including keepalive comments. + Use it as a liveness signal; keep it fast. + """ + return Subscription(self, (), on_activity) + + def of_type(self, *types: type[_T]) -> Subscription[_T]: + """Shortcut for ``subscribe().of_type(*types)``.""" + return self.subscribe().of_type(*types) + + def for_session(self, session_id: str) -> Subscription[Event]: + """Shortcut for ``subscribe().for_session(session_id)``.""" + return self.subscribe().for_session(session_id) + + def __iter__(self) -> Iterator[Event]: + """Iterates a new subscription. Prefer ``with hub.subscribe() as events`` to release it deterministically.""" + return iter(self.subscribe()) + + def _join(self, subscriber: _SyncSubscriber) -> _SyncConnection: + with self._lock: + connection = self._current + fresh = connection is None + if connection is None: + connection = self._current = _SyncConnection() + connection.subscribers.add(subscriber) + if connection.connected is not None: + subscriber.push(connection.connected) + if fresh: + threading.Thread(target=self._run, args=(connection,), name="opencode-events", daemon=True).start() + return connection + + def _leave(self, connection: _SyncConnection, subscriber: _SyncSubscriber) -> None: + with self._lock: + connection.subscribers.discard(subscriber) + if not connection.subscribers: + self._stop(connection) + + def _stop(self, connection: _SyncConnection) -> None: + # Called with the lock held. + connection.stopped.set() + if self._current is connection: + self._current = None + if connection.stream is not None: + connection.stream.interrupt() + + def _publish(self, connection: _SyncConnection, event: Event) -> None: + for subscriber in list(connection.subscribers): + if not subscriber.push(event): + self._leave(connection, subscriber) + + def _run(self, connection: _SyncConnection) -> None: + error: BaseException | None = None + attempt = 0 + try: + while not connection.stopped.is_set(): + error = None + connection.delivered = False + try: + self._read(connection) + except Exception as exc: + error = exc + if connection.stopped.is_set(): + error = None + break + if connection.delivered: + attempt = 0 + if self._reconnect is None or not _retryable(error) or not self._reconnect.allows(attempt): + break + delay = self._reconnect.delay(attempt) + log.info("Event stream ended (%r); reconnecting in %.2fs", error, delay) + attempt += 1 + if connection.stopped.wait(delay): + error = None + break + finally: + with self._lock: + self._stop(connection) + subscribers = list(connection.subscribers) + completion = _DONE if error is None else error + for subscriber in subscribers: + subscriber.finish(completion, discard=False) + + def _read(self, connection: _SyncConnection) -> None: + """Reads one connection until it ends, setting ``connection.delivered`` once an event arrives.""" + stream = self._client.event.subscribe() + with self._lock: + connection.stream = stream + stopped = connection.stopped.is_set() + try: + if stopped: + return + stream.on_activity = connection.activity + # A fresh connection replays nothing yet; reconnects must not repeat server.connected. + seen_connected = connection.connected is not None + for event in stream: + if connection.stopped.is_set(): + return + connection.delivered = True + if isinstance(event, ServerConnectedEvent): + connection.connected = event + if seen_connected: + continue + seen_connected = True + self._publish(connection, event) + finally: + connection.stream = None + stream.close() + + +class Subscription(_Filters[_E]): + """One consumer of a shared event stream. Iterate it with ``for``; close it, or use ``with``, to leave. + + The subscription joins the shared connection when entered as a context manager or first iterated, so filters + can be chained before that:: + + for event in client.events.for_session(session_id).of_type(SessionIdleEvent): + ... + + Iteration ends when the connection ends, and raises the connection's error if it failed, after delivering the + events already queued. ``close()`` may be called from any thread to stop a blocked iteration. + """ + + def __init__( + self, + hub: EventHub, + predicates: tuple[Callable[[Any], bool], ...], + on_activity: Callable[[], object] | None, + ) -> None: + self._hub = hub + self._predicates = predicates + self._on_activity = on_activity + self._subscriber: _SyncSubscriber | None = None + self._connection: _SyncConnection | None = None + self._closed = False + + def _narrowed(self, predicate: Callable[[Any], bool]) -> Any: + if self._subscriber is not None: + raise RuntimeError("Filters must be added before the subscription is iterated") + return Subscription[Any](self._hub, (*self._predicates, predicate), self._on_activity) + + def of_type(self, *types: type[_T]) -> Subscription[_T]: + """Returns a subscription that only yields events that are instances of ``types``. + + Type checkers narrow the yielded type, so ``of_type(SessionIdleEvent)`` yields ``SessionIdleEvent``. + """ + if not types: + raise TypeError("of_type() needs at least one event type") + return cast("Subscription[_T]", self._narrowed(lambda event: isinstance(event, types))) + + def for_session(self, session_id: str) -> Subscription[_E]: + """Returns a subscription that only yields events of one session (see :func:`session_id_of`).""" + return cast("Subscription[_E]", self._narrowed(lambda event: session_id_of(event) == session_id)) + + def where(self, predicate: Callable[[_E], bool]) -> Subscription[_E]: + """Returns a subscription that only yields events for which ``predicate`` returns true. + + The predicate runs on the reader thread, after the other filters of this subscription. + """ + return cast("Subscription[_E]", self._narrowed(predicate)) + + def __iter__(self) -> Self: + return self + + def __next__(self) -> _E: + return self.get() + + def get(self, timeout: float | None = None) -> _E: + """Returns the next event, waiting up to ``timeout`` seconds (forever when ``None``). + + Raises: + StopIteration: The connection ended or the subscription was closed. + TimeoutError: No event arrived in time. The subscription stays usable. + EventOverflowError: This subscriber fell too far behind. + """ + subscriber = self._start() + if subscriber is None: + raise StopIteration + with subscriber.cond: + if not subscriber.cond.wait_for(lambda: subscriber.queue or subscriber.completion is not _OPEN, timeout): + raise TimeoutError("No event arrived in time") + if subscriber.queue: + return cast("_E", subscriber.queue.popleft()) + completion = subscriber.completion + if isinstance(completion, BaseException): + raise completion + raise StopIteration + + def _start(self) -> _SyncSubscriber | None: + if self._subscriber is None and not self._closed: + self._subscriber = _SyncSubscriber(self._accepts, self._on_activity) + self._connection = self._hub._join(self._subscriber) # pyright: ignore[reportPrivateUsage] + return self._subscriber + + def close(self) -> None: + """Leaves the shared stream. Pending and future iterations end. Safe to call more than once.""" + self._closed = True + subscriber, connection = self._subscriber, self._connection + if subscriber is not None and connection is not None: + subscriber.finish(_DONE, discard=True) + self._hub._leave(connection, subscriber) # pyright: ignore[reportPrivateUsage] + + def __enter__(self) -> Self: + """Joins the shared stream now, so every event published after this point is delivered.""" + self._start() + return self + + def __exit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + self.close() + + +class _AsyncSubscriber: + def __init__(self, accepts: Callable[[object], bool], on_activity: Callable[[], object] | None) -> None: + self.accepts = accepts + self.on_activity = on_activity + self.queue: collections.deque[object] = collections.deque() + self.completion: object = _OPEN + self.wakeup: anyio.Event | None = None + + def _wake(self) -> None: + if self.wakeup is not None: + self.wakeup.set() + self.wakeup = None + + def push(self, event: object) -> bool: + if self.completion is not _OPEN or not self.accepts(event): + return True + if len(self.queue) >= SUBSCRIBER_CAPACITY: + self.finish(EventOverflowError(f"Event subscriber exceeded its {SUBSCRIBER_CAPACITY}-event capacity")) + return False + self.queue.append(event) + self._wake() + return True + + def finish(self, completion: object, *, discard: bool = True) -> None: + if self.completion is not _OPEN: + return + self.completion = completion + if discard: + self.queue.clear() + self._wake() + + +class _AsyncConnection: + def __init__(self) -> None: + self.subscribers: set[_AsyncSubscriber] = set() + self.connected: ServerConnectedEvent | None = None + self.stopped = False + self.scope: anyio.CancelScope | None = None + self.delivered = False + + def is_stopped(self) -> bool: + # A method, so type checkers do not narrow ``stopped`` across awaits that may change it. + return self.stopped + + def activity(self) -> None: + for subscriber in list(self.subscribers): + if subscriber.on_activity is not None: + subscriber.on_activity() + + +_background: set[object] = set() + + +def _spawn(run: Callable[[], Awaitable[None]]) -> None: + """Starts ``run`` as a detached task of the running event loop (asyncio or trio).""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + import trio # only reachable under trio + + trio.lowlevel.spawn_system_task(run, name="opencode-events") + return + task = loop.create_task(cast("Any", run()), name="opencode-events") + _background.add(task) + task.add_done_callback(_background.discard) + + +class AsyncEventHub: + """One shared event connection for any number of asynchronous subscribers. Works with asyncio and trio. + + ``client.events`` is the hub of a client. Create another hub with different options, such as reconnection, with + ``AsyncEventHub(client, reconnect=True)``. The connection is read by a detached task of the running event loop, + and a hub must only be used from one event loop. + + Example:: + + async with client.events.of_type(SessionIdleEvent).for_session(session.id) as idle: + await client.session.prompt(session.id, text="Hello") + await idle.get() + + Args: + client: The client that opens the connection. + reconnect: ``True`` or a :class:`Reconnect` policy to reopen the connection after it drops. Off by default. + """ + + def __init__(self, client: AsyncOpencode, *, reconnect: bool | Reconnect = False) -> None: + self._client = client + self._reconnect = _policy(reconnect) + self._current: _AsyncConnection | None = None + + def subscribe(self, *, on_activity: Callable[[], object] | None = None) -> AsyncSubscription[Event]: + """Returns a subscription to every event. It joins the shared connection when first iterated. + + Args: + on_activity: Called for every chunk received, including keepalive comments. Use it as a liveness + signal; keep it fast and non-blocking. + """ + return AsyncSubscription(self, (), on_activity) + + def of_type(self, *types: type[_T]) -> AsyncSubscription[_T]: + """Shortcut for ``subscribe().of_type(*types)``.""" + return self.subscribe().of_type(*types) + + def for_session(self, session_id: str) -> AsyncSubscription[Event]: + """Shortcut for ``subscribe().for_session(session_id)``.""" + return self.subscribe().for_session(session_id) + + def __aiter__(self) -> AsyncIterator[Event]: + """Iterates a new subscription. Prefer ``async with hub.subscribe() as events`` to release it.""" + return self.subscribe() + + def _join(self, subscriber: _AsyncSubscriber) -> _AsyncConnection: + connection = self._current + if connection is None: + connection = self._current = _AsyncConnection() + connection.subscribers.add(subscriber) + _spawn(lambda: self._run(connection)) + else: + connection.subscribers.add(subscriber) + if connection.connected is not None: + subscriber.push(connection.connected) + return connection + + def _leave(self, connection: _AsyncConnection, subscriber: _AsyncSubscriber) -> None: + connection.subscribers.discard(subscriber) + if not connection.subscribers: + self._stop(connection) + + def _stop(self, connection: _AsyncConnection) -> None: + connection.stopped = True + if self._current is connection: + self._current = None + if connection.scope is not None: + connection.scope.cancel() + + def _publish(self, connection: _AsyncConnection, event: Event) -> None: + for subscriber in list(connection.subscribers): + if not subscriber.push(event): + self._leave(connection, subscriber) + + async def _run(self, connection: _AsyncConnection) -> None: + error: BaseException | None = None + attempt = 0 + try: + with anyio.CancelScope() as scope: + connection.scope = scope + if connection.stopped: + scope.cancel() + while not connection.stopped: + error = None + connection.delivered = False + try: + await self._read(connection) + except Exception as exc: + error = exc + if connection.is_stopped(): + break + if connection.delivered: + attempt = 0 + if self._reconnect is None or not _retryable(error) or not self._reconnect.allows(attempt): + break + delay = self._reconnect.delay(attempt) + log.info("Event stream ended (%r); reconnecting in %.2fs", error, delay) + attempt += 1 + await anyio.sleep(delay) + finally: + if connection.stopped: + error = None + self._stop(connection) + completion = _DONE if error is None else error + for subscriber in list(connection.subscribers): + subscriber.finish(completion, discard=False) + + async def _read(self, connection: _AsyncConnection) -> None: + """Reads one connection until it ends, setting ``connection.delivered`` once an event arrives.""" + stream: AsyncStream[Event] = await self._client.event.subscribe() + try: + stream.on_activity = connection.activity + seen_connected = connection.connected is not None + async for event in stream: + connection.delivered = True + if isinstance(event, ServerConnectedEvent): + connection.connected = event + if seen_connected: + continue + seen_connected = True + self._publish(connection, event) + finally: + with anyio.CancelScope(shield=True): + await stream.close() + + +class AsyncSubscription(_Filters[_E]): + """One consumer of a shared event stream. Iterate it with ``async for``; close it, or use ``async with``. + + The subscription joins the shared connection when entered as a context manager or first iterated, so filters + can be chained before that:: + + async for event in client.events.for_session(session_id).of_type(SessionIdleEvent): + ... + + Iteration ends when the connection ends, and raises the connection's error if it failed, after delivering the + events already queued. Cancelling a task that waits for an event (``anyio.fail_after``, ``asyncio.timeout``, + trio cancel scopes) leaves the subscription usable; close it to leave the stream. + """ + + def __init__( + self, + hub: AsyncEventHub, + predicates: tuple[Callable[[Any], bool], ...], + on_activity: Callable[[], object] | None, + ) -> None: + self._hub = hub + self._predicates = predicates + self._on_activity = on_activity + self._subscriber: _AsyncSubscriber | None = None + self._connection: _AsyncConnection | None = None + self._closed = False + + def _narrowed(self, predicate: Callable[[Any], bool]) -> Any: + if self._subscriber is not None: + raise RuntimeError("Filters must be added before the subscription is iterated") + return AsyncSubscription[Any](self._hub, (*self._predicates, predicate), self._on_activity) + + def of_type(self, *types: type[_T]) -> AsyncSubscription[_T]: + """Returns a subscription that only yields events that are instances of ``types``. + + Type checkers narrow the yielded type, so ``of_type(SessionIdleEvent)`` yields ``SessionIdleEvent``. + """ + if not types: + raise TypeError("of_type() needs at least one event type") + return cast("AsyncSubscription[_T]", self._narrowed(lambda event: isinstance(event, types))) + + def for_session(self, session_id: str) -> AsyncSubscription[_E]: + """Returns a subscription that only yields events of one session (see :func:`session_id_of`).""" + return cast("AsyncSubscription[_E]", self._narrowed(lambda event: session_id_of(event) == session_id)) + + def where(self, predicate: Callable[[_E], bool]) -> AsyncSubscription[_E]: + """Returns a subscription that only yields events for which ``predicate`` returns true.""" + return cast("AsyncSubscription[_E]", self._narrowed(predicate)) + + def __aiter__(self) -> Self: + return self + + async def __anext__(self) -> _E: + try: + return await self.get() + except _End: + raise StopAsyncIteration from None + + async def get(self) -> _E: + """Returns the next event. Bound the wait with ``anyio.fail_after`` or your event loop's timeout. + + Raises: + EndOfStream: The connection ended or the subscription was closed. + EventOverflowError: This subscriber fell too far behind. + """ + subscriber = self._start() + if subscriber is None: + raise _End + while not subscriber.queue and subscriber.completion is _OPEN: + wakeup = subscriber.wakeup = subscriber.wakeup or anyio.Event() + await wakeup.wait() + if subscriber.queue: + return cast("_E", subscriber.queue.popleft()) + if isinstance(subscriber.completion, BaseException): + raise subscriber.completion + raise _End + + def _start(self) -> _AsyncSubscriber | None: + if self._subscriber is None and not self._closed: + self._subscriber = _AsyncSubscriber(self._accepts, self._on_activity) + self._connection = self._hub._join(self._subscriber) # pyright: ignore[reportPrivateUsage] + return self._subscriber + + async def close(self) -> None: + """Leaves the shared stream. Pending and future iterations end. Safe to call more than once.""" + self.close_nowait() + await anyio.lowlevel.checkpoint() + + def close_nowait(self) -> None: + """Synchronous form of :meth:`close`, for callbacks and ``finally`` blocks that cannot await.""" + self._closed = True + subscriber, connection = self._subscriber, self._connection + if subscriber is not None and connection is not None: + subscriber.finish(_DONE) + self._hub._leave(connection, subscriber) # pyright: ignore[reportPrivateUsage] + + async def __aenter__(self) -> Self: + """Joins the shared stream now, so every event published after this point is delivered.""" + self._start() + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + self.close_nowait() diff --git a/src/opencode/_rpc.py b/src/opencode/_rpc.py new file mode 100644 index 0000000..f9b111b --- /dev/null +++ b/src/opencode/_rpc.py @@ -0,0 +1,482 @@ +"""Typed calls to plugin RPCs. + +A plugin registers an RPC with an id, methods and events. Declare the same shape once as a class, then call it +through ``client.rpc(Definition)``, the counterpart of ``client.rpc(definition)`` in the official TypeScript client:: + + from pydantic import BaseModel + from opencode import RpcDefinition, rpc_event, rpc_method + + + class AddInput(BaseModel): + amount: int + + + class Total(BaseModel): + total: int + + + class Counter(RpcDefinition, id="counter"): + add = rpc_method(AddInput, Total) + reset = rpc_method(None, Total) + changed = rpc_event(Total) + + + counter = client.rpc(Counter) + print(counter.call(Counter.add, AddInput(amount=2)).total) + with counter.subscribe(Counter.changed) as changes: + for change in changes: + print(change.data.total) + +Inputs are validated and serialized, and outputs and event data validated, with pydantic, so any type pydantic +accepts works: models, dataclasses, ``TypedDict``, ``list[int]``. RPC failures raise :class:`RpcCallError`. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast, overload + +import pydantic +from typing_extensions import Self + +from ._base_client import type_adapter +from ._errors import RpcError, RpcInternalError +from ._exceptions import OpencodeError, validation_summary +from ._types import NOT_GIVEN, NotGiven +from .resources.rpc import AsyncRpcResource, RpcResource +from .types import RpcEvent + +if TYPE_CHECKING: + from collections.abc import Mapping + from types import TracebackType + + from ._client import AsyncOpencode, Opencode + from ._events import AsyncSubscription, Subscription + from ._types import Headers, Timeout + from .types import Location + +__all__ = [ + "AsyncRpcClient", + "AsyncRpcEvents", + "RpcCallError", + "RpcCallInternalError", + "RpcClient", + "RpcDefinition", + "RpcEventDefinition", + "RpcEventPayload", + "RpcEvents", + "RpcMethod", + "rpc_event", + "rpc_method", +] + +_D = TypeVar("_D", bound="RpcDefinition") +_I = TypeVar("_I") +_O = TypeVar("_O") +_T = TypeVar("_T") + + +class RpcCallError(OpencodeError): + """A plugin RPC method failed (the server answered with ``RpcError``). + + Attributes: + rpc_id: The RPC that was called. + method: The method that was called. + type: The error type the plugin reported, for example ``"not_found"``. + message: The error message. + data: Error details, validated against the method's ``errors`` declaration when it names this type. + """ + + def __init__(self, rpc_id: str, method: str, type: str, message: str, data: object = None) -> None: + super().__init__(f"{rpc_id}.{method} failed with {type}: {message}") + self.rpc_id = rpc_id + self.method = method + self.type = type + self.message = message + self.data = data + + +class RpcCallInternalError(RpcCallError): + """The plugin RPC crashed or returned output that does not match its declaration. + + ``type`` is ``"rpc.internal"`` or ``"rpc.invalid_output"``. Output this client cannot validate against the + method's declared output type is also reported as ``"rpc.invalid_output"``. + """ + + +class RpcMethod(Generic[_I, _O]): + """One method of an :class:`RpcDefinition`. Create it with :func:`rpc_method`.""" + + name: str + + def __init__(self, input: Any, output: Any, *, name: str | None, errors: Mapping[str, Any] | None) -> None: + self.input = input + self.output = output + self.errors = dict(errors or {}) + if name is not None: + self.name = name + + def __set_name__(self, owner: type, name: str) -> None: + if not hasattr(self, "name"): + self.name = name + + def encode(self, value: _I) -> object: + """Validates ``value`` against the input type and returns its JSON form.""" + if self.input is None: + return NOT_GIVEN if value is None else value + adapter = type_adapter(self.input) + return adapter.dump_python(adapter.validate_python(value), mode="json", by_alias=True) + + def decode(self, rpc_id: str, value: object) -> _O: + """Validates a method's output.""" + if self.output is None: + return cast("_O", None) + try: + return cast("_O", type_adapter(self.output).validate_python(value)) + except pydantic.ValidationError as exc: + raise RpcCallInternalError(rpc_id, self.name, "rpc.invalid_output", validation_summary(exc), value) from exc + + def error(self, rpc_id: str, exc: RpcError | RpcInternalError) -> RpcCallError: + """Maps a transport error to :class:`RpcCallError`, decoding declared error data.""" + body = exc.error + data: object = body.data + if body.type in self.errors and data is not None: + try: + data = type_adapter(self.errors[body.type]).validate_python(data) + except pydantic.ValidationError: + data = body.data + cls = RpcCallInternalError if isinstance(exc, RpcInternalError) else RpcCallError + return cls(rpc_id, self.name, body.type, body.message, data) + + +class RpcEventDefinition(Generic[_T]): + """One event of an :class:`RpcDefinition`. Create it with :func:`rpc_event`.""" + + name: str + + def __init__(self, data: Any, *, name: str | None) -> None: + self.data = data + if name is not None: + self.name = name + + def __set_name__(self, owner: type, name: str) -> None: + if not hasattr(self, "name"): + self.name = name + + def parse(self, event: RpcEvent) -> RpcEventPayload[_T]: + """Validates the data of a raw RPC event of this kind. + + Raises: + pydantic.ValidationError: The data does not match the declared type. + """ + data = cast("_T", type_adapter(self.data).validate_python(event.data)) + return RpcEventPayload(name=self.name, data=data, event=event) + + +@overload +def rpc_method( + input: None, output: type[_O], *, name: str | None = None, errors: Mapping[str, Any] | None = None +) -> RpcMethod[None, _O]: ... +@overload +def rpc_method( + input: type[_I], output: type[_O], *, name: str | None = None, errors: Mapping[str, Any] | None = None +) -> RpcMethod[_I, _O]: ... +def rpc_method( + input: Any, output: Any, *, name: str | None = None, errors: Mapping[str, Any] | None = None +) -> RpcMethod[Any, Any]: + """Declares an RPC method on an :class:`RpcDefinition`. + + Args: + input: The input type, or ``None`` for a method without input. + output: The output type. Use ``type(None)`` for a method that returns nothing. + name: The method name on the server. Defaults to the attribute name. + errors: Data types of the errors the method declares, by error type. Matching error data is validated + into these types on :attr:`RpcCallError.data`. + """ + return RpcMethod(input, output, name=name, errors=errors) + + +def rpc_event(data: type[_T], *, name: str | None = None) -> RpcEventDefinition[_T]: + """Declares an RPC event on an :class:`RpcDefinition`. + + Args: + data: The type of the event's ``data``. + name: The event name on the server. Defaults to the attribute name. + """ + return RpcEventDefinition(data, name=name) + + +class RpcDefinition: + """Base class for plugin RPC declarations. Pass the RPC id as a class keyword:: + + class Counter(RpcDefinition, id="counter"): + add = rpc_method(AddInput, Total) + changed = rpc_event(Total) + + The class is only a declaration; bind it to a client with ``client.rpc(Counter)``. + """ + + id: ClassVar[str] + methods: ClassVar[dict[str, RpcMethod[Any, Any]]] + events: ClassVar[dict[str, RpcEventDefinition[Any]]] + + def __init_subclass__(cls, *, id: str | None = None, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if id is not None: + cls.id = id + elif not hasattr(cls, "id"): + raise TypeError(f"{cls.__name__} needs an RPC id: class {cls.__name__}(RpcDefinition, id='...')") + methods: dict[str, RpcMethod[Any, Any]] = {} + events: dict[str, RpcEventDefinition[Any]] = {} + for base in reversed(cls.__mro__): + for value in vars(base).values(): + if isinstance(value, RpcMethod): + methods[value.name] = value + elif isinstance(value, RpcEventDefinition): + events[value.name] = value + cls.methods = methods + cls.events = events + + @classmethod + def event_type(cls, event: RpcEventDefinition[Any] | str) -> str: + """The wire ``type`` of an event of this RPC: ``rpc..``.""" + return f"rpc.{cls.id}.{event if isinstance(event, str) else event.name}" + + +@dataclass(frozen=True) +class RpcEventPayload(Generic[_T]): + """A plugin RPC event with validated data. + + Attributes: + name: The event name, without the ``rpc..`` prefix. + data: The validated event data. + event: The raw event, with its ``id``, ``created`` time and ``location``. + """ + + name: str + data: _T + event: RpcEvent + + +def _check_method(definition: type[RpcDefinition], method: RpcMethod[Any, Any]) -> None: + if not any(known is method for known in definition.methods.values()): + raise ValueError(f"{getattr(method, 'name', method)!r} is not a method of {definition.__name__}") + + +def _check_event(definition: type[RpcDefinition], event: RpcEventDefinition[Any]) -> None: + if not any(known is event for known in definition.events.values()): + raise ValueError(f"{getattr(event, 'name', event)!r} is not an event of {definition.__name__}") + + +class RpcClient(Generic[_D]): + """A plugin RPC bound to a synchronous client. Get one with ``client.rpc(Definition)``.""" + + def __init__(self, client: Opencode, definition: type[_D]) -> None: + self._client = client + self.definition = definition + + @overload + def call( + self, + method: RpcMethod[None, _O], + input: None = None, + *, + location: Location | NotGiven | None = NOT_GIVEN, + extra_headers: Headers | None = None, + timeout: Timeout | NotGiven = NOT_GIVEN, + ) -> _O: ... + @overload + def call( + self, + method: RpcMethod[_I, _O], + input: _I, + *, + location: Location | NotGiven | None = NOT_GIVEN, + extra_headers: Headers | None = None, + timeout: Timeout | NotGiven = NOT_GIVEN, + ) -> _O: ... + def call( + self, + method: RpcMethod[Any, Any], + input: Any = None, + *, + location: Location | NotGiven | None = NOT_GIVEN, + extra_headers: Headers | None = None, + timeout: Timeout | NotGiven = NOT_GIVEN, + ) -> Any: + """Calls a method of this RPC and returns its validated output. + + Raises: + RpcCallError: The method failed. + RpcCallInternalError: The method crashed or returned invalid output. + """ + _check_method(self.definition, method) + rpc_id = self.definition.id + try: + result = self._client.rpc.call( + rpc_id, + method.name, + input=method.encode(input), + location=location, + extra_headers=extra_headers, + timeout=timeout, + ) + except (RpcError, RpcInternalError) as exc: + raise method.error(rpc_id, exc) from exc + return method.decode(rpc_id, result.output) + + def events(self) -> Subscription[RpcEvent]: + """Subscribes to every event of this RPC, as raw events, on the client's shared event stream.""" + prefix = f"rpc.{self.definition.id}." + return self._client.events.of_type(RpcEvent).where(lambda event: event.type.startswith(prefix)) + + def subscribe(self, event: RpcEventDefinition[_T]) -> RpcEvents[_T]: + """Subscribes to one event of this RPC, with validated data, on the client's shared event stream.""" + _check_event(self.definition, event) + wire = self.definition.event_type(event) + return RpcEvents(self._client.events.of_type(RpcEvent).where(lambda raw: raw.type == wire), event) + + +class AsyncRpcClient(Generic[_D]): + """A plugin RPC bound to an asynchronous client. Get one with ``client.rpc(Definition)``.""" + + def __init__(self, client: AsyncOpencode, definition: type[_D]) -> None: + self._client = client + self.definition = definition + + @overload + async def call( + self, + method: RpcMethod[None, _O], + input: None = None, + *, + location: Location | NotGiven | None = NOT_GIVEN, + extra_headers: Headers | None = None, + timeout: Timeout | NotGiven = NOT_GIVEN, + ) -> _O: ... + @overload + async def call( + self, + method: RpcMethod[_I, _O], + input: _I, + *, + location: Location | NotGiven | None = NOT_GIVEN, + extra_headers: Headers | None = None, + timeout: Timeout | NotGiven = NOT_GIVEN, + ) -> _O: ... + async def call( + self, + method: RpcMethod[Any, Any], + input: Any = None, + *, + location: Location | NotGiven | None = NOT_GIVEN, + extra_headers: Headers | None = None, + timeout: Timeout | NotGiven = NOT_GIVEN, + ) -> Any: + """Calls a method of this RPC and returns its validated output. + + Raises: + RpcCallError: The method failed. + RpcCallInternalError: The method crashed or returned invalid output. + """ + _check_method(self.definition, method) + rpc_id = self.definition.id + try: + result = await self._client.rpc.call( + rpc_id, + method.name, + input=method.encode(input), + location=location, + extra_headers=extra_headers, + timeout=timeout, + ) + except (RpcError, RpcInternalError) as exc: + raise method.error(rpc_id, exc) from exc + return method.decode(rpc_id, result.output) + + def events(self) -> AsyncSubscription[RpcEvent]: + """Subscribes to every event of this RPC, as raw events, on the client's shared event stream.""" + prefix = f"rpc.{self.definition.id}." + return self._client.events.of_type(RpcEvent).where(lambda event: event.type.startswith(prefix)) + + def subscribe(self, event: RpcEventDefinition[_T]) -> AsyncRpcEvents[_T]: + """Subscribes to one event of this RPC, with validated data, on the client's shared event stream.""" + _check_event(self.definition, event) + wire = self.definition.event_type(event) + return AsyncRpcEvents(self._client.events.of_type(RpcEvent).where(lambda raw: raw.type == wire), event) + + +class RpcEvents(Generic[_T]): + """Events of one kind from a plugin RPC. Iterate with ``for``; close it, or use ``with``, to leave.""" + + def __init__(self, source: Subscription[RpcEvent], event: RpcEventDefinition[_T]) -> None: + self._source = source + self._event = event + + def __iter__(self) -> Self: + return self + + def __next__(self) -> RpcEventPayload[_T]: + return self._event.parse(next(self._source)) + + def get(self, timeout: float | None = None) -> RpcEventPayload[_T]: + """Returns the next event, waiting up to ``timeout`` seconds. See :meth:`Subscription.get`.""" + return self._event.parse(self._source.get(timeout)) + + def close(self) -> None: + """Leaves the shared event stream.""" + self._source.close() + + def __enter__(self) -> Self: + return self + + def __exit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + self.close() + + +class AsyncRpcEvents(Generic[_T]): + """Events of one kind from a plugin RPC. Iterate with ``async for``; close it, or use ``async with``.""" + + def __init__(self, source: AsyncSubscription[RpcEvent], event: RpcEventDefinition[_T]) -> None: + self._source = source + self._event = event + + def __aiter__(self) -> Self: + return self + + async def __anext__(self) -> RpcEventPayload[_T]: + return self._event.parse(await self._source.__anext__()) + + async def get(self) -> RpcEventPayload[_T]: + """Returns the next event. See :meth:`AsyncSubscription.get`.""" + return self._event.parse(await self._source.get()) + + async def close(self) -> None: + """Leaves the shared event stream.""" + await self._source.close() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + await self.close() + + +class TypedRpcResource(RpcResource): + """``client.rpc``: the generated ``call`` plus ``client.rpc(Definition)`` for typed access.""" + + def __call__(self, definition: type[_D]) -> RpcClient[_D]: + """Binds a plugin RPC declaration to this client.""" + return RpcClient(self._client, definition) + + +class AsyncTypedRpcResource(AsyncRpcResource): + """``client.rpc``: the generated ``call`` plus ``client.rpc(Definition)`` for typed access.""" + + def __call__(self, definition: type[_D]) -> AsyncRpcClient[_D]: + """Binds a plugin RPC declaration to this client.""" + return AsyncRpcClient(self._client, definition) diff --git a/src/opencode/_streaming.py b/src/opencode/_streaming.py index b00a106..53ee463 100644 --- a/src/opencode/_streaming.py +++ b/src/opencode/_streaming.py @@ -2,8 +2,10 @@ from __future__ import annotations +import contextlib import json import re +import socket from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Generic, TypeVar @@ -14,7 +16,7 @@ from ._exceptions import APIConnectionError, APIResponseValidationError, validation_summary if TYPE_CHECKING: - from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator + from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator from types import TracebackType __all__ = ["AsyncStream", "ServerSentEvent", "Stream"] @@ -111,6 +113,12 @@ class _StreamBase(Generic[_T]): def __init__(self, response: httpx.Response, adapter: pydantic.TypeAdapter[_T]) -> None: self.response = response self._adapter = adapter + self.on_activity: Callable[[], object] | None = None + """Called for every chunk the stream receives, including keepalive comments that carry no event.""" + + def _activity(self) -> None: + if self.on_activity is not None: + self.on_activity() def _decode(self, event: ServerSentEvent) -> _T: try: @@ -148,6 +156,7 @@ def _events(self) -> Generator[_T, None, None]: try: splitter = _LineSplitter() for chunk in self.response.iter_text(): + self._activity() for line in splitter.feed(chunk): if (event := decoder.feed_line(line)) is not None: yield self._decode(event) @@ -166,6 +175,17 @@ def close(self) -> None: self._iterator.close() self.response.close() + def interrupt(self) -> None: + """Unblocks a read in progress on another thread, which then ends with :class:`APIConnectionError`. + + ``close`` must run on the iterating thread. Call this from any other thread first to wake it. + """ + network_stream: Any = self.response.extensions.get("network_stream") + sock: Any = network_stream.get_extra_info("socket") if network_stream is not None else None + if isinstance(sock, socket.socket): + with contextlib.suppress(OSError): + sock.shutdown(socket.SHUT_RDWR) + def __enter__(self) -> Self: return self @@ -200,6 +220,7 @@ async def _events(self) -> AsyncGenerator[_T, None]: try: splitter = _LineSplitter() async for chunk in self.response.aiter_text(): + self._activity() for line in splitter.feed(chunk): if (event := decoder.feed_line(line)) is not None: yield self._decode(event) diff --git a/tests/test_events.py b/tests/test_events.py new file mode 100644 index 0000000..75d112a --- /dev/null +++ b/tests/test_events.py @@ -0,0 +1,396 @@ +"""The shared event hub: multiplexing, overflow, replay, reconnect, filters and cancellation.""" + +from __future__ import annotations + +import collections +import json +import queue +import threading +from collections.abc import AsyncIterator, Iterator + +import anyio +import httpx +import pytest +import respx + +from opencode import ( + AsyncEventHub, + AsyncOpencode, + EventHub, + EventOverflowError, + Opencode, + Reconnect, + UnauthorizedError, + session_id_of, +) +from opencode.types import Event, ServerConnectedEvent, SessionIdleEvent, UnknownEvent + +from .conftest import BASE_URL, PASSWORD + +SSE = {"content-type": "text/event-stream"} + + +def frame(payload: dict[str, object]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def connected(n: int = 1) -> bytes: + return frame({"type": "server.connected", "id": f"evt_c{n}", "data": {}}) + + +def idle(session_id: str, n: int = 1) -> bytes: + return frame({"type": "session.idle", "id": f"evt_i{n}", "created": 1, "data": {"sessionID": session_id}}) + + +def other(n: int = 1) -> bytes: + return frame({"type": "brand.new", "id": f"evt_o{n}", "created": 1}) + + +END = b"" + + +class SyncFeed: + """One fake SSE connection whose chunks the test pushes. ``END`` closes it.""" + + def __init__(self) -> None: + self.chunks: queue.Queue[bytes] = queue.Queue() + self.closed = threading.Event() + + def push(self, *chunks: bytes) -> None: + for chunk in chunks: + self.chunks.put(chunk) + + def __iter__(self) -> Iterator[bytes]: + try: + while (chunk := self.chunks.get(timeout=5)) != END: + yield chunk + finally: + self.closed.set() + + +class Server: + """Hands out a new feed (or a prepared response) for each connection.""" + + def __init__(self, mock_api: respx.MockRouter, *, sync: bool) -> None: + self.sync = sync + self.feeds: list[SyncFeed | AsyncFeed] = [] + self.prepared: list[httpx.Response] = [] + self.opened = threading.Semaphore(0) + mock_api.get("/api/event").mock(side_effect=self._respond) + + def _respond(self, request: httpx.Request) -> httpx.Response: + if self.prepared: + response = self.prepared.pop(0) + else: + feed: SyncFeed | AsyncFeed = SyncFeed() if self.sync else AsyncFeed() + self.feeds.append(feed) + response = httpx.Response(200, headers=SSE, content=feed) + self.opened.release() + return response + + def wait_open(self, count: int) -> SyncFeed: + while len(self.feeds) < count: + assert self.opened.acquire(timeout=5) + feed = self.feeds[count - 1] + assert isinstance(feed, SyncFeed) + return feed + + +class AsyncFeed: + """The async counterpart of ``SyncFeed``. It polls, since ``anyio.sleep`` is patched in tests.""" + + def __init__(self) -> None: + self.chunks: collections.deque[bytes] = collections.deque() + + def push(self, *chunks: bytes) -> None: + self.chunks.extend(chunks) + + async def __aiter__(self) -> AsyncIterator[bytes]: + while True: + while not self.chunks: + await anyio.lowlevel.checkpoint() + chunk = self.chunks.popleft() + if chunk == END: + return + yield chunk + + +# ---------------------------------------------------------------- sync + + +def test_subscribers_share_one_connection(client: Opencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=True) + with client.events.subscribe() as first, client.events.subscribe() as second: + feed_ready = threading.Thread(target=lambda: server.wait_open(1).push(connected())) + feed_ready.start() + assert isinstance(first.get(5), ServerConnectedEvent) + assert isinstance(second.get(5), ServerConnectedEvent) + feed_ready.join() + server.feeds[0].push(idle("ses_1")) + assert isinstance(first.get(5), SessionIdleEvent) + assert isinstance(second.get(5), SessionIdleEvent) + assert len(server.feeds) == 1 + server.feeds[0].push(END) + + +def test_late_subscriber_gets_server_connected_replayed(client: Opencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=True) + with client.events.subscribe() as first: + threading.Thread(target=lambda: server.wait_open(1).push(connected(), other())).start() + assert isinstance(first.get(5), ServerConnectedEvent) + assert isinstance(first.get(5), UnknownEvent) + with client.events.subscribe() as late: + assert isinstance(late.get(5), ServerConnectedEvent) + server.wait_open(1).push(idle("ses_1")) + assert isinstance(late.get(5), SessionIdleEvent) + server.feeds[0].push(END) + + +def test_slow_subscriber_overflows_alone( + client: Opencode, mock_api: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("opencode._events.SUBSCRIBER_CAPACITY", 3) + server = Server(mock_api, sync=True) + with client.events.subscribe() as slow, client.events.subscribe() as fast: + # Both join on first use; the slow one never reads after that. + threading.Thread(target=lambda: server.wait_open(1).push(connected())).start() + assert isinstance(slow.get(5), ServerConnectedEvent) + assert isinstance(fast.get(5), ServerConnectedEvent) + feed = server.wait_open(1) + for n in range(5): + feed.push(idle("ses_1", n)) + assert isinstance(fast.get(5), SessionIdleEvent) + with pytest.raises(EventOverflowError, match="3-event capacity"): + slow.get(5) + feed.push(other()) + assert isinstance(fast.get(5), UnknownEvent) + feed.push(END) + + +def test_last_subscriber_leaving_closes_the_connection(client: Opencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=True) + subscription = client.events.subscribe() + threading.Thread(target=lambda: server.wait_open(1).push(connected())).start() + assert isinstance(subscription.get(5), ServerConnectedEvent) + subscription.close() + with pytest.raises(StopIteration): + subscription.get(1) + feed = server.wait_open(1) + feed.push(other()) # wakes the reader, which sees the hub stopped + assert feed.closed.wait(5) + assert client.events._current is None + + +def test_close_from_another_thread_unblocks_iteration(client: Opencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=True) + subscription = client.events.subscribe() + received: list[Event] = [] + + def consume() -> None: + received.extend(subscription) + + worker = threading.Thread(target=consume) + worker.start() + server.wait_open(1).push(connected()) + while not received: + threading.Event().wait(0.01) + subscription.close() + worker.join(5) + assert not worker.is_alive() + server.feeds[0].push(END) + + +def test_stream_end_finishes_after_queued_events(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.get("/api/event").respond(headers=SSE, content=connected() + idle("ses_1")) + events = list(client.events) + assert [event.type for event in events] == ["server.connected", "session.idle"] + + +def test_filters_narrow_by_type_and_session(client: Opencode, mock_api: respx.MockRouter) -> None: + body = connected() + idle("ses_1", 1) + other() + idle("ses_2", 2) + idle("ses_1", 3) + mock_api.get("/api/event").respond(headers=SSE, content=body) + with client.events.of_type(SessionIdleEvent).for_session("ses_1") as events: + ids = [event.id for event in events] + assert ids == ["evt_i1", "evt_i3"] + with client.events.for_session("ses_2") as events: + assert [event.id for event in events] == ["evt_i2"] + with client.events.subscribe().where(lambda event: event.type == "brand.new") as events: + assert [event.id for event in events] == ["evt_o1"] + + +def test_filters_cannot_be_added_after_iteration_starts(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.get("/api/event").respond(headers=SSE, content=connected()) + with client.events.subscribe() as events: + next(events) + with pytest.raises(RuntimeError, match="before"): + events.of_type(SessionIdleEvent) + with pytest.raises(TypeError): + client.events.of_type() + + +def test_on_activity_fires_on_keepalive_comments(client: Opencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=True) + beats = threading.Semaphore(0) + with client.events.subscribe(on_activity=beats.release) as events: + threading.Thread(target=lambda: server.wait_open(1).push(b": keepalive\n\n")).start() + with pytest.raises(TimeoutError): + events.get(0.2) + assert beats.acquire(timeout=5) + server.feeds[0].push(END) + + +def test_reconnect_dedupes_server_connected(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.get("/api/event").mock( + side_effect=[ + httpx.Response(200, headers=SSE, content=connected(1) + idle("ses_1", 1)), + httpx.ConnectError("refused"), + httpx.Response(503), + httpx.Response(200, headers=SSE, content=connected(2) + idle("ses_1", 2)), + httpx.Response(401, json={"_tag": "UnauthorizedError", "message": "no"}), + ] + ) + hub = EventHub(client.with_options(max_retries=0), reconnect=Reconnect(initial_delay=0, max_delay=0)) + received: list[Event] = [] + with pytest.raises(UnauthorizedError), hub.subscribe() as events: + received.extend(events) + assert [event.id for event in received] == ["evt_c1", "evt_i1", "evt_i2"] + with hub.subscribe() as late: # the hub keeps working after a failure + assert late._hub is hub + + +def test_reconnect_gives_up_after_max_attempts(client: Opencode, mock_api: respx.MockRouter) -> None: + route = mock_api.get("/api/event").mock(side_effect=httpx.ConnectError("refused")) + hub = EventHub(client.with_options(max_retries=0), reconnect=Reconnect(initial_delay=0, max_attempts=2)) + with pytest.raises(Exception, match="refused"), hub.subscribe() as events: + next(events) + assert route.call_count == 3 + + +def test_without_reconnect_errors_end_every_subscription(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.get("/api/event").mock(side_effect=httpx.ConnectError("refused")) + with pytest.raises(Exception, match="refused"), client.with_options(max_retries=0).events.subscribe() as events: + next(events) + + +def test_session_id_of() -> None: + assert session_id_of(ServerConnectedEvent(id="e", data={})) is None + assert session_id_of(UnknownEvent.model_validate({"type": "x", "data": {"sessionID": "ses_9"}})) == "ses_9" + + +def test_generated_subscribe_still_works(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.get("/api/event").respond(headers=SSE, content=connected()) + with client.event.subscribe() as stream: + assert [event.type for event in stream] == ["server.connected"] + + +# ---------------------------------------------------------------- async + + +@pytest.fixture(params=["asyncio", "trio"]) +def anyio_backend(request: pytest.FixtureRequest) -> str: + return str(request.param) + + +async def _wait_feed(server: Server, count: int) -> AsyncFeed: + with anyio.fail_after(5): + while len(server.feeds) < count: + await anyio.lowlevel.checkpoint() # anyio.sleep is patched to return at once + feed = server.feeds[count - 1] + assert isinstance(feed, AsyncFeed) + return feed + + +@pytest.fixture +async def aclient(mock_api: respx.MockRouter) -> AsyncIterator[AsyncOpencode]: + async with AsyncOpencode(BASE_URL, password=PASSWORD, max_retries=0) as instance: + yield instance + + +@pytest.mark.anyio +async def test_async_multiplexing_and_replay(aclient: AsyncOpencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=False) + async with aclient.events.subscribe() as first, aclient.events.subscribe() as second: + with anyio.fail_after(5): + feed = await _wait_feed(server, 1) + feed.push(connected()) + assert isinstance(await first.get(), ServerConnectedEvent) + assert isinstance(await second.get(), ServerConnectedEvent) + feed.push(idle("ses_1")) + assert isinstance(await first.get(), SessionIdleEvent) + assert isinstance(await second.get(), SessionIdleEvent) + async with aclient.events.subscribe() as late: + assert isinstance(await late.get(), ServerConnectedEvent) + assert len(server.feeds) == 1 + assert aclient.events._current is None + + +@pytest.mark.anyio +async def test_async_cancellation_keeps_subscription_usable(aclient: AsyncOpencode, mock_api: respx.MockRouter) -> None: + server = Server(mock_api, sync=False) + async with aclient.events.of_type(SessionIdleEvent) as idle_events: + with anyio.move_on_after(0.05) as scope: + await idle_events.get() + assert scope.cancelled_caught + feed = await _wait_feed(server, 1) + feed.push(connected(), idle("ses_1")) + with anyio.fail_after(5): + assert (await idle_events.get()).data.session_id == "ses_1" + assert aclient.events._current is None + + +@pytest.mark.anyio +async def test_async_overflow_and_end( + aclient: AsyncOpencode, mock_api: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("opencode._events.SUBSCRIBER_CAPACITY", 2) + server = Server(mock_api, sync=False) + async with aclient.events.subscribe() as slow, aclient.events.for_session("ses_1") as fast: + feed = await _wait_feed(server, 1) + feed.push(connected()) + with anyio.fail_after(5): + for n in (1, 2, 3): + feed.push(idle("ses_1", n)) + assert (await fast.get()).id == f"evt_i{n}" + feed.push(END) + assert [event async for event in fast] == [] + with pytest.raises(EventOverflowError): + await slow.get() + + +@pytest.mark.anyio +async def test_async_reconnect_and_activity(aclient: AsyncOpencode, mock_api: respx.MockRouter) -> None: + mock_api.get("/api/event").mock( + side_effect=[ + httpx.Response(200, headers=SSE, content=b": hi\n\n" + connected(1)), + httpx.Response(502), + httpx.Response(200, headers=SSE, content=connected(2) + other(1)), + httpx.Response(401, json={"_tag": "UnauthorizedError", "message": "no"}), + ] + ) + beats: list[None] = [] + hub = AsyncEventHub(aclient, reconnect=Reconnect(initial_delay=0)) + received: list[Event] = [] + async with hub.subscribe(on_activity=lambda: beats.append(None)) as events: + with anyio.fail_after(5): + received.extend([await events.get(), await events.get()]) + with pytest.raises(UnauthorizedError): + await events.get() + assert [event.id for event in received] == ["evt_c1", "evt_o1"] + assert len(beats) >= 2 + + +@pytest.mark.anyio +async def test_async_request_cancellation_leaves_client_usable( + aclient: AsyncOpencode, mock_api: respx.MockRouter +) -> None: + async def slow(request: httpx.Request) -> httpx.Response: + await anyio.Event().wait() + raise AssertionError("unreachable") + + mock_api.get("/api/agent").mock(side_effect=slow) + with anyio.move_on_after(0.05) as scope: + await aclient.agent.list() + assert scope.cancelled_caught + mock_api.get("/api/event").respond(headers=SSE, content=connected()) + with anyio.fail_after(5): + assert [event.type async for event in aclient.events] == ["server.connected"] diff --git a/tests/test_live.py b/tests/test_live.py index 77b0ab0..a7108e8 100644 --- a/tests/test_live.py +++ b/tests/test_live.py @@ -9,9 +9,19 @@ from collections.abc import Iterator from pathlib import Path +import anyio import pytest -from opencode import AsyncOpencodeServer, Opencode, OpencodeServer, SessionNotFoundError +from opencode import ( + AsyncOpencodeServer, + Opencode, + OpencodeServer, + RpcCallError, + RpcDefinition, + SessionNotFoundError, + rpc_method, + session_id_of, +) from opencode.types import ServerConnectedEvent pytestmark = [ @@ -103,3 +113,41 @@ def test_pty_attach_echo(client: Opencode) -> None: output += chunk finally: client.pty.remove(pty.id) + + +def test_shared_events_multiplex_and_filter_by_session(client: Opencode) -> None: + beats: list[None] = [] + with client.events.subscribe(on_activity=lambda: beats.append(None)) as everything: + assert isinstance(everything.get(10), ServerConnectedEvent) + session = client.session.create(title="events") + with client.events.for_session(session.id) as only_session: + client.session.prompt(session.id, text="hello", resume=False) + event = only_session.get(10) + assert session_id_of(event) == session.id + assert client.events._current is not None + client.session.remove(session.id) + assert beats + assert client.events._current is None + + +class Missing(RpcDefinition, id="opencode-python-sdk-missing"): + ping = rpc_method(None, str) + + +def test_typed_rpc_maps_errors(client: Opencode) -> None: + with pytest.raises(RpcCallError) as caught: + client.rpc(Missing).call(Missing.ping) + assert caught.value.rpc_id == "opencode-python-sdk-missing" + assert caught.value.method == "ping" + + +@pytest.mark.anyio +async def test_async_shared_events(tmp_path: Path) -> None: + running = AsyncOpencodeServer(cwd=tmp_path, isolated=True, env=_sandbox_env(tmp_path)) + async with running, running.client() as client: + async with client.events.subscribe() as first, client.events.subscribe() as second: + with anyio.fail_after(10): + assert isinstance(await first.get(), ServerConnectedEvent) + assert isinstance(await second.get(), ServerConnectedEvent) + with pytest.raises(RpcCallError): + await client.rpc(Missing).call(Missing.ping) diff --git a/tests/test_rpc.py b/tests/test_rpc.py new file mode 100644 index 0000000..59fe9fc --- /dev/null +++ b/tests/test_rpc.py @@ -0,0 +1,154 @@ +"""Typed plugin RPC calls and events.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator + +import pytest +import respx +from pydantic import BaseModel + +from opencode import ( + AsyncOpencode, + Opencode, + RpcCallError, + RpcCallInternalError, + RpcDefinition, + RpcError, + rpc_event, + rpc_method, +) +from opencode.types import Location + +from .conftest import BASE_URL, PASSWORD + +SSE = {"content-type": "text/event-stream"} + + +class AddInput(BaseModel): + amount: int + + +class Total(BaseModel): + total: int + + +class Missing(BaseModel): + key: str + + +class Counter(RpcDefinition, id="counter"): + add = rpc_method(AddInput, Total, errors={"not_found": Missing}) + reset = rpc_method(None, Total, name="reset-all") + changed = rpc_event(Total) + cleared = rpc_event(dict[str, object]) + + +def rpc_frame(event_type: str, data: object, n: int) -> bytes: + payload = {"type": event_type, "id": f"evt_{n}", "created": 1, "location": {"directory": "/p"}, "data": data} + return f"data: {json.dumps(payload)}\n\n".encode() + + +def test_definition_collects_methods_and_events() -> None: + assert Counter.id == "counter" + assert set(Counter.methods) == {"add", "reset-all"} + assert set(Counter.events) == {"changed", "cleared"} + assert Counter.event_type(Counter.changed) == "rpc.counter.changed" + with pytest.raises(TypeError, match="needs an RPC id"): + type("Nameless", (RpcDefinition,), {}) + + +def test_call_validates_input_and_output(client: Opencode, mock_api: respx.MockRouter) -> None: + route = mock_api.post("/api/rpc/counter/add").respond(json={"output": {"total": 5}}) + result = client.rpc(Counter).call(Counter.add, AddInput(amount=2), location=Location(directory="/p")) + assert result == Total(total=5) + request = route.calls.last.request + assert json.loads(request.content) == {"input": {"amount": 2}} + assert "location" in str(request.url) + + +def test_call_without_input_uses_declared_name(client: Opencode, mock_api: respx.MockRouter) -> None: + route = mock_api.post("/api/rpc/counter/reset-all").respond(json={"output": {"total": 0}}) + assert client.rpc(Counter).call(Counter.reset).total == 0 + assert json.loads(route.calls.last.request.content) == {} + + +def test_generated_call_still_available(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.post("/api/rpc/counter/add").respond(json={"output": 1}) + assert client.rpc.call("counter", "add", input={"amount": 1}).output == 1 + + +def test_rpc_error_maps_to_rpc_call_error_with_typed_data(client: Opencode, mock_api: respx.MockRouter) -> None: + mock_api.post("/api/rpc/counter/add").respond( + 400, json={"_tag": "RpcError", "type": "not_found", "message": "no counter", "data": {"key": "k"}} + ) + with pytest.raises(RpcCallError) as caught: + client.rpc(Counter).call(Counter.add, AddInput(amount=1)) + error = caught.value + assert (error.rpc_id, error.method, error.type, error.message) == ("counter", "add", "not_found", "no counter") + assert error.data == Missing(key="k") + assert isinstance(error.__cause__, RpcError) + assert not isinstance(error, RpcCallInternalError) + + +def test_internal_error_and_invalid_output(client: Opencode, mock_api: respx.MockRouter) -> None: + route = mock_api.post("/api/rpc/counter/add") + route.respond(500, json={"_tag": "RpcInternalError", "type": "rpc.internal", "message": "boom"}) + with pytest.raises(RpcCallInternalError, match=r"rpc\.internal: boom"): + client.with_options(max_retries=0).rpc(Counter).call(Counter.add, AddInput(amount=1)) + route.respond(json={"output": {"total": "many"}}) + with pytest.raises(RpcCallInternalError) as caught: + client.rpc(Counter).call(Counter.add, AddInput(amount=1)) + assert caught.value.type == "rpc.invalid_output" + + +def test_call_rejects_foreign_method(client: Opencode) -> None: + class Other(RpcDefinition, id="other"): + add = rpc_method(AddInput, Total) + + with pytest.raises(ValueError, match="not a method of Counter"): + client.rpc(Counter).call(Other.add, AddInput(amount=1)) + with pytest.raises(ValueError, match="not an event of Counter"): + client.rpc(Counter).subscribe(rpc_event(Total)) + + +def test_subscribe_filters_and_validates_events(client: Opencode, mock_api: respx.MockRouter) -> None: + body = ( + rpc_frame("rpc.counter.changed", {"total": 1}, 1) + + rpc_frame("rpc.other.changed", {"total": 2}, 2) + + rpc_frame("rpc.counter.cleared", {}, 3) + + rpc_frame("rpc.counter.changed", {"total": 3}, 4) + ) + mock_api.get("/api/event").respond(headers=SSE, content=body) + counter = client.rpc(Counter) + with counter.subscribe(Counter.changed) as changes: + received = list(changes) + assert [change.data.total for change in received] == [1, 3] + assert received[0].name == "changed" + assert received[0].event.id == "evt_1" + with counter.events() as raw: + assert [event.type for event in raw] == ["rpc.counter.changed", "rpc.counter.cleared", "rpc.counter.changed"] + + +@pytest.fixture +async def aclient(mock_api: respx.MockRouter) -> AsyncIterator[AsyncOpencode]: + async with AsyncOpencode(BASE_URL, password=PASSWORD, max_retries=0) as instance: + yield instance + + +@pytest.mark.anyio +async def test_async_call_and_subscribe(aclient: AsyncOpencode, mock_api: respx.MockRouter) -> None: + mock_api.post("/api/rpc/counter/add").respond(json={"output": {"total": 7}}) + counter = aclient.rpc(Counter) + assert (await counter.call(Counter.add, AddInput(amount=7))).total == 7 + mock_api.post("/api/rpc/counter/reset-all").respond( + 400, json={"_tag": "RpcError", "type": "locked", "message": "busy"} + ) + with pytest.raises(RpcCallError, match="locked"): + await counter.call(Counter.reset) + mock_api.get("/api/event").respond(headers=SSE, content=rpc_frame("rpc.counter.changed", {"total": 9}, 1)) + async with counter.subscribe(Counter.changed) as changes: + assert [change.data.total async for change in changes] == [9] + async with counter.events() as raw: + assert [event.id async for event in raw] == ["evt_1"]