diff --git a/libs/langgraph/langgraph/_internal/_runnable.py b/libs/langgraph/langgraph/_internal/_runnable.py index 0c110b96f..40dc5b7d3 100644 --- a/libs/langgraph/langgraph/_internal/_runnable.py +++ b/libs/langgraph/langgraph/_internal/_runnable.py @@ -342,6 +342,17 @@ class RunnableCallable(Runnable): # If the kwarg is accepted by the function, store the key / runtime attribute to inject self.func_accepts[kw] = (runtime_key, default) + # True when the function expects a "state" or "input" first arg. + # False only when ALL non-VAR_POSITIONAL params are injected kwargs, e.g. + # `def handler(error: NodeError) -> T` — passing input positionally would + # conflict with the kwarg injection. + injected_names = set(self.func_accepts) + self.takes_input: bool = not self.explode_args and any( + p.kind == inspect.Parameter.VAR_POSITIONAL + or (p.kind in VALID_KINDS and p.name not in injected_names) + for p in params.values() + ) + def __repr__(self) -> str: repr_args = { k: v @@ -364,9 +375,12 @@ class RunnableCallable(Runnable): if self.explode_args: args, _kwargs = input kwargs = {**self.kwargs, **_kwargs, **kwargs} - else: + elif self.takes_input: args = (input,) kwargs = {**self.kwargs, **kwargs} + else: + args = () + kwargs = {**self.kwargs, **kwargs} runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME) @@ -438,9 +452,12 @@ class RunnableCallable(Runnable): if self.explode_args: args, _kwargs = input kwargs = {**self.kwargs, **_kwargs, **kwargs} - else: + elif self.takes_input: args = (input,) kwargs = {**self.kwargs, **kwargs} + else: + args = () + kwargs = {**self.kwargs, **kwargs} runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index be310f0f8..d8fa8c087 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -65,6 +65,7 @@ class _TaskFunction(Generic[P, T]): cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None, timeout: TimeoutPolicy | None = None, name: str | None = None, + error_handler: Callable[..., Any] | None = None, ) -> None: if name is not None: if hasattr(func, "__func__"): @@ -81,11 +82,14 @@ class _TaskFunction(Generic[P, T]): self.retry_policy = retry_policy self.cache_policy = cache_policy self.timeout = timeout + self._raw_error_handler = error_handler + self.error_handler = None # not used for push tasks; see __call__ functools.update_wrapper(self, func) def __call__(self, *args: P.args, **kwargs: P.kwargs) -> SyncAsyncFuture[T]: + func = self._wrap_with_error_handler(self.func) return _call_with_options( - self.func, + func, args, kwargs, retry_policy=self.retry_policy, @@ -93,6 +97,42 @@ class _TaskFunction(Generic[P, T]): timeout=self.timeout, ) + def _wrap_with_error_handler( + self, func: Callable[P, Awaitable[T]] | Callable[P, T] + ) -> Callable[P, Awaitable[T]] | Callable[P, T]: + if self._raw_error_handler is None: + return func + handler = self._raw_error_handler + task_name = getattr(func, "__name__", "task") + if is_async_callable(func): + + async def awrapped(*args: Any, **kwargs: Any) -> Any: + try: + return await func(*args, **kwargs) # type: ignore[misc] + except Exception as exc: + from langgraph.errors import NodeError + + node_error = NodeError(node=task_name, error=exc) + if is_async_callable(handler): + return await handler(*args, error=node_error, **kwargs) # type: ignore[misc] + return handler(*args, error=node_error, **kwargs) # type: ignore[misc] + + awrapped.__name__ = task_name # type: ignore[attr-defined] + return awrapped # type: ignore[return-value] + else: + + def wrapped(*args: Any, **kwargs: Any) -> Any: + try: + return func(*args, **kwargs) # type: ignore[operator] + except Exception as exc: + from langgraph.errors import NodeError + + node_error = NodeError(node=task_name, error=exc) + return handler(*args, error=node_error, **kwargs) # type: ignore[misc] + + wrapped.__name__ = task_name # type: ignore[attr-defined] + return wrapped # type: ignore[return-value] + def clear_cache(self, cache: BaseCache) -> None: """Clear the cache for this task.""" if self.cache_policy is not None: @@ -136,6 +176,7 @@ def task( retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None, timeout: float | timedelta | TimeoutPolicy | None = None, + error_handler: Callable[..., Any] | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> ( Callable[[Callable[P, Awaitable[T]] | Callable[P, T]], _TaskFunction[P, T]] @@ -243,6 +284,7 @@ def task( cache_policy=cache_policy, timeout=timeout_policy, name=name, + error_handler=error_handler, ) if __func_or_none__ is not None: diff --git a/libs/langgraph/langgraph/graph/_node.py b/libs/langgraph/langgraph/graph/_node.py index d8238c2b3..9ffa97d18 100644 --- a/libs/langgraph/langgraph/graph/_node.py +++ b/libs/langgraph/langgraph/graph/_node.py @@ -8,6 +8,7 @@ from langchain_core.runnables import Runnable, RunnableConfig from langgraph.store.base import BaseStore from langgraph._internal._typing import EMPTY_SEQ +from langgraph.errors import NodeError from langgraph.runtime import Runtime from langgraph.types import CachePolicy, RetryPolicy, StreamWriter, TimeoutPolicy from langgraph.typing import ContextT, NodeInputT, NodeInputT_contra @@ -64,6 +65,22 @@ class _NodeWithRuntime(Protocol[NodeInputT_contra, ContextT]): ) -> Any: ... +class _NodeWithNodeError(Protocol[NodeInputT_contra]): + def __call__(self, state: NodeInputT_contra, *, error: NodeError) -> Any: ... + + +class _NodeWithConfigNodeError(Protocol[NodeInputT_contra]): + def __call__( + self, state: NodeInputT_contra, *, config: RunnableConfig, error: NodeError + ) -> Any: ... + + +class _NodeWithRuntimeNodeError(Protocol[NodeInputT_contra, ContextT]): + def __call__( + self, state: NodeInputT_contra, *, runtime: Runtime[ContextT], error: NodeError + ) -> Any: ... + + # TODO: we probably don't want to explicitly support the config / store signatures once # we move to adding a context arg. Maybe what we do is we add support for kwargs with param spec # this is purely for typing purposes though, so can easily change in the coming weeks. @@ -80,6 +97,13 @@ StateNode: TypeAlias = ( | Runnable[NodeInputT, Any] ) +ErrorHandlerNode: TypeAlias = ( + StateNode[NodeInputT, ContextT] + | _NodeWithNodeError[NodeInputT] + | _NodeWithConfigNodeError[NodeInputT] + | _NodeWithRuntimeNodeError[NodeInputT, ContextT] +) + @dataclass(slots=True) class StateNodeSpec(Generic[NodeInputT, ContextT]): @@ -88,8 +112,7 @@ class StateNodeSpec(Generic[NodeInputT, ContextT]): input_schema: type[NodeInputT] retry_policy: RetryPolicy | Sequence[RetryPolicy] | None cache_policy: CachePolicy | None - is_error_handler: bool = False - error_handler_node: str | None = None + error_handler: Runnable[Any, Any] | None = None ends: tuple[str, ...] | dict[str, str] | None = EMPTY_SEQ defer: bool = False timeout: TimeoutPolicy | None = None diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 6f31431c9..275ab0519 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -65,7 +65,7 @@ from langgraph.errors import ( create_error_message, ) from langgraph.graph._branch import BranchSpec -from langgraph.graph._node import StateNode, StateNodeSpec +from langgraph.graph._node import ErrorHandlerNode, StateNode, StateNodeSpec from langgraph.managed.base import ( ManagedValueSpec, is_managed_value, @@ -772,24 +772,15 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): if destinations is not None: ends = destinations - resolved_input_schema: type[Any] = ( - input_schema or inferred_input_schema or self.state_schema - ) - handler_node_name: str | None = None - if error_handler is not None: - handler_node_name = f"__error_handler__{node}" - if handler_node_name in self.nodes: - raise ValueError( - f"Auto-generated error handler node `{handler_node_name}` already exists." - ) - self.nodes[handler_node_name] = StateNodeSpec[Any, ContextT]( - coerce_to_runnable(error_handler, name=handler_node_name, trace=False), # type: ignore[arg-type] - metadata=None, - input_schema=resolved_input_schema, - retry_policy=None, - cache_policy=None, - is_error_handler=True, + coerced_error_handler: Runnable[Any, Any] | None = ( + coerce_to_runnable( # type: ignore[arg-type] + error_handler, + name=f"__error_handler__{node}", + trace=False, ) + if error_handler is not None + else None + ) if input_schema is not None: self.nodes[node] = StateNodeSpec[NodeInputT, ContextT]( @@ -798,7 +789,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): input_schema=input_schema, retry_policy=retry_policy, cache_policy=cache_policy, - error_handler_node=handler_node_name, + error_handler=coerced_error_handler, ends=ends, defer=defer, timeout=timeout, @@ -810,7 +801,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): input_schema=inferred_input_schema, retry_policy=retry_policy, cache_policy=cache_policy, - error_handler_node=handler_node_name, + error_handler=coerced_error_handler, ends=ends, defer=defer, timeout=timeout, @@ -822,7 +813,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): input_schema=self.state_schema, retry_policy=retry_policy, cache_policy=cache_policy, - error_handler_node=handler_node_name, + error_handler=coerced_error_handler, ends=ends, defer=defer, timeout=timeout, @@ -1079,7 +1070,17 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): if interrupt: for node in interrupt: if node not in self.nodes: - raise ValueError(f"Interrupt node `{node}` not found") + # __error_handler__ is a valid virtual task name when the + # base node has an error_handler configured. + if node.startswith("__error_handler__"): + base = node[len("__error_handler__"):] + if ( + base not in self.nodes + or self.nodes[base].error_handler is None + ): + raise ValueError(f"Interrupt node `{node}` not found") + else: + raise ValueError(f"Interrupt node `{node}` not found") self.compiled = True return self @@ -1094,6 +1095,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): debug: bool = False, name: str | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, + error_handler: ErrorHandlerNode[Any, ContextT] | None = None, ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: """Compiles the `StateGraph` into a `CompiledStateGraph` object. @@ -1193,11 +1195,15 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): key for key, val in self.channels.items() if not is_managed_value(val) ] ) - 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 - } + error_handler: Runnable[Any, Any] | None = ( + coerce_to_runnable( # type: ignore[arg-type] + error_handler, + name="__graph_error_handler__", + trace=False, + ) + if error_handler is not None + else None + ) compiled = CompiledStateGraph[StateT, ContextT, InputT, OutputT]( builder=self, @@ -1220,7 +1226,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): debug=debug, store=store, cache=cache, - node_error_handler_map=node_error_handler_map, + error_handler=error_handler, name=name or "LangGraph", stream_transformers=transformers, ) @@ -1395,8 +1401,7 @@ class CompiledStateGraph( metadata=node.metadata, retry_policy=node.retry_policy, cache_policy=node.cache_policy, - is_error_handler=node.is_error_handler, - error_handler_node=node.error_handler_node, + error_handler=node.error_handler, bound=node.runnable, # type: ignore[arg-type] timeout=node.timeout, ) diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 103f6cce0..11ea31b4a 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -20,6 +20,7 @@ from typing import ( from langchain_core.callbacks import Callbacks from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager +from langchain_core.runnables import Runnable from langchain_core.runnables.config import RunnableConfig from langgraph.checkpoint.base import ( BaseCheckpointSaver, @@ -32,6 +33,7 @@ from langgraph.store.base import BaseStore from xxhash import xxh3_128_hexdigest from langgraph._internal._config import merge_configs, patch_config +from langgraph._internal._runnable import RunnableSeq from langgraph._internal._constants import ( CACHE_NS_WRITES, CONF, @@ -71,6 +73,7 @@ from langgraph.constants import TAG_HIDDEN from langgraph.errors import NodeError from langgraph.managed.base import ManagedValueMapping from langgraph.pregel._call import get_runnable_for_task, identifier +from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry from langgraph.pregel._io import read_channels from langgraph.pregel._log import logger from langgraph.pregel._read import INPUT_CACHE_KEY_TYPE, PregelNode @@ -407,6 +410,7 @@ def prepare_next_tasks( updated_channels: set[str] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + error_handler: Runnable[Any, Any] | None = None, ) -> dict[str, PregelTask] | dict[str, PregelExecutableTask]: """Prepare the set of tasks that will make up the next Pregel step. @@ -462,6 +466,7 @@ def prepare_next_tasks( input_cache=input_cache, cache_policy=cache_policy, retry_policy=retry_policy, + error_handler=error_handler, ): tasks.append(task) @@ -508,6 +513,7 @@ def prepare_next_tasks( input_cache=input_cache, cache_policy=cache_policy, retry_policy=retry_policy, + error_handler=error_handler, ): tasks.append(task) return {t.id: t for t in tasks} @@ -542,6 +548,7 @@ def prepare_single_task( input_cache: dict[INPUT_CACHE_KEY_TYPE, Any] | None = None, cache_policy: CachePolicy | None = None, retry_policy: Sequence[RetryPolicy] = (), + error_handler: Runnable[Any, Any] | None = None, ) -> None | PregelTask | PregelExecutableTask: """Prepares a single task for the next Pregel step, given a task path, which uniquely identifies a PUSH or PULL task within the graph.""" @@ -756,6 +763,7 @@ def prepare_single_task( writers=proc.flat_writers, subgraphs=proc.subgraphs, timeout=proc.timeout, + error_handler=proc.error_handler or error_handler, ) else: return PregelTask(task_id, name, task_path[:3]) @@ -1111,10 +1119,10 @@ def prepare_node_error_handler_task( failed_task: PregelExecutableTask, *, handler_node_name: str, + handler: Runnable, failed_error: BaseException, checkpoint: Checkpoint, pending_writes: list[PendingWrite], - processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, config: RunnableConfig, @@ -1125,15 +1133,12 @@ def prepare_node_error_handler_task( manager: None | ParentRunManager | AsyncParentRunManager = None, cache_policy: CachePolicy | None = None, retry_policy: Sequence[RetryPolicy] = (), -) -> PregelExecutableTask | None: - """Prepare an immediate node-level error handler task for a failed task.""" - if handler_node_name not in processes: - return None - proc = processes[handler_node_name] - proc_node = proc.node - if proc_node is None: - return None +) -> PregelExecutableTask: + """Prepare an error handler task for a failed task. + The handler borrows the failed task's write pipeline (same state channels), + so no separate node registration is needed. + """ checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", "")) task_id_func = _xxhash_str if checkpoint["v"] > 1 else _uuid5_str configurable = config.get(CONF, {}) @@ -1159,27 +1164,17 @@ def prepare_node_error_handler_task( "langgraph_path": translated_task_path, "langgraph_checkpoint_ns": task_checkpoint_ns, } - if proc.metadata: - metadata.update(proc.metadata) writes: deque[tuple[str, Any]] = deque() - - effective_retry_policy = proc.retry_policy or retry_policy - effective_cache_policy = proc.cache_policy or cache_policy - if effective_cache_policy: - args_key = effective_cache_policy.key_func(failed_task.input) - cache_key = CacheKey( - ( - CACHE_NS_WRITES, - (identifier(proc) or "__dynamic__"), - handler_node_name, - ), - xxh3_128_hexdigest( - args_key.encode() if isinstance(args_key, str) else args_key - ), - effective_cache_policy.ttl, - ) + # Mirror how regular node procs are built: combine handler with a write pipeline + # so run_with_retry invokes the full pipeline in one shot. + # - PULL node tasks: writers are in failed_task.writers → reuse them + # - PUSH functional tasks: writers are embedded in proc (empty writers list) → + # add a RETURN write so the handler's result becomes the future's value. + handler_writers = failed_task.writers + if handler_writers: + handler_proc: Runnable = RunnableSeq(handler, *handler_writers) else: - cache_key = None + handler_proc = RunnableSeq(handler, ChannelWrite([ChannelWriteEntry(RETURN)])) scratchpad = _scratchpad( config[CONF].get(CONFIG_KEY_SCRATCHPAD), @@ -1194,14 +1189,11 @@ def prepare_node_error_handler_task( runtime = runtime.override( store=store, previous=checkpoint["channel_values"].get(PREVIOUS, None) ) - additional_config: RunnableConfig = { - "metadata": metadata, - "tags": proc.tags, - } + additional_config: RunnableConfig = {"metadata": metadata} return PregelExecutableTask( handler_node_name, failed_task.input, - proc_node, + handler_proc, writes, patch_config( merge_configs(config, additional_config), @@ -1239,12 +1231,11 @@ def prepare_node_error_handler_task( }, ), PUSH_TRIGGER, - effective_retry_policy, - cache_key, + retry_policy, + None, # handlers don't cache task_id, translated_task_path, - writers=proc.flat_writers, - subgraphs=proc.subgraphs, + writers=handler_writers, # for ParentCommand / subgraph routing ) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 9949a62d9..f964f4ff7 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -22,7 +22,7 @@ from typing import ( ) from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager -from langchain_core.runnables import RunnableConfig +from langchain_core.runnables import Runnable, RunnableConfig from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( WRITES_IDX_MAP, @@ -279,6 +279,7 @@ class PregelLoop: migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + error_handler: Runnable[Any, Any] | None = None, has_graph_lifecycle_callbacks: bool = False, ) -> None: self.stream = stream @@ -303,6 +304,7 @@ class PregelLoop: self.trigger_to_nodes = trigger_to_nodes self.retry_policy = retry_policy self.cache_policy = cache_policy + self.error_handler = error_handler self.durability = durability self._has_graph_lifecycle_callbacks = has_graph_lifecycle_callbacks self._graph_lifecycle_events = deque() @@ -545,6 +547,7 @@ class PregelLoop: manager=self.manager, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + error_handler=self.error_handler, ), ): # produce debug output @@ -597,6 +600,7 @@ class PregelLoop: updated_channels=self.updated_channels, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + error_handler=self.error_handler, ) # produce debug output @@ -1368,6 +1372,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + error_handler: Runnable[Any, Any] | None = None, has_graph_lifecycle_callbacks: bool = False, ) -> None: super().__init__( @@ -1389,6 +1394,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): trigger_to_nodes=trigger_to_nodes, retry_policy=retry_policy, cache_policy=cache_policy, + error_handler=error_handler, durability=durability, has_graph_lifecycle_callbacks=has_graph_lifecycle_callbacks, ) @@ -1451,22 +1457,20 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): def schedule_error_handler( self, failed_task: PregelExecutableTask, error: BaseException ) -> PregelExecutableTask | None: - handler_node = self.nodes[failed_task.name].error_handler_node - if not handler_node: + handler = failed_task.error_handler or self.error_handler + if handler is None: return None + handler_node_name = f"__error_handler__{failed_task.name}" writes = list(failed_task.writes) writes.append((ERROR_SOURCE_NODE, failed_task.name)) - self.put_writes( - failed_task.id, - writes, - ) + self.put_writes(failed_task.id, writes) handler_task = prepare_node_error_handler_task( failed_task, - handler_node_name=handler_node, + handler_node_name=handler_node_name, + handler=handler, failed_error=error, checkpoint=self.checkpoint, pending_writes=self.checkpoint_pending_writes, - processes=self.nodes, channels=self.channels, managed=self.managed, config=failed_task.config, @@ -1478,8 +1482,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): retry_policy=self.retry_policy, cache_policy=self.cache_policy, ) - if handler_task is None: - return None self.tasks[handler_task.id] = handler_task if not self.is_replaying: self._match_writes({handler_task.id: handler_task}) @@ -1621,6 +1623,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + error_handler: Runnable[Any, Any] | None = None, has_graph_lifecycle_callbacks: bool = False, ) -> None: super().__init__( @@ -1642,6 +1645,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): trigger_to_nodes=trigger_to_nodes, retry_policy=retry_policy, cache_policy=cache_policy, + error_handler=error_handler, durability=durability, has_graph_lifecycle_callbacks=has_graph_lifecycle_callbacks, ) @@ -1706,22 +1710,20 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): async def aschedule_error_handler( self, failed_task: PregelExecutableTask, error: BaseException ) -> PregelExecutableTask | None: - handler_node = self.nodes[failed_task.name].error_handler_node - if not handler_node: + handler = failed_task.error_handler or self.error_handler + if handler is None: return None + handler_node_name = f"__error_handler__{failed_task.name}" writes = list(failed_task.writes) writes.append((ERROR_SOURCE_NODE, failed_task.name)) - self.put_writes( - failed_task.id, - writes, - ) + self.put_writes(failed_task.id, writes) handler_task = prepare_node_error_handler_task( failed_task, - handler_node_name=handler_node, + handler_node_name=handler_node_name, + handler=handler, failed_error=error, checkpoint=self.checkpoint, pending_writes=self.checkpoint_pending_writes, - processes=self.nodes, channels=self.channels, managed=self.managed, config=failed_task.config, @@ -1733,8 +1735,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): retry_policy=self.retry_policy, cache_policy=self.cache_policy, ) - if handler_task is None: - return None self.tasks[handler_task.id] = handler_task if not self.is_replaying: self._match_writes({handler_task.id: handler_task}) diff --git a/libs/langgraph/langgraph/pregel/_read.py b/libs/langgraph/langgraph/pregel/_read.py index 7d49fce85..ee0ca0bde 100644 --- a/libs/langgraph/langgraph/pregel/_read.py +++ b/libs/langgraph/langgraph/pregel/_read.py @@ -138,11 +138,8 @@ class PregelNode: metadata: Mapping[str, Any] | None """Metadata to attach to the node for tracing.""" - is_error_handler: bool - """Whether this node is registered as an error handler node.""" - - error_handler_node: str | None - """Optional handler node name for failures from this node.""" + error_handler: Runnable[Any, Any] | None + """Callable invoked after retries are exhausted; receives same input as the node.""" subgraphs: Sequence[PregelProtocol] """Subgraphs used by the node.""" @@ -159,8 +156,7 @@ class PregelNode: bound: Runnable[Any, Any] | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, - is_error_handler: bool = False, - error_handler_node: str | None = None, + error_handler: Runnable[Any, Any] | None = None, subgraphs: Sequence[PregelProtocol] | None = None, timeout: float | timedelta | TimeoutPolicy | None = None, ) -> None: @@ -177,8 +173,7 @@ class PregelNode: self.timeout = coerce_timeout_policy(timeout) self.tags = tags self.metadata = metadata - self.is_error_handler = is_error_handler - self.error_handler_node = error_handler_node + self.error_handler = error_handler if subgraphs is not None: self.subgraphs = subgraphs elif self.bound is not DEFAULT_BOUND: diff --git a/libs/langgraph/langgraph/pregel/_runner.py b/libs/langgraph/langgraph/pregel/_runner.py index 979935c9a..52c8b4481 100644 --- a/libs/langgraph/langgraph/pregel/_runner.py +++ b/libs/langgraph/langgraph/pregel/_runner.py @@ -13,7 +13,6 @@ from collections.abc import ( Collection, Iterable, Iterator, - Mapping, Sequence, ) from functools import partial @@ -143,7 +142,6 @@ class PregelRunner: put_writes: weakref.ref[Callable[[str, Sequence[tuple[str, Any]]], None]], use_astream: bool = False, node_finished: Callable[[str], None] | None = None, - node_error_handler_map: Mapping[str, str] | None = None, schedule_error_handler: Callable[ [PregelExecutableTask, BaseException], PregelExecutableTask | None ] @@ -158,19 +156,12 @@ class PregelRunner: self.put_writes = put_writes self.use_astream = use_astream self.node_finished = node_finished - self.node_error_handler_map = dict(node_error_handler_map or {}) - self.error_handler_nodes = set(self.node_error_handler_map.values()) self.schedule_error_handler = schedule_error_handler self.aschedule_error_handler = aschedule_error_handler - # Exception object ids that are already routed to graph-level error handler. - # These ids are consulted by stop/panic checks to avoid re-raising handled - # exceptions via the normal fatal path in the same run. self._handled_exception_ids: set[int] = set() def _should_route_to_error_handler(self, task: PregelExecutableTask) -> bool: - if task.name in self.error_handler_nodes: - return False - return task.name in self.node_error_handler_map + return task.error_handler is not None def tick( self, diff --git a/libs/langgraph/langgraph/pregel/_validate.py b/libs/langgraph/langgraph/pregel/_validate.py index fcfb54c9a..3b29c108b 100644 --- a/libs/langgraph/langgraph/pregel/_validate.py +++ b/libs/langgraph/langgraph/pregel/_validate.py @@ -99,14 +99,21 @@ def validate_graph( if interrupt_after_nodes != "*": for n in interrupt_after_nodes: - if n not in nodes: + if n not in nodes and not _is_valid_error_handler_interrupt(n, nodes): raise ValueError(f"Node {n} not in nodes") if interrupt_before_nodes != "*": for n in interrupt_before_nodes: - if n not in nodes: + if n not in nodes and not _is_valid_error_handler_interrupt(n, nodes): raise ValueError(f"Node {n} not in nodes") +def _is_valid_error_handler_interrupt(name: str, nodes: Mapping[str, PregelNode]) -> bool: + if not name.startswith("__error_handler__"): + return False + base = name[len("__error_handler__"):] + return base in nodes and nodes[base].error_handler is not None + + def validate_keys( keys: str | Sequence[str] | None, channels: Mapping[str, Any], diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1550e5a92..1accb474b 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -33,6 +33,7 @@ from uuid import UUID, uuid5 from langchain_core._api import beta from langchain_core.globals import get_debug from langchain_core.runnables import ( + Runnable, RunnableSequence, ) from langchain_core.runnables.base import Input, Output @@ -751,7 +752,7 @@ class Pregel( name: str = "LangGraph" trigger_to_nodes: Mapping[str, Sequence[str]] - node_error_handler_map: Mapping[str, str] + error_handler: Runnable[Any, Any] | None def __init__( self, @@ -776,7 +777,7 @@ class Pregel( context_schema: type[ContextT] | None = None, config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, - node_error_handler_map: Mapping[str, str] | None = None, + error_handler: Runnable[Any, Any] | None = None, name: str = "LangGraph", stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, **deprecated_kwargs: Unpack[DeprecatedKwargs], @@ -824,7 +825,7 @@ class Pregel( self.context_schema = context_schema self.config = config self.trigger_to_nodes = trigger_to_nodes or {} - self.node_error_handler_map = node_error_handler_map or {} + self.error_handler = error_handler self.name = name self.stream_transformers: tuple[Callable[[tuple[str, ...]], Any], ...] = tuple( stream_transformers or () @@ -2885,6 +2886,7 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + error_handler=self.error_handler, has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers), ) as loop: emit_graph_lifecycle_events(loop) @@ -2895,7 +2897,6 @@ class Pregel( ), put_writes=weakref.WeakMethod(loop.put_writes), node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), - node_error_handler_map=self.node_error_handler_map, schedule_error_handler=loop.schedule_error_handler, ) # enable subgraph streaming @@ -3337,6 +3338,7 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + error_handler=self.error_handler, has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers), ) as loop: await aemit_graph_lifecycle_events(loop) @@ -3348,7 +3350,6 @@ class Pregel( put_writes=weakref.WeakMethod(loop.put_writes), use_astream=do_stream, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), - node_error_handler_map=self.node_error_handler_map, aschedule_error_handler=loop.aschedule_error_handler, ) # enable subgraph streaming diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index fa0bdc685..528108925 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -628,6 +628,7 @@ class PregelExecutableTask: writers: Sequence[Runnable] = () subgraphs: Sequence[PregelProtocol] = () timeout: TimeoutPolicy | None = None + error_handler: Runnable | None = None class StateSnapshot(NamedTuple): diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index f5d4d74e0..184d7db7b 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -2280,3 +2280,183 @@ def test_node_without_error_handler_still_fails_run(): with pytest.raises(ValueError, match="no handler"): graph.invoke({"foo": ""}) + + +# --------------------------------------------------------------------------- +# Structural invariants from the policy-style refactor +# --------------------------------------------------------------------------- + + +def test_error_handler_not_registered_as_node(): + """After compile, no hidden __error_handler__* nodes should exist in the graph.""" + + class State(TypedDict): + foo: str + + def failing_node(state: State) -> State: + raise ValueError("boom") + + def handler(state: State, error: NodeError) -> State: + return {"foo": "handled"} + + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, error_handler=handler) + .add_edge(START, "failing_node") + .compile() + ) + + hidden = [k for k in graph.nodes if k.startswith("__error_handler__")] + assert hidden == [], f"unexpected hidden nodes: {hidden}" + + +def test_error_handler_stored_on_pregel_node(): + """The error_handler callable should be a Runnable field on PregelNode, not a name pointer.""" + + class State(TypedDict): + foo: str + + def failing_node(state: State) -> State: + raise ValueError("boom") + + def handler(state: State) -> State: + return {"foo": "handled"} + + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, error_handler=handler) + .add_edge(START, "failing_node") + .compile() + ) + + pregel_node = graph.nodes["failing_node"] + assert pregel_node.error_handler is not None, "error_handler should be set on PregelNode" + assert not hasattr(pregel_node, "error_handler_node"), "old string-pointer field should be gone" + assert not hasattr(pregel_node, "is_error_handler"), "is_error_handler flag should be gone" + + +def test_error_handler_dispatched_from_task_field(): + """error_handler on PregelExecutableTask drives dispatch — no node-map lookup needed.""" + + class State(TypedDict): + foo: str + + def failing_node(state: State) -> State: + raise ValueError("boom") + + def handler(state: State) -> State: + return {"foo": "handled"} + + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, error_handler=handler) + .add_edge(START, "failing_node") + .compile() + ) + result = graph.invoke({"foo": ""}) + assert result["foo"] == "handled" + + +# --------------------------------------------------------------------------- +# Graph-level error handler +# --------------------------------------------------------------------------- + + +def test_graph_level_error_handler_used_when_no_per_node_handler(): + """compile(error_handler=fallback) should catch failures from nodes without their own handler.""" + + class State(TypedDict): + foo: str + + def failing_node(state: State) -> State: + raise RuntimeError("node failed") + + def graph_handler(state: State, error: NodeError) -> State: + return {"foo": f"graph_handler_caught:{error.node}"} + + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node) + .add_edge(START, "failing_node") + .compile(error_handler=graph_handler) + ) + + result = graph.invoke({"foo": ""}) + assert result["foo"] == "graph_handler_caught:failing_node" + + +def test_per_node_handler_takes_precedence_over_graph_level(): + """When a node has its own error_handler, it should win over the graph-level fallback.""" + + class State(TypedDict): + foo: str + + def failing_node(state: State) -> State: + raise RuntimeError("node failed") + + def node_handler(state: State, error: NodeError) -> State: + return {"foo": "node_handler"} + + def graph_handler(state: State, error: NodeError) -> State: + return {"foo": "graph_handler"} + + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, error_handler=node_handler) + .add_edge(START, "failing_node") + .compile(error_handler=graph_handler) + ) + + result = graph.invoke({"foo": ""}) + assert result["foo"] == "node_handler" + + +# --------------------------------------------------------------------------- +# Functional API +# --------------------------------------------------------------------------- + + +def test_task_error_handler_catches_task_failure(): + """@task(error_handler=handler) should catch failures and return the handler's value.""" + + class State(TypedDict): + foo: str + + def handler(x: str, error: NodeError) -> str: + return f"recovered:{error.node}:{x}" + + @task(error_handler=handler) + def failing_task(x: str) -> str: + raise ValueError("task failed") + + @entrypoint() + def wf(state: State) -> State: + result = failing_task(state["foo"]).result() + return {"foo": result} + + result = wf.invoke({"foo": "input"}) + assert result["foo"] == "recovered:failing_task:input" + + +@NEEDS_CONTEXTVARS +def test_task_error_handler_async(): + """@task(error_handler=handler) works for async tasks.""" + import asyncio + + class State(TypedDict): + foo: str + + def handler(x: str, error: NodeError) -> str: + return f"async_recovered:{error.node}:{x}" + + @task(error_handler=handler) + async def failing_task(x: str) -> str: + raise ValueError("async task failed") + + @entrypoint() + async def wf(state: State) -> State: + result = await failing_task(state["foo"]) + return {"foo": result} + + result = asyncio.run(wf.ainvoke({"foo": "input"})) + assert result["foo"] == "async_recovered:failing_task:input"