mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-29 10:21:29 +01:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
19fd4105ff | ||
|
|
9ebc35a477 | ||
|
|
4af5142b14 | ||
|
|
0470c45f8e | ||
|
|
347397d4bc |
No files matched your search
@@ -2148,6 +2148,6 @@ yocto-queue@^0.1.0:
|
||||
integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==
|
||||
|
||||
"zod@^3.25.76 || ^4":
|
||||
version "4.6.7"
|
||||
resolved "https://registry.yarnpkg.com/zod/-/zod-4.6.7.tgz#e806cf64a7f4e37635ce9f2f52c04c3766f75e5b"
|
||||
integrity sha512-g8W2D2kU6a8untc053fX1uVBKzM7yqAxW/3zwk83+6RyF2+SKOuV+c0wBDhUwzV1PaGapmZH+rmcUzPf+3CGwg==
|
||||
version "4.3.6"
|
||||
resolved "https://registry.yarnpkg.com/zod/-/zod-4.3.6.tgz#89c56e0aa7d2b05107d894412227087885ab112a"
|
||||
integrity sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==
|
||||
@@ -22,7 +22,7 @@ from langchain_core.messages import (
|
||||
)
|
||||
from typing_extensions import TypedDict, deprecated
|
||||
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.pregel._messages import (
|
||||
_messages_delta_reducer as _messages_delta_reducer,
|
||||
@@ -343,7 +343,7 @@ def push_message(
|
||||
)
|
||||
|
||||
from langgraph.config import get_config
|
||||
from langgraph.pregel._messages import StreamMessagesHandler
|
||||
from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
|
||||
|
||||
config = get_config()
|
||||
message = next(x for x in convert_to_messages([message]))
|
||||
@@ -359,15 +359,12 @@ def push_message(
|
||||
):
|
||||
handlers = config["callbacks"]
|
||||
|
||||
if stream_handler := next(
|
||||
(x for x in handlers if isinstance(x, StreamMessagesHandler)), None
|
||||
):
|
||||
metadata = config["metadata"]
|
||||
message_meta = (
|
||||
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
|
||||
metadata,
|
||||
)
|
||||
stream_handler._emit(message_meta, message, dedupe=False)
|
||||
metadata = config["metadata"]
|
||||
for handler in handlers:
|
||||
if isinstance(handler, StreamMessagesHandler):
|
||||
ns = graph_namespace(metadata)
|
||||
if handler.subgraphs or not ns or ns == handler.parent_ns:
|
||||
handler._emit((ns, metadata), message, dedupe=False)
|
||||
|
||||
if state_key:
|
||||
config[CONF][CONFIG_KEY_SEND]([(state_key, message)])
|
||||
|
||||
@@ -23,7 +23,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResul
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph._internal._config import filter_to_user_tags
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph._internal._constants import NS_SEP, PUSH
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
@@ -55,6 +55,23 @@ def _state_values(obj: Any) -> Sequence[Any]:
|
||||
return ()
|
||||
|
||||
|
||||
def graph_namespace(metadata: Mapping[str, Any]) -> tuple[str, ...]:
|
||||
"""Namespace of the graph a task runs in, from the task's run metadata.
|
||||
|
||||
Drops the task's own segment, plus one per enclosing `@task` call: a call
|
||||
runs in its caller's graph but nests its segment under the caller's. A
|
||||
call's `langgraph_path` is `(PUSH, <caller's path>, <index>, True)`, while
|
||||
a `Send` task's ends in `False`.
|
||||
"""
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
|
||||
path = metadata.get("langgraph_path", ())
|
||||
drop = 1
|
||||
while len(path) == 4 and path[0] == PUSH and path[-1] is True:
|
||||
drop += 1
|
||||
path = path[1]
|
||||
return ns[:-drop]
|
||||
|
||||
|
||||
class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
"""A callback handler that implements stream_mode=messages.
|
||||
|
||||
@@ -148,9 +165,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
ns = graph_namespace(metadata)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
if (filtered_tags := filter_to_user_tags(tags)) is not None:
|
||||
@@ -213,9 +228,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
and kwargs.get("name") == metadata.get("langgraph_node")
|
||||
and (not tags or TAG_HIDDEN not in tags)
|
||||
):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
ns = graph_namespace(metadata)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
|
||||
@@ -7,9 +7,9 @@ from uuid import UUID
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.constants import TAG_NOSTREAM
|
||||
from langgraph.errors import GraphDrained, GraphInterrupt
|
||||
from langgraph.pregel._messages import graph_namespace
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
|
||||
try:
|
||||
@@ -92,12 +92,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
) -> tuple[str, ...] | None:
|
||||
"""Resolve the namespace this tool call should emit at, or `None` to skip.
|
||||
|
||||
Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
|
||||
derivation: parses `langgraph_checkpoint_ns` (which ends with
|
||||
the `node_name:task_id` of the calling node), drops that
|
||||
trailing segment, and returns the containing subgraph's own
|
||||
namespace. Returns `None` when the call should be silently
|
||||
suppressed:
|
||||
Mirrors `StreamMessagesHandler.on_chat_model_start`: the namespace
|
||||
is that of the graph the tool runs in (`graph_namespace`). Returns
|
||||
`None` when the call should be silently suppressed:
|
||||
|
||||
- `metadata` is missing — handler is attached to a context
|
||||
without Pregel routing info.
|
||||
@@ -110,11 +107,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
return None
|
||||
if tags and TAG_NOSTREAM in tags:
|
||||
return None
|
||||
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||
if not nskey:
|
||||
ns: tuple[str, ...] = ()
|
||||
else:
|
||||
ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
|
||||
ns = (
|
||||
graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
|
||||
)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return None
|
||||
return ns
|
||||
|
||||
@@ -2880,6 +2880,18 @@ class Pregel(
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
if CONFIG_KEY_TASK_ID not in config[CONF]:
|
||||
# Not run as a task (a child invoked with its own thread, say),
|
||||
# this graph's namespace starts at the root, so a caller's
|
||||
# handler would take this graph's events for its own graph's.
|
||||
callback_manager.handlers = [
|
||||
h for h in callback_manager.handlers if not _is_stream_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_handler(h)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
@@ -2898,9 +2910,9 @@ class Pregel(
|
||||
)
|
||||
# set up subgraph checkpointing
|
||||
config = self._own_checkpoint_config(config)
|
||||
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
use_stream_messages_v2 = bool(
|
||||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||
)
|
||||
@@ -2909,24 +2921,12 @@ class Pregel(
|
||||
if use_stream_messages_v2
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
messages_handler_cls(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
|
||||
|
||||
# set up tools stream mode
|
||||
if "tools" in stream_modes:
|
||||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamToolCallHandler(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
|
||||
run_manager.inheritable_handlers.extend(stream_handlers)
|
||||
|
||||
# set up custom stream mode
|
||||
if "custom" in stream_modes:
|
||||
@@ -3009,6 +3009,10 @@ class Pregel(
|
||||
cache_policy=self.cache_policy,
|
||||
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
|
||||
) as loop:
|
||||
# Known only now: the loop numbers the namespace of every
|
||||
# subgraph a task runs after its first.
|
||||
for handler in stream_handlers:
|
||||
handler.parent_ns = loop.checkpoint_ns
|
||||
emit_graph_lifecycle_events(loop)
|
||||
# create runner
|
||||
runner = PregelRunner(
|
||||
@@ -3291,6 +3295,18 @@ class Pregel(
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
if CONFIG_KEY_TASK_ID not in config[CONF]:
|
||||
# Not run as a task (a child invoked with its own thread, say),
|
||||
# this graph's namespace starts at the root, so a caller's
|
||||
# handler would take this graph's events for its own graph's.
|
||||
callback_manager.handlers = [
|
||||
h for h in callback_manager.handlers if not _is_stream_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_handler(h)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
@@ -3325,10 +3341,9 @@ class Pregel(
|
||||
)
|
||||
# set up subgraph checkpointing
|
||||
config = self._own_checkpoint_config(config)
|
||||
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
# namespace can be None in a root level graph?
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
use_stream_messages_v2 = bool(
|
||||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||
)
|
||||
@@ -3337,24 +3352,12 @@ class Pregel(
|
||||
if use_stream_messages_v2
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
messages_handler_cls(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
|
||||
|
||||
# set up tools stream mode
|
||||
if "tools" in stream_modes:
|
||||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamToolCallHandler(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
|
||||
run_manager.inheritable_handlers.extend(stream_handlers)
|
||||
|
||||
# set up custom stream mode
|
||||
def stream_writer(c: Any) -> None:
|
||||
@@ -3463,6 +3466,10 @@ class Pregel(
|
||||
cache_policy=self.cache_policy,
|
||||
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
|
||||
) as loop:
|
||||
# Known only now: the loop numbers the namespace of every
|
||||
# subgraph a task runs after its first.
|
||||
for handler in stream_handlers:
|
||||
handler.parent_ns = loop.checkpoint_ns
|
||||
await aemit_graph_lifecycle_events(loop)
|
||||
# create runner
|
||||
runner = PregelRunner(
|
||||
@@ -4460,6 +4467,10 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
|
||||
payload["values"] = mapper(payload["values"])
|
||||
|
||||
|
||||
def _is_stream_handler(handler: object) -> bool:
|
||||
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
|
||||
|
||||
def _resolve_parent_ns(
|
||||
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
|
||||
) -> tuple[str, ...]:
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import START, MessagesState, StateGraph
|
||||
from langgraph.graph.message import push_message
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
|
||||
|
||||
@tool
|
||||
def echo(text: str) -> str:
|
||||
"""Echo the text."""
|
||||
return text
|
||||
|
||||
|
||||
def _one_node_graph(node: Any) -> CompiledStateGraph:
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node(node)
|
||||
.add_edge(START, node.__name__)
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _chat_and_echo(state: MessagesState) -> dict:
|
||||
echo.invoke({"text": "done"})
|
||||
return {"messages": [GenericFakeChatModel(messages=iter(["one two"])).invoke("hi")]}
|
||||
|
||||
|
||||
async def _achat_and_echo(state: MessagesState) -> dict:
|
||||
await echo.ainvoke({"text": "done"})
|
||||
model = GenericFakeChatModel(messages=iter(["one two"]))
|
||||
return {"messages": [await model.ainvoke("hi")]}
|
||||
|
||||
|
||||
def _segment_names(ns: tuple[str, ...]) -> tuple[str, ...]:
|
||||
return tuple(segment.split(":")[0] for segment in ns)
|
||||
|
||||
|
||||
def _summary(part: dict) -> tuple[tuple[str, ...], str]:
|
||||
if part["type"] == "messages":
|
||||
return _segment_names(part["ns"]), part["data"][0].content
|
||||
return _segment_names(part["ns"]), part["data"]["event"]
|
||||
|
||||
|
||||
_CHAT_AND_ECHO_EVENTS = [
|
||||
((), "tool-started"),
|
||||
((), "tool-finished"),
|
||||
((), "one"),
|
||||
((), " "),
|
||||
((), "two"),
|
||||
]
|
||||
|
||||
_CHILD_CONFIGS = pytest.mark.parametrize(
|
||||
"child_config",
|
||||
[
|
||||
None,
|
||||
{"configurable": {"thread_id": "child"}},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
],
|
||||
ids=["no_config", "own_thread", "caller_thread"],
|
||||
)
|
||||
|
||||
|
||||
_TASK_CALLERS = pytest.mark.parametrize(
|
||||
"nested", [False, True], ids=["from_entrypoint", "from_task_in_task"]
|
||||
)
|
||||
|
||||
|
||||
@_TASK_CALLERS
|
||||
@pytest.mark.parametrize("subgraphs", [False, True])
|
||||
def test_task_events_stream_under_the_graph_namespace(
|
||||
nested: bool, subgraphs: bool
|
||||
) -> None:
|
||||
@task
|
||||
def chat_and_echo() -> None:
|
||||
_chat_and_echo({"messages": []})
|
||||
|
||||
@task
|
||||
def outer() -> None:
|
||||
chat_and_echo().result()
|
||||
|
||||
@entrypoint()
|
||||
def workflow(inputs: dict) -> None:
|
||||
(outer() if nested else chat_and_echo()).result()
|
||||
|
||||
parts = workflow.stream(
|
||||
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
|
||||
)
|
||||
|
||||
assert [_summary(part) for part in parts] == _CHAT_AND_ECHO_EVENTS
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@_TASK_CALLERS
|
||||
@pytest.mark.parametrize("subgraphs", [False, True])
|
||||
async def test_task_events_stream_under_the_graph_namespace_async(
|
||||
nested: bool, subgraphs: bool
|
||||
) -> None:
|
||||
@task
|
||||
async def chat_and_echo() -> None:
|
||||
await _achat_and_echo({"messages": []})
|
||||
|
||||
@task
|
||||
async def outer() -> None:
|
||||
await chat_and_echo()
|
||||
|
||||
@entrypoint()
|
||||
async def workflow(inputs: dict) -> None:
|
||||
await (outer() if nested else chat_and_echo())
|
||||
|
||||
parts = workflow.astream(
|
||||
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
|
||||
)
|
||||
|
||||
assert [_summary(part) async for part in parts] == _CHAT_AND_ECHO_EVENTS
|
||||
|
||||
|
||||
def test_task_events_stream_under_the_namespace_of_their_subgraph() -> None:
|
||||
@task
|
||||
def chat_and_echo() -> None:
|
||||
_chat_and_echo({"messages": []})
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
chat_and_echo().result()
|
||||
return {}
|
||||
|
||||
sub = _one_node_graph(node)
|
||||
parent = StateGraph(MessagesState).add_node("sub", sub).add_edge(START, "sub")
|
||||
|
||||
parts = parent.compile().stream(
|
||||
{"messages": []},
|
||||
stream_mode=["messages", "tools"],
|
||||
subgraphs=True,
|
||||
version="v2",
|
||||
)
|
||||
|
||||
assert [_summary(part) for part in parts] == [
|
||||
(("sub",), event) for _, event in _CHAT_AND_ECHO_EVENTS
|
||||
]
|
||||
|
||||
|
||||
@_CHILD_CONFIGS
|
||||
def test_child_graph_streams_only_to_the_node_streaming_it(
|
||||
child_config: dict | None,
|
||||
) -> None:
|
||||
child = _one_node_graph(_chat_and_echo)
|
||||
received = []
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
received.extend(
|
||||
child.stream(
|
||||
{"messages": []}, child_config, stream_mode=["messages", "tools"]
|
||||
)
|
||||
)
|
||||
return {}
|
||||
|
||||
parts = _one_node_graph(node).stream(
|
||||
{"messages": []},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
stream_mode=["messages", "tools"],
|
||||
)
|
||||
|
||||
assert list(parts) == []
|
||||
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@_CHILD_CONFIGS
|
||||
async def test_child_graph_streams_only_to_the_node_streaming_it_async(
|
||||
child_config: dict | None,
|
||||
) -> None:
|
||||
child = _one_node_graph(_achat_and_echo)
|
||||
received = []
|
||||
|
||||
async def node(state: MessagesState) -> dict:
|
||||
stream = child.astream(
|
||||
{"messages": []}, child_config, stream_mode=["messages", "tools"]
|
||||
)
|
||||
received.extend([part async for part in stream])
|
||||
return {}
|
||||
|
||||
parts = _one_node_graph(node).astream(
|
||||
{"messages": []},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
stream_mode=["messages", "tools"],
|
||||
)
|
||||
|
||||
assert [part async for part in parts] == []
|
||||
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"child_config",
|
||||
[
|
||||
{"configurable": {"thread_id": "child"}},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
],
|
||||
ids=["own_thread", "caller_thread"],
|
||||
)
|
||||
def test_child_graph_on_its_own_thread_skips_a_caller_streaming_subgraphs(
|
||||
child_config: dict,
|
||||
) -> None:
|
||||
child = _one_node_graph(_chat_and_echo)
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
child.invoke({"messages": []}, child_config)
|
||||
return {}
|
||||
|
||||
parts = _one_node_graph(node).stream(
|
||||
{"messages": []},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
stream_mode=["messages", "tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
|
||||
assert list(parts) == []
|
||||
|
||||
|
||||
def test_node_gets_the_events_of_every_subgraph_it_streams() -> None:
|
||||
sub = _one_node_graph(_chat_and_echo)
|
||||
counts = []
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
for _ in range(2):
|
||||
stream = sub.stream({"messages": []}, stream_mode=["messages", "tools"])
|
||||
counts.append(len(list(stream)))
|
||||
return {}
|
||||
|
||||
_one_node_graph(node).invoke({"messages": []})
|
||||
|
||||
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_node_gets_the_events_of_every_subgraph_it_streams_async() -> None:
|
||||
sub = _one_node_graph(_achat_and_echo)
|
||||
counts = []
|
||||
|
||||
async def node(state: MessagesState) -> dict:
|
||||
for _ in range(2):
|
||||
stream = sub.astream({"messages": []}, stream_mode=["messages", "tools"])
|
||||
counts.append(len([part async for part in stream]))
|
||||
return {}
|
||||
|
||||
await _one_node_graph(node).ainvoke({"messages": []})
|
||||
|
||||
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
|
||||
|
||||
|
||||
def _push(state: MessagesState) -> dict:
|
||||
push_message(AIMessage(id="pushed", content="hi"), state_key=None)
|
||||
return {}
|
||||
|
||||
|
||||
def _namespaces(graph: CompiledStateGraph, subgraphs: bool) -> list[tuple[str, ...]]:
|
||||
parts = graph.stream(
|
||||
{"messages": []}, stream_mode="messages", subgraphs=subgraphs, version="v2"
|
||||
)
|
||||
return [_segment_names(part["ns"]) for part in parts]
|
||||
|
||||
|
||||
def _streams_pushing_graph(received: list) -> Any:
|
||||
pushing = _one_node_graph(_push)
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
for message, _ in pushing.stream({"messages": []}, stream_mode="messages"):
|
||||
received.append(message.id)
|
||||
return {}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def test_pushed_message_streams_under_the_namespace_of_its_graph() -> None:
|
||||
pushing = _one_node_graph(_push)
|
||||
parent = StateGraph(MessagesState).add_node("sub", pushing).add_edge(START, "sub")
|
||||
|
||||
assert _namespaces(pushing, subgraphs=True) == [()]
|
||||
assert _namespaces(parent.compile(), subgraphs=False) == []
|
||||
assert _namespaces(parent.compile(), subgraphs=True) == [("sub",)]
|
||||
|
||||
|
||||
def test_pushed_message_streams_to_the_node_that_streams_its_graph() -> None:
|
||||
received: list[str] = []
|
||||
graph = _one_node_graph(_streams_pushing_graph(received))
|
||||
|
||||
assert _namespaces(graph, subgraphs=False) == []
|
||||
assert received == ["pushed"]
|
||||
|
||||
|
||||
def test_pushed_message_streams_to_every_stream_that_takes_it() -> None:
|
||||
received: list[str] = []
|
||||
graph = _one_node_graph(_streams_pushing_graph(received))
|
||||
|
||||
assert _namespaces(graph, subgraphs=True) == [("node",)]
|
||||
assert received == ["pushed"]
|
||||
Generated
+3
-3
@@ -3419,11 +3419,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "84.1.0"
|
||||
version = "83.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/73/eb/7e5f628ee2eb168fecbf8026de7856b2543d6d609cc85f086b9a91ac93d6/setuptools-84.1.0.tar.gz", hash = "sha256:23657205b63d4d31395ab917815247f694aafad28cd21316bf5fd45a51048ea0", size = 1171783, upload-time = "2026-10-19T20:11:43.489Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/18/02/b4bbb90f8f29e334898ff45c6d9b1fffaa6f378017052b3023a06262b5f7/setuptools-84.1.0-py3-none-any.whl", hash = "sha256:5231689c6ac8156cd155cfb98c3fd6fb252c3e7e9bd3d78b64d3c55f132624c9", size = 819380, upload-time = "2026-10-19T20:11:41.684Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -14,14 +14,9 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_araise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.errors import _araise_for_status_typed
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, aiter_lines_raw
|
||||
|
||||
@@ -151,8 +146,7 @@ class HttpClient:
|
||||
on_response: Callable[[httpx.Response], None] | None = None,
|
||||
reconnect_limit: int = 5,
|
||||
) -> Any:
|
||||
"""Send a request that reconnects to its Location header if the response is
|
||||
cut off, with backoff between attempts."""
|
||||
"""Send a request that automatically reconnects to Location header."""
|
||||
request_headers, content = await _aencode_json(json)
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
@@ -182,27 +176,13 @@ class HttpClient:
|
||||
stacklevel=2,
|
||||
)
|
||||
await r.aclose()
|
||||
for attempt in range(reconnect_limit):
|
||||
await asyncio.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
|
||||
# A restart or a deploy looks like a refused connection, a
|
||||
# dropped body, a 429 or a 5xx; other errors won't go away.
|
||||
try:
|
||||
return await self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=0,
|
||||
)
|
||||
except httpx.HTTPStatusError as err:
|
||||
status = err.response.status_code
|
||||
if attempt == reconnect_limit - 1 or (
|
||||
status != 429 and status < 500
|
||||
):
|
||||
raise
|
||||
except httpx.TransportError:
|
||||
if attempt == reconnect_limit - 1:
|
||||
raise
|
||||
return await self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=reconnect_limit - 1,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
@@ -245,96 +225,70 @@ class HttpClient:
|
||||
}
|
||||
if last_event_id:
|
||||
current_headers["Last-Event-ID"] = last_event_id
|
||||
if reconnect_path is not None and not any(
|
||||
key.lower() == "last-event-id" and value
|
||||
for key, value in current_headers.items()
|
||||
):
|
||||
warnings.warn(
|
||||
"Reconnecting without a Last-Event-ID, so events sent while "
|
||||
"disconnected are not replayed. Only runs created with "
|
||||
"stream_resumable=True send the event IDs a reconnect resumes "
|
||||
"from.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
current_method = method if reconnect_path is None else "GET"
|
||||
current_content = content if reconnect_path is None else None
|
||||
current_params = params if reconnect_path is None else None
|
||||
|
||||
reconnect_cause: httpx.HTTPError | None = None
|
||||
try:
|
||||
async with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
await _araise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
retry = False
|
||||
async with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
await _araise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
|
||||
# parse SSE
|
||||
decoder = SSEDecoder(last_event_id=last_event_id)
|
||||
try:
|
||||
async for line in aiter_lines_raw(res):
|
||||
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or decoder.has_data:
|
||||
# Progress since the last drop earns a fresh
|
||||
# set of reconnect attempts.
|
||||
reconnect_attempts = 0
|
||||
yield sse
|
||||
except httpx.HTTPError as err:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
# parse SSE
|
||||
decoder = SSEDecoder()
|
||||
try:
|
||||
async for line in aiter_lines_raw(res):
|
||||
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or decoder.has_data:
|
||||
# decoder.decode(b"") flushes the in-flight event and may
|
||||
# return an empty placeholder when there is no pending
|
||||
# message. Skip these no-op events so the stream doesn't
|
||||
# emit a trailing blank item after reconnects.
|
||||
if sse.event or sse.data is not None:
|
||||
yield sse
|
||||
except (httpx.TransportError, RateLimitError, InternalServerError) as err:
|
||||
# A reconnect that can't open (connection refused, 429, 5xx or a
|
||||
# non-SSE body) usually means the server is restarting or
|
||||
# mid-deploy, so retry it like a drop. Before the server sends a
|
||||
# Location there is nothing to reconnect to, so raise.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
if reconnect_cause is not None:
|
||||
except httpx.HTTPError:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
retry = True
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
# decoder.decode(b"") flushes the in-flight event and may
|
||||
# return an empty placeholder when there is no pending
|
||||
# message. Skip these no-op events so the stream doesn't
|
||||
# emit a trailing blank item after reconnects.
|
||||
yield sse
|
||||
if retry:
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts > max_reconnect_attempts:
|
||||
raise httpx.TransportError(
|
||||
"Exceeded maximum SSE reconnection attempts"
|
||||
) from reconnect_cause
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(reconnect_attempts - 1, base=1.0, cap=5.0)
|
||||
)
|
||||
)
|
||||
continue
|
||||
break
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal, TypedDict, cast
|
||||
@@ -24,7 +25,7 @@ from langchain_core.language_models.chat_model_stream import AsyncChatModelStrea
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._async.http import HttpClient
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param, _reconnect_delay
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param
|
||||
from langgraph_sdk.schema import LangSmithTracing, QueryParamTypes
|
||||
from langgraph_sdk.stream.controller import _SeenEventIds
|
||||
from langgraph_sdk.stream.decoders import (
|
||||
@@ -1690,13 +1691,11 @@ class AsyncThreadStream:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt `attempt`."""
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._shared_reconnect_backoff_base,
|
||||
cap=self._shared_reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
base = self._shared_reconnect_backoff_base
|
||||
cap = self._shared_reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
await asyncio.sleep(delay + jitter)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a post-ready transport drop.
|
||||
|
||||
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import tzinfo
|
||||
@@ -195,15 +194,6 @@ def _validate_reconnect_location(base_url: httpx.URL, location: str) -> str:
|
||||
return location
|
||||
|
||||
|
||||
def _reconnect_delay(attempt: int, *, base: float, cap: float) -> float:
|
||||
"""Seconds to wait before reconnect attempt `attempt` (0-based).
|
||||
|
||||
Exponential backoff capped at `cap`, plus up to 25% jitter.
|
||||
"""
|
||||
delay = min(cap, base * (2**attempt))
|
||||
return delay + random.uniform(0, delay * 0.25)
|
||||
|
||||
|
||||
def _provided_vals(d: Mapping[str, Any]) -> dict[str, Any]:
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
|
||||
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from typing import Any, cast
|
||||
@@ -14,14 +13,9 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_raise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.errors import _raise_for_status_typed
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, iter_lines_raw
|
||||
|
||||
@@ -152,8 +146,7 @@ class SyncHttpClient:
|
||||
on_response: Callable[[httpx.Response], None] | None = None,
|
||||
reconnect_limit: int = 5,
|
||||
) -> Any:
|
||||
"""Send a request that reconnects to its Location header if the response is
|
||||
cut off, with backoff between attempts."""
|
||||
"""Send a request that automatically reconnects to Location header."""
|
||||
request_headers, content = _encode_json(json)
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
@@ -183,27 +176,13 @@ class SyncHttpClient:
|
||||
stacklevel=2,
|
||||
)
|
||||
r.close()
|
||||
for attempt in range(reconnect_limit):
|
||||
time.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
|
||||
# A restart or a deploy looks like a refused connection, a
|
||||
# dropped body, a 429 or a 5xx; other errors won't go away.
|
||||
try:
|
||||
return self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=0,
|
||||
)
|
||||
except httpx.HTTPStatusError as err:
|
||||
status = err.response.status_code
|
||||
if attempt == reconnect_limit - 1 or (
|
||||
status != 429 and status < 500
|
||||
):
|
||||
raise
|
||||
except httpx.TransportError:
|
||||
if attempt == reconnect_limit - 1:
|
||||
raise
|
||||
return self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=reconnect_limit - 1,
|
||||
)
|
||||
|
||||
def stream(
|
||||
self,
|
||||
@@ -248,91 +227,67 @@ class SyncHttpClient:
|
||||
}
|
||||
if last_event_id:
|
||||
current_headers["Last-Event-ID"] = last_event_id
|
||||
if reconnect_path is not None and not any(
|
||||
key.lower() == "last-event-id" and value
|
||||
for key, value in current_headers.items()
|
||||
):
|
||||
warnings.warn(
|
||||
"Reconnecting without a Last-Event-ID, so events sent while "
|
||||
"disconnected are not replayed. Only runs created with "
|
||||
"stream_resumable=True send the event IDs a reconnect resumes "
|
||||
"from.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
current_method = method if reconnect_path is None else "GET"
|
||||
current_content = content if reconnect_path is None else None
|
||||
current_params = params if reconnect_path is None else None
|
||||
|
||||
reconnect_cause: httpx.HTTPError | None = None
|
||||
try:
|
||||
with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
_raise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
retry = False
|
||||
with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
_raise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
|
||||
decoder = SSEDecoder(last_event_id=last_event_id)
|
||||
try:
|
||||
for line in iter_lines_raw(res):
|
||||
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or decoder.has_data:
|
||||
# Progress since the last drop earns a fresh
|
||||
# set of reconnect attempts.
|
||||
reconnect_attempts = 0
|
||||
yield sse
|
||||
except httpx.HTTPError as err:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
decoder = SSEDecoder()
|
||||
try:
|
||||
for line in iter_lines_raw(res):
|
||||
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or decoder.has_data:
|
||||
# See async stream implementation for rationale on
|
||||
# skipping empty flush events.
|
||||
if sse.event or sse.data is not None:
|
||||
yield sse
|
||||
except (httpx.TransportError, RateLimitError, InternalServerError) as err:
|
||||
# A reconnect that can't open (connection refused, 429, 5xx or a
|
||||
# non-SSE body) usually means the server is restarting or
|
||||
# mid-deploy, so retry it like a drop. Before the server sends a
|
||||
# Location there is nothing to reconnect to, so raise.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
if reconnect_cause is not None:
|
||||
except httpx.HTTPError:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
retry = True
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
# See async stream implementation for rationale on
|
||||
# skipping empty flush events.
|
||||
yield sse
|
||||
if retry:
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts > max_reconnect_attempts:
|
||||
raise httpx.TransportError(
|
||||
"Exceeded maximum SSE reconnection attempts"
|
||||
) from reconnect_cause
|
||||
time.sleep(_reconnect_delay(reconnect_attempts - 1, base=1.0, cap=5.0))
|
||||
)
|
||||
continue
|
||||
break
|
||||
|
||||
|
||||
@@ -30,14 +30,18 @@ class BytesLineDecoder:
|
||||
# See https://docs.python.org/3/glossary.html#term-universal-newlines
|
||||
NEWLINE_CHARS = b"\n\r"
|
||||
|
||||
if not text:
|
||||
return []
|
||||
if self.trailing_cr and text.startswith(b"\n"):
|
||||
text = text[1:]
|
||||
self.trailing_cr = text.endswith(b"\r")
|
||||
# We always push a trailing `\r` into the next decode iteration.
|
||||
if self.trailing_cr:
|
||||
text = b"\r" + text
|
||||
self.trailing_cr = False
|
||||
if text.endswith(b"\r"):
|
||||
self.trailing_cr = True
|
||||
text = text[:-1]
|
||||
|
||||
if not text:
|
||||
return []
|
||||
# NOTE: the edge case input of empty text doesn't occur in practice,
|
||||
# because other httpx internals filter out this value
|
||||
return [] # pragma: no cover
|
||||
|
||||
trailing_newline = text[-1] in NEWLINE_CHARS
|
||||
lines = cast(list[BytesLike], text.splitlines())
|
||||
@@ -62,7 +66,7 @@ class BytesLineDecoder:
|
||||
return lines
|
||||
|
||||
def flush(self) -> list[BytesLike]:
|
||||
if not self.buffer:
|
||||
if not self.buffer and not self.trailing_cr:
|
||||
return []
|
||||
|
||||
lines: list[BytesLike] = [self.buffer]
|
||||
@@ -72,25 +76,18 @@ class BytesLineDecoder:
|
||||
|
||||
|
||||
class SSEDecoder:
|
||||
def __init__(self, *, last_event_id: str | None = None) -> None:
|
||||
def __init__(self) -> None:
|
||||
self._event = ""
|
||||
self._data = bytearray()
|
||||
self._last_event_id = last_event_id
|
||||
self._last_event_id: str | None = None
|
||||
self._retry: int | None = None
|
||||
self._at_stream_start = True
|
||||
self._has_data = False
|
||||
self._id_changed = False
|
||||
|
||||
@property
|
||||
def has_data(self) -> bool:
|
||||
"""Whether the last decoded block contained a data field."""
|
||||
return self._has_data
|
||||
|
||||
@property
|
||||
def last_event_id(self) -> str | None:
|
||||
"""Return the last event identifier that was seen.
|
||||
|
||||
`None` until an ID is supplied or received, and `""` after an empty
|
||||
`None` until the stream sends an `id` field, and `""` after an empty
|
||||
one, which resets the reconnect cursor.
|
||||
"""
|
||||
|
||||
@@ -104,26 +101,17 @@ class SSEDecoder:
|
||||
line = line.removeprefix(b"\xef\xbb\xbf")
|
||||
|
||||
if not line:
|
||||
self._has_data = bool(self._data)
|
||||
if (
|
||||
not self._event
|
||||
and not self._data
|
||||
and not self._id_changed
|
||||
and self._last_event_id is None
|
||||
and self._retry is None
|
||||
):
|
||||
return None
|
||||
|
||||
body = self._data[:-1]
|
||||
data = None
|
||||
if body and not body.isspace():
|
||||
try:
|
||||
data = orjson.loads(body)
|
||||
except orjson.JSONDecodeError:
|
||||
data = orjson.loads(body.decode(errors="replace"))
|
||||
|
||||
sse = StreamPart(
|
||||
event=self._event,
|
||||
data=data, # ty: ignore[invalid-argument-type]
|
||||
data=orjson.loads(self._data) if self._data else None, # ty: ignore[invalid-argument-type]
|
||||
id=self._last_event_id or None,
|
||||
)
|
||||
|
||||
@@ -131,7 +119,6 @@ class SSEDecoder:
|
||||
self._event = ""
|
||||
self._data = bytearray()
|
||||
self._retry = None
|
||||
self._id_changed = False
|
||||
|
||||
return sse
|
||||
|
||||
@@ -144,16 +131,16 @@ class SSEDecoder:
|
||||
value = value[1:]
|
||||
|
||||
if fieldname == b"event":
|
||||
self._event = value.decode(errors="replace")
|
||||
self._event = value.decode()
|
||||
elif fieldname == b"data":
|
||||
if self._data:
|
||||
self._data.extend(b"\n")
|
||||
self._data.extend(value)
|
||||
self._data.extend(b"\n")
|
||||
elif fieldname == b"id":
|
||||
if b"\0" in value:
|
||||
pass
|
||||
else:
|
||||
self._last_event_id = value.decode(errors="replace")
|
||||
self._id_changed = True
|
||||
self._last_event_id = value.decode()
|
||||
elif fieldname == b"retry":
|
||||
with contextlib.suppress(TypeError, ValueError):
|
||||
self._retry = int(value)
|
||||
|
||||
@@ -15,6 +15,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
from collections import OrderedDict
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
@@ -22,7 +23,6 @@ from typing import Any
|
||||
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._shared.utilities import _reconnect_delay
|
||||
from langgraph_sdk.stream.transport import AsyncProtocolTransport, EventStreamHandle
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -275,13 +275,11 @@ class StreamController:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt *attempt*."""
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
base = self._reconnect_backoff_base
|
||||
cap = self._reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
await asyncio.sleep(delay + jitter)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
@@ -12,7 +13,6 @@ from typing import Any
|
||||
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._shared.utilities import _reconnect_delay
|
||||
from langgraph_sdk.stream.errors import (
|
||||
is_non_retryable_stream_error,
|
||||
stream_error_frame,
|
||||
@@ -282,13 +282,11 @@ class SyncStreamController:
|
||||
|
||||
def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
|
||||
time.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
base = self._reconnect_backoff_base
|
||||
cap = self._reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
time.sleep(delay + jitter)
|
||||
|
||||
def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -1,17 +0,0 @@
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_sleep(monkeypatch):
|
||||
"""Record backoff sleeps without actually waiting."""
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_async_sleep(d):
|
||||
sleeps.append(d)
|
||||
|
||||
monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
|
||||
monkeypatch.setattr(time, "sleep", lambda d: sleeps.append(d))
|
||||
return sleeps
|
||||
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from collections.abc import Iterator, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -13,7 +13,6 @@ from langgraph_sdk._async.runs import _wrap_stream_v2
|
||||
from langgraph_sdk._shared.utilities import _sse_to_v2_dict
|
||||
from langgraph_sdk._sync.runs import _wrap_stream_v2_sync
|
||||
from langgraph_sdk.client import HttpClient, SyncHttpClient
|
||||
from langgraph_sdk.errors import InternalServerError, NotFoundError
|
||||
from langgraph_sdk.schema import (
|
||||
CheckpointPayload,
|
||||
CheckpointsStreamPart,
|
||||
@@ -131,110 +130,6 @@ def test_sse_decoder_strips_a_leading_bom():
|
||||
assert decoder.decode(b"") == StreamPart(event="values", data=1, id=None) # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
def test_control_blocks_do_not_repeat_after_a_cursor_update():
|
||||
decoder = SSEDecoder()
|
||||
decoder.decode(b"id: cursor")
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=None, id="cursor") # ty: ignore[invalid-argument-type]
|
||||
decoder.decode(b": heartbeat")
|
||||
assert decoder.decode(b"") is None
|
||||
assert decoder.decode(b"") is None
|
||||
decoder.decode(b"id:")
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=None, id=None) # ty: ignore[invalid-argument-type]
|
||||
assert decoder.last_event_id == ""
|
||||
assert decoder.decode(b"") is None
|
||||
|
||||
|
||||
def test_cr_lines_are_emitted_without_waiting_for_another_chunk():
|
||||
decoder = BytesLineDecoder()
|
||||
assert decoder.decode(b"data: 1\r\r") == [b"data: 1", b""]
|
||||
assert decoder.decode(b"\n") == []
|
||||
assert decoder.decode(b"data: 2\r") == [b"data: 2"]
|
||||
assert decoder.decode(b"") == []
|
||||
assert decoder.decode(b"\n\r") == [b""]
|
||||
assert decoder.flush() == []
|
||||
|
||||
|
||||
def test_empty_data_lines_preserve_their_newlines():
|
||||
decoder = SSEDecoder()
|
||||
for line in (b"data", b"data: 1", b"data:"):
|
||||
decoder.decode(line)
|
||||
assert decoder._data == b"\n1\n\n"
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=1, id=None) # ty: ignore[invalid-argument-type]
|
||||
decoder.decode(b"data:")
|
||||
decoder.decode(b"data:")
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=None, id=None) # ty: ignore[invalid-argument-type]
|
||||
assert decoder.has_data
|
||||
assert decoder.decode(b"") is None
|
||||
assert not decoder.has_data
|
||||
|
||||
|
||||
def test_malformed_utf8_is_replaced_in_fields_and_json_strings():
|
||||
decoder = SSEDecoder()
|
||||
for line in (b"event: \xff", b"id: \xff", b'data: "\xff"'):
|
||||
decoder.decode(line)
|
||||
assert decoder.decode(b"") == StreamPart(event="\ufffd", data="\ufffd", id="\ufffd") # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_null_and_empty_data_are_not_dropped(asynchronous: bool):
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
content=b"data: null\n\ndata:\n\n",
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = [part async for part in HttpClient(client).stream("/stream", "GET")]
|
||||
else:
|
||||
with httpx.Client(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = list(SyncHttpClient(client).stream("/stream", "GET"))
|
||||
assert parts == [StreamPart(event="", data=None), StreamPart(event="", data=None)] # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_cursor_is_inherited_after_reconnect(asynchronous: bool):
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if len(requests) == 1:
|
||||
stream_cls = AsyncListByteStream if asynchronous else ListByteStream
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream", "Location": "/reconnect"},
|
||||
stream=stream_cls(
|
||||
[b"id: cursor\nevent: values\ndata: 1\n\n"],
|
||||
httpx.RemoteProtocolError("incomplete chunked read"),
|
||||
),
|
||||
)
|
||||
assert (b"Last-Event-ID", b"cursor") in request.headers.raw
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
content=b"event: values\ndata: 2\n\n",
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = [part async for part in HttpClient(client).stream("/stream", "GET")]
|
||||
else:
|
||||
with httpx.Client(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = list(SyncHttpClient(client).stream("/stream", "GET"))
|
||||
assert [part.id for part in parts] == ["cursor", "cursor"]
|
||||
|
||||
|
||||
# --- HTTP client streaming ---
|
||||
|
||||
|
||||
@@ -281,7 +176,6 @@ def test_sync_http_client_stream_flushes_trailing_event():
|
||||
assert parts == [StreamPart(event="foo", data={"bar": 1})]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_http_client_stream_recovers_after_disconnect():
|
||||
reconnect_path = "/reconnect"
|
||||
first_chunks = [
|
||||
@@ -347,7 +241,6 @@ def test_sync_http_client_stream_recovers_after_disconnect():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_http_client_stream_recovers_after_disconnect():
|
||||
reconnect_path = "/reconnect"
|
||||
first_chunks = [
|
||||
@@ -482,33 +375,28 @@ _CALLER_CURSORS = pytest.mark.parametrize(
|
||||
|
||||
@_EMPTY_IDS
|
||||
@_CALLER_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
empty_id: bytes, headers: dict[str, str] | None
|
||||
):
|
||||
with pytest.warns(UserWarning, match="Last-Event-ID"):
|
||||
reconnect = _sync_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
reconnect = _sync_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
|
||||
assert "last-event-id" not in reconnect.headers
|
||||
|
||||
|
||||
@_EMPTY_IDS
|
||||
@_CALLER_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
empty_id: bytes, headers: dict[str, str] | None
|
||||
):
|
||||
with pytest.warns(UserWarning, match="Last-Event-ID"):
|
||||
reconnect = await _async_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
reconnect = await _async_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
|
||||
assert "last-event-id" not in reconnect.headers
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
reconnect = _sync_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n"], {"last-event-id": "5"}
|
||||
@@ -517,7 +405,6 @@ def test_sync_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
assert reconnect.headers.get_list("last-event-id") == ["1"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
reconnect = await _async_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n"], {"last-event-id": "5"}
|
||||
@@ -526,14 +413,12 @@ async def test_async_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
assert reconnect.headers.get_list("last-event-id") == ["1"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id():
|
||||
reconnect = _sync_reconnect([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"})
|
||||
|
||||
assert reconnect.headers.get_list("last-event-id") == ["5"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id():
|
||||
reconnect = await _async_reconnect(
|
||||
[b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}
|
||||
@@ -542,275 +427,6 @@ async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id
|
||||
assert reconnect.headers.get_list("last-event-id") == ["5"]
|
||||
|
||||
|
||||
# A reconnect that sends no Last-Event-ID can't replay what was sent while it was
|
||||
# disconnected, so it warns.
|
||||
_RECONNECT_CURSORS = pytest.mark.parametrize(
|
||||
("chunks", "headers", "warns"),
|
||||
[
|
||||
([b"event: values\ndata: 1\n\n"], None, True),
|
||||
([b"id: 1\nevent: values\ndata: 1\n\n"], None, False),
|
||||
([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}, False),
|
||||
],
|
||||
ids=["no-cursor", "stream-cursor", "caller-cursor"],
|
||||
)
|
||||
|
||||
|
||||
@_RECONNECT_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_reconnect_warns_only_without_a_cursor(
|
||||
chunks: list[bytes],
|
||||
headers: dict[str, str] | None,
|
||||
warns: bool,
|
||||
recwarn: pytest.WarningsRecorder,
|
||||
):
|
||||
_sync_reconnect(chunks, headers)
|
||||
|
||||
assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
|
||||
|
||||
|
||||
@_RECONNECT_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_reconnect_warns_only_without_a_cursor(
|
||||
chunks: list[bytes],
|
||||
headers: dict[str, str] | None,
|
||||
warns: bool,
|
||||
recwarn: pytest.WarningsRecorder,
|
||||
):
|
||||
await _async_reconnect(chunks, headers)
|
||||
|
||||
assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
|
||||
|
||||
|
||||
# --- reconnects survive a server restart or a deploy ---
|
||||
|
||||
SSE = {"Content-Type": "text/event-stream"}
|
||||
DROP = httpx.RemoteProtocolError("incomplete chunked read")
|
||||
|
||||
|
||||
def _event(i: int) -> bytes:
|
||||
return f"id: {i}\nevent: values\ndata: {i}\n\n".encode()
|
||||
|
||||
|
||||
def _scripted(steps: list[Callable[[httpx.Request], httpx.Response]]):
|
||||
"""Handler that plays one step per request and records each request."""
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
if len(seen) > len(steps):
|
||||
raise AssertionError("unexpected request")
|
||||
return steps[len(seen) - 1](request)
|
||||
|
||||
return handler, seen
|
||||
|
||||
|
||||
def _ok(stream_cls, chunks, *, drop=False, location=None):
|
||||
headers = dict(SSE)
|
||||
if location:
|
||||
headers["Location"] = location
|
||||
|
||||
def step(_request):
|
||||
return httpx.Response(
|
||||
200, headers=headers, stream=stream_cls(chunks, DROP if drop else None)
|
||||
)
|
||||
|
||||
return step
|
||||
|
||||
|
||||
def _status(code):
|
||||
return lambda request: httpx.Response(
|
||||
code, json={"detail": "upstream"}, request=request
|
||||
)
|
||||
|
||||
|
||||
def _connect_error(request):
|
||||
raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
|
||||
|
||||
|
||||
def _not_sse(request):
|
||||
return httpx.Response(
|
||||
200, headers={"Content-Type": "text/html"}, text="<html>", request=request
|
||||
)
|
||||
|
||||
|
||||
def _transient_steps(stream_cls, failure):
|
||||
"""POST drops after event 1; first reconnect fails; second succeeds."""
|
||||
return [
|
||||
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
|
||||
failure,
|
||||
_ok(stream_cls, [_event(2), b"event: end\ndata: null\n\n"]),
|
||||
]
|
||||
|
||||
|
||||
def _long_stream_steps(stream_cls, drops):
|
||||
"""POST + `drops` reconnects, each delivering one event then dropping."""
|
||||
steps = [_ok(stream_cls, [_event(0)], drop=True, location="/reconnect")]
|
||||
for i in range(1, drops):
|
||||
steps.append(_ok(stream_cls, [_event(i)], drop=True))
|
||||
steps.append(_ok(stream_cls, [_event(drops), b"event: end\ndata: null\n\n"]))
|
||||
return steps
|
||||
|
||||
|
||||
def _run_sync(steps):
|
||||
handler, seen = _scripted(steps)
|
||||
with httpx.Client(
|
||||
transport=httpx.MockTransport(handler), base_url="https://example.com"
|
||||
) as client:
|
||||
parts = list(SyncHttpClient(client).stream("/stream", "POST", json={}))
|
||||
return parts, seen
|
||||
|
||||
|
||||
async def _run_async(steps):
|
||||
sync_handler, seen = _scripted(steps)
|
||||
|
||||
async def handler(request):
|
||||
return sync_handler(request)
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(handler), base_url="https://example.com"
|
||||
) as client:
|
||||
parts = [p async for p in HttpClient(client).stream("/stream", "POST", json={})]
|
||||
return parts, seen
|
||||
|
||||
|
||||
FAILURES = {
|
||||
"429": _status(429),
|
||||
"502": _status(502),
|
||||
"503": _status(503),
|
||||
"connect_error": _connect_error,
|
||||
"not_sse": _not_sse,
|
||||
}
|
||||
|
||||
|
||||
# A reconnect that hits a 429 / 502 / 503 / connect error / non-SSE body, then
|
||||
# succeeds, should recover.
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("failure", FAILURES.values(), ids=FAILURES.keys())
|
||||
def test_sync_reconnect_recovers_from_transient_failure(failure):
|
||||
parts, seen = _run_sync(_transient_steps(ListByteStream, failure))
|
||||
assert [p.event for p in parts] == ["values", "values", "end"]
|
||||
assert len(seen) == 3
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("failure", FAILURES.values(), ids=FAILURES.keys())
|
||||
async def test_async_reconnect_recovers_from_transient_failure(failure):
|
||||
parts, seen = await _run_async(_transient_steps(AsyncListByteStream, failure))
|
||||
assert [p.event for p in parts] == ["values", "values", "end"]
|
||||
assert len(seen) == 3
|
||||
|
||||
|
||||
# Any other 4xx on a reconnect won't go away by retrying, so it raises.
|
||||
|
||||
|
||||
def _client_error_steps(stream_cls):
|
||||
return [
|
||||
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
|
||||
_status(404),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_reconnect_does_not_retry_a_client_error():
|
||||
with pytest.raises(NotFoundError):
|
||||
_run_sync(_client_error_steps(ListByteStream))
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_reconnect_does_not_retry_a_client_error():
|
||||
with pytest.raises(NotFoundError):
|
||||
await _run_async(_client_error_steps(AsyncListByteStream))
|
||||
|
||||
|
||||
# The first request isn't a reconnect, so it raises even when it names a Location.
|
||||
|
||||
|
||||
def _first_request_503(request):
|
||||
return httpx.Response(
|
||||
503,
|
||||
headers={"Location": "/reconnect"},
|
||||
json={"detail": "upstream"},
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_first_request_failure_is_not_retried():
|
||||
with pytest.raises(InternalServerError):
|
||||
_run_sync([_first_request_503])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_first_request_failure_is_not_retried():
|
||||
with pytest.raises(InternalServerError):
|
||||
await _run_async([_first_request_503])
|
||||
|
||||
|
||||
# A server that stays down gets five reconnects, backing off 1, 2, 4, 5 and 5s
|
||||
# (each plus up to 25% jitter), so about 20s, before the stream gives up. The
|
||||
# error it gives up with is caused by the last failure, so callers see why.
|
||||
|
||||
GIVE_UP = "Exceeded maximum SSE reconnection attempts"
|
||||
|
||||
|
||||
def _server_down_steps(stream_cls):
|
||||
return [
|
||||
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
|
||||
*[_status(502)] * 5,
|
||||
]
|
||||
|
||||
|
||||
def _assert_backed_off_for_about_20s(sleeps: list[float]) -> None:
|
||||
assert len(sleeps) == 5
|
||||
for slept, delay in zip(sleeps, [1, 2, 4, 5, 5], strict=True):
|
||||
assert delay <= slept <= delay * 1.25
|
||||
|
||||
|
||||
def test_sync_reconnect_gives_up_after_five_backed_off_attempts(no_sleep):
|
||||
with pytest.raises(httpx.TransportError, match=GIVE_UP) as exc_info:
|
||||
_run_sync(_server_down_steps(ListByteStream))
|
||||
assert isinstance(exc_info.value.__cause__, InternalServerError)
|
||||
_assert_backed_off_for_about_20s(no_sleep)
|
||||
|
||||
|
||||
async def test_async_reconnect_gives_up_after_five_backed_off_attempts(no_sleep):
|
||||
with pytest.raises(httpx.TransportError, match=GIVE_UP) as exc_info:
|
||||
await _run_async(_server_down_steps(AsyncListByteStream))
|
||||
assert isinstance(exc_info.value.__cause__, InternalServerError)
|
||||
_assert_backed_off_for_about_20s(no_sleep)
|
||||
|
||||
|
||||
# A long stream that drops 6 times, each reconnect succeeding and making
|
||||
# progress, should not run out of attempts.
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_long_stream_survives_many_spaced_out_drops():
|
||||
parts, _ = _run_sync(_long_stream_steps(ListByteStream, drops=6))
|
||||
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_long_stream_survives_many_spaced_out_drops():
|
||||
parts, _ = await _run_async(_long_stream_steps(AsyncListByteStream, drops=6))
|
||||
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
|
||||
|
||||
|
||||
# Reconnects should back off rather than hammer the server.
|
||||
|
||||
|
||||
def test_sync_reconnect_backs_off(no_sleep):
|
||||
_run_sync(_long_stream_steps(ListByteStream, drops=2))
|
||||
assert len(no_sleep) == 2 and all(d > 0 for d in no_sleep)
|
||||
|
||||
|
||||
async def test_async_reconnect_backs_off(no_sleep):
|
||||
await _run_async(_long_stream_steps(AsyncListByteStream, drops=2))
|
||||
assert len(no_sleep) == 2 and all(d > 0 for d in no_sleep)
|
||||
|
||||
|
||||
# --- _sse_to_v2_dict conversion ---
|
||||
|
||||
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
||||
from typing import Any
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -25,16 +24,7 @@ class _BodyDropsMidRead(httpx.SyncByteStream, httpx.AsyncByteStream):
|
||||
raise httpx.ReadError("connection dropped")
|
||||
|
||||
|
||||
def _done(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json={"done": True})
|
||||
|
||||
|
||||
def _wait_handler(
|
||||
requests: list[httpx.Request],
|
||||
reconnects: Sequence[Callable[[httpx.Request], httpx.Response]] = (_done,),
|
||||
):
|
||||
"""POST drops mid-body; the reconnects answer from `reconnects` in order."""
|
||||
|
||||
def _wait_handler(requests: list[httpx.Request]):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if request.method == "POST":
|
||||
@@ -43,7 +33,7 @@ def _wait_handler(
|
||||
headers={"Location": "/threads/t1/runs/r1/join"},
|
||||
stream=_BodyDropsMidRead(),
|
||||
)
|
||||
return reconnects[len(requests) - 2](request)
|
||||
return httpx.Response(200, json={"done": True})
|
||||
|
||||
return handler
|
||||
|
||||
@@ -56,7 +46,6 @@ def _assert_bodyless_get(request: httpx.Request) -> None:
|
||||
assert request.read() == b""
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_wait_reconnects_with_a_bodyless_get():
|
||||
requests: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(_wait_handler(requests))
|
||||
@@ -69,7 +58,6 @@ def test_sync_wait_reconnects_with_a_bodyless_get():
|
||||
_assert_bodyless_get(requests[1])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_wait_reconnects_with_a_bodyless_get():
|
||||
requests: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(_wait_handler(requests))
|
||||
@@ -84,89 +72,6 @@ async def test_async_wait_reconnects_with_a_bodyless_get():
|
||||
_assert_bodyless_get(requests[1])
|
||||
|
||||
|
||||
def _connect_error(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
|
||||
|
||||
|
||||
async def _wait(
|
||||
handler: Callable[[httpx.Request], httpx.Response], *, asynchronous: bool
|
||||
) -> Any:
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
return await RunsClient(HttpClient(client)).wait("t1", "agent", input={})
|
||||
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||
return SyncRunsClient(SyncHttpClient(client)).wait("t1", "agent", input={})
|
||||
|
||||
|
||||
# What a server restart or a deploy looks like to the reconnect GET: each one is
|
||||
# retried after a backoff.
|
||||
RECONNECT_FAILURES = {
|
||||
"429": lambda request: httpx.Response(429, request=request),
|
||||
"503": lambda request: httpx.Response(503, request=request),
|
||||
"connect_error": _connect_error,
|
||||
"body_drops": lambda _request: httpx.Response(
|
||||
200,
|
||||
headers={"Location": "/threads/t1/runs/r1/join"},
|
||||
stream=_BodyDropsMidRead(),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"failure", RECONNECT_FAILURES.values(), ids=RECONNECT_FAILURES.keys()
|
||||
)
|
||||
async def test_wait_retries_a_failed_reconnect_after_a_backoff(
|
||||
failure, asynchronous, no_sleep
|
||||
):
|
||||
requests: list[httpx.Request] = []
|
||||
with pytest.warns(UserWarning, match="attempting reconnect"):
|
||||
result = await _wait(
|
||||
_wait_handler(requests, [failure, _done]), asynchronous=asynchronous
|
||||
)
|
||||
|
||||
assert result == {"done": True}
|
||||
assert len(requests) == 3
|
||||
assert len(no_sleep) == 2 and min(no_sleep) >= 1, no_sleep
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_wait_raises_after_five_failed_reconnects(asynchronous):
|
||||
requests: list[httpx.Request] = []
|
||||
with (
|
||||
pytest.warns(UserWarning, match="attempting reconnect"),
|
||||
pytest.raises(httpx.HTTPStatusError, match="503"),
|
||||
):
|
||||
await _wait(
|
||||
_wait_handler(requests, [RECONNECT_FAILURES["503"]] * 5),
|
||||
asynchronous=asynchronous,
|
||||
)
|
||||
|
||||
assert len(requests) == 6
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_wait_does_not_retry_a_client_error(asynchronous):
|
||||
requests: list[httpx.Request] = []
|
||||
with (
|
||||
pytest.warns(UserWarning, match="attempting reconnect"),
|
||||
pytest.raises(httpx.HTTPStatusError, match="404"),
|
||||
):
|
||||
await _wait(
|
||||
_wait_handler(
|
||||
requests, [lambda request: httpx.Response(404, request=request)]
|
||||
),
|
||||
asynchronous=asynchronous,
|
||||
)
|
||||
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
def test_sync_wait_raises_the_run_error():
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
|
||||
Reference in new issue
Block a user