From 3fff7bc9283245abc5a8c864dea3979db80728a5 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Mon, 4 May 2026 08:52:52 -0400 Subject: [PATCH] feat(langgraph): forward kwargs through stream_events(version="v3") (#7696) --- libs/langgraph/langgraph/pregel/main.py | 49 +++++++- .../test_stream_events_v3_kwarg_forwarding.py | 108 ++++++++++++++++++ 2 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 libs/langgraph/tests/test_stream_events_v3_kwarg_forwarding.py diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index fdc84c808..1550e5a92 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -374,6 +374,25 @@ class NodeBuilder: ) +# Kwargs that ``stream_events(version="v3")`` / ``astream_events(version="v3")`` +# manage internally and must not be overridden by callers. ``stream_mode`` is +# derived from the transformer mux; ``subgraphs`` is forced True so nested +# namespaces flow through scoped muxes. Forwarding either to the inner +# ``stream(...)`` would silently break v3's invariants, so we raise instead. +_V3_INVARIANT_KWARGS: tuple[str, ...] = ("stream_mode", "subgraphs") + + +def _reject_v3_invariant_kwargs(kwargs: dict[str, Any]) -> None: + collisions = [k for k in _V3_INVARIANT_KWARGS if k in kwargs] + if collisions: + raise TypeError( + "stream_events(version='v3') / astream_events(version='v3') do " + f"not accept {', '.join(collisions)}; v3 owns these " + "(stream_mode is built from the transformer mux, subgraphs is " + "forced True so nested namespaces flow through scoped muxes)." + ) + + def _collect_stream_modes(mux: Any) -> list[StreamMode]: """Return the union of `required_stream_modes` across registered transformers. @@ -3459,9 +3478,16 @@ class Pregel( interrupt_after: All | Sequence[str] | None = None, control: RunControl | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, + **kwargs: Any, ) -> Any: """Internal v3 sync streaming implementation. Public entry: stream_events(version='v3'). + Extra keyword arguments are forwarded to the underlying ``stream(...)`` + call. The dispatcher in ``stream_events`` rejects ``stream_mode`` and + ``subgraphs`` since v3 owns them (``stream_mode`` is derived from the + transformer mux; ``subgraphs`` is always True so nested namespaces + flow through scoped muxes). + !!! warning The v3 streaming protocol is experimental and may change. @@ -3493,6 +3519,7 @@ class Pregel( interrupt_before=interrupt_before, interrupt_after=interrupt_after, control=control, + **kwargs, ) ) return GraphRunStream(graph_iter, mux) @@ -3507,9 +3534,16 @@ class Pregel( interrupt_after: All | Sequence[str] | None = None, control: RunControl | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, + **kwargs: Any, ) -> Any: """Internal v3 async streaming implementation. Public entry: astream_events(version='v3'). + Extra keyword arguments are forwarded to the underlying ``astream(...)`` + call. The dispatcher in ``astream_events`` rejects ``stream_mode`` and + ``subgraphs`` since v3 owns them (``stream_mode`` is derived from the + transformer mux; ``subgraphs`` is always True so nested namespaces + flow through scoped muxes). + !!! warning The v3 streaming protocol is experimental and may change. @@ -3540,6 +3574,7 @@ class Pregel( interrupt_before=interrupt_before, interrupt_after=interrupt_after, control=control, + **kwargs, ).__aiter__() return AsyncGraphRunStream(graph_aiter, mux) @@ -3564,6 +3599,7 @@ class Pregel( interrupt_after: All | Sequence[str] | None = None, control: RunControl | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, + **kwargs: Any, ) -> Any: ... def stream_events( @@ -3622,13 +3658,20 @@ class Pregel( `stream_transformers`. Factories are called as `factory(scope)` so they can propagate to subgraph scopes. Only used for `version="v3"`. - **kwargs: Forwarded to the v1/v2 path. + **kwargs: For `version="v1"`/`"v2"`, forwarded to + `Runnable.stream_events`. For `version="v3"`, forwarded + to the underlying `stream(...)` call (e.g. `context`, + `durability`, `output_keys`, `print_mode`, `debug`). + `stream_mode` and `subgraphs` are not accepted under + `version="v3"` and raise `TypeError` if supplied; v3 + owns them. Returns: For `version="v3"`, a `GraphRunStream` the caller iterates to drive the run. Otherwise an `Iterator[StreamEvent]`. """ if version == "v3": + _reject_v3_invariant_kwargs(kwargs) return self._pregel_stream_v3( input, config, @@ -3636,6 +3679,7 @@ class Pregel( interrupt_after=interrupt_after, control=control, transformers=transformers, + **kwargs, ) return super().stream_events(input, config, version=version, **kwargs) @@ -3660,6 +3704,7 @@ class Pregel( interrupt_after: All | Sequence[str] | None = None, control: RunControl | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, + **kwargs: Any, ) -> Awaitable[Any]: ... def astream_events( @@ -3689,6 +3734,7 @@ class Pregel( See `stream_events` for full argument and return documentation. """ if version == "v3": + _reject_v3_invariant_kwargs(kwargs) return self._apregel_stream_v3( input, config, @@ -3696,6 +3742,7 @@ class Pregel( interrupt_after=interrupt_after, control=control, transformers=transformers, + **kwargs, ) return super().astream_events(input, config, version=version, **kwargs) diff --git a/libs/langgraph/tests/test_stream_events_v3_kwarg_forwarding.py b/libs/langgraph/tests/test_stream_events_v3_kwarg_forwarding.py new file mode 100644 index 000000000..a5f353ad3 --- /dev/null +++ b/libs/langgraph/tests/test_stream_events_v3_kwarg_forwarding.py @@ -0,0 +1,108 @@ +"""Tests that ``(a)stream_events(version="v3")`` forwards extra kwargs to the +underlying ``(a)stream`` call, and rejects the kwargs v3 owns internally. + +Background: prior to this change the v3 dispatcher silently dropped ``**kwargs`` +on the v3 branch while forwarding them on v1/v2, so callers passing e.g. +``context=...`` saw their value disappear with no error. v3 now forwards +caller kwargs to the inner ``(a)stream`` call but rejects ``stream_mode`` and +``subgraphs`` since v3 owns them (``stream_mode`` is built from the +transformer mux; ``subgraphs`` is forced True so nested namespaces flow +through scoped muxes). +""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from typing import Any + +import pytest +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.runtime import Runtime + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +@dataclass +class _Ctx: + api_key: str + + +class _State(TypedDict): + message: str + + +def _build_context_reading_graph(): + def read_context(state: _State, runtime: Runtime[_Ctx]) -> dict[str, Any]: + return {"message": f"api key: {runtime.context.api_key}"} + + builder = StateGraph(state_schema=_State, context_schema=_Ctx) + builder.add_node("read_context", read_context) + builder.add_edge(START, "read_context") + builder.add_edge("read_context", END) + return builder.compile() + + +class TestKwargForwardingSync: + def test_context_reaches_node(self) -> None: + run = _build_context_reading_graph().stream_events( + {"message": "hello"}, + version="v3", + context=_Ctx(api_key="sk_sync"), + ) + assert run.output == {"message": "api key: sk_sync"} + + def test_rejects_stream_mode(self) -> None: + graph = _build_context_reading_graph() + with pytest.raises(TypeError, match="stream_mode"): + graph.stream_events( + {"message": "hello"}, + version="v3", + stream_mode=["values"], + ) + + def test_rejects_subgraphs(self) -> None: + graph = _build_context_reading_graph() + with pytest.raises(TypeError, match="subgraphs"): + graph.stream_events( + {"message": "hello"}, + version="v3", + subgraphs=False, + ) + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +class TestKwargForwardingAsync: + async def test_context_reaches_node(self) -> None: + run = await _build_context_reading_graph().astream_events( + {"message": "hello"}, + version="v3", + context=_Ctx(api_key="sk_async"), + ) + output = await run.output() + assert output == {"message": "api key: sk_async"} + + async def test_rejects_stream_mode(self) -> None: + graph = _build_context_reading_graph() + with pytest.raises(TypeError, match="stream_mode"): + await graph.astream_events( + {"message": "hello"}, + version="v3", + stream_mode=["values"], + ) + + async def test_rejects_subgraphs(self) -> None: + graph = _build_context_reading_graph() + with pytest.raises(TypeError, match="subgraphs"): + await graph.astream_events( + {"message": "hello"}, + version="v3", + subgraphs=False, + )