mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 03:09:45 +02:00
syachamaneni-lcandGitHub
f25a0f4f4c
- Distinguish between Node cancellations and other cancellations - Use python 3.11+ feature where `cancelling() == 0` when it is the node cancelling - Bubble up the node cancellation example so the client can take care of it instead of silently failing without reporting it. Read the full contributing guidelines: https://docs.langchain.com/oss/python/contributing/overview > **All contributions must be in English.** See the [language policy](https://docs.langchain.com/oss/python/contributing/overview#language-policy). If you paste a large clearly AI generated description here your PR may be IGNORED or CLOSED! Thank you for contributing to LangGraph! Follow these steps to have your pull request considered as ready for review. 1. PR title: Should follow the format: TYPE(SCOPE): DESCRIPTION - feat(langgraph): add multi-tenant support - Allowed TYPE and SCOPE values: https://github.com/langchain-ai/langgraph/blob/main/.github/workflows/pr_lint.yml#L19-L43 2. PR description: - Write 1-2 sentences summarizing the change. - The `Fixes #xx` line at the top is **required** for external contributions — update the issue number and keep the keyword. This links your PR to the approved issue and auto-closes it on merge. - If there are any breaking changes, please clearly describe them. - If this PR depends on another PR being merged first, please include "Depends on #PR_NUMBER" in the description. 3. Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified. - We will not consider a PR unless these three are passing in CI. 4. How did you verify your code works? - Existing unit tests pass - New unit tests added - Used langgraph deployment to make sure the feature is working as expected. Additional guidelines: - All external PRs must link to an issue or discussion where a solution has been approved by a maintainer, and you must be assigned to that issue. PRs without prior approval will be closed. - PRs should not touch more than one package unless absolutely necessary. - Do not update the `uv.lock` files or add dependencies to `pyproject.toml` files (even optional ones) unless you have explicit permission to do so by a maintainer. ## Social handles (optional) <!-- If you'd like a shoutout on release, add your socials below --> Twitter: @ LinkedIn: https://linkedin.com/in/
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"
|