Compare commits

...
Author SHA1 Message Date
John Kennedyandopen-swe[bot] <open-swe@users.noreply.github.com> a7cd539c94 fix: preserve full root messages inbox
Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
2026-07-09 04:29:18 +00:00
Langster Security Agent d0524a4473 Fix root messages inbox bounds 2026-07-09 04:08:19 +00:00
4 changed files with 167 additions and 6 deletions
+32 -3
View File
@@ -911,6 +911,34 @@ 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:
if root_inbox.maxsize > 0 and root_inbox.qsize() >= root_inbox.maxsize - 1:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while buffering "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
)
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:
try:
root_inbox.put_nowait(None)
except asyncio.QueueFull as exc:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while closing "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
) from exc
async def _subgraphs_iter(self) -> AsyncGenerator[ScopedStreamHandle, None]:
if self._thread._transport is None:
raise RuntimeError("AsyncThreadStream not entered - use `async with`.")
@@ -943,7 +971,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 +989,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 +1363,8 @@ class AsyncThreadStream:
that arrive at namespace `[]` before `thread.messages` has subscribed.
"""
if self._root_messages_inbox is None:
self._root_messages_inbox = asyncio.Queue()
maxsize = self._max_queue_size + 1 if self._max_queue_size > 0 else 0
self._root_messages_inbox = asyncio.Queue(maxsize=maxsize)
return self._root_messages_inbox
def _register_active_message_stream(self, stream: AsyncChatModelStream) -> None:
+31 -3
View File
@@ -954,6 +954,34 @@ 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:
if root_inbox.maxsize > 0 and root_inbox.qsize() >= root_inbox.maxsize - 1:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while buffering "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
)
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:
try:
root_inbox.put_nowait(None)
except queue.Full as exc:
raise RuntimeError(
"Root messages inbox exceeded max_queue_size while closing "
"root-scope messages. Iterate thread.messages concurrently "
"or increase max_queue_size."
) from exc
def _subgraphs_iter(self) -> Iterator[SyncScopedStreamHandle]:
if self._thread._transport is None:
raise RuntimeError("SyncThreadStream not entered — use `with`.")
@@ -986,7 +1014,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 +1035,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 +1250,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=1025)
return self._root_messages_inbox
def _register_active_message_stream(self, stream: ChatModelStream) -> None:
@@ -2,7 +2,13 @@
from __future__ import annotations
import asyncio
from typing import cast
from unittest.mock import MagicMock
import httpx
import pytest
from langchain_protocol import Event
from langgraph_sdk._async.http import HttpClient
from langgraph_sdk._async.threads import ThreadsClient
@@ -447,6 +453,56 @@ 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 reserve one terminal sentinel slot."""
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 == 17
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[Event | None] = asyncio.Queue(maxsize=2)
inbox.put_nowait(cast(Event, message_start_event(seq=1, message_id="msg-1")))
with pytest.raises(RuntimeError, match="Root messages inbox exceeded"):
_SubgraphsProjection._put_root_message(
inbox,
cast(
Event,
message_text_delta_event(seq=2, text="overflow", message_id="msg-1"),
),
)
def test_subgraphs_root_message_close_preserves_full_inbox():
"""Closing a full root messages inbox must not evict buffered events."""
from langgraph_sdk._async.stream import _SubgraphsProjection
inbox: asyncio.Queue[Event | None] = asyncio.Queue(maxsize=3)
first = cast(Event, message_start_event(seq=1, message_id="msg-1"))
second = cast(Event, message_finish_event(seq=2, message_id="msg-1"))
inbox.put_nowait(first)
inbox.put_nowait(second)
_SubgraphsProjection._signal_root_inbox_closed(inbox)
assert inbox.get_nowait() is first
assert inbox.get_nowait() is second
assert inbox.get_nowait() is None
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."""
@@ -2,6 +2,7 @@
from __future__ import annotations
import queue
from typing import Any, cast
import httpx
@@ -398,6 +399,53 @@ def test_sync_tool_calls_run_error_fails_active_handle():
# ---------------------------------------------------------------------------
def test_sync_root_messages_inbox_is_bounded():
"""Root messages inbox must reserve one terminal sentinel slot."""
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 == 1025
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=2)
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_subgraphs_root_message_close_preserves_full_inbox():
"""Closing a full root messages inbox must not evict buffered events."""
from langgraph_sdk._sync.stream import _SyncSubgraphsProjection
inbox: queue.Queue[Event | None] = queue.Queue(maxsize=3)
first = cast(Event, message_start_event(seq=1, message_id="msg-1"))
second = cast(Event, message_finish_event(seq=2, message_id="msg-1"))
inbox.put_nowait(first)
inbox.put_nowait(second)
_SyncSubgraphsProjection._signal_root_inbox_closed(inbox)
assert inbox.get_nowait() is first
assert inbox.get_nowait() is second
assert inbox.get_nowait() is None
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()