codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from typing import TYPE_CHECKING, Any, Literal, cast5 6from langchain_core.language_models._compat_bridge import message_to_events7from langchain_core.language_models.chat_model_stream import (8    AsyncChatModelStream,9    ChatModelStream,10)11from langchain_core.messages import AIMessageChunk, BaseMessage12from langchain_protocol.protocol import MessagesData13from typing_extensions import NotRequired, TypedDict14 15from langgraph.errors import GraphDrained, GraphInterrupt16from langgraph.stream._types import ProtocolEvent, StreamTransformer17from langgraph.stream.run_stream import AsyncSubgraphRunStream, SubgraphRunStream18from langgraph.stream.stream_channel import StreamChannel19 20if TYPE_CHECKING:21    from collections.abc import Awaitable, Callable22 23    from langgraph.stream._mux import StreamMux24 25_logger = logging.getLogger(__name__)26 27 28class ValuesTransformer(StreamTransformer):29    """Capture values events as a drainable stream of state snapshots.30 31    Provides the `run.values` projection. `run.output`,32    `run.interrupted` and `run.interrupts` are tracked directly33    by the run stream and do not depend on this transformer.34 35    Native transformer — projection keys are exposed as direct36    attributes on the run stream (e.g. `run.values`).37 38    Only values events at the run's own level are captured; snapshots39    from deeper subgraphs are left in the main event log but excluded40    from the projection. "Own level" is defined by `scope`, which41    `stream_events(version="v3")` / `astream_events(version="v3")` populate from the caller's42    checkpoint namespace so that a nested `stream_events(version="v3")` call still43    sees its own root snapshots.44    """45 46    _native = True47    required_stream_modes = ("values",)48 49    def __init__(self, scope: tuple[str, ...] = ()) -> None:50        super().__init__(scope)51        self._log: StreamChannel[dict[str, Any]] = StreamChannel()52        self._latest: dict[str, Any] | None = None53        self._interrupted = False54        self._interrupts: list[Any] = []55        # Cached as a list once for cheap equality with the protocol56        # event's `namespace` field, which is `list[str]`.57        self._scope_list: list[str] = list(scope)58 59    def init(self) -> dict[str, Any]:60        return {"values": self._log}61 62    @property63    def error(self) -> BaseException | None:64        """The error that ended the run, or `None` if it succeeded.65 66        Set by the mux when it auto-fails the projection log.67        """68        return self._log._error69 70    def process(self, event: ProtocolEvent) -> bool:71        if event["method"] != "values":72            return True73        params = event["params"]74        if params["namespace"] != self._scope_list:75            return True76        self._latest = params["data"]77        interrupts = params.get("interrupts", ())78        if interrupts:79            self._interrupted = True80            self._interrupts.extend(interrupts)81        self._log.push(params["data"])82        return True83 84 85class CustomTransformer(StreamTransformer):86    """Capture custom events as a drainable stream of arbitrary payloads.87 88    Nodes emit custom data via `get_stream_writer()`. This transformer89    surfaces those events on `run.custom` as a `StreamChannel[Any]`,90    preserving payloads in arrival order.91 92    Only events at the run's own scope are captured; custom data from93    deeper subgraphs is available on the respective subgraph handle's94    `.custom` projection.95 96    Native transformer — `run.custom` is a direct attribute.97    """98 99    _native = True100    required_stream_modes = ("custom",)101 102    def __init__(self, scope: tuple[str, ...] = ()) -> None:103        super().__init__(scope)104        self._log: StreamChannel[Any] = StreamChannel()105        self._scope_list: list[str] = list(scope)106 107    def init(self) -> dict[str, Any]:108        return {"custom": self._log}109 110    def process(self, event: ProtocolEvent) -> bool:111        if event["method"] != "custom":112            return True113        params = event["params"]114        if params["namespace"] != self._scope_list:115            return True116        self._log.push(params["data"])117        return True118 119 120class UpdatesTransformer(StreamTransformer):121    """Capture updates events as a drainable stream of node outputs.122 123    Surfaces `stream_mode="updates"` data on `run.updates` as a124    `StreamChannel[dict[str, Any]]`. Each item is a dict mapping a node125    (or task) name to the update it returned after a step.126 127    Only events at the run's own scope are captured; updates from deeper128    subgraphs are available on the respective subgraph handle's129    `.updates` projection.130 131    Native transformer — `run.updates` is a direct attribute.132    """133 134    _native = True135    required_stream_modes = ("updates",)136 137    def __init__(self, scope: tuple[str, ...] = ()) -> None:138        super().__init__(scope)139        self._log: StreamChannel[dict[str, Any]] = StreamChannel()140        self._scope_list: list[str] = list(scope)141 142    def init(self) -> dict[str, Any]:143        return {"updates": self._log}144 145    def process(self, event: ProtocolEvent) -> bool:146        if event["method"] != "updates":147            return True148        params = event["params"]149        if params["namespace"] != self._scope_list:150            return True151        self._log.push(params["data"])152        return True153 154 155class MessagesTransformer(StreamTransformer):156    """Capture messages events as ChatModelStream objects.157 158    The messages projection yields one `ChatModelStream` (or159    `AsyncChatModelStream`) per LLM call. Consumers iterate160    `run.messages` to get stream handles, then use each handle's typed161    projections (`.text`, `.reasoning`, `.tool_calls`, `.usage`,162    `.output`) for per-message content.163 164    Two input shapes are handled (via `params["data"] = (payload,165    metadata)` from `StreamMessagesHandler`):166 167    1. Protocol event (dict with `"event"` key) — emitted by168       `stream_events(version="v3")` / `astream_events(version="v3")` via the `on_stream_event`169       callback. Routed to an existing `ChatModelStream` by170       `metadata["run_id"]`. A `message-start` event creates a new171       stream; `message-finish` closes it.172    2. Whole `AIMessage` — emitted from `on_chain_end` when a node173       returns a finalized message. Replayed as a synthetic protocol174       event lifecycle via `message_to_events`, then the175       already-complete stream is pushed to the log.176 177    V1 `AIMessageChunk` tuples (from `on_llm_new_token`) are not178    streamed into this projection: chat models that want to populate179    `run.messages` with content-block streaming must use180    `stream_events(version="v3")` / `astream_events(version="v3")`. Models called via the legacy181    `stream()` method still surface their final `AIMessage` via182    `on_chain_end` when a node returns it as state.183 184    Only events at the run's own level are projected; tokens from185    deeper subgraphs are left in the main event log but excluded from186    `.messages`. "Own level" is defined by `scope`, which187    `stream_events(version="v3")` / `astream_events(version="v3")` populate from the caller's checkpoint188    namespace so that a `stream_events(version="v3")` call inside a node still sees its189    own root chat model streams on `.messages`. Consumers that need190    subgraph tokens should iterate the raw event stream or register a191    custom transformer.192 193    Native transformer — the `messages` projection is exposed as a194    direct attribute on the run stream.195    """196 197    _native = True198    required_stream_modes = ("messages",)199 200    def __init__(self, scope: tuple[str, ...] = ()) -> None:201        super().__init__(scope)202        self._log: StreamChannel[ChatModelStream] = StreamChannel()203        # Correlate protocol events back to a ChatModelStream by run_id204        # (attached to the event's metadata by StreamMessagesHandler).205        self._by_run: dict[str, ChatModelStream] = {}206        self._pump_fn: Callable[[], bool] | None = None207        self._apump_fn: Callable[[], Awaitable[bool]] | None = None208        # Cached as a list once for cheap equality with the protocol209        # event's `namespace` field, which is `list[str]`.210        self._scope_list: list[str] = list(scope)211 212    def init(self) -> dict[str, Any]:213        return {"messages": self._log}214 215    def _bind_pump(self, fn: Callable[[], bool]) -> None:216        """Wire the sync pull callback. Called by GraphRunStream._wire_request_more."""217        self._pump_fn = fn218 219    def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:220        """Wire the async pull callback.221 222        Called by `AsyncGraphRunStream._wire_arequest_more` so each223        `AsyncChatModelStream` this transformer creates can drive the224        shared graph pump from its projection cursors.225        """226        self._apump_fn = fn227 228    def _make_stream(229        self,230        *,231        namespace: list[str],232        node: str | None,233        message_id: str | None,234    ) -> ChatModelStream:235        """Create a ChatModelStream (sync) or AsyncChatModelStream (async).236 237        Wires whichever pump is bound. Prefers the async pump so nested238        iteration under `AsyncGraphRunStream` drives the graph forward239        without a background task. The unwired fallback (no pump bound)240        is used by unit tests that dispatch events manually.241        """242        if self._apump_fn is not None:243            astream = AsyncChatModelStream(244                namespace=namespace,245                node=node,246                message_id=message_id,247            )248            astream.set_arequest_more(self._apump_fn)249            return astream250        if self._pump_fn is not None:251            stream: ChatModelStream = ChatModelStream(252                namespace=namespace,253                node=node,254                message_id=message_id,255            )256            stream.set_request_more(self._pump_fn)257            return stream258        return AsyncChatModelStream(259            namespace=namespace,260            node=node,261            message_id=message_id,262        )263 264    def process(self, event: ProtocolEvent) -> bool:265        if event["method"] != "messages":266            return True267        params = event["params"]268        if params["namespace"] != self._scope_list:269            return True270 271        payload, metadata = params["data"]272        node: str | None = metadata.get("langgraph_node")273        run_id = str(metadata.get("run_id", "")) if metadata else ""274 275        if isinstance(payload, dict) and "event" in payload:276            self._route_protocol_event(277                cast("MessagesData", payload), run_id=run_id, node=node278            )279        elif isinstance(payload, BaseMessage) and not isinstance(280            payload, AIMessageChunk281        ):282            self._route_whole_message(payload, node=node)283        # Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored;284        # v1 streaming callers must switch to stream_events(version="v3") to populate this285        # projection.286 287        return True288 289    def _route_protocol_event(290        self,291        event: MessagesData,292        *,293        run_id: str,294        node: str | None,295    ) -> None:296        event_type = event.get("event")297        if event_type == "message-start":298            message_id = event.get("message_id")299            stream = self._make_stream(300                namespace=[],301                node=node,302                message_id=str(message_id) if message_id is not None else None,303            )304            self._by_run[run_id] = stream305            self._log.push(stream)306            stream.dispatch(event)307        elif run_id in self._by_run:308            stream = self._by_run[run_id]309            stream.dispatch(event)310            if event_type == "message-finish":311                del self._by_run[run_id]312 313    def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None:314        stream = self._make_stream(namespace=[], node=node, message_id=message.id)315        for evt in message_to_events(message, message_id=message.id):316            stream.dispatch(evt)317        self._log.push(stream)318 319    def finalize(self) -> None:320        """Clear any routing state — streams close themselves via `message-finish`."""321        self._by_run.clear()322 323    def fail(self, err: BaseException) -> None:324        """Propagate run error to any streams still open when the graph fails."""325        for stream in list(self._by_run.values()):326            stream.fail(err)327        self._by_run.clear()328 329 330SubgraphStatus = Literal["started", "completed", "failed", "interrupted", "drained"]331 332 333def _parse_ns_segment(segment: str) -> tuple[str, str | None]:334    """Split a namespace segment into `(graph_name, trigger_call_id)`.335 336    Segments are formatted `node_name:task_id` by `prepare_next_tasks`.337    Returns `(segment, None)` if no `:` is present.338    """339    name, sep, task_id = segment.partition(":")340    return name, task_id if sep else None341 342 343class LifecyclePayload(TypedDict, total=False):344    """Payload of a lifecycle event surfaced on the `lifecycle` channel.345 346    Auto-forwarded as `lifecycle` protocol events (no `custom:` prefix347    because `LifecycleTransformer` is a native transformer) so remote348    SDK clients receive the same data in-process consumers see via349    `run.lifecycle`.350    """351 352    event: SubgraphStatus353    namespace: list[str]354    graph_name: NotRequired[str]355    trigger_call_id: NotRequired[str]356    error: NotRequired[str]357 358 359class _TasksLifecycleBase(StreamTransformer):360    """Shared bookkeeping for `tasks`-event-driven lifecycle inference.361 362    Both `LifecycleTransformer` (wire-serializable channel) and363    `SubgraphTransformer` (in-process navigation handles) discover364    subgraphs by watching the same `tasks` stream — `started` on the365    first event at a tracked namespace, terminal status when the366    parent's `TaskResultPayload` arrives. Centralizing the dispatch367    + open-set bookkeeping here keeps the inference rules from368    drifting between the two surfaces.369 370    Subclasses provide three template-method hooks:371 372    - `_should_track(ns)` — scope filter (e.g. multi-depth vs373      direct-children-only).374    - `_on_started(ns, graph_name, trigger_call_id)` — first sighting375      action (push payload / build handle / etc.). Called once per376      discovered namespace.377    - `_on_terminal(ns, status, error)` — terminal action (push378      terminal payload / mark handle status). Called once per379      tracked namespace at result time, or via `finalize` / `fail`380      sweeps if no parent result arrived.381 382    Tasks events are suppressed from the main event log (`process`383    returns False) — they're folded into whichever projection the384    subclass populates; consumers iterating the raw protocol stream385    see the higher-level view.386    """387 388    required_stream_modes = ("tasks",)389 390    def __init__(self, scope: tuple[str, ...] = ()) -> None:391        super().__init__(scope)392        self._seen: set[tuple[str, ...]] = set()393        # Maps tracked namespace -> task_id of the parent task whose394        # `TaskResultPayload` will close it.395        self._open: dict[tuple[str, ...], str] = {}396 397    # --- Template-method hooks (subclass overrides) ---398 399    def _should_track(self, ns: tuple[str, ...]) -> bool:400        """Scope filter — return True iff `ns` is in this transformer's region."""401        raise NotImplementedError402 403    def _on_started(404        self,405        ns: tuple[str, ...],406        graph_name: str | None,407        trigger_call_id: str | None,408    ) -> None:409        """Fired once per discovered namespace (first observed task event)."""410        raise NotImplementedError411 412    def _on_terminal(413        self,414        ns: tuple[str, ...],415        status: SubgraphStatus,416        error: str | None,417    ) -> None:418        """Fired once per tracked namespace when its parent's result arrives,419        or via finalize/fail safety-net sweeps.420        """421        raise NotImplementedError422 423    # --- Dispatch + bookkeeping (shared) ---424 425    def process(self, event: ProtocolEvent) -> bool:426        if event["method"] != "tasks":427            return True428        ns = tuple(event["params"]["namespace"])429        data = event["params"]["data"]430        if "result" in data:431            self._handle_task_result(ns, data)432        else:433            self._handle_task_start(ns)434        # Tasks events are folded into the synthesized projections;435        # suppress from the main event log so iterators don't double-see436        # the same information in two shapes.437        return False438 439    def _handle_task_start(self, ns: tuple[str, ...]) -> None:440        if not self._should_track(ns) or ns in self._seen:441            return442        self._seen.add(ns)443        graph_name, trigger_call_id = _parse_ns_segment(ns[-1])444        self._on_started(ns, graph_name or None, trigger_call_id)445        if trigger_call_id is not None:446            self._open[ns] = trigger_call_id447 448    def _pop_terminal_transitions(449        self, ns: tuple[str, ...], data: dict[str, Any]450    ) -> list[tuple[tuple[str, ...], SubgraphStatus, str | None]]:451        """Return and remove tracked children closed by this task result."""452        result_id = data.get("id")453        if not result_id:454            return []455        transitions: list[tuple[tuple[str, ...], SubgraphStatus, str | None]] = []456        for child_ns, parent_task_id in list(self._open.items()):457            if child_ns[:-1] != ns or parent_task_id != result_id:458                continue459            status, error = _terminal_from_result(data)460            transitions.append((child_ns, status, error))461            del self._open[child_ns]462        return transitions463 464    def _handle_task_result(self, ns: tuple[str, ...], data: dict[str, Any]) -> None:465        for child_ns, status, error in self._pop_terminal_transitions(ns, data):466            self._on_terminal(child_ns, status, error)467 468    def finalize(self) -> None:469        """Emit `completed` for any tracked namespace still open at run end."""470        for ns in list(self._open):471            self._on_terminal(ns, "completed", None)472        self._open.clear()473 474    def fail(self, err: BaseException) -> None:475        """Emit terminal status for any tracked namespace still open."""476        status, error_str = _status_from_exception(err)477        for ns in list(self._open):478            self._on_terminal(ns, status, error_str)479        self._open.clear()480 481 482def _status_from_exception(err: BaseException) -> tuple[SubgraphStatus, str | None]:483    """Map a run exception to a subgraph terminal status and error string."""484    if isinstance(err, GraphDrained):485        return "drained", None486    if isinstance(err, GraphInterrupt):487        return "interrupted", None488    return "failed", str(err)489 490 491def _terminal_from_result(492    payload: dict[str, Any],493) -> tuple[SubgraphStatus, str | None]:494    """Map a `TaskResultPayload` to a `(status, error)` pair.495 496    Order matters: a result with both `error` and `interrupts` prefers497    the interrupt classification, since `GraphInterrupt` manifests as498    a populated `interrupts` list, not as `error`.499    """500    if payload.get("interrupts"):501        return "interrupted", None502    error = payload.get("error")503    if error:504        return "failed", str(error)505    return "completed", None506 507 508class LifecycleTransformer(_TasksLifecycleBase):509    """Surface subgraph lifecycle as `lifecycle` protocol events.510 511    Pushes `LifecyclePayload` to a `StreamChannel` named `lifecycle`.512    The channel is auto-forwarded by the mux so payloads land in the513    main event log under `method = "lifecycle"` (native transformer —514    no `custom:` prefix) — visible to remote SDK clients over the515    wire and to in-process consumers via `run.lifecycle`.516 517    Tracks subgraphs at every depth strictly below the transformer's518    scope, so a graph → subgraph → subgraph chain produces lifecycle519    events for both nested levels in a flat stream.520 521    Native transformer — projection key `lifecycle` is exposed as522    `run.lifecycle`.523    """524 525    _native = True526 527    def __init__(self, scope: tuple[str, ...] = ()) -> None:528        super().__init__(scope)529        self._channel: StreamChannel[LifecyclePayload] = StreamChannel("lifecycle")530 531    def init(self) -> dict[str, Any]:532        return {"lifecycle": self._channel}533 534    def _should_track(self, ns: tuple[str, ...]) -> bool:535        depth = len(self.scope)536        return len(ns) > depth and ns[:depth] == self.scope537 538    def _on_started(539        self,540        ns: tuple[str, ...],541        graph_name: str | None,542        trigger_call_id: str | None,543    ) -> None:544        if trigger_call_id is None:545            # Without a task id we can't correlate a parent-result546            # event back to this namespace — skip the started payload547            # and rely on finalize/fail to close.548            return549        payload: LifecyclePayload = {"event": "started", "namespace": list(ns)}550        if graph_name:551            payload["graph_name"] = graph_name552        payload["trigger_call_id"] = trigger_call_id553        self._channel.push(payload)554 555    def _on_terminal(556        self,557        ns: tuple[str, ...],558        status: SubgraphStatus,559        error: str | None,560    ) -> None:561        payload: LifecyclePayload = {"event": status, "namespace": list(ns)}562        if error is not None:563            payload["error"] = error564        self._channel.push(payload)565 566 567class SubgraphTransformer(_TasksLifecycleBase):568    """Discover subgraph invocations as in-process navigation handles.569 570    Per discovered direct-child subgraph, builds a `SubgraphRunStream`571    (or `AsyncSubgraphRunStream`) wrapping a child mini-mux scoped to572    the subgraph's namespace. Consumers iterate `run.subgraphs` to573    receive handles, then drill into `handle.values` / `handle.messages`574    / `handle.subgraphs` (recursive grandchildren) / `handle.lifecycle`.575 576    Each mini-mux owns its own scope and uses its own577    `SubgraphTransformer` to discover its direct children, so578    grandchildren live on the child handle — never on the root's579    `subgraphs` log. Forwarding events into the matching child mini-mux580    is what keeps the child's projections populated.581 582    Native transformer — `subgraphs` is exposed as `run.subgraphs`.583    """584 585    _native = True586    supports_sync = True587 588    def __init__(self, scope: tuple[str, ...] = ()) -> None:589        super().__init__(scope)590        self._log: StreamChannel[SubgraphRunStream | AsyncSubgraphRunStream] = (591            StreamChannel()592        )593        self._handles: dict[594            tuple[str, ...], SubgraphRunStream | AsyncSubgraphRunStream595        ] = {}596        self._mux: StreamMux | None = None597 598    def init(self) -> dict[str, Any]:599        return {"subgraphs": self._log}600 601    def _on_register(self, mux: Any) -> None:602        self._mux = mux603 604    def _should_track(self, ns: tuple[str, ...]) -> bool:605        # Direct children only — grandchildren are picked up by the606        # child mini-mux's own SubgraphTransformer.607        depth = len(self.scope)608        return len(ns) == depth + 1 and ns[:depth] == self.scope609 610    def _on_started(611        self,612        ns: tuple[str, ...],613        graph_name: str | None,614        trigger_call_id: str | None,615    ) -> None:616        if self._mux is None:617            return618        try:619            child_mux = self._mux._make_child(ns)620        except RuntimeError:621            return622        handle_cls = AsyncSubgraphRunStream if child_mux.is_async else SubgraphRunStream623        handle = handle_cls(624            mux=child_mux,625            path=ns,626            graph_name=graph_name,627            trigger_call_id=trigger_call_id,628        )629        self._handles[ns] = handle630        self._log.push(handle)631 632    def _on_terminal(633        self,634        ns: tuple[str, ...],635        status: SubgraphStatus,636        error: str | None,637    ) -> None:638        handle = self._handles.get(ns)639        if handle is None or not self._mark_terminal(handle, status, error):640            return641        self._close_or_fail_handle(handle, status, error)642 643    async def _aon_terminal(644        self,645        ns: tuple[str, ...],646        status: SubgraphStatus,647        error: str | None,648    ) -> None:649        handle = self._handles.get(ns)650        if handle is None or not self._mark_terminal(handle, status, error):651            return652        await self._aclose_or_fail_handle(handle, status, error)653 654    def _mark_terminal(655        self,656        handle: SubgraphRunStream | AsyncSubgraphRunStream,657        status: SubgraphStatus,658        error: str | None,659    ) -> bool:660        """Mark a handle terminal once. Returns True on first transition."""661        if handle._seen_terminal:662            return False663        handle.status = status664        if error is not None and handle.error is None:665            handle.error = error666        handle._seen_terminal = True667        return True668 669    def _close_or_fail_handle(670        self,671        handle: SubgraphRunStream | AsyncSubgraphRunStream,672        status: SubgraphStatus,673        error: str | None,674    ) -> None:675        if handle._mux is None or handle._mux._events._closed:676            return677        if status == "failed":678            handle._mux.fail(RuntimeError(error or "Subgraph failed"))679        else:680            handle._mux.close()681 682    async def _aclose_or_fail_handle(683        self,684        handle: SubgraphRunStream | AsyncSubgraphRunStream,685        status: SubgraphStatus,686        error: str | None,687    ) -> None:688        if handle._mux is None or handle._mux._events._closed:689            return690        if status == "failed":691            await handle._mux.afail(RuntimeError(error or "Subgraph failed"))692        else:693            await handle._mux.aclose()694 695    def _handle_for_event(696        self, event: ProtocolEvent697    ) -> SubgraphRunStream | AsyncSubgraphRunStream | None:698        ns = tuple(event["params"]["namespace"])699        depth = len(self.scope)700        if len(ns) < depth + 1:701            return None702        handle = self._handles.get(ns[: depth + 1])703        if handle is None or handle._mux is None or handle._mux._events._closed:704            return None705        return handle706 707    def process(self, event: ProtocolEvent) -> bool:708        # Run tasks bookkeeping first so a `started` handle exists709        # by the time we forward the event to the child mini-mux.710        keep = super().process(event)711        handle = self._handle_for_event(event)712        if handle is not None:713            handle._observe_event(event)714            handle._mux.push(event)715        return keep716 717    async def aprocess(self, event: ProtocolEvent) -> bool:718        # Async counterpart: repeats the tasks bookkeeping here so719        # child mini-muxes receive events through their async lane.720        if event["method"] == "tasks":721            ns = tuple(event["params"]["namespace"])722            data = event["params"]["data"]723            if "result" in data:724                for child_ns, status, error in self._pop_terminal_transitions(ns, data):725                    await self._aon_terminal(child_ns, status, error)726            else:727                self._handle_task_start(ns)728            keep = False729        else:730            keep = True731        handle = self._handle_for_event(event)732        if handle is not None:733            handle._observe_event(event)734            await handle._mux.apush(event)735        return keep736 737    def _complete_open_handles(self) -> BaseException | None:738        first_error: BaseException | None = None739        for ns in list(self._open):740            try:741                self._on_terminal(ns, "completed", None)742            except BaseException as e:743                if first_error is None:744                    first_error = e745        self._open.clear()746        for handle in self._handles.values():747            if self._mark_terminal(handle, "completed", None):748                try:749                    self._close_or_fail_handle(handle, "completed", None)750                except BaseException as e:751                    if first_error is None:752                        first_error = e753        return first_error754 755    async def _acomplete_open_handles(self) -> BaseException | None:756        first_error: BaseException | None = None757        for ns in list(self._open):758            try:759                await self._aon_terminal(ns, "completed", None)760            except BaseException as e:761                if first_error is None:762                    first_error = e763        self._open.clear()764        for handle in self._handles.values():765            if self._mark_terminal(handle, "completed", None):766                try:767                    await self._aclose_or_fail_handle(handle, "completed", None)768                except BaseException as e:769                    if first_error is None:770                        first_error = e771        return first_error772 773    def finalize(self) -> None:774        first_error = self._complete_open_handles()775        if first_error is not None:776            raise first_error777 778    async def afinalize(self) -> None:779        first_error = await self._acomplete_open_handles()780        if first_error is not None:781            raise first_error782 783    def fail(self, err: BaseException) -> None:784        status, error_str = _status_from_exception(err)785        self._open.clear()786        for handle in self._handles.values():787            self._mark_terminal(handle, status, error_str)788            if handle._mux is not None and not handle._mux._events._closed:789                try:790                    handle._mux.fail(err)791                except Exception:792                    _logger.warning(793                        "Error failing subgraph mini-mux at %s; "794                        "subscribers may not see the terminal error.",795                        handle.path,796                        exc_info=True,797                    )798 799    async def afail(self, err: BaseException) -> None:800        status, error_str = _status_from_exception(err)801        self._open.clear()802        for handle in self._handles.values():803            self._mark_terminal(handle, status, error_str)804            if handle._mux is not None and not handle._mux._events._closed:805                try:806                    await handle._mux.afail(err)807                except Exception:808                    _logger.warning(809                        "Error failing subgraph mini-mux at %s; "810                        "subscribers may not see the terminal error.",811                        handle.path,812                        exc_info=True,813                    )814 815 816class CheckpointsTransformer(StreamTransformer):817    """Capture checkpoint events as a drainable stream.818 819    Surfaces `stream_mode="checkpoints"` data on `run.checkpoints` as820    a `StreamChannel[dict[str, Any]]`. Each item is in the same format821    as returned by `get_state()`.822 823    Checkpoint events are only emitted when a checkpointer is configured824    on the graph. When no checkpointer is present, the projection exists825    but receives no events.826 827    Only events at the run's own scope are captured; checkpoint data from828    deeper subgraphs is available on the respective subgraph handle's829    `.checkpoints` projection.830 831    Native transformer — `run.checkpoints` is a direct attribute.832    """833 834    _native = True835    required_stream_modes = ("checkpoints",)836 837    def __init__(self, scope: tuple[str, ...] = ()) -> None:838        super().__init__(scope)839        self._log: StreamChannel[dict[str, Any]] = StreamChannel()840        self._scope_list: list[str] = list(scope)841 842    def init(self) -> dict[str, Any]:843        return {"checkpoints": self._log}844 845    def process(self, event: ProtocolEvent) -> bool:846        if event["method"] != "checkpoints":847            return True848        params = event["params"]849        if params["namespace"] != self._scope_list:850            return True851        self._log.push(params["data"])852        return True853 854 855class DebugTransformer(StreamTransformer):856    """Capture debug events as a drainable stream.857 858    Surfaces `stream_mode="debug"` data on `run.debug` as a859    `StreamChannel[dict[str, Any]]`. Each item is a debug event with860    step-level detail (checkpoint snapshots, task payloads, and861    task results wrapped with step number and timestamp).862 863    Only events at the run's own scope are captured; debug data from864    deeper subgraphs is available on the respective subgraph handle's865    `.debug` projection.866 867    Native transformer — `run.debug` is a direct attribute.868    """869 870    _native = True871    required_stream_modes = ("debug",)872 873    def __init__(self, scope: tuple[str, ...] = ()) -> None:874        super().__init__(scope)875        self._log: StreamChannel[dict[str, Any]] = StreamChannel()876        self._scope_list: list[str] = list(scope)877 878    def init(self) -> dict[str, Any]:879        return {"debug": self._log}880 881    def process(self, event: ProtocolEvent) -> bool:882        if event["method"] != "debug":883            return True884        params = event["params"]885        if params["namespace"] != self._scope_list:886            return True887        self._log.push(params["data"])888        return True889 890 891class TasksTransformer(StreamTransformer):892    """Capture raw task events as a drainable stream.893 894    Surfaces `stream_mode="tasks"` data on `run.tasks` as a895    `StreamChannel[dict[str, Any]]`. Each item is a task payload896    (start or result).897 898    `LifecycleTransformer` and `SubgraphTransformer` also consume899    `tasks` events for subgraph discovery and lifecycle tracking.900    This transformer captures the raw payloads independently for901    consumers who need task-level detail.902 903    Only events at the run's own scope are captured; task data from904    deeper subgraphs is available on the respective subgraph handle's905    `.tasks` projection.906 907    Native transformer — `run.tasks` is a direct attribute.908    """909 910    _native = True911    required_stream_modes = ("tasks",)912 913    def __init__(self, scope: tuple[str, ...] = ()) -> None:914        super().__init__(scope)915        self._log: StreamChannel[dict[str, Any]] = StreamChannel()916        self._scope_list: list[str] = list(scope)917 918    def init(self) -> dict[str, Any]:919        return {"tasks": self._log}920 921    def process(self, event: ProtocolEvent) -> bool:922        if event["method"] != "tasks":923            return True924        params = event["params"]925        if params["namespace"] != self._scope_list:926            return True927        self._log.push(params["data"])928        return True929