mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-28 18:01:30 +01:00
Compare commits
8
Commits
sdk==0.4.10
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a3e5dd739 | ||
|
|
eb8ec111b3 | ||
|
|
8c177f7388 | ||
|
|
55013146ae | ||
|
|
858891aaa1 | ||
|
|
1526b0e82d | ||
|
|
c2c73b4f91 | ||
|
|
03e73ff51b |
No files matched your search
@@ -228,6 +228,10 @@ def draw_graph(
|
||||
metadata["__interrupt"] = "before"
|
||||
elif name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "after"
|
||||
# Duck-typed so core doesn't import ToolNode from langgraph-prebuilt.
|
||||
tools_by_name = getattr(node.bound, "tools_by_name", None)
|
||||
if isinstance(tools_by_name, Mapping) and tools_by_name:
|
||||
metadata["__tools"] = list(tools_by_name)
|
||||
graph.add_node(node.bound, name, metadata=metadata or None)
|
||||
# add start node
|
||||
if START not in nodes:
|
||||
|
||||
@@ -216,7 +216,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
if not self.subgraphs and len(ns) > 0:
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
for value in _state_values(inputs):
|
||||
|
||||
@@ -34,6 +34,7 @@ from langgraph._internal._constants import (
|
||||
CONFIG_KEY_THREAD_ID,
|
||||
CONFIG_KEY_TIMED_ATTEMPT_OBSERVER,
|
||||
NS_SEP,
|
||||
RETURN,
|
||||
)
|
||||
from langgraph._internal._runnable import create_task_in_config_context
|
||||
from langgraph._internal._timeout import sync_timeout_unsupported
|
||||
@@ -733,15 +734,14 @@ async def arun_with_retry(
|
||||
if stream:
|
||||
async for _ in task.proc.astream(task.input, config):
|
||||
pass
|
||||
break
|
||||
return _return_value(task)
|
||||
return await task.proc.ainvoke(task.input, config)
|
||||
result = await _arun_with_timeout(
|
||||
task, config, resolved_timeout, attempt_ctx, stream=stream
|
||||
)
|
||||
_finish_timed_attempt(config, attempt_ctx)
|
||||
if stream:
|
||||
# if successful, end
|
||||
break
|
||||
return _return_value(task)
|
||||
return result
|
||||
except ParentCommand as exc:
|
||||
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
|
||||
@@ -831,6 +831,12 @@ async def arun_with_retry(
|
||||
config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
|
||||
|
||||
|
||||
def _return_value(task: PregelExecutableTask) -> Any:
|
||||
# Streaming drains the task instead of returning its result, so a caller
|
||||
# awaiting a `@task` needs the value the task wrote.
|
||||
return next((value for channel, value in task.writes if channel == RETURN), None)
|
||||
|
||||
|
||||
def _should_retry_on(retry_policy: RetryPolicy, exc: Exception) -> bool:
|
||||
"""Check if the given exception should be retried based on the retry policy."""
|
||||
if isinstance(retry_policy.retry_on, Sequence):
|
||||
|
||||
@@ -9,6 +9,7 @@ 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.protocol import StreamChunk
|
||||
|
||||
try:
|
||||
@@ -188,6 +189,10 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
return
|
||||
ns, tool_call_id, token = info
|
||||
self._reset_writer(token)
|
||||
# An interrupted or drained tool is paused, not failed: it runs again
|
||||
# on resume and reports its start and end then.
|
||||
if isinstance(error, (GraphInterrupt, GraphDrained)):
|
||||
return
|
||||
self.stream(
|
||||
(
|
||||
ns,
|
||||
|
||||
@@ -3310,7 +3310,9 @@ class Pregel(
|
||||
True
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
and not isinstance(h, StreamMessagesHandler)
|
||||
and not isinstance(
|
||||
h, (StreamMessagesHandler, StreamToolCallHandler)
|
||||
)
|
||||
),
|
||||
False,
|
||||
)
|
||||
|
||||
@@ -847,6 +847,11 @@ class RemoteGraph(PregelProtocol):
|
||||
# filter for what was actually requested
|
||||
if mode not in requested:
|
||||
continue
|
||||
# subgraph events may have been requested only for the parent stream.
|
||||
# messages are exempt: the server dedupes them across namespaces, so
|
||||
# dropping them here would also drop the root-level message
|
||||
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
|
||||
continue
|
||||
|
||||
if chunk.event.startswith("messages"):
|
||||
chunk = chunk._replace(data=tuple(chunk.data))
|
||||
@@ -997,6 +1002,11 @@ class RemoteGraph(PregelProtocol):
|
||||
# filter for what was actually requested
|
||||
if mode not in requested:
|
||||
continue
|
||||
# subgraph events may have been requested only for the parent stream.
|
||||
# messages are exempt: the server dedupes them across namespaces, so
|
||||
# dropping them here would also drop the root-level message
|
||||
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
|
||||
continue
|
||||
|
||||
if chunk.event.startswith("messages"):
|
||||
chunk = chunk._replace(data=tuple(chunk.data))
|
||||
|
||||
@@ -133,6 +133,11 @@
|
||||
"ToolNode"
|
||||
],
|
||||
"name": "tools"
|
||||
},
|
||||
"metadata": {
|
||||
"__tools": [
|
||||
"search_api"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -221,6 +226,11 @@
|
||||
"ToolNode"
|
||||
],
|
||||
"name": "tools"
|
||||
},
|
||||
"metadata": {
|
||||
"__tools": [
|
||||
"search_api"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
|
||||
@@ -8950,6 +8950,39 @@ def test_subgraph_streaming_sync() -> None:
|
||||
assert result["num_chunks"] == 9
|
||||
|
||||
|
||||
def test_subgraph_streaming_node_output_messages_sync() -> None:
|
||||
"""Messages returned by a node (not streamed by an LLM) are emitted when a
|
||||
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
|
||||
|
||||
def respond(state: MessagesState) -> dict:
|
||||
return {"messages": [AIMessage("hi there", id="ai1")]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("respond", respond)
|
||||
.add_edge(START, "respond")
|
||||
.compile()
|
||||
)
|
||||
msgs = {"messages": [("user", "hello")]}
|
||||
|
||||
top_level = [msg for msg, _ in subgraph.stream(msgs, stream_mode="messages")]
|
||||
|
||||
def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
|
||||
nonlocal in_node
|
||||
in_node = [
|
||||
msg for msg, _ in subgraph.stream(msgs, config, stream_mode="messages")
|
||||
]
|
||||
return {}
|
||||
|
||||
in_node: list = []
|
||||
StateGraph(MessagesState).add_node("parent", parent_node).add_edge(
|
||||
START, "parent"
|
||||
).compile().invoke(msgs)
|
||||
|
||||
assert [m.content for m in top_level] == ["hi there"]
|
||||
assert [m.content for m in in_node] == ["hi there"]
|
||||
|
||||
|
||||
def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None:
|
||||
class State(TypedDict):
|
||||
messages: list[str]
|
||||
|
||||
@@ -6834,6 +6834,31 @@ async def test_multiple_updates() -> None:
|
||||
]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_task_in_subgraph_node_returns_its_value_with_tools_mode() -> None:
|
||||
class State(TypedDict):
|
||||
out: int | None
|
||||
|
||||
@task
|
||||
async def double(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
async def node(state: State) -> State:
|
||||
return {"out": await double(21)}
|
||||
|
||||
inner = StateGraph(State).add_node("node", node).add_edge(START, "node").compile()
|
||||
graph = (
|
||||
StateGraph(State).add_node("inner", inner).add_edge(START, "inner").compile()
|
||||
)
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in graph.astream({"out": None}, stream_mode=["values", "tools"])
|
||||
]
|
||||
|
||||
assert chunks[-1] == ("values", {"out": 42})
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_falsy_return_from_task(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
"""Test with a falsy return from a task."""
|
||||
@@ -7598,6 +7623,40 @@ async def test_entrypoint_from_async_generator() -> None:
|
||||
yield "b"
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("timeout", [None, 30])
|
||||
@pytest.mark.parametrize("version", ["v1", "v2"])
|
||||
async def test_awaited_task_returns_its_value_under_astream_events(
|
||||
version: Literal["v1", "v2"], timeout: int | None
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
doubled: int
|
||||
zeroed: int
|
||||
|
||||
@task(timeout=timeout)
|
||||
async def double(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
@task(timeout=timeout)
|
||||
async def zero() -> int:
|
||||
return 0
|
||||
|
||||
async def node(state: State) -> State:
|
||||
return {"doubled": await double(21), "zeroed": await zero()}
|
||||
|
||||
graph = StateGraph(State).add_node("node", node).add_edge(START, "node").compile()
|
||||
|
||||
outputs = [
|
||||
event["data"]["output"]
|
||||
async for event in graph.astream_events(
|
||||
{"doubled": 0, "zeroed": 1}, version=version
|
||||
)
|
||||
if event["event"] == "on_chain_end" and event["name"] == "node"
|
||||
]
|
||||
|
||||
assert outputs == [{"doubled": 42, "zeroed": 0}]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_named_tasks_functional() -> None:
|
||||
class Foo:
|
||||
@@ -9444,6 +9503,44 @@ async def test_subgraph_streaming_async() -> None:
|
||||
assert result["num_chunks"] == 9
|
||||
|
||||
|
||||
async def test_subgraph_streaming_node_output_messages_async() -> None:
|
||||
"""Messages returned by a node (not streamed by an LLM) are emitted when a
|
||||
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
|
||||
|
||||
async def respond(state: MessagesState) -> dict:
|
||||
return {"messages": [AIMessage("hi there", id="ai1")]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("respond", respond)
|
||||
.add_edge(START, "respond")
|
||||
.compile()
|
||||
)
|
||||
msgs = {"messages": [("user", "hello")]}
|
||||
|
||||
top_level = [msg async for msg, _ in subgraph.astream(msgs, stream_mode="messages")]
|
||||
|
||||
async def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
|
||||
nonlocal in_node
|
||||
in_node = [
|
||||
msg
|
||||
async for msg, _ in subgraph.astream(msgs, config, stream_mode="messages")
|
||||
]
|
||||
return {}
|
||||
|
||||
in_node: list = []
|
||||
await (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("parent", parent_node)
|
||||
.add_edge(START, "parent")
|
||||
.compile()
|
||||
.ainvoke(msgs)
|
||||
)
|
||||
|
||||
assert [m.content for m in top_level] == ["hi there"]
|
||||
assert [m.content for m in in_node] == ["hi there"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_null_resume_disallowed_with_multiple_interrupts(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
|
||||
@@ -16,9 +16,11 @@ from langgraph_sdk.schema import StreamPart
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph._internal._constants import CONFIG_KEY_STREAM
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph, add_messages
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel.protocol import StreamProtocol
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
from langgraph.types import Interrupt, StateSnapshot
|
||||
from tests.any_str import AnyStr
|
||||
@@ -863,6 +865,87 @@ async def test_astream_subgraphs_keeps_caller_namespace():
|
||||
assert stream_parts == _CALLER_PREFIXED
|
||||
|
||||
|
||||
# Subgraph events requested only on behalf of a streaming parent graph
|
||||
_PARENT_STREAMED_PARTS = [
|
||||
StreamPart(event="updates|sub:abc", data={"inner": {"chunk": "inner"}}),
|
||||
StreamPart(event="messages|sub:abc", data=[{"content": "token"}, {}]),
|
||||
StreamPart(event="updates", data={"sub": {"chunk": "root"}}),
|
||||
]
|
||||
|
||||
|
||||
def _parent_streamed_config(parent_parts: list) -> RunnableConfig:
|
||||
return {
|
||||
"configurable": {
|
||||
"thread_id": "thread_1",
|
||||
"checkpoint_ns": "outer:123",
|
||||
CONFIG_KEY_STREAM: StreamProtocol(
|
||||
parent_parts.append, {"updates", "messages"}
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_stream_without_subgraphs_skips_parent_requested_subgraph_events():
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = _PARENT_STREAMED_PARTS
|
||||
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
|
||||
|
||||
for stream_mode, expected in [
|
||||
("updates", [{"sub": {"chunk": "root"}}]),
|
||||
# messages are not filtered: the server dedupes them across namespaces
|
||||
("messages", [({"content": "token"}, {})]),
|
||||
]:
|
||||
parent_parts: list = []
|
||||
stream_parts = list(
|
||||
remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
_parent_streamed_config(parent_parts),
|
||||
stream_mode=stream_mode,
|
||||
)
|
||||
)
|
||||
|
||||
assert stream_parts == expected
|
||||
# subgraph events are still requested and forwarded to the parent
|
||||
assert mock_sync_client.runs.stream.call_args.kwargs["stream_subgraphs"]
|
||||
assert [ns for ns, _, _ in parent_parts] == [
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123",),
|
||||
]
|
||||
|
||||
|
||||
async def test_astream_without_subgraphs_skips_parent_requested_subgraph_events():
|
||||
mock_async_client = MagicMock()
|
||||
async_iter = MagicMock()
|
||||
async_iter.__aiter__.return_value = _PARENT_STREAMED_PARTS
|
||||
mock_async_client.runs.stream.return_value = async_iter
|
||||
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
|
||||
|
||||
for stream_mode, expected in [
|
||||
("updates", [{"sub": {"chunk": "root"}}]),
|
||||
# messages are not filtered: the server dedupes them across namespaces
|
||||
("messages", [({"content": "token"}, {})]),
|
||||
]:
|
||||
parent_parts: list = []
|
||||
stream_parts = [
|
||||
part
|
||||
async for part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
_parent_streamed_config(parent_parts),
|
||||
stream_mode=stream_mode,
|
||||
)
|
||||
]
|
||||
|
||||
assert stream_parts == expected
|
||||
# subgraph events are still requested and forwarded to the parent
|
||||
assert mock_async_client.runs.stream.call_args.kwargs["stream_subgraphs"]
|
||||
assert [ns for ns, _, _ in parent_parts] == [
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123",),
|
||||
]
|
||||
|
||||
|
||||
def test_invoke():
|
||||
# set up test
|
||||
mock_sync_client = MagicMock()
|
||||
|
||||
@@ -13,13 +13,17 @@ from typing import Annotated, Any
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.prebuilt import ToolNode, ToolRuntime
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import GraphDrained
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.pregel._tools import _tool_call_writer
|
||||
from langgraph.runtime import RunControl
|
||||
from langgraph.types import Command, interrupt
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
@@ -54,14 +58,14 @@ def _caller_async(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1")
|
||||
return caller
|
||||
|
||||
|
||||
def _build_graph(caller, tools) -> Any:
|
||||
def _build_graph(caller, tools, checkpointer=None) -> Any:
|
||||
sg = StateGraph(_State)
|
||||
sg.add_node("caller", caller)
|
||||
sg.add_node("tools", ToolNode(tools))
|
||||
sg.add_edge(START, "caller")
|
||||
sg.add_edge("caller", "tools")
|
||||
sg.add_edge("tools", END)
|
||||
return sg.compile()
|
||||
return sg.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _tool_events(stream) -> list[tuple[tuple[str, ...], dict]]:
|
||||
@@ -146,6 +150,105 @@ class TestSyncGraphSyncTool:
|
||||
assert kinds == ["tool-started", "tool-error"]
|
||||
assert events[1][1]["message"] == "nope"
|
||||
|
||||
def test_interrupt_in_tool_is_not_a_tool_error(self) -> None:
|
||||
@tool
|
||||
def ask() -> str:
|
||||
"""asks."""
|
||||
return interrupt("ok?")
|
||||
|
||||
graph = _build_graph(_caller_sync("ask", {}), [ask], InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
paused = _tool_events(
|
||||
graph.stream(
|
||||
{"messages": []}, config, stream_mode=["tools"], subgraphs=True
|
||||
)
|
||||
)
|
||||
resumed = _tool_events(
|
||||
graph.stream(
|
||||
Command(resume="yes"), config, stream_mode=["tools"], subgraphs=True
|
||||
)
|
||||
)
|
||||
|
||||
assert [p["event"] for _, p in paused] == ["tool-started"]
|
||||
assert [(p["event"], p["tool_call_id"]) for _, p in resumed] == [
|
||||
("tool-started", "tc1"),
|
||||
("tool-finished", "tc1"),
|
||||
]
|
||||
|
||||
def test_tool_handing_off_to_the_parent_still_ends_its_call(self) -> None:
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
done: bool
|
||||
|
||||
def hand_off(state: State) -> Command:
|
||||
return Command(graph=Command.PARENT, goto="finish", update={"done": True})
|
||||
|
||||
inner = StateGraph(State).add_node("hand_off", hand_off)
|
||||
inner = inner.add_edge(START, "hand_off").compile()
|
||||
|
||||
@tool
|
||||
def handoff() -> str:
|
||||
"""hands off."""
|
||||
inner.invoke({"messages": [], "done": False})
|
||||
return "unreachable"
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("caller", _caller_sync("handoff", {}))
|
||||
.add_node("tools", ToolNode([handoff]))
|
||||
.add_node("finish", lambda state: {"done": True})
|
||||
.add_edge(START, "caller")
|
||||
.add_edge("caller", "tools")
|
||||
.add_edge("tools", END)
|
||||
.add_edge("finish", END)
|
||||
.compile()
|
||||
)
|
||||
events = [
|
||||
payload["event"]
|
||||
for _, payload in _tool_events(
|
||||
graph.stream(
|
||||
{"messages": [], "done": False},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
assert events[0] == "tool-started"
|
||||
assert events[-1] == "tool-error", events
|
||||
|
||||
def test_drained_tool_is_not_a_tool_error(self) -> None:
|
||||
control = RunControl()
|
||||
|
||||
def request_drain(state: _State) -> dict:
|
||||
control.request_drain()
|
||||
return {}
|
||||
|
||||
inner = (
|
||||
StateGraph(_State)
|
||||
.add_node("request_drain", request_drain)
|
||||
.add_node("after", lambda state: {})
|
||||
.add_edge(START, "request_drain")
|
||||
.add_edge("request_drain", "after")
|
||||
.compile()
|
||||
)
|
||||
|
||||
@tool
|
||||
def drained() -> str:
|
||||
"""drains."""
|
||||
inner.invoke({"messages": []}, control=control)
|
||||
return "unreachable"
|
||||
|
||||
graph = _build_graph(_caller_sync("drained", {}), [drained])
|
||||
chunks: list = []
|
||||
with pytest.raises(GraphDrained):
|
||||
for chunk in graph.stream(
|
||||
{"messages": []}, stream_mode=["tools"], subgraphs=True
|
||||
):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert [p["event"] for _, p in _tool_events(chunks)] == ["tool-started"]
|
||||
|
||||
def test_writer_unset_outside_tool(self) -> None:
|
||||
# Outside any tool body the ContextVar that ToolRuntime reads
|
||||
# is unset — emitting from there would be a no-op.
|
||||
@@ -196,6 +299,25 @@ class TestAsyncGraphAsyncTool:
|
||||
assert kinds == ["tool-started", "tool-output-delta", "tool-finished"]
|
||||
assert events[1][1]["delta"] == "hi"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupt_in_async_tool_is_not_a_tool_error(self) -> None:
|
||||
@tool
|
||||
async def aask() -> str:
|
||||
"""asks."""
|
||||
return interrupt("ok?")
|
||||
|
||||
graph = _build_graph(_caller_async("aask", {}), [aask], InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
events = []
|
||||
for graph_input in ({"messages": []}, Command(resume="yes")):
|
||||
async for ns, mode, payload in graph.astream(
|
||||
graph_input, config, stream_mode=["tools"], subgraphs=True
|
||||
):
|
||||
if mode == "tools":
|
||||
events.append(payload["event"])
|
||||
|
||||
assert events == ["tool-started", "tool-started", "tool-finished"]
|
||||
|
||||
|
||||
class TestConcurrentToolCalls:
|
||||
def test_parallel_tool_calls_do_not_bleed(self) -> None:
|
||||
|
||||
@@ -14,9 +14,14 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import _araise_for_status_typed
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_araise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, aiter_lines_raw
|
||||
|
||||
@@ -217,6 +222,12 @@ class HttpClient:
|
||||
current_headers = dict(
|
||||
request_headers if reconnect_path is None else reconnect_headers
|
||||
)
|
||||
if last_event_id is not None:
|
||||
current_headers = {
|
||||
key: value
|
||||
for key, value in current_headers.items()
|
||||
if key.lower() != "last-event-id"
|
||||
}
|
||||
if last_event_id:
|
||||
current_headers["Last-Event-ID"] = last_event_id
|
||||
|
||||
@@ -224,65 +235,80 @@ class HttpClient:
|
||||
current_content = content if reconnect_path is None else None
|
||||
current_params = params if reconnect_path is None else None
|
||||
|
||||
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_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}"
|
||||
)
|
||||
|
||||
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()
|
||||
try:
|
||||
async for line in aiter_lines_raw(res):
|
||||
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
# 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 sse.data is not None:
|
||||
# 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""):
|
||||
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
|
||||
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:
|
||||
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:
|
||||
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,7 +16,6 @@ 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
|
||||
@@ -25,7 +24,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
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param, _reconnect_delay
|
||||
from langgraph_sdk.schema import LangSmithTracing, QueryParamTypes
|
||||
from langgraph_sdk.stream.controller import _SeenEventIds
|
||||
from langgraph_sdk.stream.decoders import (
|
||||
@@ -1691,11 +1690,13 @@ class AsyncThreadStream:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt `attempt`."""
|
||||
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)
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._shared_reconnect_backoff_base,
|
||||
cap=self._shared_reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a post-ready transport drop.
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import tzinfo
|
||||
@@ -194,6 +195,15 @@ 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,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from typing import Any, cast
|
||||
@@ -13,9 +14,14 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import _raise_for_status_typed
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_raise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, iter_lines_raw
|
||||
|
||||
@@ -219,6 +225,12 @@ class SyncHttpClient:
|
||||
current_headers = dict(
|
||||
request_headers if reconnect_path is None else reconnect_headers
|
||||
)
|
||||
if last_event_id is not None:
|
||||
current_headers = {
|
||||
key: value
|
||||
for key, value in current_headers.items()
|
||||
if key.lower() != "last-event-id"
|
||||
}
|
||||
if last_event_id:
|
||||
current_headers["Last-Event-ID"] = last_event_id
|
||||
|
||||
@@ -226,62 +238,75 @@ class SyncHttpClient:
|
||||
current_content = content if reconnect_path is None else None
|
||||
current_params = params if reconnect_path is None else None
|
||||
|
||||
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_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}"
|
||||
)
|
||||
|
||||
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()
|
||||
try:
|
||||
for line in iter_lines_raw(res):
|
||||
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
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 sse.data is not None:
|
||||
# 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""):
|
||||
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
|
||||
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:
|
||||
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:
|
||||
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
|
||||
|
||||
|
||||
@@ -811,7 +811,7 @@ class _On:
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
resources: str | Sequence[str],
|
||||
resources: str | Sequence[str] | None = None,
|
||||
actions: str | Sequence[str] | None = None,
|
||||
) -> Callable[[AHO], AHO]: ...
|
||||
|
||||
|
||||
@@ -104,7 +104,7 @@ class SSEDecoder:
|
||||
if (
|
||||
not self._event
|
||||
and not self._data
|
||||
and not self._last_event_id
|
||||
and self._last_event_id is None
|
||||
and self._retry is None
|
||||
):
|
||||
return None
|
||||
|
||||
@@ -15,7 +15,6 @@ 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
|
||||
@@ -23,6 +22,7 @@ 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,11 +275,13 @@ class StreamController:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt *attempt*."""
|
||||
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)
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
@@ -13,6 +12,7 @@ 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,11 +282,13 @@ class SyncStreamController:
|
||||
|
||||
def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
|
||||
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)
|
||||
time.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -1,8 +1,49 @@
|
||||
import ast
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph_sdk import Auth
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[{}, {"actions": "read"}, {"resources": None}],
|
||||
)
|
||||
def test_global_handler_factory_accepts_optional_resources(kwargs) -> None:
|
||||
auth = Auth()
|
||||
tree = ast.parse(inspect.getsource(type(auth.on)))
|
||||
cls = tree.body[0]
|
||||
assert isinstance(cls, ast.ClassDef)
|
||||
factory_overload = next(
|
||||
node
|
||||
for node in cls.body
|
||||
if isinstance(node, ast.FunctionDef)
|
||||
and node.name == "__call__"
|
||||
and node.decorator_list
|
||||
and any(arg.arg == "resources" for arg in node.args.kwonlyargs)
|
||||
)
|
||||
resources = next(
|
||||
arg for arg in factory_overload.args.kwonlyargs if arg.arg == "resources"
|
||||
)
|
||||
assert resources.annotation is not None
|
||||
assert ast.unparse(resources.annotation) == "str | Sequence[str] | None"
|
||||
default = factory_overload.args.kw_defaults[0]
|
||||
assert isinstance(default, ast.Constant) and default.value is None
|
||||
|
||||
@auth.on(**kwargs)
|
||||
async def handler(ctx, value):
|
||||
del ctx, value
|
||||
return None
|
||||
|
||||
if kwargs.get("actions"):
|
||||
assert auth._handlers == {("*", "read"): [handler]}
|
||||
assert auth._global_handlers == []
|
||||
else:
|
||||
assert auth._global_handlers == [handler]
|
||||
assert auth._handlers == {}
|
||||
|
||||
|
||||
def test_handler_multiple_resources_and_actions() -> None:
|
||||
auth = Auth()
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator, Sequence
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -13,6 +15,7 @@ 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,
|
||||
@@ -67,6 +70,19 @@ class ListByteStream(httpx.ByteStream):
|
||||
return None
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
def iter_lines_raw(payload: list[bytes]) -> Iterator[BytesLike]:
|
||||
decoder = BytesLineDecoder()
|
||||
for part in payload:
|
||||
@@ -176,6 +192,7 @@ 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 = [
|
||||
@@ -241,6 +258,7 @@ 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 = [
|
||||
@@ -312,18 +330,14 @@ async def test_http_client_stream_recovers_after_disconnect():
|
||||
]
|
||||
|
||||
|
||||
def _drop_after_an_empty_id(stream_cls: type, reconnects: list[httpx.Request]):
|
||||
def _drop_after(chunks: list[bytes], stream_cls: type, reconnects: list[httpx.Request]):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
if request.method == "POST":
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream", "Location": "/reconnect"},
|
||||
stream=stream_cls(
|
||||
[
|
||||
b"id: 1\nevent: values\ndata: 1\n\n",
|
||||
b"id:\nevent: values\ndata: 2\n\n",
|
||||
],
|
||||
httpx.RemoteProtocolError("incomplete chunked read"),
|
||||
chunks, httpx.RemoteProtocolError("incomplete chunked read")
|
||||
),
|
||||
)
|
||||
reconnects.append(request)
|
||||
@@ -336,26 +350,335 @@ def _drop_after_an_empty_id(stream_cls: type, reconnects: list[httpx.Request]):
|
||||
return handler
|
||||
|
||||
|
||||
def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id():
|
||||
def _sync_reconnect(
|
||||
chunks: list[bytes], headers: dict[str, str] | None
|
||||
) -> httpx.Request:
|
||||
reconnects: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(_drop_after_an_empty_id(ListByteStream, reconnects))
|
||||
transport = httpx.MockTransport(_drop_after(chunks, ListByteStream, reconnects))
|
||||
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||
list(SyncHttpClient(client).stream("/stream", "POST", json={}))
|
||||
|
||||
assert "last-event-id" not in reconnects[0].headers
|
||||
list(SyncHttpClient(client).stream("/stream", "POST", json={}, headers=headers))
|
||||
return reconnects[0]
|
||||
|
||||
|
||||
async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id():
|
||||
async def _async_reconnect(
|
||||
chunks: list[bytes], headers: dict[str, str] | None
|
||||
) -> httpx.Request:
|
||||
reconnects: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(
|
||||
_drop_after_an_empty_id(AsyncListByteStream, reconnects)
|
||||
_drop_after(chunks, AsyncListByteStream, reconnects)
|
||||
)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
[part async for part in HttpClient(client).stream("/stream", "POST", json={})]
|
||||
[
|
||||
part
|
||||
async for part in HttpClient(client).stream(
|
||||
"/stream", "POST", json={}, headers=headers
|
||||
)
|
||||
]
|
||||
return reconnects[0]
|
||||
|
||||
assert "last-event-id" not in reconnects[0].headers
|
||||
|
||||
_EMPTY_IDS = pytest.mark.parametrize(
|
||||
"empty_id",
|
||||
[b"id:\nevent: values\ndata: 2\n\n", b"id:\n\n"],
|
||||
ids=["with-an-event", "alone"],
|
||||
)
|
||||
_CALLER_CURSORS = pytest.mark.parametrize(
|
||||
"headers",
|
||||
[None, {"Last-Event-ID": "5"}, {"last-event-id": "5"}],
|
||||
ids=["no-caller-cursor", "caller-cursor", "lowercase-caller-cursor"],
|
||||
)
|
||||
|
||||
|
||||
@_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
|
||||
):
|
||||
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
|
||||
):
|
||||
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"}
|
||||
)
|
||||
|
||||
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"}
|
||||
)
|
||||
|
||||
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"}
|
||||
)
|
||||
|
||||
assert reconnect.headers.get_list("last-event-id") == ["5"]
|
||||
|
||||
|
||||
# --- 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 ---
|
||||
|
||||
Reference in new issue
Block a user