diff --git a/libs/langgraph/langgraph/pregel/_remote_run_stream.py b/libs/langgraph/langgraph/pregel/_remote_run_stream.py index 06d8efd09..2e6cd1ef5 100644 --- a/libs/langgraph/langgraph/pregel/_remote_run_stream.py +++ b/libs/langgraph/langgraph/pregel/_remote_run_stream.py @@ -11,6 +11,7 @@ from langchain_core.runnables import RunnableConfig from langgraph_sdk._async.stream import AsyncThreadStream from langgraph_sdk._sync.stream import SyncThreadStream from langgraph_sdk.client import LangGraphClient, SyncLangGraphClient +from langgraph_sdk.stream.decoders import DataDecoder from langgraph.types import Command @@ -43,9 +44,10 @@ def _translate_command_input(input: Any) -> Any: class _ChannelProjection: """Decoded projection for a wire channel the SDK doesn't type natively. - Subscribes to `channel` and yields each event's `params["data"]` — the same - item shape the SDK's typed projections yield (`_ValuesProjection` etc.) and - that local's `UpdatesTransformer` / `CheckpointsTransformer` / + Subscribes to `channel` and decodes each event's `params["data"]` through the + SDK's `DataDecoder` — the same decoder the SDK's own plain-payload projections + (`values` / `updates` / `checkpoints` / `tasks`) use, which yields the item + shape that local's `UpdatesTransformer` / `CheckpointsTransformer` / `TasksTransformer` / `CustomTransformer` push, so iterating this matches the corresponding local projection. Iterate with `for` against a sync stream and `async for` against an async stream (matching the underlying SDK). Opening @@ -56,30 +58,23 @@ class _ChannelProjection: self._sdk = sdk self._channel = channel - @staticmethod - def _data(event: Any) -> Any: - """Extract `params.data` from a protocol event, tolerating odd shapes.""" - params = event.get("params") if isinstance(event, dict) else None - return params.get("data") if isinstance(params, dict) else None - def __iter__(self) -> Iterator[Any]: # Sync lane: the sync adapter's SDK returns a sync iterator here. + decoder = DataDecoder(self._channel) events = cast(Iterator[Any], self._sdk.subscribe([self._channel])) for event in events: - data = self._data(event) - if data is not None: - yield data + yield from decoder.feed(event) def __aiter__(self) -> AsyncIterator[Any]: return self._aiter() async def _aiter(self) -> AsyncIterator[Any]: # Async lane: the async adapter's SDK returns an async iterator here. + decoder = DataDecoder(self._channel) events = cast(AsyncIterator[Any], self._sdk.subscribe([self._channel])) async for event in events: - data = self._data(event) - if data is not None: - yield data + for item in decoder.feed(event): + yield item class _ProjectionRegistry(Mapping[str, Any]): @@ -243,7 +238,7 @@ class _RemoteGraphRunStream: return self._events_iter def interleave(self, *names: str) -> Iterator[tuple[str, Any]]: - raise NotImplementedError + yield from self._sdk.interleave_projections(list(names)) class _AsyncRemoteGraphRunStream: diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index d7fe15d98..23c40397e 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ dependencies = [ "langchain-core>=1.4.0,<2", "langgraph-checkpoint>=4.1.0,<5.0.0", - "langgraph-sdk>=0.4.0,<0.5.0", + "langgraph-sdk>=0.4.1,<0.5.0", "langgraph-prebuilt>=1.1.0,<1.2.0", "xxhash>=3.5.0", "pydantic>=2.7.4", diff --git a/libs/langgraph/tests/test_remote_graph_v3.py b/libs/langgraph/tests/test_remote_graph_v3.py index db7f53e56..c85594c71 100644 --- a/libs/langgraph/tests/test_remote_graph_v3.py +++ b/libs/langgraph/tests/test_remote_graph_v3.py @@ -139,14 +139,15 @@ def test_projection_registry_typed_decoded_and_custom(): def test_channel_projection_decodes_params_data(): sdk = MagicMock() - # Two events with data + one malformed/dataless event that must be skipped. + # Real wire events carry `method`; the SDK `DataDecoder` yields matching + # events' `params.data` and skips dataless and off-channel ones. sdk.subscribe = MagicMock( return_value=iter( [ - {"params": {"data": {"n": 1}}}, - {"params": {}}, # no data -> skipped - {"params": {"data": {"n": 2}}}, - {"unexpected": "shape"}, # not a params dict -> skipped + {"method": "checkpoints", "params": {"data": {"n": 1}}}, + {"method": "checkpoints", "params": {}}, # no data -> skipped + {"method": "checkpoints", "params": {"data": {"n": 2}}}, + {"method": "lifecycle", "params": {"data": {"n": 3}}}, # other channel ] ) ) @@ -155,6 +156,39 @@ def test_channel_projection_decodes_params_data(): sdk.subscribe.assert_called_once_with(["checkpoints"]) +@pytest.mark.anyio +async def test_channel_projection_decodes_params_data_async(): + """Async lane mirrors the sync lane: `async for` over the SDK's async + subscription, decoded through the same `DataDecoder`.""" + + class _FakeAsyncEvents: + def __init__(self, items): + self._items = list(items) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._items: + raise StopAsyncIteration + return self._items.pop(0) + + sdk = MagicMock() + sdk.subscribe = MagicMock( + return_value=_FakeAsyncEvents( + [ + {"method": "checkpoints", "params": {"data": {"n": 1}}}, + {"method": "checkpoints", "params": {}}, # no data -> skipped + {"method": "checkpoints", "params": {"data": {"n": 2}}}, + {"method": "lifecycle", "params": {"data": {"n": 3}}}, # other channel + ] + ) + ) + proj = _ChannelProjection(sdk, "checkpoints") + assert [item async for item in proj] == [{"n": 1}, {"n": 2}] + sdk.subscribe.assert_called_once_with(["checkpoints"]) + + def test_sync_adapter_translates_command_input(): sync_client = MagicMock() sdk_thread = MagicMock() @@ -213,11 +247,14 @@ def test_abort_swallows_cancel_failure_and_still_closes(): sdk_thread.close.assert_called_once() -def test_sync_interleave_raises_not_implemented(): - adapter, _, _ = _make_sync_adapter() +def test_sync_interleave_delegates_to_interleave_projections(): + adapter, _, sdk_thread = _make_sync_adapter() + pairs = [("values", {"x": 1}), ("messages", object())] + sdk_thread.interleave_projections.return_value = pairs with adapter as stream: - with pytest.raises(NotImplementedError): - list(stream.interleave("messages")) + result = list(stream.interleave("values", "messages")) + assert result == pairs + sdk_thread.interleave_projections.assert_called_once_with(["values", "messages"]) def test_async_adapter_has_no_interleave():