feat(client): add event handler implementation for websockets

Author: stainless-app[bot]

src/dedalus_sdk/_event_handler.py

@@ -0,0 +1,72 @@
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Callable
+
+EventHandler = Callable[..., Any]
+
+
+class EventHandlerRegistry:
+    """Thread-safe (optional) registry of event handlers."""
+
+    def __init__(self, *, use_lock: bool = False) -> None:
+        self._handlers: dict[str, list[EventHandler]] = {}
+        self._once_ids: set[int] = set()
+        self._lock: threading.Lock | None = threading.Lock() if use_lock else None
+
+    def _acquire(self) -> None:
+        if self._lock is not None:
+            self._lock.acquire()
+
+    def _release(self) -> None:
+        if self._lock is not None:
+            self._lock.release()
+
+    def add(self, event_type: str, handler: EventHandler, *, once: bool = False) -> None:
+        self._acquire()
+        try:
+            handlers = self._handlers.setdefault(event_type, [])
+            handlers.append(handler)
+            if once:
+                self._once_ids.add(id(handler))
+        finally:
+            self._release()
+
+    def remove(self, event_type: str, handler: EventHandler) -> None:
+        self._acquire()
+        try:
+            handlers = self._handlers.get(event_type)
+            if handlers is not None:
+                try:
+                    handlers.remove(handler)
+                except ValueError:
+                    pass
+                self._once_ids.discard(id(handler))
+        finally:
+            self._release()
+
+    def get_handlers(self, event_type: str) -> list[EventHandler]:
+        """Return a snapshot of handlers for the given event type, removing once-handlers."""
+        self._acquire()
+        try:
+            handlers = self._handlers.get(event_type)
+            if not handlers:
+                return []
+            result = list(handlers)
+            to_remove = [h for h in result if id(h) in self._once_ids]
+            for h in to_remove:
+                handlers.remove(h)
+                self._once_ids.discard(id(h))
+            return result
+        finally:
+            self._release()
+
+    def has_handlers(self, event_type: str) -> bool:
+        self._acquire()
+        try:
+            handlers = self._handlers.get(event_type)
+            return bool(handlers)
+        finally:
+            self._release()

src/dedalus_sdk/resources/machines/terminals.py

@@ -7,7 +7,7 @@
 import random
 import logging
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Dict, Callable, Iterator, Awaitable, cast
+from typing import TYPE_CHECKING, Any, Dict, Union, Callable, Iterator, Awaitable, cast
 from typing_extensions import AsyncIterator
 
 import httpx
@@ -27,10 +27,12 @@
 from ...pagination import SyncCursorPage, AsyncCursorPage
 from ..._exceptions import DedalusError
 from ..._base_client import AsyncPaginator, _merge_mappings, make_request_options
+from ..._event_handler import EventHandlerRegistry
 from ...types.machines import terminal_list_params, terminal_create_params
 from ...types.machines.terminal import Terminal
 from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close
 from ...types.websocket_connection_options import WebSocketConnectionOptions
+from ...types.machines.terminal_error_event import TerminalErrorEvent
 from ...types.machines.terminal_client_event import TerminalClientEvent
 from ...types.machines.terminal_server_event import TerminalServerEvent
 from ...types.machines.terminal_client_event_param import TerminalClientEventParam
@@ -606,6 +608,7 @@ def __init__(
         self._extra_query = extra_query
         self._extra_headers = extra_headers
         self._intentionally_closed = False
+        self._event_handler_registry = EventHandlerRegistry(use_lock=False)
 
     async def __aiter__(self) -> AsyncIterator[TerminalServerEvent]:
         """
@@ -736,6 +739,86 @@ async def _reconnect(self, exc: Exception) -> bool:
 
         return False
 
+    def on(
+        self, event_type: str, handler: Callable[..., Any] | None = None
+    ) -> Union[AsyncTerminalsResourceConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]:
+        """Adds the handler to the end of the handlers list for the given event type.
+
+        No checks are made to see if the handler has already been added. Multiple calls
+        passing the same combination of event type and handler will result in the handler
+        being added, and called, multiple times.
+
+        Can be used as a method (returns ``self`` for chaining)::
+
+            connection.on("output", my_handler)
+
+        Or as a decorator::
+
+            @connection.on("output")
+            async def my_handler(event): ...
+        """
+        if handler is not None:
+            self._event_handler_registry.add(event_type, handler)
+            return self
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            self._event_handler_registry.add(event_type, fn)
+            return fn
+
+        return decorator
+
+    def off(self, event_type: str, handler: Callable[..., Any]) -> AsyncTerminalsResourceConnection:
+        """Remove a previously registered event handler."""
+        self._event_handler_registry.remove(event_type, handler)
+        return self
+
+    def once(
+        self, event_type: str, handler: Callable[..., Any] | None = None
+    ) -> Union[AsyncTerminalsResourceConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]:
+        """Register a one-time event handler.
+
+        Automatically removed after first invocation.
+        """
+        if handler is not None:
+            self._event_handler_registry.add(event_type, handler, once=True)
+            return self
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            self._event_handler_registry.add(event_type, fn, once=True)
+            return fn
+
+        return decorator
+
+    async def dispatch_events(self) -> None:
+        """Run the event loop, dispatching received events to registered handlers.
+
+        Blocks until the connection is closed. This is the push-based
+        alternative to iterating with ``async for event in connection``.
+
+        If an ``"error"`` event arrives and no handler is registered for
+        ``"error"`` or ``"event"``, an ``DedalusError`` is raised.
+        """
+        import asyncio
+
+        async for event in self:
+            event_type = event.type
+            specific = self._event_handler_registry.get_handlers(event_type)
+            generic = self._event_handler_registry.get_handlers("event")
+
+            if event_type == "error" and not specific and not generic:
+                if isinstance(event, TerminalErrorEvent):
+                    raise DedalusError(f"WebSocket error: {event}")
+
+            for handler in specific:
+                result = handler(event)
+                if asyncio.iscoroutine(result):
+                    await result
+
+            for handler in generic:
+                result = handler(event)
+                if asyncio.iscoroutine(result):
+                    await result
+
 
 class AsyncTerminalsResourceConnectionManager:
     """
@@ -785,7 +868,7 @@ def __init__(
 
     async def __aenter__(self) -> AsyncTerminalsResourceConnection:
         """
-        👋 If your application doesn't work well with the context manager approach then you
+        If your application doesn't work well with the context manager approach then you
         can call this method directly to initiate a connection.
 
         **Warning**: You must remember to close the connection with `.close()`.
@@ -893,6 +976,7 @@ def __init__(
         self._extra_query = extra_query
         self._extra_headers = extra_headers
         self._intentionally_closed = False
+        self._event_handler_registry = EventHandlerRegistry(use_lock=True)
 
     def __iter__(self) -> Iterator[TerminalServerEvent]:
         """
@@ -1021,6 +1105,80 @@ def _reconnect(self, exc: Exception) -> bool:
 
         return False
 
+    def on(
+        self, event_type: str, handler: Callable[..., Any] | None = None
+    ) -> Union[TerminalsResourceConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]:
+        """Adds the handler to the end of the handlers list for the given event type.
+
+        No checks are made to see if the handler has already been added. Multiple calls
+        passing the same combination of event type and handler will result in the handler
+        being added, and called, multiple times.
+
+        Can be used as a method (returns ``self`` for chaining)::
+
+            connection.on("output", my_handler)
+
+        Or as a decorator::
+
+            @connection.on("output")
+            def my_handler(event): ...
+        """
+        if handler is not None:
+            self._event_handler_registry.add(event_type, handler)
+            return self
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            self._event_handler_registry.add(event_type, fn)
+            return fn
+
+        return decorator
+
+    def off(self, event_type: str, handler: Callable[..., Any]) -> TerminalsResourceConnection:
+        """Remove a previously registered event handler."""
+        self._event_handler_registry.remove(event_type, handler)
+        return self
+
+    def once(
+        self, event_type: str, handler: Callable[..., Any] | None = None
+    ) -> Union[TerminalsResourceConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]:
+        """Register a one-time event handler.
+
+        Automatically removed after first invocation.
+        """
+        if handler is not None:
+            self._event_handler_registry.add(event_type, handler, once=True)
+            return self
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            self._event_handler_registry.add(event_type, fn, once=True)
+            return fn
+
+        return decorator
+
+    def dispatch_events(self) -> None:
+        """Run the event loop, dispatching received events to registered handlers.
+
+        Blocks the current thread until the connection is closed. This is the push-based
+        alternative to iterating with ``for event in connection``.
+
+        If an ``"error"`` event arrives and no handler is registered for
+        ``"error"`` or ``"event"``, an ``DedalusError`` is raised.
+        """
+        for event in self:
+            event_type = event.type
+            specific = self._event_handler_registry.get_handlers(event_type)
+            generic = self._event_handler_registry.get_handlers("event")
+
+            if event_type == "error" and not specific and not generic:
+                if isinstance(event, TerminalErrorEvent):
+                    raise DedalusError(f"WebSocket error: {event}")
+
+            for handler in specific:
+                handler(event)
+
+            for handler in generic:
+                handler(event)
+
 
 class TerminalsResourceConnectionManager:
     """
@@ -1070,7 +1228,7 @@ def __init__(
 
     def __enter__(self) -> TerminalsResourceConnection:
         """
-        👋 If your application doesn't work well with the context manager approach then you
+        If your application doesn't work well with the context manager approach then you
         can call this method directly to initiate a connection.
 
         **Warning**: You must remember to close the connection with `.close()`.