From 3db82d5b6d6ff02069d99ab60d77d347976f50d2 Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Mon, 11 May 2026 16:37:16 -0700 Subject: [PATCH] feat(langgraph): add set_node_defaults() to StateGraph (#7747) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `StateGraph.set_node_defaults()` — a fluent builder method for setting graph-wide node policies in one place. Per-node values from `add_node()` always take precedence. Defaults are applied at `compile()` time. ```python graph = ( StateGraph(State) .set_node_defaults( retry_policy=RetryPolicy(max_attempts=3), error_handler=my_fallback_handler, timeout=TimeoutPolicy(run_timeout=30), ) .add_node("a", node_a) .add_node("b", node_b, retry_policy=custom) # overrides default .add_edge(START, "a") .compile() ) ``` Error handlers are never invoked on error-handler nodes themselves — handler failures fail the run. Not inherited by subgraphs. ## Supported defaults | `set_node_defaults()` kwarg | Fallback for `add_node(...)` kwarg | Applies to error-handler nodes? | |---|---|---| | `retry_policy` | `retry_policy` | Yes | | `cache_policy` | `cache_policy` | No — caching handler results is unsafe | | `error_handler` | `error_handler` | No — handlers must never catch themselves | | `timeout` | `timeout` | Yes | ## Changes - `libs/langgraph/langgraph/graph/state.py` — new `_NodeDefaults` dataclass, `set_node_defaults()` method on `StateGraph`; `compile()` applies builder defaults to node specs with per-branch rules for which defaults apply to error-handler nodes. - `libs/langgraph/tests/test_retry.py` — 14 new tests covering all four policy types, per-node override precedence, chaining, combined retry+handler, handler exclusion, `RunnableConfig` injection, and name collision. ## Verification `make format`, `make lint`, `make test` all passing in `libs/langgraph`. --------- Co-authored-by: Sydney Runkle --- libs/langgraph/langgraph/graph/state.py | 135 +++++++- libs/langgraph/tests/test_retry.py | 389 +++++++++++++++++++++++- 2 files changed, 521 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 6f31431c9..deb0e9e94 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -6,7 +6,7 @@ import typing import warnings from collections import defaultdict from collections.abc import Awaitable, Callable, Hashable, Sequence -from dataclasses import is_dataclass +from dataclasses import dataclass, is_dataclass from datetime import timedelta from functools import partial from inspect import isclass, isfunction, ismethod, signature @@ -95,6 +95,17 @@ __all__ = ("StateGraph", "CompiledStateGraph") logger = logging.getLogger(__name__) _CHANNEL_BRANCH_TO = "branch:to:{}" +_DEFAULT_ERROR_HANDLER_NODE = "__default_error_handler__" + + +@dataclass(slots=True) +class _NodeDefaults: + """Default node policies applied to every node at compile time.""" + + retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None + cache_policy: CachePolicy | None = None + error_handler: StateNode[Any, Any] | None = None + timeout: TimeoutPolicy | None = None def _warn_invalid_state_schema(schema: type[Any] | Any) -> None: @@ -251,10 +262,77 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): self.output_schema = cast(type[OutputT], output_schema or state_schema) self.context_schema = context_schema + self._node_defaults: _NodeDefaults = _NodeDefaults() + self._add_schema(self.state_schema) self._add_schema(self.input_schema, allow_managed=False) self._add_schema(self.output_schema, allow_managed=False) + def set_node_defaults( + self, + *, + retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, + cache_policy: CachePolicy | None = None, + error_handler: StateNode[Any, ContextT] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, + ) -> Self: + """Set default node policies that apply to every node in this graph. + + Per-node values passed to `add_node` always take precedence over these + defaults. Defaults are applied at `compile()` time. Policies set here + are **not** inherited by subgraphs. + + `retry_policy` and `timeout` defaults apply to **all** nodes, + including error-handler nodes. `cache_policy` and `error_handler` + defaults only apply to regular nodes -- caching error-handler results + is unsafe, and handlers must never catch themselves. + + Args: + retry_policy: Default retry policy for nodes that don't specify + their own via `add_node(..., retry_policy=...)`. Also applies + to error-handler nodes. + cache_policy: Default cache policy for nodes that don't specify + their own via `add_node(..., cache_policy=...)`. Does **not** + apply to error-handler nodes. + error_handler: Default error handler invoked when any regular node + raises and does not have its own `error_handler` set via + `add_node`. The handler is **not** invoked when an + error-handler node itself raises -- handler failures fail the + run. + timeout: Default timeout policy for nodes that don't specify their + own via `add_node(..., timeout=...)`. Also applies to + error-handler nodes. Accepts a `TimeoutPolicy`, a number of + seconds (`float`), or a `timedelta`. + + Returns: + Self: The builder instance, for chaining. + + Example: + ```python + graph = ( + StateGraph(State) + .set_node_defaults( + retry_policy=RetryPolicy(max_attempts=3), + error_handler=my_fallback_handler, + ) + .add_node("a", node_a) + .add_node("b", node_b, retry_policy=custom_retry) # overrides default + .add_edge(START, "a") + .compile() + ) + ``` + """ + defaults = self._node_defaults + if retry_policy is not None: + defaults.retry_policy = retry_policy + if cache_policy is not None: + defaults.cache_policy = cache_policy + if error_handler is not None: + defaults.error_handler = error_handler + if timeout is not None: + defaults.timeout = coerce_timeout_policy(timeout) + return self + @property def _all_edges(self) -> set[tuple[str, str]]: return self.edges | { @@ -1193,10 +1271,63 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): key for key, val in self.channels.items() if not is_managed_value(val) ] ) + # Apply builder defaults to node specs. Per-node values always win. + # Error-handler routing and cache_policy are only assigned to regular + # nodes. Retry and timeout defaults also apply to error-handler nodes. + defaults = self._node_defaults + default_handler_name: str | None = None + if defaults.error_handler is not None: + if _DEFAULT_ERROR_HANDLER_NODE in self.nodes: + raise ValueError( + f"Auto-generated default error handler node " + f"`{_DEFAULT_ERROR_HANDLER_NODE}` already exists." + ) + default_handler_name = _DEFAULT_ERROR_HANDLER_NODE + self.nodes[default_handler_name] = StateNodeSpec[Any, ContextT]( + coerce_to_runnable( + defaults.error_handler, # type: ignore[arg-type] + name=default_handler_name, + trace=False, + ), + metadata=None, + input_schema=self.state_schema, + retry_policy=None, + cache_policy=None, + is_error_handler=True, + ) + + # Apply builder defaults to node specs. Per-node values always win. + for spec in self.nodes.values(): + # error_handler: regular nodes only — handlers must never + # catch themselves or other handlers. + if ( + not spec.is_error_handler + and default_handler_name is not None + and spec.error_handler_node is None + ): + spec.error_handler_node = default_handler_name + # retry: all nodes — handlers should be retried on transient + # failures just like regular nodes. + if defaults.retry_policy is not None and spec.retry_policy is None: + spec.retry_policy = defaults.retry_policy + # cache: regular nodes only — caching an error-handler result + # is unsafe because the input (failed-node state) may differ + # across failures even when the cache key matches. + if ( + not spec.is_error_handler + and defaults.cache_policy is not None + and spec.cache_policy is None + ): + spec.cache_policy = defaults.cache_policy + # timeout: all nodes — a stuck handler should be cancelled the + # same way a stuck regular node would be. + if defaults.timeout is not None and spec.timeout is None: + spec.timeout = defaults.timeout + node_error_handler_map = { node_name: spec.error_handler_node for node_name, spec in self.nodes.items() - if spec.error_handler_node is not None + if not spec.is_error_handler and spec.error_handler_node is not None } compiled = CompiledStateGraph[StateT, ContextT, InputT, OutputT]( diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index f5d4d74e0..be57e9360 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -15,7 +15,7 @@ from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, BaseCallback 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 RunnableLambda, RunnableParallel +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 @@ -2280,3 +2280,390 @@ def test_node_without_error_handler_still_fails_run(): 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"