diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index 870ce2556..b05ab40b5 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -911,6 +911,29 @@ class _SubgraphsProjection: def __aiter__(self) -> AsyncIterator[ScopedStreamHandle]: return self._subgraphs_iter() + @staticmethod + def _put_root_message( + root_inbox: asyncio.Queue[Event | None], item: Event + ) -> None: + try: + root_inbox.put_nowait(item) + except asyncio.QueueFull as exc: + raise RuntimeError( + "Root messages inbox exceeded max_queue_size while buffering " + "root-scope messages. Iterate thread.messages concurrently " + "or increase max_queue_size." + ) from exc + + @staticmethod + def _signal_root_inbox_closed(root_inbox: asyncio.Queue[Event | None]) -> None: + while True: + try: + root_inbox.put_nowait(None) + return + except asyncio.QueueFull: + with contextlib.suppress(asyncio.QueueEmpty): + root_inbox.get_nowait() + async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]: if self._thread._transport is None: raise RuntimeError("AsyncThreadStream not entered - use `async with`.") @@ -943,7 +966,7 @@ class _SubgraphsProjection: and item.get("method") == "messages" and tuple(_event_namespace(params_field)) == self._scope ): - root_inbox.put_nowait(item) + self._put_root_message(root_inbox, item) for handle in decoder.feed(item): yield handle finally: @@ -961,7 +984,7 @@ class _SubgraphsProjection: handle._finish(terminal_status) self._thread._unregister_subscription(sub.id) if root_inbox is not None: - root_inbox.put_nowait(None) + self._signal_root_inbox_closed(root_inbox) class ToolCallHandle: @@ -1335,7 +1358,7 @@ class AsyncThreadStream: that arrive at namespace `[]` before `thread.messages` has subscribed. """ if self._root_messages_inbox is None: - self._root_messages_inbox = asyncio.Queue() + self._root_messages_inbox = asyncio.Queue(maxsize=self._max_queue_size) return self._root_messages_inbox def _register_active_message_stream(self, stream: AsyncChatModelStream) -> None: diff --git a/libs/sdk-py/langgraph_sdk/_sync/stream.py b/libs/sdk-py/langgraph_sdk/_sync/stream.py index 9405e7f83..141cf01e8 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/stream.py +++ b/libs/sdk-py/langgraph_sdk/_sync/stream.py @@ -954,6 +954,27 @@ class _SyncSubgraphsProjection: def __iter__(self) -> Iterator[SyncScopedStreamHandle]: return self._subgraphs_iter() + @staticmethod + def _put_root_message(root_inbox: queue.Queue[Event | None], item: Event) -> None: + try: + root_inbox.put_nowait(item) + except queue.Full as exc: + raise RuntimeError( + "Root messages inbox exceeded max_queue_size while buffering " + "root-scope messages. Iterate thread.messages concurrently " + "or increase max_queue_size." + ) from exc + + @staticmethod + def _signal_root_inbox_closed(root_inbox: queue.Queue[Event | None]) -> None: + while True: + try: + root_inbox.put_nowait(None) + return + except queue.Full: + with contextlib.suppress(queue.Empty): + root_inbox.get_nowait() + def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]: if self._thread._transport is None: raise RuntimeError("SyncThreadStream not entered — use `with`.") @@ -986,7 +1007,7 @@ class _SyncSubgraphsProjection: and item.get("method") == "messages" and tuple(_event_namespace(params_field)) == self._scope ): - root_inbox.put_nowait(item) + self._put_root_message(root_inbox, item) for handle in decoder.feed(cast(dict[str, Any], item)): yield handle finally: @@ -1007,7 +1028,7 @@ class _SyncSubgraphsProjection: handle._finish(terminal_status) self._thread._unregister_subscription(sub.id) if root_inbox is not None: - root_inbox.put_nowait(None) + self._signal_root_inbox_closed(root_inbox) class _SyncExtensionsProjection: @@ -1222,7 +1243,7 @@ class SyncThreadStream: def _activate_root_messages_inbox(self) -> queue.Queue[Event | None]: if self._root_messages_inbox is None: - self._root_messages_inbox = queue.Queue() + self._root_messages_inbox = queue.Queue(maxsize=1024) return self._root_messages_inbox def _register_active_message_stream(self, stream: ChatModelStream) -> None: diff --git a/libs/sdk-py/tests/streaming/test_scoped_handles.py b/libs/sdk-py/tests/streaming/test_scoped_handles.py index 13955b225..18841f9de 100644 --- a/libs/sdk-py/tests/streaming/test_scoped_handles.py +++ b/libs/sdk-py/tests/streaming/test_scoped_handles.py @@ -2,7 +2,11 @@ from __future__ import annotations +import asyncio +from unittest.mock import MagicMock + import httpx +import pytest from langgraph_sdk._async.http import HttpClient from langgraph_sdk._async.threads import ThreadsClient @@ -447,6 +451,36 @@ def test_scoped_handle_inboxes_bounded_by_max_queue_size(): assert handle._tasks_inbox.maxsize == 16 +def test_root_messages_inbox_bounded_by_max_queue_size(): + """Root messages inbox must use the stream queue bound.""" + from langgraph_sdk._async.stream import AsyncThreadStream + + thread = AsyncThreadStream( + http=MagicMock(), + thread_id="t-1", + assistant_id="agent", + max_queue_size=16, + ) + + inbox = thread._activate_root_messages_inbox() + + assert inbox.maxsize == 16 + + +def test_subgraphs_root_message_overflow_raises_runtime_error(): + """Overflowing the root messages inbox must fail explicitly.""" + from langgraph_sdk._async.stream import _SubgraphsProjection + + inbox = asyncio.Queue(maxsize=1) + inbox.put_nowait(message_start_event(seq=1, message_id="msg-1")) + + with pytest.raises(RuntimeError, match="Root messages inbox exceeded"): + _SubgraphsProjection._put_root_message( + inbox, + message_text_delta_event(seq=2, text="overflow", message_id="msg-1"), + ) + + async def test_child_handle_inherits_max_queue_size_from_parent(): """Grandchild ScopedStreamHandles created by _HandleSubgraphsProjection inherit the parent's max_queue_size so all queues are consistently bounded.""" diff --git a/libs/sdk-py/tests/streaming/test_sync_projections.py b/libs/sdk-py/tests/streaming/test_sync_projections.py index 07708793e..38749efdf 100644 --- a/libs/sdk-py/tests/streaming/test_sync_projections.py +++ b/libs/sdk-py/tests/streaming/test_sync_projections.py @@ -2,6 +2,7 @@ from __future__ import annotations +import queue from typing import Any, cast import httpx @@ -398,6 +399,36 @@ def test_sync_tool_calls_run_error_fails_active_handle(): # --------------------------------------------------------------------------- +def test_sync_root_messages_inbox_is_bounded(): + """Root messages inbox must have an explicit maximum size.""" + fake = SyncFakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + fake.set_state({}) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + inbox = thread._activate_root_messages_inbox() + + assert inbox.maxsize == 1024 + + +def test_sync_subgraphs_root_message_overflow_raises_runtime_error(): + """Overflowing the root messages inbox must fail explicitly.""" + from langgraph_sdk._sync.stream import _SyncSubgraphsProjection + + inbox: queue.Queue[Event | None] = queue.Queue(maxsize=1) + inbox.put_nowait(cast(Event, message_start_event(seq=1, message_id="msg-1"))) + + with pytest.raises(RuntimeError, match="Root messages inbox exceeded"): + _SyncSubgraphsProjection._put_root_message( + inbox, + cast( + Event, + message_text_delta_event(seq=2, text="overflow", message_id="msg-1"), + ), + ) + + def test_sync_drain_messages_inbox_pre_dispatches_before_yield(): """When draining the root inbox, str(message.text) must work immediately on yield.""" fake = SyncFakeServer()