coco-framework 0.1.0__py3-none-any.whl

coco_graph/stream/__init__.py

@@ -0,0 +1,45 @@
+"""Streaming infrastructure for CocoGraph.
+
+Compile a graph with `transformers=[...]` and call `graph.stream_events(version="v3")` /
+`graph.astream_events(version="v3")` to drive a transformer pipeline that projects the
+graph's raw events into ergonomic per-channel streams.
+"""
+
+from coco_graph.stream._types import ProtocolEvent, StreamTransformer
+from coco_graph.stream.run_stream import (
+    AsyncGraphRunStream,
+    AsyncSubgraphRunStream,
+    GraphRunStream,
+    SubgraphRunStream,
+)
+from coco_graph.stream.stream_channel import StreamChannel
+from coco_graph.stream.transformers import (
+    CheckpointsTransformer,
+    CustomTransformer,
+    DebugTransformer,
+    LifecyclePayload,
+    LifecycleTransformer,
+    SubgraphStatus,
+    SubgraphTransformer,
+    TasksTransformer,
+    UpdatesTransformer,
+)
+
+__all__ = [
+    "AsyncGraphRunStream",
+    "AsyncSubgraphRunStream",
+    "CheckpointsTransformer",
+    "CustomTransformer",
+    "DebugTransformer",
+    "GraphRunStream",
+    "LifecyclePayload",
+    "LifecycleTransformer",
+    "ProtocolEvent",
+    "StreamChannel",
+    "StreamTransformer",
+    "SubgraphRunStream",
+    "SubgraphStatus",
+    "SubgraphTransformer",
+    "TasksTransformer",
+    "UpdatesTransformer",
+]

coco_graph/stream/_convert.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import time
+from typing import Any, cast
+
+from coco_graph.stream._types import ProtocolEvent, _ProtocolEventParams
+from coco_graph.types import StreamPart
+
+
+def convert_to_protocol_event(part: StreamPart) -> ProtocolEvent:
+    """Convert a v2 StreamPart to a ProtocolEvent.
+
+    Args:
+        part: A stream part with keys `type`, `ns`, `data`, and
+            optionally `interrupts` (present on values events).
+
+    Returns:
+        The equivalent ProtocolEvent.
+    """
+    part_dict = cast(dict[str, Any], part)
+    params: _ProtocolEventParams = {
+        "namespace": list(part_dict["ns"]),
+        "timestamp": int(time.time() * 1000),
+        "data": part_dict["data"],
+    }
+    if "interrupts" in part_dict:
+        params["interrupts"] = part_dict["interrupts"]
+    return {
+        "type": "event",
+        "method": part_dict["type"],
+        "params": params,
+    }

coco_graph/stream/_mux.py

@@ -0,0 +1,523 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from coco_graph.stream._types import (
+    ProtocolEvent,
+    StreamTransformer,
+    transformer_requires_async,
+)
+from coco_graph.stream.stream_channel import StreamChannel
+
+TransformerFactory = Callable[["tuple[str, ...]"], StreamTransformer]
+"""Factory that builds a scoped transformer for a mux.
+
+Called once per `StreamMux` with the mux's scope (typically `()` for
+the root). Standard transformer classes accept a single positional
+scope argument, so the class itself is a valid factory. User
+transformers can close over their config:
+`lambda scope: MyTransformer(scope, foo=...)`.
+"""
+
+
+class StreamMux:
+    """Central event dispatcher for the streaming infrastructure.
+
+    Owns the main event log and routes events through a transformer
+    pipeline. StreamChannels with a name discovered in transformer
+    projections are auto-wired so that every `push()` also injects a
+    `ProtocolEvent` into the main log. StreamChannels without a name
+    are local-only.
+
+    Pass `is_async=True` when the mux will be consumed via async
+    iteration (`handler.astream()`). All StreamChannel instances
+    discovered during registration are automatically bound to the
+    matching mode.
+
+    Attributes:
+        extensions: Merged projection dict across all registered
+            transformers. Treat as read-only — mutations won't be
+            reflected back in individual transformers' state.
+        native_keys: Projection keys contributed by transformers with
+            `_native = True`.
+    """
+
+    def __init__(
+        self,
+        transformers: list[StreamTransformer] | None = None,
+        *,
+        is_async: bool = False,
+        factories: list[TransformerFactory] | None = None,
+        scope: tuple[str, ...] = (),
+        _assign_seq: bool = True,
+    ) -> None:
+        """Initialize the mux and register transformers in order.
+
+        Callers pass either `transformers` (pre-built instances) or
+        `factories` (callables producing fresh instances per mux). Each
+        transformer's `init()` is called, projections are merged into
+        `extensions`, `_native` keys are recorded in `native_keys`, and
+        any StreamChannel instances are bound and (if named) wired.
+
+        Transformers with `StreamTransformer.before_builtins = True` are
+        registered ahead of the rest, preserving relative order within
+        each lane. This lets content-mutating transformers (PII
+        redaction, content filters, etc.) run before built-ins like
+        `MessagesTransformer` that eagerly snapshot text fields into
+        their projections. See `StreamTransformer.before_builtins` for
+        the contract and foot-guns.
+
+        Args:
+            transformers: Already-built transformer instances. Registered
+                only on this mux — they are NOT cloned into child
+                mini-muxes built by `_make_child`. Use `factories` for
+                transformers that should propagate to nested scopes.
+            is_async: True for async dispatch (`apush` / `aclose` /
+                `afail`), False for the sync path.
+            factories: One-argument callables `(scope) -> StreamTransformer`.
+                Called once with this mux's `scope` here, and cloned
+                again per child scope by `_make_child` so each
+                sub-mux gets fresh instances.
+            scope: The namespace the mux operates within. The root mux
+                is `()`.
+            _assign_seq: Internal flag for child muxes. Root muxes assign
+                monotonic `seq` numbers when appending to their main event
+                log; child muxes share forwarded event objects and must not
+                mutate their envelopes.
+
+        Raises:
+            RuntimeError: If any transformer requires an async run but
+                the mux is in sync mode.
+            TypeError: If a transformer's `init()` doesn't return a dict.
+            ValueError: If transformers' projection keys collide.
+        """
+        self.is_async = is_async
+        self.scope: tuple[str, ...] = scope
+        self._assign_seq = _assign_seq
+        self._events: StreamChannel[ProtocolEvent] = StreamChannel()
+        self._events._bind(is_async=is_async)
+        self._events._bind_mux(self)
+        self._transformers: list[StreamTransformer] = []
+        self._channels: list[StreamChannel[Any]] = []
+        self._seq = 0
+        self._push_seq = 0
+
+        self.extensions: dict[str, Any] = {}
+        self.native_keys: set[str] = set()
+        self._projection_owners: dict[str, str] = {}
+        self._transformer_by_key: dict[str, StreamTransformer] = {}
+
+        # Stored only when constructed from factories — used by
+        # `_make_child` to clone the transformer pipeline at a deeper
+        # scope. Pre-built transformers can't be cloned, so a mux
+        # built with `transformers=` rejects child construction.
+        self._factories: list[TransformerFactory] | None = (
+            list(factories) if factories is not None else None
+        )
+        self._pump_fn: Callable[[], bool] | None = None
+        self._apump_fn: Callable[[], Awaitable[bool]] | None = None
+
+        # Factories run first (they propagate to child mini-muxes via
+        # `_make_child`), then any pre-built `transformers=` instances
+        # are registered as root-only — they aren't cloned for child
+        # scopes. Within each group, transformers with
+        # `before_builtins = True` are registered ahead of the rest so
+        # they observe (and may mutate) events before built-ins like
+        # `MessagesTransformer`. The order *within* each lane matches
+        # the supplied sequence.
+        pre: list[StreamTransformer] = []
+        rest: list[StreamTransformer] = []
+        if factories is not None:
+            for factory in factories:
+                transformer = factory(scope)
+                (
+                    pre if getattr(transformer, "before_builtins", False) else rest
+                ).append(transformer)
+            for transformer in (*pre, *rest):
+                self._register(transformer)
+            pre.clear()
+            rest.clear()
+        for transformer in transformers or ():
+            (pre if getattr(transformer, "before_builtins", False) else rest).append(
+                transformer
+            )
+        for transformer in (*pre, *rest):
+            self._register(transformer)
+
+    def transformer_by_key(self, key: str) -> StreamTransformer | None:
+        """Return the transformer that contributed `key` to the projection."""
+        return self._transformer_by_key.get(key)
+
+    def _next_push_seq(self) -> int:
+        self._push_seq += 1
+        return self._push_seq
+
+    # ------------------------------------------------------------------
+    # Pump wiring + mini-mux nesting
+    # ------------------------------------------------------------------
+
+    def bind_pump(self, fn: Callable[[], bool]) -> None:
+        """Wire the sync pull callback onto every projection in this mux.
+
+        Records the pump on the mux so child mini-muxes built by
+        `_make_child` can inherit it. Propagates to:
+        - the main event log (`self._events`)
+        - every projection StreamChannel in `extensions`
+        - any registered transformer that exposes `_bind_pump` (e.g.
+          `MessagesTransformer` so `ChatModelStream` instances drive the
+          shared pump from their cursors)
+        """
+        self._pump_fn = fn
+        self._events._request_more = fn
+        for ch in self._channels:
+            ch._request_more = fn
+        for transformer in self._transformers:
+            bind = getattr(transformer, "_bind_pump", None)
+            if bind is not None:
+                bind(fn)
+
+    def bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
+        """Async counterpart to `bind_pump`."""
+        self._apump_fn = fn
+        self._events._arequest_more = fn
+        for ch in self._channels:
+            ch._arequest_more = fn
+        for transformer in self._transformers:
+            abind = getattr(transformer, "_bind_apump", None)
+            if abind is not None:
+                abind(fn)
+
+    def _make_child(self, scope: tuple[str, ...]) -> StreamMux:
+        """Build a mini-mux with the same factories scoped to `scope`.
+
+        Used by `SubgraphTransformer` to attach a fresh transformer
+        pipeline to each discovered subgraph handle. The child mux
+        inherits the current pump bindings (so cursors on its
+        projection logs drive the root pump), carries the same factory
+        list forward to any grandchild subgraphs, and does not assign
+        `seq` numbers so forwarded events can be shared without
+        mutating their envelope.
+
+        Raises:
+            RuntimeError: If the mux was not constructed with
+                `factories=`. Mini-muxes require factories so each scope
+                gets its own fresh transformer instances.
+        """
+        if self._factories is None:
+            raise RuntimeError(
+                "StreamMux._make_child requires the mux to be constructed "
+                "with `factories=`; pre-built transformers can't be "
+                "cloned to a new scope."
+            )
+        child = StreamMux(
+            factories=self._factories,
+            is_async=self.is_async,
+            scope=scope,
+            _assign_seq=False,
+        )
+        if self._pump_fn is not None:
+            child.bind_pump(self._pump_fn)
+        if self._apump_fn is not None:
+            child.bind_apump(self._apump_fn)
+        return child
+
+    def _register(self, transformer: StreamTransformer) -> None:
+        """Register a single transformer.
+
+        Calls `transformer.init()`, stores the transformer for event
+        processing, binds any StreamChannel instances in the projection,
+        and merges the projection into `extensions`.
+        """
+        if transformer_requires_async(transformer) and not self.is_async:
+            raise RuntimeError(
+                f"{type(transformer).__name__} requires an async run — "
+                "it overrides aprocess/afinalize/afail or sets "
+                "requires_async=True. Use astream(), not stream()."
+            )
+        projection = transformer.init()
+        if not isinstance(projection, dict):
+            raise TypeError(
+                f"StreamTransformer.init() must return a dict, "
+                f"got {type(projection).__name__}"
+            )
+        conflicts = set(projection) & set(self.extensions)
+        if conflicts:
+            attributions = ", ".join(
+                f"{key!r} (owned by {self._projection_owners[key]})"
+                for key in sorted(conflicts)
+            )
+            raise ValueError(
+                f"Transformer {type(transformer).__name__} returned "
+                f"projection keys that conflict with already-registered "
+                f"keys: {attributions}"
+            )
+        is_native = bool(getattr(transformer, "_native", False))
+        self._transformers.append(transformer)
+        self._bind_and_wire(projection, native=is_native)
+        self.extensions.update(projection)
+        owner_name = type(transformer).__name__
+        for key in projection:
+            self._projection_owners[key] = owner_name
+            self._transformer_by_key[key] = transformer
+        if is_native:
+            self.native_keys.update(projection.keys())
+        transformer._on_register(self)
+
+    def push(self, event: ProtocolEvent) -> None:
+        """Route an event through all transformers, then append to the main log.
+
+        Each transformer's `process()` is called in registration order.
+        If any transformer returns False, the event is suppressed from
+        the main log, but transformers that already saw it keep their
+        side effects.
+
+        On the root mux, `seq` is assigned right before an event enters
+        the main log, not before the transformer pipeline runs. This
+        ensures that events auto-forwarded from StreamChannels during
+        `process()` get earlier seq numbers than the original event,
+        preserving monotonic ordering in the root log. Child muxes do
+        not assign `seq`, so subgraph forwarding can share event objects
+        without mutating their envelopes.
+
+        Args:
+            event: The protocol event to dispatch.
+        """
+        keep = True
+        for transformer in self._transformers:
+            if not transformer.process(event):
+                keep = False
+        if keep:
+            if self._assign_seq:
+                self._seq += 1
+                event["seq"] = self._seq
+            self._events.push(event)
+
+    def close(self) -> None:
+        """Finalize all transformers, close all projections and the main log.
+
+        StreamChannels discovered in transformer projections are
+        auto-closed after `finalize()` runs — transformers don't need
+        to close them manually. If any transformer's `finalize()` raises,
+        the remaining transformers, projections, and the main log are
+        still closed; the first error is re-raised after cleanup
+        completes.
+
+        Raises:
+            BaseException: The first error raised by a transformer's
+                `finalize()`, re-raised after cleanup finishes.
+        """
+        first_error: BaseException | None = None
+        for transformer in self._transformers:
+            try:
+                transformer.finalize()
+            except BaseException as e:
+                if first_error is None:
+                    first_error = e
+        for ch in self._channels:
+            if not ch._closed:
+                ch.close()
+        self._events.close()
+        if first_error is not None:
+            raise first_error
+
+    def fail(self, err: BaseException) -> None:
+        """Fail all transformers, projections, and the main log.
+
+        StreamChannels discovered in transformer projections are
+        auto-failed — transformers don't need to fail them manually.
+        If any transformer's `fail()` raises, the remaining
+        transformers, projections, and the main log are still failed.
+
+        Args:
+            err: The exception that ended the run.
+        """
+        for transformer in self._transformers:
+            try:
+                transformer.fail(err)
+            except BaseException:
+                pass
+        for ch in self._channels:
+            if not ch._closed:
+                ch.fail(err)
+        self._events.fail(err)
+
+    # ------------------------------------------------------------------
+    # Async dispatch
+    # ------------------------------------------------------------------
+
+    async def apush(self, event: ProtocolEvent) -> None:
+        """Dispatch an event on the async lane.
+
+        Awaits each transformer's `aprocess` in registration order
+        before appending to the main log. A slow `aprocess` serializes
+        the pipeline by design — that's the guarantee that lets a later
+        transformer (or a synchronous consumer) see the result of the
+        async work. For decoupled work, use `schedule()` from inside
+        `process` / `aprocess` instead.
+
+        The main log append is a non-blocking `push` — matching v1's
+        `put_nowait` shape. The root mux assigns `seq`; child muxes do
+        not, so forwarded subgraph events can be shared without copying.
+        Memory is bounded by caller pace via the caller-driven pump; see
+        `StreamChannel` for the full tradeoff story.
+
+        Args:
+            event: The protocol event to dispatch.
+        """
+        keep = True
+        for transformer in self._transformers:
+            if not await transformer.aprocess(event):
+                keep = False
+        if keep:
+            if self._assign_seq:
+                self._seq += 1
+                event["seq"] = self._seq
+            self._events.push(event)
+
+    async def aclose(self) -> None:
+        """Finalize on the async lane.
+
+        Awaits every task started via `StreamTransformer.schedule()`
+        across all transformers, then calls `afinalize()` on each,
+        then auto-closes channels and the main event log.
+
+        If any scheduled task raised under `on_error="raise"`, or any
+        transformer's `afinalize` raises, the exception propagates.
+        The caller (the pump) handles it by routing into `afail`.
+
+        Raises:
+            BaseException: The first scheduled-task or `afinalize`
+                error, re-raised after cleanup.
+        """
+        pending = self._collect_scheduled_tasks()
+        if pending:
+            results = await asyncio.gather(*pending, return_exceptions=True)
+            first_err = next(
+                (
+                    r
+                    for r in results
+                    if isinstance(r, BaseException)
+                    and not isinstance(r, asyncio.CancelledError)
+                ),
+                None,
+            )
+            if first_err is not None:
+                raise first_err
+
+        first_error: BaseException | None = None
+        for transformer in self._transformers:
+            try:
+                await transformer.afinalize()
+            except BaseException as e:
+                if first_error is None:
+                    first_error = e
+        for ch in self._channels:
+            if not ch._closed:
+                ch.close()
+        self._events.close()
+        if first_error is not None:
+            raise first_error
+
+    async def afail(self, err: BaseException) -> None:
+        """Fail on the async lane.
+
+        Cancels every scheduled task across all transformers, awaits
+        them to completion, then runs each transformer's `afail` hook
+        and auto-fails channels and the main event log.
+
+        Args:
+            err: The exception that ended the run.
+        """
+        pending = self._collect_scheduled_tasks()
+        for task in pending:
+            task.cancel()
+        if pending:
+            await asyncio.gather(*pending, return_exceptions=True)
+
+        for transformer in self._transformers:
+            try:
+                await transformer.afail(err)
+            except BaseException:
+                pass
+        for ch in self._channels:
+            if not ch._closed:
+                ch.fail(err)
+        if not self._events._closed:
+            self._events.fail(err)
+
+    def _collect_scheduled_tasks(self) -> list[asyncio.Task[Any]]:
+        """Return a snapshot of in-flight tasks scheduled via transformers."""
+        return [
+            task
+            for transformer in self._transformers
+            for task in getattr(transformer, "_stream_scheduled_tasks", ())
+            if not task.done()
+        ]
+
+    # ------------------------------------------------------------------
+    # Binding and StreamChannel auto-wiring
+    # ------------------------------------------------------------------
+
+    def _bind_and_wire(
+        self, projection: dict[str, Any], *, native: bool = False
+    ) -> None:
+        """Bind and optionally wire StreamChannel instances in a projection.
+
+        All StreamChannels are bound and tracked. Channels with a name
+        are additionally wired for protocol auto-forwarding.
+
+        Args:
+            projection: The projection dict returned by a transformer's
+                `init()`.
+            native: True when the owning transformer is `_native`.
+                Named channels owned by a native transformer use the
+                channel name directly as the protocol method;
+                user-defined channels are prefixed with `custom:`.
+        """
+        for value in projection.values():
+            if isinstance(value, StreamChannel):
+                value._bind(is_async=self.is_async)
+                value._bind_mux(self)
+                self._channels.append(value)
+                if value.name is not None:
+                    method = value.name if native else f"custom:{value.name}"
+
+                    def _make_forward(method_name: str) -> Callable[[Any], None]:
+                        def _forward(item: Any) -> None:
+                            self._forward(method_name, item)
+
+                        return _forward
+
+                    value._wire(_make_forward(method))
+
+    def _forward(self, method: str, item: Any) -> None:
+        """Inject a ProtocolEvent for a StreamChannel push.
+
+        Forwarded events bypass the transformer pipeline to avoid
+        infinite recursion (a transformer that pushes to a channel
+        during `process()` would re-trigger itself). These events are
+        visible in this mux's main event log but are not passed through
+        transformers' `process()` methods. Only the root mux assigns
+        `seq` to forwarded channel events.
+
+        Args:
+            method: The full protocol method (already with or without
+                the `custom:` prefix; resolved by `_bind_and_wire`).
+            item: The payload pushed onto the channel.
+        """
+        event: ProtocolEvent = {
+            "type": "event",
+            "method": method,
+            "params": {
+                "namespace": [],
+                "timestamp": int(time.time() * 1000),
+                "data": item,
+            },
+        }
+        if self._assign_seq:
+            self._seq += 1
+            event["seq"] = self._seq
+        self._events.push(event)