simplify: drop _should_route_to_error_handler, remove functional API error handler, clean up prepare_node_error_handler_task signature

This commit is contained in:
Sydney Runkle
2026-05-11 16:22:10 -07:00
parent 2cd7ecc81e
commit a55365f3f9
5 changed files with 10 additions and 103 deletions
+1 -43
View File
@@ -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:
+1 -2
View File
@@ -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, {})
+3 -6
View File
@@ -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:
+5 -8
View File
@@ -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.
-44
View File
@@ -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"