From a55365f3f94d4c2f1536ac9dbcb9c5a8e6a18980 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Mon, 11 May 2026 16:22:10 -0700 Subject: [PATCH] simplify: drop _should_route_to_error_handler, remove functional API error handler, clean up prepare_node_error_handler_task signature --- libs/langgraph/langgraph/func/__init__.py | 44 +--------------------- libs/langgraph/langgraph/pregel/_algo.py | 3 +- libs/langgraph/langgraph/pregel/_loop.py | 9 ++--- libs/langgraph/langgraph/pregel/_runner.py | 13 +++---- libs/langgraph/tests/test_retry.py | 44 ---------------------- 5 files changed, 10 insertions(+), 103 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index d8fa8c087..be310f0f8 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -65,7 +65,6 @@ 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__"): @@ -82,14 +81,11 @@ 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( - func, + self.func, args, kwargs, retry_policy=self.retry_policy, @@ -97,42 +93,6 @@ 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: @@ -176,7 +136,6 @@ 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]] @@ -284,7 +243,6 @@ 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/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 11ea31b4a..2e3b59a6a 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -1118,7 +1118,6 @@ def prepare_push_task_send( def prepare_node_error_handler_task( failed_task: PregelExecutableTask, *, - handler_node_name: str, handler: Runnable, failed_error: BaseException, checkpoint: Checkpoint, @@ -1131,7 +1130,6 @@ def prepare_node_error_handler_task( store: BaseStore | None = None, checkpointer: BaseCheckpointSaver | None = None, manager: None | ParentRunManager | AsyncParentRunManager = None, - cache_policy: CachePolicy | None = None, retry_policy: Sequence[RetryPolicy] = (), ) -> PregelExecutableTask: """Prepare an error handler task for a failed task. @@ -1139,6 +1137,7 @@ def prepare_node_error_handler_task( The handler borrows the failed task's write pipeline (same state channels), so no separate node registration is needed. """ + handler_node_name = f"__error_handler__{failed_task.name}" checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", "")) task_id_func = _xxhash_str if checkpoint["v"] > 1 else _uuid5_str configurable = config.get(CONF, {}) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index f964f4ff7..444562ced 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -23,6 +23,7 @@ from typing import ( from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager from langchain_core.runnables import Runnable, RunnableConfig + from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( WRITES_IDX_MAP, @@ -1460,13 +1461,11 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): 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) handler_task = prepare_node_error_handler_task( failed_task, - handler_node_name=handler_node_name, handler=handler, failed_error=error, checkpoint=self.checkpoint, @@ -1480,7 +1479,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): checkpointer=self.checkpointer, manager=self.manager, retry_policy=self.retry_policy, - cache_policy=self.cache_policy, ) self.tasks[handler_task.id] = handler_task if not self.is_replaying: @@ -1489,6 +1487,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): self.output_writes(task.id, task.writes, cached=True) return handler_task + + def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" super().put_writes(task_id, writes) @@ -1713,13 +1713,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): 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) handler_task = prepare_node_error_handler_task( failed_task, - handler_node_name=handler_node_name, handler=handler, failed_error=error, checkpoint=self.checkpoint, @@ -1733,7 +1731,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): checkpointer=self.checkpointer, manager=self.manager, retry_policy=self.retry_policy, - cache_policy=self.cache_policy, ) self.tasks[handler_task.id] = handler_task if not self.is_replaying: diff --git a/libs/langgraph/langgraph/pregel/_runner.py b/libs/langgraph/langgraph/pregel/_runner.py index 52c8b4481..886909e72 100644 --- a/libs/langgraph/langgraph/pregel/_runner.py +++ b/libs/langgraph/langgraph/pregel/_runner.py @@ -160,9 +160,6 @@ class PregelRunner: self.aschedule_error_handler = aschedule_error_handler self._handled_exception_ids: set[int] = set() - def _should_route_to_error_handler(self, task: PregelExecutableTask) -> bool: - return task.error_handler is not None - def tick( self, tasks: Iterable[PregelExecutableTask], @@ -213,7 +210,7 @@ class PregelRunner: self.commit(t, exc) if ( not isinstance(exc, GraphBubbleUp) - and self._should_route_to_error_handler(t) + and t.error_handler is not None and self.schedule_error_handler is not None ): self._handled_exception_ids.add(id(exc)) @@ -286,7 +283,7 @@ class PregelRunner: futures[get_waiter()] = None elif ( (task_exc := _exception(fut)) - and self._should_route_to_error_handler(task) + and task.error_handler is not None and not isinstance(task_exc, GraphBubbleUp) ): self._handled_exception_ids.add(id(task_exc)) @@ -405,7 +402,7 @@ class PregelRunner: self.commit(t, exc) if ( not isinstance(exc, GraphBubbleUp) - and self._should_route_to_error_handler(t) + and t.error_handler is not None and self.aschedule_error_handler is not None ): self._handled_exception_ids.add(id(exc)) @@ -485,7 +482,7 @@ class PregelRunner: futures[get_waiter()] = None elif ( (task_exc := _exception(fut)) - and self._should_route_to_error_handler(task) + and task.error_handler is not None and not isinstance(task_exc, GraphBubbleUp) ): self._handled_exception_ids.add(id(task_exc)) @@ -585,7 +582,7 @@ class PregelRunner: else: # save error to checkpointer task.writes.append((ERROR, exception)) - if self._should_route_to_error_handler(task) and not isinstance( + if task.error_handler is not None and not isinstance( exception, GraphBubbleUp ): # Mark early in commit path; loop-side routing may happen later. diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index 184d7db7b..7afc5e433 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -2416,47 +2416,3 @@ def test_per_node_handler_takes_precedence_over_graph_level(): # --------------------------------------------------------------------------- -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"