Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 26 additions & 3 deletions libs/sdk-py/langgraph_sdk/_async/stream.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
"""Async thread-centric streaming surface for the v3 protocol.

`AsyncThreadStream` is an async context manager that owns a
Expand Down Expand Up @@ -911,6 +911,29 @@
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()
Comment thread
open-swe[bot] marked this conversation as resolved.
Outdated

async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]:
if self._thread._transport is None:
raise RuntimeError("AsyncThreadStream not entered - use `async with`.")
Expand Down Expand Up @@ -943,7 +966,7 @@
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:
Expand All @@ -961,7 +984,7 @@
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:
Expand Down Expand Up @@ -1335,7 +1358,7 @@
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:
Expand Down
27 changes: 24 additions & 3 deletions libs/sdk-py/langgraph_sdk/_sync/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -954,6 +954,27 @@ def __init__(self, thread: SyncThreadStream, scope: tuple[str, ...] = ()) -> Non
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`.")
Expand Down Expand Up @@ -986,7 +1007,7 @@ def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]:
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:
Expand All @@ -1007,7 +1028,7 @@ def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]:
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:
Expand Down Expand Up @@ -1222,7 +1243,7 @@ def _reconcile_stream(self, candidate_filter: SubscribeParams) -> None:

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:
Expand Down
34 changes: 34 additions & 0 deletions libs/sdk-py/tests/streaming/test_scoped_handles.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."""
Expand Down
31 changes: 31 additions & 0 deletions libs/sdk-py/tests/streaming/test_sync_projections.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from __future__ import annotations

import queue
from typing import Any, cast

import httpx
Expand Down Expand Up @@ -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()
Expand Down
Loading