diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 37062c64d..08ba69260 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NoReturn from langchain_core._api import beta @@ -33,6 +33,21 @@ async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None: pass +def _missing_projection(run: object, name: str) -> NoReturn: + """Raise for a projection name that was never registered. + + Reads the mux out of `__dict__` because `__getattr__` also fires before + `__init__` assigns `_mux` and for dunder probes, where `run._mux` would + recurse until the stack is exhausted. + """ + mux = run.__dict__.get("_mux") + registered = sorted(mux.native_keys) if mux is not None else [] + raise AttributeError( + f"{type(run).__name__!r} object has no attribute {name!r} " + f"(registered projections: {', '.join(registered) or 'none'})" + ) + + @beta(message="The v3 streaming protocol on Pregel is experimental.") class GraphRunStream: """Sync run stream with caller-driven pumping. @@ -54,15 +69,32 @@ class GraphRunStream: experimental and may change. """ - # Native projections always registered by `stream_events(version="v3")`. - # Attached dynamically by the `setattr` loop in `__init__`; declared here - # so type checkers see them. Opt-in native projections (`updates`, - # `custom`, `checkpoints`, `debug`, `tasks`) are only present when their - # transformer is registered, so they are reached via `extensions[...]`. + # Native projections, attached dynamically by the `setattr` loop in + # `__init__` and declared here so type checkers see them. + # + # Always registered by `stream_events(version="v3")`: values: StreamChannel[dict[str, Any]] messages: StreamChannel[ChatModelStream] lifecycle: StreamChannel[LifecyclePayload] subgraphs: StreamChannel[SubgraphRunStream] + # Registered on demand via `compile(transformers=...)` or + # `stream_events(transformers=...)`; reading one whose transformer was not + # registered raises AttributeError. Projections contributed by transformers + # outside this package are covered by `__getattr__` instead. + updates: StreamChannel[dict[str, Any]] + custom: StreamChannel[Any] + checkpoints: StreamChannel[dict[str, Any]] + debug: StreamChannel[dict[str, Any]] + tasks: StreamChannel[dict[str, Any]] + + def __getattr__(self, name: str) -> StreamChannel[Any]: + """Type the projections of transformers declared outside this package. + + Projection names come from a registry, so no annotation here can name + them all. The cost is that a misspelling type-checks too, and fails at + runtime instead. + """ + _missing_projection(self, name) def __init__( self, @@ -345,15 +377,24 @@ class AsyncGraphRunStream: experimental and may change. """ - # Native projections always registered by `astream_events(version="v3")`. - # Attached dynamically by the `setattr` loop in `__init__`; declared here - # so type checkers see them. Opt-in native projections (`updates`, - # `custom`, `checkpoints`, `debug`, `tasks`) are only present when their - # transformer is registered, so they are reached via `extensions[...]`. + # Native projections, attached dynamically by the `setattr` loop in + # `__init__` and declared here so type checkers see them. + # + # Always registered by `astream_events(version="v3")`: values: StreamChannel[dict[str, Any]] messages: StreamChannel[AsyncChatModelStream] lifecycle: StreamChannel[LifecyclePayload] subgraphs: StreamChannel[AsyncSubgraphRunStream] + # Registered on demand; see `GraphRunStream`. + updates: StreamChannel[dict[str, Any]] + custom: StreamChannel[Any] + checkpoints: StreamChannel[dict[str, Any]] + debug: StreamChannel[dict[str, Any]] + tasks: StreamChannel[dict[str, Any]] + + def __getattr__(self, name: str) -> StreamChannel[Any]: + """Type projections declared elsewhere. See `GraphRunStream`.""" + _missing_projection(self, name) def __init__( self, diff --git a/libs/langgraph/tests/test_stream_events_v3.py b/libs/langgraph/tests/test_stream_events_v3.py index 71d81325e..fff1ded3d 100644 --- a/libs/langgraph/tests/test_stream_events_v3.py +++ b/libs/langgraph/tests/test_stream_events_v3.py @@ -6,6 +6,7 @@ Type-narrowing is validated via `assert_type` calls in `_check_type_narrowing`. from __future__ import annotations +import copy import operator import sys from dataclasses import dataclass @@ -34,8 +35,10 @@ from langgraph.stream import ( GraphRunStream, LifecyclePayload, StreamChannel, + StreamTransformer, SubgraphRunStream, ) +from langgraph.stream._types import ProtocolEvent from langgraph.types import ( CheckpointPayload, CheckpointStreamPart, @@ -1199,6 +1202,29 @@ def _check_type_narrowing(part: StreamPart[_StateT, _OutputT]) -> None: # type and the always-registered native projections. +class _MarkerTransformer(StreamTransformer): + """Native transformer contributing a key this module doesn't declare. + + Stands in for any transformer defined outside this package — projections + whose names `GraphRunStream` can't enumerate, so they resolve through + `__getattr__` instead of a class annotation. + """ + + _native = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: StreamChannel[str] = StreamChannel() + + def init(self) -> dict[str, Any]: + return {"marker": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._log.push("saw_values") + return True + + def _check_stream_events_v3_typing() -> None: """Compile-time checks for sync v3 typing — never called at runtime.""" graph = _make_simple_graph().compile() @@ -1208,6 +1234,16 @@ def _check_stream_events_v3_typing() -> None: assert_type(run.messages, StreamChannel[ChatModelStream]) assert_type(run.lifecycle, StreamChannel[LifecyclePayload]) assert_type(run.subgraphs, StreamChannel[SubgraphRunStream]) + # Opt-in projections from transformers this package ships carry their real + # item type even though they are only present once registered. + assert_type(run.updates, StreamChannel[dict[str, Any]]) + assert_type(run.custom, StreamChannel[Any]) + assert_type(run.checkpoints, StreamChannel[dict[str, Any]]) + assert_type(run.debug, StreamChannel[dict[str, Any]]) + assert_type(run.tasks, StreamChannel[dict[str, Any]]) + # Projections this module can't enumerate resolve through `__getattr__` + # as `StreamChannel[Any]` rather than failing with attr-defined. + assert_type(run.marker, StreamChannel[Any]) async def _check_astream_events_v3_typing() -> None: @@ -1219,3 +1255,66 @@ async def _check_astream_events_v3_typing() -> None: assert_type(run.messages, StreamChannel[AsyncChatModelStream]) assert_type(run.lifecycle, StreamChannel[LifecyclePayload]) assert_type(run.subgraphs, StreamChannel[AsyncSubgraphRunStream]) + assert_type(run.updates, StreamChannel[dict[str, Any]]) + assert_type(run.custom, StreamChannel[Any]) + assert_type(run.checkpoints, StreamChannel[dict[str, Any]]) + assert_type(run.debug, StreamChannel[dict[str, Any]]) + assert_type(run.tasks, StreamChannel[dict[str, Any]]) + assert_type(run.marker, StreamChannel[Any]) + + +def test_undeclared_native_projection_is_attached() -> None: + """A native projection this module doesn't declare still works at runtime. + + `__getattr__` is a type-checker fallback only — it must not shadow the + `setattr` loop that attaches registered native projections. + """ + graph = _make_simple_graph().compile() + run = graph.stream_events( + _SIMPLE_INPUT, version="v3", transformers=[_MarkerTransformer] + ) + + marker_iter = iter(run.marker) + assert run.output is not None + assert run.marker is run.extensions["marker"] + assert "saw_values" in list(marker_iter) + + +def test_unregistered_projection_raises_attribute_error() -> None: + """An unregistered projection name still fails at runtime. + + The `__getattr__` fallback exists to satisfy type checkers; it must not + make unknown names resolve to anything. The message lists what *is* + registered so a typo is diagnosable from the traceback alone. + """ + graph = _make_simple_graph().compile() + run = graph.stream_events(_SIMPLE_INPUT, version="v3") + + # `marker` type-checks via `__getattr__` but was never registered here. + with pytest.raises(AttributeError) as exc_info: + run.marker + + message = str(exc_info.value) + assert "marker" in message + # The always-registered natives are listed as the alternatives. + assert "messages" in message + + # Registered projections still resolve, and the run is unaffected. + assert isinstance(run.messages, StreamChannel) + assert run.output is not None + + +def test_getattr_fallback_does_not_recurse_before_init() -> None: + """`__getattr__` reads the mux from `__dict__`, so it is safe pre-init. + + `self._mux` would re-enter `__getattr__` and overflow the stack when the + attribute is missing, which is reachable via `hasattr` / `copy` / pickle + probing on a partially constructed instance. + """ + bare = GraphRunStream.__new__(GraphRunStream) + + with pytest.raises(AttributeError, match="anything"): + bare.anything + + assert hasattr(bare, "anything") is False + assert copy.copy(bare) is not None