wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
2944 lines
94 KiB
Python
2944 lines
94 KiB
Python
import asyncio
|
||
import contextlib
|
||
import operator
|
||
import sys
|
||
import threading
|
||
import time
|
||
from collections import deque
|
||
from collections.abc import AsyncIterator
|
||
from datetime import datetime, timedelta
|
||
from typing import Annotated, Any
|
||
from unittest.mock import Mock, patch
|
||
from uuid import uuid4
|
||
|
||
import pytest
|
||
from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, BaseCallbackHandler
|
||
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
||
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage
|
||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
|
||
from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnableParallel
|
||
from langgraph.checkpoint.memory import InMemorySaver, MemorySaver
|
||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||
from typing_extensions import TypedDict
|
||
|
||
from langgraph._internal._constants import (
|
||
CONF,
|
||
CONFIG_KEY_CHECKPOINT_ID,
|
||
CONFIG_KEY_CHECKPOINT_NS,
|
||
CONFIG_KEY_RUNTIME,
|
||
CONFIG_KEY_SEND,
|
||
CONFIG_KEY_STREAM,
|
||
CONFIG_KEY_TASK_ID,
|
||
CONFIG_KEY_THREAD_ID,
|
||
CONFIG_KEY_TIMED_ATTEMPT_OBSERVER,
|
||
)
|
||
from langgraph._internal._runnable import RunnableCallable
|
||
from langgraph._internal._timeout import coerce_timeout_policy
|
||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||
from langgraph.channels.last_value import LastValue
|
||
from langgraph.errors import (
|
||
GraphInterrupt,
|
||
NodeCancelledError,
|
||
NodeError,
|
||
NodeTimeoutError,
|
||
ParentCommand,
|
||
)
|
||
from langgraph.func import entrypoint, task
|
||
from langgraph.graph import END, START, StateGraph, add_messages
|
||
from langgraph.pregel import NodeBuilder, Pregel
|
||
from langgraph.pregel._read import PregelNode
|
||
from langgraph.pregel._retry import (
|
||
_checkpoint_ns_for_parent_command,
|
||
_ensure_execution_info,
|
||
_should_retry_on,
|
||
_TimedAttemptScope,
|
||
arun_with_retry,
|
||
run_with_retry,
|
||
)
|
||
from langgraph.pregel.protocol import StreamProtocol
|
||
from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime
|
||
from langgraph.types import (
|
||
Command,
|
||
PregelExecutableTask,
|
||
RetryPolicy,
|
||
Send,
|
||
TimeoutPolicy,
|
||
)
|
||
|
||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||
sys.version_info < (3, 11),
|
||
reason="Python 3.11+ is required for async contextvars support",
|
||
)
|
||
|
||
# `asyncio.Task.cancelling()` is Python 3.11+. The LSD-1507 fix in
|
||
# `langgraph/pregel/_retry.py` falls back to a no-op on 3.10 (preserves the
|
||
# existing CancelledError-as-silent-tear-down behaviour) because there is no
|
||
# reliable way to distinguish user-raised from framework-initiated
|
||
# cancellation without that API. Tests for the converted behaviour gate on
|
||
# the same Python version boundary.
|
||
NEEDS_TASK_CANCELLING = pytest.mark.skipif(
|
||
sys.version_info < (3, 11),
|
||
reason="LSD-1507 user-cancellation conversion requires Python 3.11+ "
|
||
"(asyncio.Task.cancelling)",
|
||
)
|
||
|
||
|
||
def test_should_retry_on_single_exception():
|
||
"""Test retry with a single exception type."""
|
||
policy = RetryPolicy(retry_on=ValueError)
|
||
|
||
# Should retry on ValueError
|
||
assert _should_retry_on(policy, ValueError("test error")) is True
|
||
|
||
# Should not retry on other exceptions
|
||
assert _should_retry_on(policy, TypeError("test error")) is False
|
||
assert _should_retry_on(policy, Exception("test error")) is False
|
||
|
||
|
||
def test_should_retry_on_sequence_of_exceptions():
|
||
"""Test retry with a sequence of exception types."""
|
||
policy = RetryPolicy(retry_on=(ValueError, KeyError))
|
||
|
||
# Should retry on listed exceptions
|
||
assert _should_retry_on(policy, ValueError("test error")) is True
|
||
assert _should_retry_on(policy, KeyError("test error")) is True
|
||
|
||
# Should not retry on other exceptions
|
||
assert _should_retry_on(policy, TypeError("test error")) is False
|
||
assert _should_retry_on(policy, Exception("test error")) is False
|
||
|
||
|
||
def test_should_retry_on_subclass_of_exception():
|
||
"""Test retry on subclass of specified exception."""
|
||
|
||
class CustomError(ValueError):
|
||
pass
|
||
|
||
policy = RetryPolicy(retry_on=ValueError)
|
||
|
||
# Should retry on subclass of specified exception
|
||
assert _should_retry_on(policy, CustomError("test error")) is True
|
||
|
||
|
||
def test_should_retry_on_callable():
|
||
"""Test retry with a callable predicate."""
|
||
|
||
# Only retry on ValueError with message containing 'retry'
|
||
def should_retry(exc: Exception) -> bool:
|
||
return isinstance(exc, ValueError) and "retry" in str(exc)
|
||
|
||
policy = RetryPolicy(retry_on=should_retry)
|
||
|
||
# Should retry when predicate returns True
|
||
assert _should_retry_on(policy, ValueError("please retry this")) is True
|
||
|
||
# Should not retry when predicate returns False
|
||
assert _should_retry_on(policy, ValueError("other error")) is False
|
||
assert _should_retry_on(policy, TypeError("please retry this")) is False
|
||
|
||
|
||
def test_should_retry_on_invalid_type():
|
||
"""Test retry with an invalid retry_on type."""
|
||
policy = RetryPolicy(retry_on=123) # type: ignore
|
||
|
||
with pytest.raises(TypeError, match="retry_on must be an Exception class"):
|
||
_should_retry_on(policy, ValueError("test error"))
|
||
|
||
|
||
def test_should_retry_on_empty_sequence():
|
||
"""Test retry with an empty sequence."""
|
||
policy = RetryPolicy(retry_on=())
|
||
|
||
# Should not retry when sequence is empty
|
||
assert _should_retry_on(policy, ValueError("test error")) is False
|
||
|
||
|
||
def test_checkpoint_ns_for_parent_command() -> None:
|
||
assert _checkpoint_ns_for_parent_command("") == ""
|
||
assert _checkpoint_ns_for_parent_command("node:1") == ""
|
||
assert _checkpoint_ns_for_parent_command("node:1|child:2") == "node:1"
|
||
assert _checkpoint_ns_for_parent_command("node:1|1|child:2") == "node:1"
|
||
assert _checkpoint_ns_for_parent_command("node:1|1|child:2|1") == "node:1"
|
||
assert (
|
||
_checkpoint_ns_for_parent_command("parent:1|1|child:1|1|node:1|1")
|
||
== "parent:1|1|child:1"
|
||
)
|
||
assert (
|
||
_checkpoint_ns_for_parent_command("parent:1|1|child:1|1|node:1")
|
||
== "parent:1|1|child:1"
|
||
)
|
||
|
||
|
||
def test_should_retry_default_retry_on():
|
||
"""Test the default retry_on function."""
|
||
import httpx
|
||
import requests
|
||
|
||
# Create a RetryPolicy with default_retry_on
|
||
policy = RetryPolicy()
|
||
|
||
# Should retry on ConnectionError
|
||
assert _should_retry_on(policy, ConnectionError("connection refused")) is True
|
||
|
||
# Should not retry on common programming errors
|
||
assert _should_retry_on(policy, ValueError("invalid value")) is False
|
||
assert _should_retry_on(policy, TypeError("invalid type")) is False
|
||
assert _should_retry_on(policy, ArithmeticError("division by zero")) is False
|
||
assert _should_retry_on(policy, ImportError("module not found")) is False
|
||
assert _should_retry_on(policy, LookupError("key not found")) is False
|
||
assert _should_retry_on(policy, NameError("name not defined")) is False
|
||
assert _should_retry_on(policy, SyntaxError("invalid syntax")) is False
|
||
assert _should_retry_on(policy, RuntimeError("runtime error")) is False
|
||
assert _should_retry_on(policy, ReferenceError("weak reference")) is False
|
||
assert _should_retry_on(policy, StopIteration()) is False
|
||
assert _should_retry_on(policy, StopAsyncIteration()) is False
|
||
assert _should_retry_on(policy, OSError("file not found")) is False
|
||
|
||
# Should retry on httpx.HTTPStatusError with 5xx status code
|
||
response_5xx = Mock()
|
||
response_5xx.status_code = 503
|
||
http_error_5xx = httpx.HTTPStatusError(
|
||
"server error", request=Mock(), response=response_5xx
|
||
)
|
||
assert _should_retry_on(policy, http_error_5xx) is True
|
||
|
||
# Should not retry on httpx.HTTPStatusError with 4xx status code
|
||
response_4xx = Mock()
|
||
response_4xx.status_code = 404
|
||
http_error_4xx = httpx.HTTPStatusError(
|
||
"not found", request=Mock(), response=response_4xx
|
||
)
|
||
assert _should_retry_on(policy, http_error_4xx) is False
|
||
|
||
# Should retry on requests.HTTPError with 5xx status code
|
||
response_req_5xx = Mock()
|
||
response_req_5xx.status_code = 502
|
||
req_error_5xx = requests.HTTPError("bad gateway")
|
||
req_error_5xx.response = response_req_5xx
|
||
assert _should_retry_on(policy, req_error_5xx) is True
|
||
|
||
# Should not retry on requests.HTTPError with 4xx status code
|
||
response_req_4xx = Mock()
|
||
response_req_4xx.status_code = 400
|
||
req_error_4xx = requests.HTTPError("bad request")
|
||
req_error_4xx.response = response_req_4xx
|
||
assert _should_retry_on(policy, req_error_4xx) is False
|
||
|
||
# Should retry on requests.HTTPError with no response
|
||
req_error_no_resp = requests.HTTPError("connection error")
|
||
req_error_no_resp.response = None
|
||
assert _should_retry_on(policy, req_error_no_resp) is True
|
||
|
||
# NodeTimeoutError should be retryable by default
|
||
assert (
|
||
_should_retry_on(
|
||
policy,
|
||
NodeTimeoutError("node", 1.0, kind="run", run_timeout=0.5),
|
||
)
|
||
is True
|
||
)
|
||
|
||
# Should retry on other exceptions by default
|
||
class CustomException(Exception):
|
||
pass
|
||
|
||
assert _should_retry_on(policy, CustomException("custom error")) is True
|
||
|
||
|
||
def test_graph_with_single_retry_policy():
|
||
"""Test a simple graph with a single RetryPolicy for a node."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempt_count = 0
|
||
attempt_numbers: list[int] = []
|
||
first_attempt_times: list[float | None] = []
|
||
|
||
def failing_node(state: State, runtime: Runtime):
|
||
nonlocal attempt_count
|
||
attempt_count += 1
|
||
assert runtime.execution_info.node_attempt == attempt_count
|
||
attempt_numbers.append(runtime.execution_info.node_attempt)
|
||
first_attempt_times.append(runtime.execution_info.node_first_attempt_time)
|
||
if attempt_count < 3: # Fail the first two attempts
|
||
raise ValueError("Intentional failure")
|
||
return {"foo": "success"}
|
||
|
||
def other_node(state: State):
|
||
return {"foo": "other_node"}
|
||
|
||
# Create a retry policy with specific parameters
|
||
retry_policy = RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.01, # Short interval for tests
|
||
backoff_factor=2.0,
|
||
jitter=False, # Disable jitter for predictable timing
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
# Create and compile the graph
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("failing_node", failing_node, retry_policy=retry_policy)
|
||
.add_node("other_node", other_node)
|
||
.add_edge(START, "failing_node")
|
||
.add_edge("failing_node", "other_node")
|
||
.compile()
|
||
)
|
||
|
||
with patch("time.sleep") as mock_sleep:
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
# Verify retry behavior
|
||
assert attempt_count == 3 # The node should have been tried 3 times
|
||
assert attempt_numbers == [1, 2, 3]
|
||
assert len(first_attempt_times) == 3
|
||
assert first_attempt_times[0] is not None
|
||
assert first_attempt_times[1] == first_attempt_times[0]
|
||
assert first_attempt_times[2] == first_attempt_times[0]
|
||
assert result["foo"] == "other_node" # Final result should be from other_node
|
||
|
||
# Verify the sleep intervals
|
||
call_args_list = [args[0][0] for args in mock_sleep.call_args_list]
|
||
assert call_args_list == [0.01, 0.02]
|
||
|
||
|
||
def test_runtime_execution_info_defaults_without_retry():
|
||
"""Test execution_info defaults when no retry and no config are provided."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
captured = {}
|
||
|
||
def node(state: State, runtime: Runtime):
|
||
captured["node_attempt"] = runtime.execution_info.node_attempt
|
||
captured["node_first_attempt_time"] = (
|
||
runtime.execution_info.node_first_attempt_time
|
||
)
|
||
return {"foo": "ok"}
|
||
|
||
graph = StateGraph(State).add_node("node", node).add_edge(START, "node").compile()
|
||
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
assert result["foo"] == "ok"
|
||
assert captured["node_attempt"] == 1
|
||
assert isinstance(captured["node_first_attempt_time"], float)
|
||
|
||
|
||
def test_graph_with_jitter_retry_policy():
|
||
"""Test a graph with a RetryPolicy that uses jitter."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempt_count = 0
|
||
|
||
def failing_node(state):
|
||
nonlocal attempt_count
|
||
attempt_count += 1
|
||
if attempt_count < 2: # Fail the first attempt
|
||
raise ValueError("Intentional failure")
|
||
return {"foo": "success"}
|
||
|
||
# Create a retry policy with jitter enabled
|
||
retry_policy = RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.01,
|
||
jitter=True, # Enable jitter for randomized backoff
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
# Create and compile the graph
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("failing_node", failing_node, retry_policy=retry_policy)
|
||
.add_edge(START, "failing_node")
|
||
.compile()
|
||
)
|
||
|
||
# Test graph execution with mocked random and sleep
|
||
with (
|
||
patch("random.uniform", return_value=0.05) as mock_random,
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
# Verify retry behavior
|
||
assert attempt_count == 2 # The node should have been tried twice
|
||
assert result["foo"] == "success"
|
||
|
||
# Verify jitter was applied
|
||
mock_random.assert_called_with(0, 1) # Jitter should use random.uniform(0, 1)
|
||
mock_sleep.assert_called_with(0.01 + 0.05) # Sleep should include jitter
|
||
|
||
|
||
def test_graph_with_multiple_retry_policies():
|
||
"""Test a graph with multiple retry policies for a node."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
error_type: str
|
||
|
||
attempt_counts = {"value_error": 0, "key_error": 0}
|
||
|
||
def failing_node(state):
|
||
error_type = state["error_type"]
|
||
|
||
if error_type == "value_error":
|
||
attempt_counts["value_error"] += 1
|
||
if attempt_counts["value_error"] < 2:
|
||
raise ValueError("Value error")
|
||
elif error_type == "key_error":
|
||
attempt_counts["key_error"] += 1
|
||
if attempt_counts["key_error"] < 3:
|
||
raise KeyError("Key error")
|
||
|
||
return {"foo": f"recovered_from_{error_type}"}
|
||
|
||
# Create multiple retry policies
|
||
value_error_policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
key_error_policy = RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.02,
|
||
jitter=False,
|
||
retry_on=KeyError,
|
||
)
|
||
|
||
# Create and compile the graph with a list of retry policies
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node(
|
||
"failing_node",
|
||
failing_node,
|
||
retry_policy=(value_error_policy, key_error_policy),
|
||
)
|
||
.add_edge(START, "failing_node")
|
||
.compile()
|
||
)
|
||
|
||
# Test ValueError scenario
|
||
with patch("time.sleep"):
|
||
result_value_error = graph.invoke({"foo": "", "error_type": "value_error"})
|
||
|
||
assert attempt_counts["value_error"] == 2
|
||
assert result_value_error["foo"] == "recovered_from_value_error"
|
||
|
||
# Reset attempt counts
|
||
attempt_counts = {"value_error": 0, "key_error": 0}
|
||
|
||
# Test KeyError scenario
|
||
with patch("time.sleep"):
|
||
result_key_error = graph.invoke({"foo": "", "error_type": "key_error"})
|
||
|
||
assert attempt_counts["key_error"] == 3
|
||
assert result_key_error["foo"] == "recovered_from_key_error"
|
||
|
||
|
||
def test_graph_with_max_attempts_exceeded():
|
||
"""Test a graph where max_attempts is exceeded."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def always_failing_node(state):
|
||
raise ValueError("Always fails")
|
||
|
||
# Create a retry policy with limited attempts
|
||
retry_policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
# Create and compile the graph
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("always_failing", always_failing_node, retry_policy=retry_policy)
|
||
.add_edge(START, "always_failing")
|
||
.compile()
|
||
)
|
||
|
||
# Test graph execution
|
||
with (
|
||
patch("time.sleep") as mock_sleep,
|
||
pytest.raises(ValueError, match="Always fails"),
|
||
):
|
||
graph.invoke({"foo": ""})
|
||
|
||
mock_sleep.assert_called_with(0.01)
|
||
|
||
|
||
def test_execution_info_identity_fields_populated_on_retry():
|
||
"""Test that thread_id, task_id, run_id, etc. are populated in execution_info during retries."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempt_count = 0
|
||
captured_infos: list[dict] = []
|
||
|
||
def failing_node(state: State, runtime: Runtime):
|
||
nonlocal attempt_count
|
||
attempt_count += 1
|
||
info = runtime.execution_info
|
||
captured_infos.append(
|
||
{
|
||
"thread_id": info.thread_id,
|
||
"run_id": info.run_id,
|
||
"node_attempt": info.node_attempt,
|
||
"node_first_attempt_time": info.node_first_attempt_time,
|
||
"checkpoint_ns": info.checkpoint_ns,
|
||
}
|
||
)
|
||
if attempt_count < 2:
|
||
raise ValueError("Intentional failure")
|
||
return {"foo": "success"}
|
||
|
||
retry_policy = RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("failing_node", failing_node, retry_policy=retry_policy)
|
||
.add_edge(START, "failing_node")
|
||
.compile(checkpointer=MemorySaver())
|
||
)
|
||
|
||
with patch("time.sleep"):
|
||
result = graph.invoke(
|
||
{"foo": ""},
|
||
config={"configurable": {"thread_id": "retry-thread"}},
|
||
)
|
||
|
||
assert result["foo"] == "success"
|
||
assert len(captured_infos) == 2
|
||
|
||
# Both attempts should have the same thread_id and first_attempt_time
|
||
assert captured_infos[0]["thread_id"] == "retry-thread"
|
||
assert captured_infos[1]["thread_id"] == "retry-thread"
|
||
assert (
|
||
captured_infos[0]["node_first_attempt_time"]
|
||
== captured_infos[1]["node_first_attempt_time"]
|
||
)
|
||
|
||
# node_attempt should increment
|
||
assert captured_infos[0]["node_attempt"] == 1
|
||
assert captured_infos[1]["node_attempt"] == 2
|
||
|
||
|
||
def test_ensure_execution_info_noop_when_already_set():
|
||
"""Test that _ensure_execution_info is a no-op when execution_info exists."""
|
||
existing_info = ExecutionInfo(
|
||
checkpoint_id="cp-1", checkpoint_ns="ns-1", task_id="task-1"
|
||
)
|
||
runtime = DEFAULT_RUNTIME.override(execution_info=existing_info)
|
||
config = {CONF: {CONFIG_KEY_THREAD_ID: "thread-1"}}
|
||
task = Mock(id="task-2")
|
||
|
||
result = _ensure_execution_info(runtime, config, task)
|
||
assert result is runtime
|
||
assert result.execution_info is existing_info
|
||
|
||
|
||
def test_ensure_execution_info_creates_from_config():
|
||
"""Test that _ensure_execution_info creates ExecutionInfo from config when missing."""
|
||
runtime = DEFAULT_RUNTIME.override(execution_info=None)
|
||
config = {
|
||
"run_id": "run-123",
|
||
CONF: {
|
||
CONFIG_KEY_CHECKPOINT_ID: "cp-42",
|
||
CONFIG_KEY_CHECKPOINT_NS: "ns-42",
|
||
CONFIG_KEY_TASK_ID: "task-42",
|
||
CONFIG_KEY_THREAD_ID: "thread-42",
|
||
},
|
||
}
|
||
task = Mock(id="fallback-task-id")
|
||
|
||
result = _ensure_execution_info(runtime, config, task)
|
||
assert result.execution_info is not None
|
||
assert result.execution_info.checkpoint_id == "cp-42"
|
||
assert result.execution_info.checkpoint_ns == "ns-42"
|
||
assert result.execution_info.task_id == "task-42"
|
||
assert result.execution_info.thread_id == "thread-42"
|
||
assert result.execution_info.run_id == "run-123"
|
||
|
||
|
||
def test_ensure_execution_info_falls_back_to_task_id():
|
||
"""Test that _ensure_execution_info uses task.id when CONFIG_KEY_TASK_ID is missing."""
|
||
runtime = DEFAULT_RUNTIME.override(execution_info=None)
|
||
config = {CONF: {}}
|
||
task = Mock(id="fallback-task-id")
|
||
|
||
result = _ensure_execution_info(runtime, config, task)
|
||
assert result.execution_info.task_id == "fallback-task-id"
|
||
assert result.execution_info.checkpoint_id == ""
|
||
assert result.execution_info.checkpoint_ns == ""
|
||
|
||
|
||
def test_run_with_retry_creates_execution_info_when_missing():
|
||
"""Test that run_with_retry works when runtime has no execution_info (distributed runtime scenario)."""
|
||
captured_infos: list[ExecutionInfo] = []
|
||
|
||
class FakeProc:
|
||
def invoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
captured_infos.append(runtime.execution_info)
|
||
return input
|
||
|
||
runtime = DEFAULT_RUNTIME.override(execution_info=None)
|
||
config = {
|
||
"run_id": "run-abc",
|
||
CONF: {
|
||
CONFIG_KEY_RUNTIME: runtime,
|
||
CONFIG_KEY_CHECKPOINT_ID: "cp-99",
|
||
CONFIG_KEY_CHECKPOINT_NS: "__start__:task123",
|
||
CONFIG_KEY_TASK_ID: "task123",
|
||
CONFIG_KEY_THREAD_ID: "thread-xyz",
|
||
},
|
||
}
|
||
|
||
task = PregelExecutableTask(
|
||
name="__start__",
|
||
input={"messages": []},
|
||
proc=FakeProc(),
|
||
writes=deque(),
|
||
config=config,
|
||
triggers=["__start__"],
|
||
retry_policy=[],
|
||
cache_key=None,
|
||
id="task123",
|
||
path=("__pregel_pull", "__start__"),
|
||
)
|
||
|
||
run_with_retry(task, retry_policy=None)
|
||
|
||
assert len(captured_infos) == 1
|
||
info = captured_infos[0]
|
||
assert info.checkpoint_id == "cp-99"
|
||
assert info.checkpoint_ns == "__start__:task123"
|
||
assert info.task_id == "task123"
|
||
assert info.thread_id == "thread-xyz"
|
||
assert info.run_id == "run-abc"
|
||
assert info.node_attempt == 1
|
||
assert info.node_first_attempt_time is not None
|
||
|
||
|
||
def _make_task(
|
||
proc,
|
||
*,
|
||
timeout=None,
|
||
retry_policy=(),
|
||
name="timed",
|
||
task_id="tid",
|
||
writers=(),
|
||
):
|
||
runtime = DEFAULT_RUNTIME.override(execution_info=None)
|
||
writes = deque()
|
||
config = {
|
||
"run_id": "run-x",
|
||
CONF: {
|
||
CONFIG_KEY_RUNTIME: runtime,
|
||
CONFIG_KEY_CHECKPOINT_ID: "cp",
|
||
CONFIG_KEY_CHECKPOINT_NS: f"{name}:{task_id}",
|
||
CONFIG_KEY_SEND: writes.extend,
|
||
CONFIG_KEY_TASK_ID: task_id,
|
||
CONFIG_KEY_THREAD_ID: "thr",
|
||
},
|
||
}
|
||
return PregelExecutableTask(
|
||
name=name,
|
||
input=None,
|
||
proc=proc,
|
||
writes=writes,
|
||
config=config,
|
||
triggers=[name],
|
||
retry_policy=retry_policy,
|
||
cache_key=None,
|
||
id=task_id,
|
||
path=("__pregel_pull", name),
|
||
writers=writers,
|
||
timeout=coerce_timeout_policy(timeout),
|
||
)
|
||
|
||
|
||
def _idle_timeout(value: float | timedelta) -> TimeoutPolicy:
|
||
return TimeoutPolicy(idle_timeout=value)
|
||
|
||
|
||
def test_coerce_timeout_policy_scalar_is_run_timeout():
|
||
assert coerce_timeout_policy(None) is None
|
||
policy = coerce_timeout_policy(timedelta(milliseconds=250))
|
||
assert policy == TimeoutPolicy(run_timeout=0.25)
|
||
assert Send("node", None, timeout=timedelta(milliseconds=250)).timeout == policy
|
||
|
||
idle_policy = coerce_timeout_policy(TimeoutPolicy(idle_timeout=1.5))
|
||
assert idle_policy == TimeoutPolicy(idle_timeout=1.5)
|
||
|
||
with pytest.raises(ValueError, match="run_timeout must be greater than 0"):
|
||
coerce_timeout_policy(0)
|
||
|
||
|
||
def test_coerce_timeout_policy_returns_same_instance_for_already_coerced():
|
||
policy = coerce_timeout_policy(TimeoutPolicy(run_timeout=1.0, idle_timeout=2.0))
|
||
assert coerce_timeout_policy(policy) is policy
|
||
assert TimeoutPolicy.coerce(policy) is policy
|
||
|
||
|
||
def test_send_timeout_round_trips_through_msgpack_serde():
|
||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||
packet = Send(
|
||
"worker",
|
||
{"x": 1},
|
||
timeout=TimeoutPolicy(run_timeout=1, idle_timeout=2),
|
||
)
|
||
|
||
assert serde.loads_typed(serde.dumps_typed(packet)) == packet
|
||
|
||
|
||
def test_send_without_timeout_round_trips_through_msgpack_serde():
|
||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||
packet = Send("worker", {"x": 1})
|
||
|
||
assert serde.loads_typed(serde.dumps_typed(packet)) == packet
|
||
|
||
|
||
def test_run_with_retry_rejects_sync_timeout_without_starting_proc():
|
||
started = False
|
||
|
||
class Proc:
|
||
def invoke(self, input, config):
|
||
nonlocal started
|
||
started = True
|
||
return input
|
||
|
||
task = _make_task(Proc(), timeout=_idle_timeout(0.05), name="sync")
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
run_with_retry(task, retry_policy=None)
|
||
assert not started
|
||
|
||
|
||
def test_run_with_retry_without_timeout_runs_sync_directly():
|
||
class FastProc:
|
||
def invoke(self, input, config):
|
||
return "ok"
|
||
|
||
task = _make_task(FastProc(), timeout=None)
|
||
assert run_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
|
||
def test_idle_timeout_guard_call_does_not_hold_scope_lock():
|
||
scope = _TimedAttemptScope()
|
||
|
||
def call():
|
||
assert not scope._lock.locked()
|
||
return "ok"
|
||
|
||
assert scope._guard_call(call)() == "ok"
|
||
|
||
|
||
def test_idle_timeout_guard_stream_does_not_hold_scope_lock():
|
||
scope = _TimedAttemptScope()
|
||
|
||
def stream(chunk):
|
||
assert not scope._lock.locked()
|
||
assert chunk == ((), "custom", "ok")
|
||
|
||
scope._guard_stream(StreamProtocol(stream, {"custom"}))(((), "custom", "ok"))
|
||
|
||
|
||
def test_idle_timeout_guard_stream_writer_does_not_hold_scope_lock():
|
||
scope = _TimedAttemptScope()
|
||
|
||
def stream_writer(chunk):
|
||
assert not scope._lock.locked()
|
||
assert chunk == "ok"
|
||
|
||
scope._guard_stream_writer(stream_writer)("ok")
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_ok_when_fast():
|
||
class FastProc:
|
||
async def ainvoke(self, input, config):
|
||
return "ok"
|
||
|
||
task = _make_task(FastProc(), timeout=_idle_timeout(1.0))
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_retries_when_retry_on_timeout():
|
||
calls: list[float] = []
|
||
|
||
class FlakyProc:
|
||
async def ainvoke(self, input, config):
|
||
calls.append(time.monotonic())
|
||
if len(calls) < 2:
|
||
await asyncio.sleep(0.5)
|
||
return "late"
|
||
return "ok"
|
||
|
||
policy = RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.0,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
)
|
||
task = _make_task(FlakyProc(), timeout=_idle_timeout(0.05), retry_policy=(policy,))
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
assert len(calls) == 2
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_entrypoint_timeout_allows_pre_timeout_child_task_to_run():
|
||
child_started = threading.Event()
|
||
|
||
@task()
|
||
def child(value: int) -> int:
|
||
child_started.set()
|
||
return value + 1
|
||
|
||
@entrypoint(timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
async def parent(value: int) -> int:
|
||
child(value)
|
||
await asyncio.sleep(0.2)
|
||
return value
|
||
|
||
with pytest.raises(NodeTimeoutError):
|
||
await parent.ainvoke(1)
|
||
assert child_started.wait(timeout=1.0)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_accepts_timedelta():
|
||
class SlowProc:
|
||
async def ainvoke(self, input, config):
|
||
await asyncio.sleep(0.5)
|
||
return input
|
||
|
||
task = _make_task(SlowProc(), timeout=_idle_timeout(timedelta(milliseconds=50)))
|
||
with pytest.raises(NodeTimeoutError):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_fires_async():
|
||
class SlowProc:
|
||
async def ainvoke(self, input, config):
|
||
await asyncio.sleep(1.0)
|
||
return input
|
||
|
||
task = _make_task(SlowProc(), timeout=_idle_timeout(0.05), name="aslow")
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.node == "aslow"
|
||
assert excinfo.value.idle_timeout == 0.05
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_run_timeout_is_not_refreshed_by_heartbeat():
|
||
class HeartbeatingProc:
|
||
async def ainvoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
while True:
|
||
runtime.heartbeat()
|
||
await asyncio.sleep(0.01)
|
||
|
||
task = _make_task(HeartbeatingProc(), timeout=0.05, name="run-timeout")
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.kind == "run"
|
||
assert excinfo.value.run_timeout == 0.05
|
||
assert excinfo.value.idle_timeout is None
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_node_timeout_error_carries_both_configured_timeouts():
|
||
"""Both `idle_timeout` and `run_timeout` reflect the configured policy
|
||
even when only one of them fires."""
|
||
|
||
class SlowProc:
|
||
async def ainvoke(self, input, config):
|
||
await asyncio.sleep(1.0)
|
||
|
||
task = _make_task(
|
||
SlowProc(),
|
||
timeout=TimeoutPolicy(run_timeout=0.05, idle_timeout=0.5),
|
||
name="both",
|
||
)
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.kind == "run"
|
||
assert excinfo.value.run_timeout == 0.05
|
||
assert excinfo.value.idle_timeout == 0.5
|
||
# `timeout` is the one that fired.
|
||
assert excinfo.value.timeout == 0.05
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_does_not_swallow_proc_asyncio_timeout():
|
||
calls = 0
|
||
|
||
class InnerTimeoutProc:
|
||
async def ainvoke(self, input, config):
|
||
nonlocal calls
|
||
calls += 1
|
||
raise asyncio.TimeoutError("inner")
|
||
|
||
# `retry_on=NodeTimeoutError` + `calls == 1` is the load-bearing assertion:
|
||
# if the proc's TimeoutError were misclassified as NodeTimeoutError it
|
||
# would be retried, and `calls` would be 2.
|
||
policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.0,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
)
|
||
task = _make_task(
|
||
InnerTimeoutProc(),
|
||
timeout=_idle_timeout(1.0),
|
||
retry_policy=(policy,),
|
||
name="parent",
|
||
)
|
||
with pytest.raises(asyncio.TimeoutError, match="inner"):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert calls == 1
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_does_not_swallow_proc_node_timeout():
|
||
child_timeout = NodeTimeoutError("child", 0.2, kind="idle", idle_timeout=0.1)
|
||
|
||
class ChildTimeoutProc:
|
||
async def ainvoke(self, input, config):
|
||
raise child_timeout
|
||
|
||
task = _make_task(ChildTimeoutProc(), timeout=_idle_timeout(1.0), name="parent")
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value is child_timeout
|
||
assert excinfo.value.node == "child"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_idle_timeout_resets_on_stream_event():
|
||
events = []
|
||
|
||
class StreamingProc:
|
||
async def ainvoke(self, input, config):
|
||
for _ in range(3):
|
||
await asyncio.sleep(0.08)
|
||
config[CONF][CONFIG_KEY_STREAM](((), "custom", "tick"))
|
||
return "ok"
|
||
|
||
task = _make_task(StreamingProc(), timeout=_idle_timeout(0.2), name="streaming")
|
||
task.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol(events.append, {"custom"})
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
assert len(events) == 3
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_idle_timeout_resets_on_runtime_stream_writer():
|
||
events = []
|
||
|
||
class WriterProc:
|
||
async def ainvoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
for _ in range(3):
|
||
await asyncio.sleep(0.08)
|
||
runtime.stream_writer("tick")
|
||
return "ok"
|
||
|
||
task = _make_task(WriterProc(), timeout=_idle_timeout(0.2), name="writer")
|
||
runtime = task.config[CONF][CONFIG_KEY_RUNTIME]
|
||
task.config[CONF][CONFIG_KEY_RUNTIME] = runtime.override(
|
||
stream_writer=events.append
|
||
)
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
assert events == ["tick", "tick", "tick"]
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_astream_with_retry_idle_timeout_resets_on_yielded_chunks():
|
||
class StreamingProc:
|
||
async def astream(self, input, config):
|
||
for i in range(3):
|
||
await asyncio.sleep(0.08)
|
||
yield i
|
||
|
||
task = _make_task(StreamingProc(), timeout=_idle_timeout(0.2), name="astream")
|
||
await arun_with_retry(task, retry_policy=None, stream=True)
|
||
|
||
|
||
class _HandlerEmittingProc:
|
||
"""Proc that fires `on_llm_new_token` on every handler attached to its config."""
|
||
|
||
def __init__(self, iterations: int = 1, sleep_s: float = 0.0) -> None:
|
||
self.iterations = iterations
|
||
self.sleep_s = sleep_s
|
||
|
||
async def ainvoke(self, input, config):
|
||
run_id = uuid4()
|
||
for _ in range(self.iterations):
|
||
if self.sleep_s:
|
||
await asyncio.sleep(self.sleep_s)
|
||
for handler in config["callbacks"]:
|
||
handler.on_llm_new_token("tok", run_id=run_id)
|
||
return "ok"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_idle_timeout_resets_on_runtime_heartbeat():
|
||
class HeartbeatProc:
|
||
async def ainvoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
for _ in range(3):
|
||
await asyncio.sleep(0.08)
|
||
runtime.heartbeat()
|
||
return "ok"
|
||
|
||
task = _make_task(HeartbeatProc(), timeout=_idle_timeout(0.15), name="heartbeat")
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_heartbeat_refresh_mode_ignores_stream_events():
|
||
events = []
|
||
|
||
class StreamingProc:
|
||
async def ainvoke(self, input, config):
|
||
while True:
|
||
await asyncio.sleep(0.01)
|
||
config[CONF][CONFIG_KEY_STREAM](((), "custom", "tick"))
|
||
|
||
task = _make_task(
|
||
StreamingProc(),
|
||
timeout=TimeoutPolicy(idle_timeout=0.05, refresh_on="heartbeat"),
|
||
name="heartbeat-only",
|
||
)
|
||
task.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol(events.append, {"custom"})
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.kind == "idle"
|
||
assert events
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_heartbeat_refresh_mode_accepts_heartbeat():
|
||
class HeartbeatProc:
|
||
async def ainvoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
for _ in range(3):
|
||
await asyncio.sleep(0.03)
|
||
runtime.heartbeat()
|
||
return "ok"
|
||
|
||
task = _make_task(
|
||
HeartbeatProc(),
|
||
timeout=TimeoutPolicy(idle_timeout=0.08, refresh_on="heartbeat"),
|
||
name="heartbeat-only",
|
||
)
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
|
||
def test_runtime_heartbeat_outside_idle_attempt_is_no_op():
|
||
DEFAULT_RUNTIME.heartbeat()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_idle_timeout_resets_on_callback_event():
|
||
task = _make_task(
|
||
_HandlerEmittingProc(iterations=3, sleep_s=0.08),
|
||
timeout=_idle_timeout(0.15),
|
||
name="cb",
|
||
)
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_idle_timeout_preserves_existing_callbacks():
|
||
seen: list[str] = []
|
||
|
||
class RecordingHandler(BaseCallbackHandler):
|
||
run_inline = True
|
||
|
||
def on_llm_new_token(self, token, *, run_id, **kwargs):
|
||
seen.append(token)
|
||
|
||
task = _make_task(_HandlerEmittingProc(), timeout=_idle_timeout(0.5), name="cb-pre")
|
||
task.config["callbacks"] = [RecordingHandler()]
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
assert seen == ["tok"]
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_discards_stale_executor_writes():
|
||
release_first_attempt = threading.Event()
|
||
|
||
class FlakyAsyncProc:
|
||
def __init__(self) -> None:
|
||
self.calls = 0
|
||
|
||
async def ainvoke(self, input, config):
|
||
self.calls += 1
|
||
if self.calls == 1:
|
||
|
||
def late_write() -> str:
|
||
release_first_attempt.wait(timeout=1.0)
|
||
config[CONF][CONFIG_KEY_SEND]([("value", "stale")])
|
||
return "late"
|
||
|
||
return await asyncio.to_thread(late_write)
|
||
release_first_attempt.set()
|
||
config[CONF][CONFIG_KEY_SEND]([("value", "fresh")])
|
||
return "ok"
|
||
|
||
policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.0,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
)
|
||
task = _make_task(
|
||
FlakyAsyncProc(), timeout=_idle_timeout(0.05), retry_policy=(policy,)
|
||
)
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
await asyncio.sleep(0.05)
|
||
assert task.writes == deque([("value", "fresh")])
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_discards_pre_timeout_writes():
|
||
class SlowAsyncWriterProc:
|
||
async def ainvoke(self, input, config):
|
||
config[CONF][CONFIG_KEY_SEND]([("value", "stale-before-idle-timeout")])
|
||
await asyncio.sleep(0.2)
|
||
return "late"
|
||
|
||
task = _make_task(
|
||
SlowAsyncWriterProc(), timeout=_idle_timeout(0.05), name="aslow-writer"
|
||
)
|
||
with pytest.raises(NodeTimeoutError):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert task.writes == deque()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_astream_with_retry_timeout_discards_pre_timeout_writes():
|
||
class SlowStreamWriterProc:
|
||
async def astream(self, input, config):
|
||
config[CONF][CONFIG_KEY_SEND]([("value", "stale-before-idle-timeout")])
|
||
await asyncio.sleep(0.2)
|
||
if False:
|
||
yield None
|
||
|
||
task = _make_task(
|
||
SlowStreamWriterProc(), timeout=_idle_timeout(0.05), name="astream-writer"
|
||
)
|
||
with pytest.raises(NodeTimeoutError):
|
||
await arun_with_retry(task, retry_policy=None, stream=True)
|
||
assert task.writes == deque()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_cannot_be_swallowed():
|
||
class StubbornProc:
|
||
async def ainvoke(self, input, config):
|
||
try:
|
||
await asyncio.sleep(1.0)
|
||
except asyncio.CancelledError:
|
||
config[CONF][CONFIG_KEY_SEND]([("value", "stale")])
|
||
await asyncio.sleep(0)
|
||
return "late"
|
||
return "ok"
|
||
|
||
task = _make_task(StubbornProc(), timeout=_idle_timeout(0.05), name="stubborn")
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.node == "stubborn"
|
||
await asyncio.sleep(0.05)
|
||
assert task.writes == deque()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_astream_with_retry_timeout_cannot_be_swallowed():
|
||
class StubbornStreamProc:
|
||
async def astream(self, input, config):
|
||
try:
|
||
await asyncio.sleep(1.0)
|
||
except asyncio.CancelledError:
|
||
config[CONF][CONFIG_KEY_SEND]([("value", "stale")])
|
||
await asyncio.sleep(0)
|
||
if False:
|
||
yield None
|
||
return
|
||
yield "ok"
|
||
|
||
task = _make_task(
|
||
StubbornStreamProc(), timeout=_idle_timeout(0.05), name="stubborn-stream"
|
||
)
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None, stream=True)
|
||
assert excinfo.value.node == "stubborn-stream"
|
||
await asyncio.sleep(0.05)
|
||
assert task.writes == deque()
|
||
|
||
|
||
class _TimeoutState(TypedDict):
|
||
x: int
|
||
|
||
|
||
def test_timeout_validation_is_eager_across_apis():
|
||
with pytest.raises(ValueError, match="greater than 0"):
|
||
task(timeout=0)
|
||
|
||
with pytest.raises(ValueError, match="greater than 0"):
|
||
entrypoint(timeout=0)
|
||
|
||
with pytest.raises(ValueError, match="greater than 0"):
|
||
NodeBuilder().set_timeout(0)
|
||
|
||
with pytest.raises(ValueError, match="greater than 0"):
|
||
PregelNode(channels="x", triggers=["x"], timeout=0)
|
||
|
||
with pytest.raises(ValueError, match="greater than 0"):
|
||
Send("slow", {}, timeout=0)
|
||
|
||
builder = StateGraph(_TimeoutState)
|
||
with pytest.raises(ValueError, match="greater than 0"):
|
||
builder.add_node("slow", lambda state: state, timeout=0)
|
||
|
||
|
||
def test_timeout_rejects_sync_functional_apis_at_declaration_time():
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
|
||
@task(timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
def sync_task(value: int) -> int:
|
||
return value
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
|
||
@entrypoint(timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
def sync_entrypoint(value: int) -> int:
|
||
return value
|
||
|
||
|
||
def test_state_graph_compile_rejects_sync_node_timeout():
|
||
def slow(state: _TimeoutState) -> _TimeoutState:
|
||
return {"x": state["x"] + 1}
|
||
|
||
builder = StateGraph(_TimeoutState)
|
||
builder.add_node("slow", slow, timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
builder.add_edge(START, "slow")
|
||
builder.add_edge("slow", END)
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
builder.compile()
|
||
|
||
|
||
def test_pregel_validate_rejects_sync_writer_timeout():
|
||
async def bound(value: int) -> int:
|
||
return value + 1
|
||
|
||
def sync_writer(value: int) -> int:
|
||
return value
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
Pregel(
|
||
nodes={
|
||
"slow": PregelNode(
|
||
channels="input",
|
||
triggers=["input"],
|
||
bound=RunnableLambda(bound),
|
||
writers=[RunnableLambda(sync_writer)],
|
||
timeout=TimeoutPolicy(run_timeout=1),
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
|
||
|
||
def test_pregel_validate_rejects_wrapped_sync_runnable_lambda_timeout():
|
||
def slow(value: int) -> int:
|
||
return value + 1
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
Pregel(
|
||
nodes={
|
||
"slow": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(RunnableLambda(slow).with_config(tags=["wrapped"]))
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
|
||
|
||
def test_pregel_validate_accepts_wrapped_async_runnable_lambda_timeout():
|
||
async def slow(value: int) -> int:
|
||
return value + 1
|
||
|
||
Pregel(
|
||
nodes={
|
||
"slow": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(RunnableLambda(slow).with_config(tags=["wrapped"]))
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
|
||
|
||
def test_pregel_validate_rejects_parallel_sync_branch_timeout():
|
||
def sync_branch(value: int) -> int:
|
||
return value + 1
|
||
|
||
async def async_branch(value: int) -> int:
|
||
return value + 1
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
Pregel(
|
||
nodes={
|
||
"parallel": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(
|
||
RunnableParallel(
|
||
sync=RunnableLambda(sync_branch),
|
||
async_=RunnableLambda(async_branch),
|
||
)
|
||
)
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(dict),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
|
||
|
||
def test_pregel_validate_rejects_sync_node_timeout():
|
||
def slow(value: int) -> int:
|
||
return value + 1
|
||
|
||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||
Pregel(
|
||
nodes={
|
||
"slow": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(slow)
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_pregel_validate_accepts_async_runnable_lambda_timeout():
|
||
async def slow(value: int) -> int:
|
||
await asyncio.sleep(0.2)
|
||
return value + 1
|
||
|
||
graph = Pregel(
|
||
nodes={
|
||
"slow": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(RunnableLambda(slow))
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
with pytest.raises(NodeTimeoutError):
|
||
await graph.ainvoke(1)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_pregel_validate_accepts_runnable_callable_with_sync_and_async_timeout():
|
||
def sync(value: int) -> int:
|
||
return value + 1
|
||
|
||
async def async_(value: int) -> int:
|
||
await asyncio.sleep(0.2)
|
||
return value + 1
|
||
|
||
graph = Pregel(
|
||
nodes={
|
||
"slow": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(RunnableCallable(sync, async_))
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
with pytest.raises(NodeTimeoutError):
|
||
await graph.ainvoke(1)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_state_graph_add_node_timeout_e2e():
|
||
async def slow(state: _TimeoutState) -> _TimeoutState:
|
||
await asyncio.sleep(1.0)
|
||
return {"x": state["x"] + 1}
|
||
|
||
builder = StateGraph(_TimeoutState)
|
||
builder.add_node("slow", slow, timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
builder.add_edge(START, "slow")
|
||
builder.add_edge("slow", END)
|
||
graph = builder.compile()
|
||
with pytest.raises(NodeTimeoutError):
|
||
await graph.ainvoke({"x": 1})
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_send_timeout_overrides_target_node_timeout():
|
||
async def slow(state: _TimeoutState) -> _TimeoutState:
|
||
await asyncio.sleep(0.2)
|
||
return {"x": state["x"] + 1}
|
||
|
||
def route(state: _TimeoutState) -> list[Send]:
|
||
return [Send("slow", state, timeout=TimeoutPolicy(idle_timeout=0.05))]
|
||
|
||
builder = StateGraph(_TimeoutState)
|
||
builder.add_node("slow", slow, timeout=TimeoutPolicy(idle_timeout=1.0))
|
||
builder.add_conditional_edges(START, route)
|
||
builder.add_edge("slow", END)
|
||
graph = builder.compile()
|
||
|
||
with pytest.raises(NodeTimeoutError) as excinfo:
|
||
await graph.ainvoke({"x": 1})
|
||
assert excinfo.value.node == "slow"
|
||
assert excinfo.value.kind == "idle"
|
||
assert excinfo.value.idle_timeout == 0.05
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_state_graph_add_node_timeout_composes_with_retry():
|
||
"""add_node(..., timeout=TimeoutPolicy(...)) retries then succeeds."""
|
||
|
||
attempts: list[int] = []
|
||
|
||
async def flaky(state: _TimeoutState) -> _TimeoutState:
|
||
attempts.append(len(attempts))
|
||
if len(attempts) < 2:
|
||
await asyncio.sleep(1.0)
|
||
return {"x": state["x"] + 1}
|
||
|
||
builder = StateGraph(_TimeoutState)
|
||
builder.add_node(
|
||
"flaky",
|
||
flaky,
|
||
timeout=TimeoutPolicy(idle_timeout=0.3),
|
||
retry_policy=RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.0,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
),
|
||
)
|
||
builder.add_edge(START, "flaky")
|
||
builder.add_edge("flaky", END)
|
||
graph = builder.compile()
|
||
result = await graph.ainvoke({"x": 0})
|
||
assert result == {"x": 1}
|
||
assert len(attempts) == 2
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
@pytest.mark.anyio
|
||
async def test_task_decorator_timeout_e2e():
|
||
@task(timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
async def slow_task(x: int) -> int:
|
||
await asyncio.sleep(0.2)
|
||
return x + 1
|
||
|
||
@entrypoint()
|
||
async def workflow(x: int) -> int:
|
||
return await slow_task(x)
|
||
|
||
with pytest.raises(NodeTimeoutError):
|
||
await workflow.ainvoke(1)
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
@pytest.mark.anyio
|
||
async def test_task_decorator_preserves_user_idle_timeout_kwarg():
|
||
@task(timeout=TimeoutPolicy(idle_timeout=1.0))
|
||
async def echo_idle_timeout(*, idle_timeout: int) -> int:
|
||
await asyncio.sleep(0)
|
||
return idle_timeout
|
||
|
||
@entrypoint()
|
||
async def workflow(x: int) -> int:
|
||
return await echo_idle_timeout(idle_timeout=x)
|
||
|
||
assert await workflow.ainvoke(5) == 5
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
@pytest.mark.anyio
|
||
async def test_task_decorator_preserves_user_timeout_kwarg():
|
||
@task(timeout=1.0)
|
||
async def echo_timeout(*, timeout: int) -> int:
|
||
await asyncio.sleep(0)
|
||
return timeout
|
||
|
||
@entrypoint()
|
||
async def workflow(x: int) -> int:
|
||
return await echo_timeout(timeout=x)
|
||
|
||
assert await workflow.ainvoke(5) == 5
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_entrypoint_timeout_e2e():
|
||
@entrypoint(timeout=TimeoutPolicy(idle_timeout=0.05))
|
||
async def slow_workflow(x: int) -> int:
|
||
await asyncio.sleep(0.2)
|
||
return x
|
||
|
||
with pytest.raises(NodeTimeoutError):
|
||
await slow_workflow.ainvoke(1)
|
||
|
||
|
||
class _MessageStreamState(TypedDict):
|
||
messages: Annotated[list[BaseMessage], add_messages]
|
||
|
||
|
||
class _SlowStreamingChatModel(GenericFakeChatModel):
|
||
async def _astream(
|
||
self,
|
||
messages: list[BaseMessage],
|
||
stop: list[str] | None = None,
|
||
run_manager: AsyncCallbackManagerForLLMRun | None = None,
|
||
**kwargs: Any,
|
||
) -> AsyncIterator[ChatGenerationChunk]:
|
||
for i in range(3):
|
||
await asyncio.sleep(0.08)
|
||
chunk = ChatGenerationChunk(
|
||
message=AIMessageChunk(
|
||
content=str(i),
|
||
chunk_position="last" if i == 2 else None,
|
||
)
|
||
)
|
||
if run_manager:
|
||
await run_manager.on_llm_new_token(str(i), chunk=chunk)
|
||
yield chunk
|
||
|
||
def _generate(
|
||
self,
|
||
messages: list[BaseMessage],
|
||
stop: list[str] | None = None,
|
||
run_manager: Any | None = None,
|
||
**kwargs: Any,
|
||
) -> ChatResult:
|
||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=""))])
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_idle_timeout_resets_on_message_stream_callbacks():
|
||
model = _SlowStreamingChatModel(messages=iter([]))
|
||
|
||
async def call_model(state: _MessageStreamState) -> _MessageStreamState:
|
||
response = await model.ainvoke(state["messages"])
|
||
return {"messages": [response]}
|
||
|
||
builder = StateGraph(_MessageStreamState)
|
||
builder.add_node(
|
||
"call_model",
|
||
call_model,
|
||
timeout=TimeoutPolicy(idle_timeout=0.15),
|
||
)
|
||
builder.add_edge(START, "call_model")
|
||
builder.add_edge("call_model", END)
|
||
graph = builder.compile()
|
||
|
||
chunks: list[str] = []
|
||
async for chunk, _metadata in graph.astream(
|
||
{"messages": [HumanMessage(content="hi")]},
|
||
stream_mode="messages",
|
||
):
|
||
chunks.append(chunk.content)
|
||
assert chunks == ["0", "1", "2"]
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_node_builder_timeout_e2e():
|
||
async def slow(value: int) -> int:
|
||
await asyncio.sleep(0.2)
|
||
return value + 1
|
||
|
||
graph = Pregel(
|
||
nodes={
|
||
"slow": (
|
||
NodeBuilder()
|
||
.subscribe_only("input")
|
||
.do(slow)
|
||
.set_timeout(TimeoutPolicy(idle_timeout=0.05))
|
||
.write_to("output")
|
||
)
|
||
},
|
||
channels={
|
||
"input": EphemeralValue(int),
|
||
"output": LastValue(int),
|
||
},
|
||
input_channels="input",
|
||
output_channels="output",
|
||
)
|
||
with pytest.raises(NodeTimeoutError):
|
||
await graph.ainvoke(1)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_observer_tracks_attempts():
|
||
events: list = []
|
||
|
||
class FlakyProc:
|
||
async def ainvoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
if runtime.execution_info.node_attempt == 1:
|
||
await asyncio.sleep(0.2)
|
||
return "ok"
|
||
|
||
policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.0,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
)
|
||
task = _make_task(
|
||
FlakyProc(),
|
||
timeout=_idle_timeout(0.05),
|
||
retry_policy=(policy,),
|
||
name="flaky",
|
||
)
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
starts = [event for event in events if event.event == "start"]
|
||
finishes = [event for event in events if event.event == "finish"]
|
||
assert [event.context.attempt for event in starts] == [1, 2]
|
||
assert [event.context.attempt for event in finishes] == [1, 2]
|
||
assert [event.status for event in finishes] == ["error", "success"]
|
||
assert starts[0].context.idle_timeout_secs == 0.05
|
||
assert starts[0].context.task_name == "flaky"
|
||
assert isinstance(starts[0].context.started_at, datetime)
|
||
assert isinstance(finishes[0].finished_at, datetime)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_observer_emits_progress_on_heartbeat():
|
||
events: list = []
|
||
|
||
# `_TimedAttemptScope.__init__` sets `_last_progress` to `time.monotonic()`,
|
||
# but the watchdog itself doesn't start running until after `wrap_config`
|
||
# and task scheduling — under CI load that gap can be large enough to eat
|
||
# the entire idle window before the task body's first await even runs. We
|
||
# defend against that by:
|
||
# 1. Using a generous idle_timeout so scheduling slack stays well within it.
|
||
# 2. Calling `runtime.heartbeat()` BEFORE the first sleep, which resets
|
||
# `_last_progress` to "now" the moment the task body actually starts.
|
||
idle_timeout_s = 1.0
|
||
|
||
class HeartbeatProc:
|
||
async def ainvoke(self, input, config):
|
||
runtime = config[CONF][CONFIG_KEY_RUNTIME]
|
||
runtime.heartbeat() # reset the idle clock at task-body entry
|
||
for _ in range(8):
|
||
await asyncio.sleep(0.05)
|
||
runtime.heartbeat()
|
||
return "ok"
|
||
|
||
task = _make_task(
|
||
HeartbeatProc(), timeout=_idle_timeout(idle_timeout_s), name="heartbeat"
|
||
)
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
by_event = [ev.event for ev in events]
|
||
assert by_event[0] == "start"
|
||
assert by_event[-1] == "finish"
|
||
progress = [ev for ev in events if ev.event == "progress"]
|
||
assert progress, "expected at least one progress event from heartbeat"
|
||
# Rate limit is `idle_timeout / 4` = 0.25s; with the task running for
|
||
# ~400ms we expect 1–2 progress events (well below the 9 heartbeats).
|
||
assert len(progress) <= len(by_event)
|
||
for ev in progress:
|
||
assert ev.context.task_name == "heartbeat"
|
||
assert ev.context.attempt == 1
|
||
assert ev.context.idle_timeout_secs == idle_timeout_s
|
||
assert isinstance(ev.progress_at, datetime)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_observer_treats_parent_command_as_non_error():
|
||
events: list = []
|
||
|
||
class ParentProc:
|
||
async def ainvoke(self, input, config):
|
||
raise ParentCommand(Command(graph=Command.PARENT))
|
||
|
||
task = _make_task(ParentProc(), timeout=_idle_timeout(0.05), name="parent")
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
with pytest.raises(ParentCommand):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
|
||
finish = next(event for event in events if event.event == "finish")
|
||
assert finish.status == "success"
|
||
assert finish.error_type is None
|
||
assert finish.error_message is None
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_observer_finishes_when_parent_writer_errors():
|
||
events: list = []
|
||
|
||
class ParentProc:
|
||
async def ainvoke(self, input, config):
|
||
raise ParentCommand(Command(graph="parent", update={"value": "updated"}))
|
||
|
||
class FailingWriter:
|
||
def invoke(self, input, config):
|
||
raise ValueError("writer failed")
|
||
|
||
task = _make_task(
|
||
ParentProc(),
|
||
timeout=_idle_timeout(0.05),
|
||
name="parent",
|
||
writers=(FailingWriter(),),
|
||
)
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
with pytest.raises(ValueError, match="writer failed"):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
|
||
finish = next(event for event in events if event.event == "finish")
|
||
assert finish.status == "error"
|
||
assert finish.error_type == "ValueError"
|
||
assert finish.error_message == "writer failed"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_timeout_observer_treats_bubble_up_as_non_error():
|
||
events: list = []
|
||
|
||
class BubbleProc:
|
||
async def ainvoke(self, input, config):
|
||
raise GraphInterrupt(())
|
||
|
||
task = _make_task(BubbleProc(), timeout=_idle_timeout(0.05), name="bubble")
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
with pytest.raises(GraphInterrupt):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
|
||
finish = next(event for event in events if event.event == "finish")
|
||
assert finish.status == "success"
|
||
assert finish.error_type is None
|
||
assert finish.error_message is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Watcher invariant: any timeout that retry/error_handler can recover from
|
||
# MUST emit `finish=error` BEFORE the in-process recovery work happens. The
|
||
# external watchdog (langgraph-api) relies on this so it only kills a worker
|
||
# when no `finish` arrives within the deadline. The tests below pin down the
|
||
# three recovery paths so a refactor that moves `_finish_timed_attempt` past
|
||
# an `await` (or past the final `raise`) trips CI.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_observer_emits_finish_before_retry_backoff():
|
||
"""`finish=error` of attempt N must arrive before the retry backoff sleep."""
|
||
timeline: list[tuple[float, Any]] = []
|
||
|
||
class TimingOutOnceProc:
|
||
def __init__(self) -> None:
|
||
self.calls = 0
|
||
|
||
async def ainvoke(self, input, config):
|
||
self.calls += 1
|
||
if self.calls == 1:
|
||
await asyncio.sleep(1.0)
|
||
return "ok"
|
||
|
||
backoff = 0.25
|
||
policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=backoff,
|
||
backoff_factor=1.0,
|
||
max_interval=backoff,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
)
|
||
task = _make_task(
|
||
TimingOutOnceProc(),
|
||
timeout=_idle_timeout(0.05),
|
||
retry_policy=(policy,),
|
||
name="backoff_watcher",
|
||
)
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = lambda ev: timeline.append(
|
||
(time.monotonic(), ev)
|
||
)
|
||
|
||
assert await arun_with_retry(task, retry_policy=None) == "ok"
|
||
|
||
starts = [(t, ev) for t, ev in timeline if ev.event == "start"]
|
||
finishes = [(t, ev) for t, ev in timeline if ev.event == "finish"]
|
||
assert [ev.context.attempt for _, ev in starts] == [1, 2]
|
||
assert [ev.status for _, ev in finishes] == ["error", "success"]
|
||
|
||
first_finish_t = finishes[0][0]
|
||
second_start_t = starts[1][0]
|
||
|
||
# The watcher relies on this gap: `finish=error` for attempt 1 must arrive
|
||
# before `arun_with_retry` enters `await asyncio.sleep(backoff)`. We give a
|
||
# generous slack to keep this stable on slow CI; the structural invariant
|
||
# is "finish lands first", not "the gap equals exactly backoff".
|
||
assert second_start_t - first_finish_t >= backoff * 0.5, (
|
||
f"finish=error appears to be emitted after retry backoff sleep; "
|
||
f"gap was {second_start_t - first_finish_t:.3f}s, expected >= {backoff * 0.5:.3f}s"
|
||
)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_state_graph_observer_emits_finish_before_error_handler_start():
|
||
"""Original task's `finish=error` must arrive before the error_handler task's `start`."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
async def slow_node(state: State) -> State:
|
||
await asyncio.sleep(1.0)
|
||
return {"foo": "should-not-happen"}
|
||
|
||
async def handler_node(state: State, error: NodeError) -> State:
|
||
return {"foo": "handled"}
|
||
|
||
events: list = []
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node(
|
||
"slow",
|
||
slow_node,
|
||
timeout=TimeoutPolicy(idle_timeout=0.05),
|
||
error_handler=handler_node,
|
||
)
|
||
.add_edge(START, "slow")
|
||
.compile()
|
||
)
|
||
|
||
result = await graph.ainvoke(
|
||
{"foo": ""},
|
||
config={
|
||
"configurable": {CONFIG_KEY_TIMED_ATTEMPT_OBSERVER: events.append},
|
||
},
|
||
)
|
||
assert result["foo"] == "handled"
|
||
|
||
# Filter to events from the failing node only — the handler node has no
|
||
# timeout configured here, so it doesn't appear in the observer stream.
|
||
slow_events = [ev for ev in events if ev.context.task_name == "slow"]
|
||
starts = [ev for ev in slow_events if ev.event == "start"]
|
||
finishes = [ev for ev in slow_events if ev.event == "finish"]
|
||
assert len(starts) == 1
|
||
assert len(finishes) == 1
|
||
assert finishes[0].status == "error"
|
||
assert finishes[0].error_type == "NodeTimeoutError"
|
||
# The slow task's finish-error event must precede every event for any
|
||
# follow-up task in the same observer stream.
|
||
slow_finish_index = events.index(finishes[0])
|
||
for ev in events[slow_finish_index + 1 :]:
|
||
assert ev.context.task_name == "slow" or ev.event == "start", (
|
||
f"unexpected event {ev.event} for {ev.context.task_name} "
|
||
f"after slow's finish=error"
|
||
)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_observer_emits_finish_before_final_raise_on_exhaustion():
|
||
"""When retry exhausts and the timeout propagates, the final `finish=error` must
|
||
be emitted before `arun_with_retry` re-raises."""
|
||
events: list = []
|
||
|
||
class AlwaysTimingOutProc:
|
||
async def ainvoke(self, input, config):
|
||
await asyncio.sleep(1.0)
|
||
return "never"
|
||
|
||
policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.0,
|
||
jitter=False,
|
||
retry_on=NodeTimeoutError,
|
||
)
|
||
task = _make_task(
|
||
AlwaysTimingOutProc(),
|
||
timeout=_idle_timeout(0.05),
|
||
retry_policy=(policy,),
|
||
name="never_succeeds",
|
||
)
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
|
||
with pytest.raises(NodeTimeoutError):
|
||
await arun_with_retry(task, retry_policy=None)
|
||
|
||
starts = [ev for ev in events if ev.event == "start"]
|
||
finishes = [ev for ev in events if ev.event == "finish"]
|
||
assert [ev.context.attempt for ev in starts] == [1, 2]
|
||
assert [ev.context.attempt for ev in finishes] == [1, 2]
|
||
assert [ev.status for ev in finishes] == ["error", "error"]
|
||
assert all(ev.error_type == "NodeTimeoutError" for ev in finishes)
|
||
# Both finish events were observed BEFORE arun_with_retry raised, otherwise
|
||
# the `with pytest.raises` block would have exited before `events` got
|
||
# populated with the second finish.
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_sync_sleep_in_async_node_bypasses_timeout_and_emits_finish_success():
|
||
"""Sync `time.sleep` inside an async node blocks the event loop so the
|
||
in-process watchdog cannot fire. We document the resulting behavior here:
|
||
|
||
1. `NodeTimeoutError` is NOT raised, even though the sync sleep exceeds
|
||
`idle_timeout`.
|
||
2. The node's normal return value flows through.
|
||
3. `finish=success` is emitted to the observer.
|
||
|
||
This is the canonical case where the in-process timeout is defeated and
|
||
the only safety net is the external watcher (langgraph-api), which
|
||
SIGKILLs the worker when no `finish` arrives within its deadline. The
|
||
catch is that with a *short* sync sleep the event loop unblocks before
|
||
the watcher's deadline expires, so the watcher legitimately does not
|
||
kill — meaning the configured `idle_timeout` is silently honored at the
|
||
process level only when the block is long enough to outlast the
|
||
watcher's grace.
|
||
|
||
This is the documented "Cooperative cancellation" caveat on
|
||
`TimeoutPolicy`. The test pins the behavior so any future change that
|
||
starts raising `NodeTimeoutError` for sync-blocked async nodes (or stops
|
||
emitting `finish=success`) is caught.
|
||
"""
|
||
events: list = []
|
||
|
||
class SyncSleepingProc:
|
||
async def ainvoke(self, input, config):
|
||
time.sleep(0.1)
|
||
return "completed_despite_timeout"
|
||
|
||
task = _make_task(
|
||
SyncSleepingProc(),
|
||
timeout=_idle_timeout(0.05),
|
||
name="sync_sleeper",
|
||
)
|
||
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
|
||
|
||
result = await arun_with_retry(task, retry_policy=None)
|
||
|
||
assert result == "completed_despite_timeout"
|
||
starts = [ev for ev in events if ev.event == "start"]
|
||
finishes = [ev for ev in events if ev.event == "finish"]
|
||
assert len(starts) == 1
|
||
assert len(finishes) == 1
|
||
assert finishes[0].status == "success"
|
||
assert finishes[0].error_type is None
|
||
|
||
|
||
def test_graph_error_handler_runs_after_retry_exhaustion():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempts = 0
|
||
captured: dict[str, object] = {}
|
||
|
||
def always_failing_node(state: State) -> State:
|
||
nonlocal attempts
|
||
attempts += 1
|
||
raise ValueError("Always fails")
|
||
|
||
def err_handler_node(state: State, error: NodeError) -> Command:
|
||
captured["from_node_name"] = error.node
|
||
captured["from_node_error"] = error.error
|
||
return Command(update={"foo": "handled"}, goto="after_handler")
|
||
|
||
def after_handler(state: State) -> State:
|
||
return {"foo": f"{state['foo']}_after"}
|
||
|
||
retry_policy = RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node(
|
||
"always_failing",
|
||
always_failing_node,
|
||
retry_policy=retry_policy,
|
||
error_handler=err_handler_node,
|
||
)
|
||
.add_node("after_handler", after_handler)
|
||
.add_edge(START, "always_failing")
|
||
.compile()
|
||
)
|
||
|
||
with patch("time.sleep"):
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
assert attempts == 2
|
||
assert result["foo"] == "handled_after"
|
||
assert captured["from_node_name"] == "always_failing"
|
||
assert isinstance(captured["from_node_error"], BaseException)
|
||
|
||
|
||
def test_graph_error_handler_can_route_with_command():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempts = 0
|
||
|
||
def always_failing_node(state: State) -> State:
|
||
nonlocal attempts
|
||
attempts += 1
|
||
raise ValueError("Always fails")
|
||
|
||
def err_handler_node(state: State) -> Command:
|
||
return Command(update={"foo": "handled"}, goto="next_node")
|
||
|
||
def next_node(state: State) -> State:
|
||
return {"foo": f"{state['foo']}_next"}
|
||
|
||
retry_policy = RetryPolicy(
|
||
max_attempts=1,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node(
|
||
"always_failing",
|
||
always_failing_node,
|
||
retry_policy=retry_policy,
|
||
error_handler=err_handler_node,
|
||
)
|
||
.add_node("next_node", next_node)
|
||
.add_edge(START, "always_failing")
|
||
.compile()
|
||
)
|
||
|
||
result = graph.invoke({"foo": ""})
|
||
assert attempts == 1
|
||
assert result["foo"] == "handled_next"
|
||
|
||
|
||
def test_graph_error_handler_failure_fails_run():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def always_failing_node(state: State) -> State:
|
||
raise ValueError("Always fails")
|
||
|
||
def err_handler_node(state: State) -> State:
|
||
raise RuntimeError("handler failed")
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("always_failing", always_failing_node, error_handler=err_handler_node)
|
||
.add_edge(START, "always_failing")
|
||
.compile()
|
||
)
|
||
|
||
with pytest.raises(RuntimeError, match="handler failed"):
|
||
graph.invoke({"foo": ""})
|
||
|
||
|
||
def test_graph_error_handler_handles_subgraph_internal_failure():
|
||
class SubState(TypedDict):
|
||
foo: str
|
||
|
||
class ParentState(TypedDict):
|
||
foo: str
|
||
|
||
parent_handler_called = False
|
||
captured: dict[str, object] = {}
|
||
|
||
def sub_fail_node(state: SubState) -> SubState:
|
||
raise ValueError("subgraph boom")
|
||
|
||
def parent_handler(state: ParentState, error: NodeError) -> ParentState:
|
||
nonlocal parent_handler_called
|
||
parent_handler_called = True
|
||
captured["from_node_name"] = error.node
|
||
captured["from_node_error"] = error.error
|
||
return {"foo": "handled_by_parent"}
|
||
|
||
subgraph = (
|
||
StateGraph(SubState)
|
||
.add_node("sub_fail_node", sub_fail_node)
|
||
.add_edge(START, "sub_fail_node")
|
||
.compile()
|
||
)
|
||
|
||
parent_graph = (
|
||
StateGraph(ParentState)
|
||
.add_node("subgraph_node", subgraph, error_handler=parent_handler)
|
||
.add_edge(START, "subgraph_node")
|
||
.compile()
|
||
)
|
||
|
||
result = parent_graph.invoke({"foo": ""})
|
||
assert result["foo"] == "handled_by_parent"
|
||
assert parent_handler_called is True
|
||
assert captured["from_node_name"] == "subgraph_node"
|
||
assert isinstance(captured["from_node_error"], BaseException)
|
||
|
||
|
||
def test_graph_error_handler_error_context_survives_checkpoint_resume():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
captured: dict[str, object] = {}
|
||
|
||
def always_failing_node(state: State) -> State:
|
||
raise RuntimeError("failed before handler")
|
||
|
||
def err_handler_node(state: State, error: NodeError) -> State:
|
||
captured["from_node_name"] = error.node
|
||
captured["from_node_error"] = error.error
|
||
return {"foo": "handled_after_resume"}
|
||
|
||
checkpointer = InMemorySaver()
|
||
config = {"configurable": {"thread_id": "graph-error-resume"}}
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("always_failing", always_failing_node, error_handler=err_handler_node)
|
||
.add_edge(START, "always_failing")
|
||
.compile(
|
||
checkpointer=checkpointer,
|
||
interrupt_before=["__error_handler__always_failing"],
|
||
)
|
||
)
|
||
|
||
# First run pauses before handler, after failure context is checkpointed.
|
||
graph.invoke({"foo": ""}, config)
|
||
# Resume should execute handler and recover serialized error context.
|
||
result = graph.invoke(None, config)
|
||
|
||
assert result["foo"] == "handled_after_resume"
|
||
assert captured["from_node_name"] == "always_failing"
|
||
assert isinstance(captured["from_node_error"], BaseException)
|
||
|
||
|
||
def test_graph_error_handler_does_not_swallow_interrupt_concurrent():
|
||
"""When a graph error handler is configured and a node calls interrupt()
|
||
concurrently with other nodes, the interrupt must still be raised — not
|
||
silently swallowed."""
|
||
from langgraph.types import interrupt
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def node_a(state: State) -> State:
|
||
# This node uses interrupt() which raises GraphInterrupt
|
||
val = interrupt("need human input")
|
||
return {"foo": f"a_{val}"}
|
||
|
||
def node_b(state: State) -> State:
|
||
return {}
|
||
|
||
def err_handler(state: State) -> State:
|
||
return {"foo": "handled"}
|
||
|
||
checkpointer = InMemorySaver()
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("node_a", node_a, error_handler=err_handler)
|
||
.add_node("node_b", node_b)
|
||
# Fan-out: both node_a and node_b run concurrently
|
||
.add_edge(START, "node_a")
|
||
.add_edge(START, "node_b")
|
||
.compile(checkpointer=checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "test-interrupt-concurrent"}}
|
||
|
||
# First invoke should pause at the interrupt, not silently complete
|
||
graph.invoke({"foo": ""}, config)
|
||
|
||
# The graph should have an interrupt pending
|
||
state = graph.get_state(config)
|
||
assert len(state.tasks) > 0
|
||
|
||
# There should be a pending interrupt from node_a
|
||
interrupts = [t for t in state.tasks if hasattr(t, "interrupts") and t.interrupts]
|
||
assert len(interrupts) > 0, (
|
||
"GraphInterrupt was swallowed — interrupt() in node_a "
|
||
"should have paused execution"
|
||
)
|
||
|
||
|
||
def test_node_error_handlers_route_to_matching_handler():
|
||
class State(TypedDict):
|
||
route: str
|
||
foo: Annotated[list[str], operator.add]
|
||
|
||
def route_node(state: State) -> State:
|
||
return {"foo": []}
|
||
|
||
def choose_node(state: State) -> str:
|
||
return state["route"]
|
||
|
||
def fail_a(state: State) -> State:
|
||
raise ValueError("a failed")
|
||
|
||
def fail_b(state: State) -> State:
|
||
raise RuntimeError("b failed")
|
||
|
||
def handler_a(state: State, error: NodeError) -> State:
|
||
assert error.node == "fail_a"
|
||
return {"foo": ["handled_a"]}
|
||
|
||
def handler_b(state: State, error: NodeError) -> State:
|
||
assert error.node == "fail_b"
|
||
return {"foo": ["handled_b"]}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("route_node", route_node)
|
||
.add_node("fail_a", fail_a, error_handler=handler_a)
|
||
.add_node("fail_b", fail_b, error_handler=handler_b)
|
||
.add_edge(START, "route_node")
|
||
.add_conditional_edges("route_node", choose_node, path_map=["fail_a", "fail_b"])
|
||
.compile()
|
||
)
|
||
|
||
result_a = graph.invoke({"route": "fail_a", "foo": []})
|
||
result_b = graph.invoke({"route": "fail_b", "foo": []})
|
||
assert result_a["foo"] == ["handled_a"]
|
||
assert result_b["foo"] == ["handled_b"]
|
||
|
||
|
||
def test_node_without_error_handler_still_fails_run():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def fail_without_handler(state: State) -> State:
|
||
raise ValueError("no handler")
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("fail_without_handler", fail_without_handler)
|
||
.add_edge(START, "fail_without_handler")
|
||
.compile()
|
||
)
|
||
|
||
with pytest.raises(ValueError, match="no handler"):
|
||
graph.invoke({"foo": ""})
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# set_node_defaults()
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_set_node_defaults_error_handler_catches_all_nodes():
|
||
class State(TypedDict):
|
||
route: str
|
||
foo: Annotated[list[str], operator.add]
|
||
|
||
def route_node(state: State) -> Command:
|
||
return Command(goto=state["route"])
|
||
|
||
def fail_a(state: State) -> State:
|
||
raise RuntimeError("a failed")
|
||
|
||
def fail_b(state: State) -> State:
|
||
raise RuntimeError("b failed")
|
||
|
||
captured: dict[str, list[str]] = {"nodes": []}
|
||
|
||
def default_handler(state: State, error: NodeError) -> State:
|
||
captured["nodes"].append(error.node)
|
||
return {"foo": [f"handled_{error.node}"]}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=default_handler)
|
||
.add_node("route_node", route_node)
|
||
.add_node("fail_a", fail_a)
|
||
.add_node("fail_b", fail_b)
|
||
.add_edge(START, "route_node")
|
||
.add_conditional_edges(
|
||
"route_node", lambda s: s["route"], path_map=["fail_a", "fail_b"]
|
||
)
|
||
.compile()
|
||
)
|
||
|
||
result_a = graph.invoke({"route": "fail_a", "foo": []})
|
||
result_b = graph.invoke({"route": "fail_b", "foo": []})
|
||
assert result_a["foo"] == ["handled_fail_a"]
|
||
assert result_b["foo"] == ["handled_fail_b"]
|
||
assert "fail_a" in captured["nodes"]
|
||
assert "fail_b" in captured["nodes"]
|
||
|
||
|
||
def test_set_node_defaults_error_handler_overridden_by_node_handler():
|
||
class State(TypedDict):
|
||
route: str
|
||
foo: Annotated[list[str], operator.add]
|
||
|
||
def route_node(state: State) -> Command:
|
||
return Command(goto=state["route"])
|
||
|
||
def fail_a(state: State) -> State:
|
||
raise RuntimeError("a failed")
|
||
|
||
def fail_b(state: State) -> State:
|
||
raise RuntimeError("b failed")
|
||
|
||
captured: dict[str, list[str]] = {"handler": []}
|
||
|
||
def node_handler(state: State, error: NodeError) -> State:
|
||
captured["handler"].append(f"node:{error.node}")
|
||
return {"foo": [f"node_handled_{error.node}"]}
|
||
|
||
def default_handler(state: State, error: NodeError) -> State:
|
||
captured["handler"].append(f"default:{error.node}")
|
||
return {"foo": [f"default_handled_{error.node}"]}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=default_handler)
|
||
.add_node("route_node", route_node)
|
||
.add_node("fail_a", fail_a, error_handler=node_handler)
|
||
.add_node("fail_b", fail_b)
|
||
.add_edge(START, "route_node")
|
||
.add_conditional_edges(
|
||
"route_node", lambda s: s["route"], path_map=["fail_a", "fail_b"]
|
||
)
|
||
.compile()
|
||
)
|
||
|
||
result_a = graph.invoke({"route": "fail_a", "foo": []})
|
||
assert result_a["foo"] == ["node_handled_fail_a"]
|
||
assert "node:fail_a" in captured["handler"]
|
||
assert "default:fail_a" not in captured["handler"]
|
||
|
||
result_b = graph.invoke({"route": "fail_b", "foo": []})
|
||
assert result_b["foo"] == ["default_handled_fail_b"]
|
||
assert "default:fail_b" in captured["handler"]
|
||
|
||
|
||
def test_set_node_defaults_error_handler_skips_per_node_handler_nodes():
|
||
"""If a per-node error handler itself raises, the default handler must NOT
|
||
catch it -- the run should fail."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def always_failing(state: State) -> State:
|
||
raise RuntimeError("node boom")
|
||
|
||
def broken_handler(state: State, error: NodeError) -> State:
|
||
raise RuntimeError("handler boom")
|
||
|
||
def default_handler(state: State, error: NodeError) -> State:
|
||
return {"foo": "default recovered"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=default_handler)
|
||
.add_node("always_failing", always_failing, error_handler=broken_handler)
|
||
.add_edge(START, "always_failing")
|
||
.compile()
|
||
)
|
||
|
||
with pytest.raises(RuntimeError, match="handler boom"):
|
||
graph.invoke({"foo": ""})
|
||
|
||
|
||
def test_set_node_defaults_error_handler_failure_fails_run():
|
||
"""When the default handler itself raises, the run fails (no infinite
|
||
recursion, no double-routing)."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def always_failing(state: State) -> State:
|
||
raise RuntimeError("node boom")
|
||
|
||
def broken_default_handler(state: State, error: NodeError) -> State:
|
||
raise RuntimeError("default handler boom")
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=broken_default_handler)
|
||
.add_node("always_failing", always_failing)
|
||
.add_edge(START, "always_failing")
|
||
.compile()
|
||
)
|
||
|
||
with pytest.raises(RuntimeError, match="default handler boom"):
|
||
graph.invoke({"foo": ""})
|
||
|
||
|
||
def test_set_node_defaults_error_handler_receives_runnable_config():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def always_failing(state: State) -> State:
|
||
raise RuntimeError("boom")
|
||
|
||
captured: dict[str, Any] = {}
|
||
|
||
def default_handler(
|
||
state: State, error: NodeError, config: RunnableConfig
|
||
) -> State:
|
||
captured["thread_id"] = config["configurable"].get("thread_id")
|
||
return {"foo": "handled"}
|
||
|
||
checkpointer = MemorySaver()
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=default_handler)
|
||
.add_node("always_failing", always_failing)
|
||
.add_edge(START, "always_failing")
|
||
.compile(checkpointer=checkpointer)
|
||
)
|
||
|
||
thread_id = str(uuid4())
|
||
result = graph.invoke(
|
||
{"foo": ""}, config={"configurable": {"thread_id": thread_id}}
|
||
)
|
||
assert result["foo"] == "handled"
|
||
assert captured["thread_id"] == thread_id
|
||
|
||
|
||
def test_set_node_defaults_error_handler_collides_with_user_node():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def default_handler(state: State, error: NodeError) -> State:
|
||
return {"foo": "handled"}
|
||
|
||
builder = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=default_handler)
|
||
.add_node("__default_error_handler__", lambda s: s)
|
||
.add_edge(START, "__default_error_handler__")
|
||
)
|
||
|
||
with pytest.raises(ValueError, match="__default_error_handler__"):
|
||
builder.compile()
|
||
|
||
|
||
def test_set_node_defaults_retry_policy():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempts = 0
|
||
|
||
def flaky_node(state: State) -> State:
|
||
nonlocal attempts
|
||
attempts += 1
|
||
if attempts < 3:
|
||
raise ValueError("not yet")
|
||
return {"foo": "ok"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(
|
||
retry_policy=RetryPolicy(
|
||
max_attempts=3, initial_interval=0.01, jitter=False, retry_on=ValueError
|
||
)
|
||
)
|
||
.add_node("flaky", flaky_node)
|
||
.add_edge(START, "flaky")
|
||
.compile()
|
||
)
|
||
|
||
with patch("time.sleep"):
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
assert result["foo"] == "ok"
|
||
assert attempts == 3
|
||
|
||
|
||
def test_set_node_defaults_retry_policy_per_node_wins():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempts = 0
|
||
|
||
def flaky_node(state: State) -> State:
|
||
nonlocal attempts
|
||
attempts += 1
|
||
if attempts < 2:
|
||
raise ValueError("not yet")
|
||
return {"foo": "ok"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(
|
||
retry_policy=RetryPolicy(
|
||
max_attempts=1, initial_interval=0.01, jitter=False, retry_on=ValueError
|
||
)
|
||
)
|
||
.add_node(
|
||
"flaky",
|
||
flaky_node,
|
||
retry_policy=RetryPolicy(
|
||
max_attempts=3,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
),
|
||
)
|
||
.add_edge(START, "flaky")
|
||
.compile()
|
||
)
|
||
|
||
with patch("time.sleep"):
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
assert result["foo"] == "ok"
|
||
assert attempts == 2
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_set_node_defaults_timeout():
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
async def slow_node(state: State) -> State:
|
||
await asyncio.sleep(10)
|
||
return {"foo": "should-not-happen"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(timeout=TimeoutPolicy(run_timeout=0.05))
|
||
.add_node("slow", slow_node)
|
||
.add_edge(START, "slow")
|
||
.compile()
|
||
)
|
||
|
||
from langgraph.errors import NodeTimeoutError
|
||
|
||
with pytest.raises(NodeTimeoutError):
|
||
await graph.ainvoke({"foo": ""})
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_set_node_defaults_timeout_per_node_wins():
|
||
"""Per-node timeout overrides the default; a generous per-node timeout
|
||
allows a node to complete even when the builder default is very short."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
async def quick_node(state: State) -> State:
|
||
await asyncio.sleep(0.05)
|
||
return {"foo": "done"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(timeout=TimeoutPolicy(run_timeout=0.01))
|
||
.add_node("quick", quick_node, timeout=TimeoutPolicy(run_timeout=5.0))
|
||
.add_edge(START, "quick")
|
||
.compile()
|
||
)
|
||
|
||
result = await graph.ainvoke({"foo": ""})
|
||
assert result["foo"] == "done"
|
||
|
||
|
||
def test_set_node_defaults_chaining():
|
||
"""set_node_defaults() is chainable and can be called in any order relative to add_node."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
def always_failing(state: State) -> State:
|
||
raise RuntimeError("boom")
|
||
|
||
def handler(state: State, error: NodeError) -> State:
|
||
return {"foo": "handled"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("a", always_failing)
|
||
.add_edge(START, "a")
|
||
.set_node_defaults(
|
||
retry_policy=RetryPolicy(
|
||
max_attempts=1, initial_interval=0.01, jitter=False
|
||
),
|
||
error_handler=handler,
|
||
)
|
||
.compile()
|
||
)
|
||
|
||
result = graph.invoke({"foo": ""})
|
||
assert result["foo"] == "handled"
|
||
|
||
|
||
def test_set_node_defaults_combined_retry_and_error_handler():
|
||
"""Retries are exhausted first, then the error handler runs."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
attempts = 0
|
||
captured: dict[str, Any] = {}
|
||
|
||
def always_failing(state: State) -> State:
|
||
nonlocal attempts
|
||
attempts += 1
|
||
raise ValueError("Always fails")
|
||
|
||
def handler(state: State, error: NodeError) -> State:
|
||
captured["error"] = str(error.error)
|
||
return {"foo": "handled"}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(
|
||
retry_policy=RetryPolicy(
|
||
max_attempts=2,
|
||
initial_interval=0.01,
|
||
jitter=False,
|
||
retry_on=ValueError,
|
||
),
|
||
error_handler=handler,
|
||
)
|
||
.add_node("fail", always_failing)
|
||
.add_edge(START, "fail")
|
||
.compile()
|
||
)
|
||
|
||
with patch("time.sleep"):
|
||
result = graph.invoke({"foo": ""})
|
||
|
||
assert result["foo"] == "handled"
|
||
assert attempts == 2
|
||
assert captured["error"] == "Always fails"
|
||
|
||
|
||
def test_error_handler_resumes_after_crash():
|
||
"""If the error handler crashes, resuming should re-schedule the handler
|
||
(not re-execute the original failed node)."""
|
||
|
||
class State(TypedDict):
|
||
foo: str
|
||
|
||
call_count = {"node": 0, "handler": 0}
|
||
captured_errors: list[NodeError] = []
|
||
|
||
def failing_node(state: State) -> State:
|
||
call_count["node"] += 1
|
||
raise RuntimeError("boom")
|
||
|
||
handler_should_fail = [True]
|
||
|
||
def handler(state: State, error: NodeError) -> State:
|
||
call_count["handler"] += 1
|
||
captured_errors.append(error)
|
||
if handler_should_fail[0]:
|
||
raise RuntimeError("handler crash")
|
||
return {"foo": "recovered"}
|
||
|
||
checkpointer = MemorySaver()
|
||
graph = (
|
||
StateGraph(State)
|
||
.set_node_defaults(error_handler=handler)
|
||
.add_node("fail", failing_node)
|
||
.add_edge(START, "fail")
|
||
.compile(checkpointer=checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "t1"}}
|
||
|
||
# First invoke: node fails -> handler runs -> handler crashes -> run fails
|
||
with pytest.raises(RuntimeError, match="handler crash"):
|
||
graph.invoke({"foo": ""}, config)
|
||
|
||
assert call_count["node"] == 1
|
||
assert call_count["handler"] == 1
|
||
assert captured_errors[0].node == "fail"
|
||
assert isinstance(captured_errors[0].error, RuntimeError)
|
||
assert str(captured_errors[0].error) == "boom"
|
||
|
||
# Resume: handler should run again, NOT the original node
|
||
handler_should_fail[0] = False
|
||
result = graph.invoke(None, config)
|
||
|
||
assert result["foo"] == "recovered"
|
||
assert call_count["node"] == 1 # NOT re-executed
|
||
assert call_count["handler"] == 2 # ran again on resume
|
||
# on resume the error was round-tripped through the checkpointer, so it
|
||
# may be deserialized as a string representation rather than the original
|
||
# exception type — verify the node name and that the error content matches.
|
||
assert captured_errors[1].node == "fail"
|
||
assert "boom" in str(captured_errors[1].error)
|
||
|
||
|
||
def test_error_handler_resumes_after_crash_multiple_nodes():
|
||
"""When multiple nodes fail in the same superstep and all have error handlers:
|
||
- error handlers start running while other nodes may still be in-flight
|
||
- resuming re-schedules each handler (not re-executes the original nodes)
|
||
"""
|
||
|
||
class State(TypedDict):
|
||
results: Annotated[list[str], operator.add]
|
||
|
||
call_count = {"a": 0, "b": 0, "handler_a": 0, "handler_b": 0}
|
||
handler_a_started = threading.Event()
|
||
|
||
def node_a(state: State) -> State:
|
||
call_count["a"] += 1
|
||
raise RuntimeError("a failed")
|
||
|
||
def node_b(state: State) -> State:
|
||
call_count["b"] += 1
|
||
# Block until handler_a has started — proves the error handler runs
|
||
# concurrently with in-flight nodes in the same superstep.
|
||
assert handler_a_started.wait(timeout=5), "handler_a never started"
|
||
raise RuntimeError("b failed")
|
||
|
||
handler_should_fail = [True]
|
||
|
||
def handler_a(state: State, error: NodeError) -> State:
|
||
call_count["handler_a"] += 1
|
||
assert error.node == "a"
|
||
assert "a failed" in str(error.error)
|
||
handler_a_started.set()
|
||
if handler_should_fail[0]:
|
||
raise RuntimeError("handler_a crash")
|
||
return {"results": [f"recovered_a:{error.node}"]}
|
||
|
||
def handler_b(state: State, error: NodeError) -> State:
|
||
call_count["handler_b"] += 1
|
||
assert error.node == "b"
|
||
assert "b failed" in str(error.error)
|
||
if handler_should_fail[0]:
|
||
raise RuntimeError("handler_b crash")
|
||
return {"results": [f"recovered_b:{error.node}"]}
|
||
|
||
checkpointer = MemorySaver()
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("a", node_a, error_handler=handler_a)
|
||
.add_node("b", node_b, error_handler=handler_b)
|
||
.add_edge(START, "a")
|
||
.add_edge(START, "b")
|
||
.compile(checkpointer=checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "t1"}}
|
||
|
||
# First invoke: node_a fails immediately -> handler_a starts (sets event) ->
|
||
# node_b unblocks and fails -> handler_b starts -> both handlers crash
|
||
with pytest.raises(RuntimeError):
|
||
graph.invoke({"results": []}, config)
|
||
|
||
assert call_count["a"] == 1
|
||
assert call_count["b"] == 1
|
||
assert call_count["handler_a"] == 1
|
||
assert call_count["handler_b"] == 1
|
||
|
||
# Resume: both handlers should run again, NOT the original nodes
|
||
handler_should_fail[0] = False
|
||
handler_a_started.clear()
|
||
result = graph.invoke(None, config)
|
||
|
||
assert call_count["a"] == 1 # NOT re-executed
|
||
assert call_count["b"] == 1 # NOT re-executed
|
||
assert call_count["handler_a"] == 2 # ran again on resume
|
||
assert call_count["handler_b"] == 2 # ran again on resume
|
||
assert "recovered_a:a" in result["results"]
|
||
assert "recovered_b:b" in result["results"]
|
||
|
||
|
||
@NEEDS_TASK_CANCELLING
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_user_raised_cancelled_becomes_node_cancelled():
|
||
class UserCancelsProc:
|
||
async def ainvoke(self, input, config):
|
||
raise asyncio.CancelledError("nope")
|
||
|
||
task = _make_task(UserCancelsProc(), name="user-cancel")
|
||
|
||
with pytest.raises(NodeCancelledError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.node == "user-cancel"
|
||
# original CancelledError chained for debugging
|
||
assert isinstance(excinfo.value.__cause__, asyncio.CancelledError)
|
||
|
||
|
||
@NEEDS_TASK_CANCELLING
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_user_raised_cancelled_with_timeout_policy():
|
||
"""The timeout path runs the node in a child task; the conversion must
|
||
still trigger for user-raised ``CancelledError``."""
|
||
|
||
class UserCancelsProc:
|
||
async def ainvoke(self, input, config):
|
||
raise asyncio.CancelledError
|
||
|
||
task = _make_task(
|
||
UserCancelsProc(), timeout=_idle_timeout(1.0), name="user-cancel-timed"
|
||
)
|
||
|
||
with pytest.raises(NodeCancelledError) as excinfo:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.node == "user-cancel-timed"
|
||
|
||
|
||
def test_run_with_retry_sync_node_raising_cancelled_becomes_node_cancelled():
|
||
class SyncUserCancelsProc:
|
||
def invoke(self, input, config):
|
||
raise asyncio.CancelledError("sync nope")
|
||
|
||
task = _make_task(SyncUserCancelsProc(), timeout=None, name="sync-user-cancel")
|
||
|
||
with pytest.raises(NodeCancelledError) as excinfo:
|
||
run_with_retry(task, retry_policy=None)
|
||
assert excinfo.value.node == "sync-user-cancel"
|
||
assert isinstance(excinfo.value.__cause__, asyncio.CancelledError)
|
||
|
||
|
||
@NEEDS_TASK_CANCELLING
|
||
@pytest.mark.anyio
|
||
async def test_arun_with_retry_external_cancel_propagates_as_cancelled():
|
||
"""When the asyncio task running ``arun_with_retry`` is cancelled from the
|
||
outside, the cancellation must still propagate as
|
||
``asyncio.CancelledError``. Converting it to ``NodeCancelledError`` would
|
||
break the runner's ability to cancel sibling tasks during cleanup."""
|
||
|
||
started = asyncio.Event()
|
||
observed: list[BaseException] = []
|
||
|
||
class SlowProc:
|
||
async def ainvoke(self, input, config):
|
||
started.set()
|
||
await asyncio.sleep(10.0)
|
||
return "never"
|
||
|
||
task = _make_task(SlowProc(), timeout=None, name="external-cancel")
|
||
|
||
async def runner():
|
||
try:
|
||
await arun_with_retry(task, retry_policy=None)
|
||
except BaseException as exc:
|
||
observed.append(exc)
|
||
raise
|
||
|
||
bg = asyncio.create_task(runner())
|
||
await started.wait()
|
||
bg.cancel()
|
||
# We expect the cancellation to surface to us as well; swallow it here so
|
||
# the test runner's own task isn't poisoned by the cancel.
|
||
with contextlib.suppress(asyncio.CancelledError):
|
||
await bg
|
||
assert observed, "runner did not observe any exception"
|
||
# Framework cancellation must remain a CancelledError, not be rewritten as
|
||
# NodeCancelledError.
|
||
assert isinstance(observed[0], asyncio.CancelledError)
|
||
assert not isinstance(observed[0], NodeCancelledError)
|
||
|
||
|
||
@NEEDS_TASK_CANCELLING
|
||
@pytest.mark.anyio
|
||
async def test_pregel_user_raised_cancellederror_fails_run():
|
||
"""End-to-end: a two-branch graph where one branch raises
|
||
``asyncio.CancelledError`` must fail the run instead of returning
|
||
a partial-success state. This is the LSD-1507 customer scenario."""
|
||
|
||
class _S(TypedDict, total=False):
|
||
vals: Annotated[list[str], operator.add]
|
||
|
||
async def ok(state: _S) -> _S:
|
||
return {"vals": ["ok"]}
|
||
|
||
async def boom(state: _S) -> _S:
|
||
raise asyncio.CancelledError("user-raised in node body")
|
||
|
||
graph = (
|
||
StateGraph(_S)
|
||
.add_node("ok", ok)
|
||
.add_node("boom", boom)
|
||
.add_edge(START, "ok")
|
||
.add_edge(START, "boom")
|
||
.add_edge("ok", END)
|
||
.add_edge("boom", END)
|
||
.compile()
|
||
)
|
||
|
||
with pytest.raises(NodeCancelledError) as excinfo:
|
||
await graph.ainvoke({"vals": []})
|
||
assert excinfo.value.node == "boom"
|