From 800071d0d45849161294b5065cc1707755facef4 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Wed, 29 Apr 2026 10:12:28 -0700 Subject: [PATCH] chore: idle timeout (#7631) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds per-node `timeout` support to async StateGraph/Pregel nodes and to the functional API (`@task` / `@entrypoint`). A timeout caps how long a single node attempt may run, either as a hard wall-clock budget (`run_timeout`), or as an idle window that resets on observable progress (`idle_timeout`), or both. When exceeded, LangGraph raises `NodeTimeoutError`, clears writes from the failed attempt, and lets the existing retry policy decide whether to retry. ## Public API A single `timeout=` kwarg on `add_node`, `@task`, `@entrypoint`, and `NodeBuilder.set_timeout`. Pass a number/`timedelta` for the simple case (treated as a hard wall-clock cap), or a `TimeoutPolicy` (in `langgraph.types`) for finer control: ```python from datetime import timedelta from langgraph.types import TimeoutPolicy # simple: hard wall-clock cap on each attempt builder.add_node("call_model", call_model, timeout=60) builder.add_node("call_model", call_model, timeout=timedelta(minutes=2)) # full control builder.add_node( "call_model", call_model, timeout=TimeoutPolicy( run_timeout=120, # hard wall-clock cap in seconds, never refreshed idle_timeout=30, # cap on time without observable progress, units in seconds refresh_on="auto", # "auto" | "heartbeat" ), ) ``` - `run_timeout`: hard wall-clock cap on a single attempt; never refreshed. - `idle_timeout`: progress-resetting cap. Refreshed by writes, stream output, yielded async stream chunks, child-task scheduling, runtime stream-writer calls, and any LangChain callback event from descendants of the node's run. `runtime.heartbeat()` is a manual signal for work that doesn't naturally emit any of these. - `refresh_on="heartbeat"` narrows the refresh source to explicit `runtime.heartbeat()` only — useful when you want a strict idle definition that isn't reset by chatty subordinates. For long-running async work that doesn't naturally emit progress: ```python async def call_model(state: State, runtime: Runtime) -> State: while still_working: ... runtime.heartbeat() return {"messages": [response]} ``` `NodeTimeoutError` subclasses `TimeoutError` and carries `node`, `timeout`, `run_timeout`, `idle_timeout`, `elapsed`, and `kind` (`"run"` or `"idle"`). If the node's `retry_policy` permits `TimeoutError` it'll retry; the timer resets per attempt. ## Why async-only Sync Python code cannot be safely cancelled in-process, so timeouts only apply to async nodes/tasks. Sync nodes with a `timeout` are rejected at compile time (covers direct nodes, wrapped runnables, sequences, and `RunnableParallel` branches); `run_with_retry` rejects them at runtime as a safety net. ## What gets cancelled and what doesn't When a watchdog fires: 1. The attempt scope is closed under a lock so any in-flight `CONFIG_KEY_SEND` / stream / child-task scheduling that races with the timeout is dropped atomically. 2. Buffered `task.writes` are cleared so pre-timeout writes from the failed attempt don't leak into the checkpoint after a retry succeeds. 3. The background `asyncio.Task` is cancelled; its eventual exception is drained via a done-callback so asyncio doesn't log it. Only the watchdog's own `TimeoutError` converts to `NodeTimeoutError`, so user-raised `asyncio.TimeoutError`, built-in `TimeoutError`, and `NodeTimeoutError` from a child node continue to propagate unchanged. Child tasks already scheduled before the timeout fired still complete — they aren't part of the cancelled task's structured cancellation surface. This is intentional and tested. ## External-watchdog hook WARNING: THIS API IS IN ALPHA AND SUBJECT TO CHANGE. `CONFIG_KEY_TIMED_ATTEMPT_OBSERVER` is a per-config callback that receives lifecycle events for each timed attempt: - `start` — fired before the proc runs, with `task_id`, `task_name`, `attempt`, `run_id`, `thread_id`, `checkpoint_ns`, `started_at`, and the configured `run_timeout_secs` / `idle_timeout_secs` / `refresh_on`. - `progress` — fired on each progress signal that resets the idle clock, rate-limited to ~4 events per `idle_timeout` window so token-rate callbacks don't flood the observer. Carries the same context plus `progress_at`. - `finish` — fired with `finished_at`, `status` (`"success"`/`"error"`), `error_type`, `error_message`. `ParentCommand` and `GraphBubbleUp` are treated as control flow, not errors. This lets an orchestrating process listen to `start` + rolling `progress` to compute its own kill deadline (`progress_at + idle_timeout_secs`) and hard-kill a worker process if the in-process cancellation deadlocks. Observer callbacks run in whatever thread fires them, and any exception they raise is logged and swallowed. ## Implementation Notes `runtime.heartbeat()` updates the progress timestamp without taking the guarded write lock. This avoids lock overhead on high-frequency callback/token paths and accepts a small timestamp race window, which is negligible for expected coarse idle-timeout configurations. ## Testing Covers timeout retry behavior, user-raised timeout propagation, stale-write suppression, stream / callback / heartbeat progress resets, sync-node rejection, `RunnableParallel` branch validation, `StateGraph.add_node` behavior, lower-level Pregel behavior, observer start/progress/finish events, and functional API compatibility. --------- Co-authored-by: Will Fu-Hinthorn Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- .../langgraph/_internal/_constants.py | 3 + .../langgraph/_internal/_runnable.py | 13 + .../langgraph/langgraph/_internal/_timeout.py | 50 + libs/langgraph/langgraph/errors.py | 58 +- libs/langgraph/langgraph/func/__init__.py | 57 +- libs/langgraph/langgraph/graph/_node.py | 3 +- libs/langgraph/langgraph/graph/state.py | 21 + libs/langgraph/langgraph/pregel/_algo.py | 16 +- libs/langgraph/langgraph/pregel/_call.py | 31 +- libs/langgraph/langgraph/pregel/_read.py | 13 +- libs/langgraph/langgraph/pregel/_retry.py | 510 +++++++- libs/langgraph/langgraph/pregel/_runner.py | 7 + libs/langgraph/langgraph/pregel/_utils.py | 77 +- libs/langgraph/langgraph/pregel/main.py | 20 +- libs/langgraph/langgraph/runtime.py | 19 + libs/langgraph/langgraph/types.py | 36 + libs/langgraph/tests/test_retry.py | 1135 ++++++++++++++++- 17 files changed, 2038 insertions(+), 31 deletions(-) create mode 100644 libs/langgraph/langgraph/_internal/_timeout.py diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index d28289053..f2c57f3ca 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -56,6 +56,8 @@ CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns") # holds the current checkpoint_ns, "" for root graph CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished") # holds a callback to be called when a node is finished +CONFIG_KEY_TIMED_ATTEMPT_OBSERVER = sys.intern("__pregel_timed_attempt_observer") +# holds a callback to be called when an idle-timed node attempt starts or finishes CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad") # holds a mutable dict for temporary storage scoped to the current task CONFIG_KEY_RUNNER_SUBMIT = sys.intern("__pregel_runner_submit") @@ -109,6 +111,7 @@ RESERVED = { CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_TIMED_ATTEMPT_OBSERVER, CONFIG_KEY_RESUME_MAP, CONFIG_KEY_STREAM_MESSAGES_V2, # other constants diff --git a/libs/langgraph/langgraph/_internal/_runnable.py b/libs/langgraph/langgraph/_internal/_runnable.py index 63e03f544..2c1a55ffa 100644 --- a/libs/langgraph/langgraph/_internal/_runnable.py +++ b/libs/langgraph/langgraph/_internal/_runnable.py @@ -117,6 +117,19 @@ def set_config_context( ctx.run(_unset_config_context, config_token, run) +def create_task_in_config_context( + coro_factory: Callable[[], Coroutine[Any, Any, Any]], config: RunnableConfig +) -> asyncio.Task[Any]: + """Create an asyncio.Task that inherits `config` as the child runnable context. + + `asyncio.create_task` snapshots the current contextvars onto the new task, + so calling `create_task` while the config context is set ensures the task + sees `config` via `var_child_runnable_config` and any tracing parent. + """ + with set_config_context(config) as context: + return context.run(lambda: asyncio.create_task(coro_factory())) + + # Before Python 3.11 native StrEnum is not available class StrEnum(str, enum.Enum): """A string enum.""" diff --git a/libs/langgraph/langgraph/_internal/_timeout.py b/libs/langgraph/langgraph/_internal/_timeout.py new file mode 100644 index 000000000..f100e36cb --- /dev/null +++ b/libs/langgraph/langgraph/_internal/_timeout.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from datetime import timedelta +from typing import Literal + +from langgraph.types import TimeoutPolicy + +_SYNC_TIMEOUT_PREFIX = ( + "Node timeouts are only supported for async nodes because sync Python " + "execution cannot be safely cancelled in-process." +) + + +def _coerce_timeout_seconds( + value: float | timedelta | None, *, field: str +) -> float | None: + if value is None: + return None + seconds = value.total_seconds() if isinstance(value, timedelta) else float(value) + if seconds <= 0: + raise ValueError(f"{field} must be greater than 0") + return seconds + + +def coerce_timeout_policy( + value: float | timedelta | TimeoutPolicy | None, +) -> TimeoutPolicy | None: + """Normalize a timeout value to positive-second policy fields.""" + if value is not None and not isinstance(value, TimeoutPolicy): + value = TimeoutPolicy(run_timeout=value) + if value is None: + return None + if value.refresh_on not in ("auto", "heartbeat"): + raise ValueError("refresh_on must be 'auto' or 'heartbeat'") + run_timeout = _coerce_timeout_seconds(value.run_timeout, field="run_timeout") + idle_timeout_s = _coerce_timeout_seconds(value.idle_timeout, field="idle_timeout") + if run_timeout is None and idle_timeout_s is None: + return None + return TimeoutPolicy( + run_timeout=run_timeout, + idle_timeout=idle_timeout_s, + refresh_on=value.refresh_on, + ) + + +def sync_timeout_unsupported( + name: str, *, kind: Literal["Node", "Task"] = "Node" +) -> ValueError: + """Build the canonical error for using `timeout` with a sync target.""" + return ValueError(f"{_SYNC_TIMEOUT_PREFIX} {kind} {name!r} is sync.") diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index fb648879e..aef1bf92f 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Sequence from enum import Enum -from typing import Any +from typing import Any, Literal from warnings import warn # EmptyChannelError is re-exported from langgraph.channels.base @@ -20,6 +20,7 @@ __all__ = ( "GraphBubbleUp", "GraphInterrupt", "NodeInterrupt", + "NodeTimeoutError", "ParentCommand", "EmptyInputError", "TaskNotFound", @@ -125,3 +126,58 @@ class TaskNotFound(Exception): """Raised when the executor is unable to find a task (for distributed mode).""" pass + + +class NodeTimeoutError(TimeoutError): + """Raised when a node invocation exceeds one of its configured timeouts. + + Subclasses the built-in `TimeoutError`, so existing `except TimeoutError` + handlers keep working. If the node has a `retry_policy` whose `retry_on` + permits `TimeoutError`, the attempt will be retried. + + Both `idle_timeout` and `run_timeout` reflect the configured policy at the + time of the failure (each is `None` if not configured). `kind` and + `timeout` identify which one fired. + """ + + node: str + timeout: float + run_timeout: float | None + idle_timeout: float | None + elapsed: float + kind: Literal["idle", "run"] + + def __init__( + self, + node: str, + elapsed: float, + *, + kind: Literal["idle", "run"], + idle_timeout: float | None = None, + run_timeout: float | None = None, + ) -> None: + if kind == "idle": + if idle_timeout is None: + raise ValueError("idle_timeout is required when kind='idle'") + message = ( + f"Node '{node}' exceeded its idle timeout of " + f"{idle_timeout:.3f}s without making progress " + f"(elapsed: {elapsed:.3f}s)." + ) + self.timeout = idle_timeout + elif kind == "run": + if run_timeout is None: + raise ValueError("run_timeout is required when kind='run'") + message = ( + f"Node '{node}' exceeded its run timeout of " + f"{run_timeout:.3f}s (elapsed: {elapsed:.3f}s)." + ) + self.timeout = run_timeout + else: + raise ValueError("kind must be 'idle' or 'run'") + super().__init__(message) + self.node = node + self.elapsed = elapsed + self.kind = kind + self.idle_timeout = idle_timeout + self.run_timeout = run_timeout diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index c7443e0a2..be310f0f8 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -5,6 +5,7 @@ import inspect import warnings from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass +from datetime import timedelta from typing import ( Any, Generic, @@ -22,6 +23,11 @@ from typing_extensions import Unpack from langgraph._internal import _serde from langgraph._internal._constants import CACHE_NS_WRITES, PREVIOUS +from langgraph._internal._runnable import is_async_callable +from langgraph._internal._timeout import ( + coerce_timeout_policy, + sync_timeout_unsupported, +) from langgraph._internal._typing import MISSING, DeprecatedKwargs from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue @@ -31,13 +37,19 @@ from langgraph.pregel._call import ( P, SyncAsyncFuture, T, - call, + _call_with_options, get_runnable_for_entrypoint, identifier, ) from langgraph.pregel._read import PregelNode from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry -from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode +from langgraph.types import ( + _DC_KWARGS, + CachePolicy, + RetryPolicy, + StreamMode, + TimeoutPolicy, +) from langgraph.typing import ContextT from langgraph.warnings import LangGraphDeprecatedSinceV05, LangGraphDeprecatedSinceV10 @@ -51,6 +63,7 @@ class _TaskFunction(Generic[P, T]): *, retry_policy: Sequence[RetryPolicy], cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None, + timeout: TimeoutPolicy | None = None, name: str | None = None, ) -> None: if name is not None: @@ -67,15 +80,17 @@ class _TaskFunction(Generic[P, T]): self.func = func self.retry_policy = retry_policy self.cache_policy = cache_policy + self.timeout = timeout functools.update_wrapper(self, func) def __call__(self, *args: P.args, **kwargs: P.kwargs) -> SyncAsyncFuture[T]: - return call( + return _call_with_options( self.func, + args, + kwargs, retry_policy=self.retry_policy, cache_policy=self.cache_policy, - *args, - **kwargs, + timeout=self.timeout, ) def clear_cache(self, cache: BaseCache) -> None: @@ -98,6 +113,7 @@ def task( name: str | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Callable[ [Callable[P, Awaitable[T]] | Callable[P, T]], @@ -119,6 +135,7 @@ def task( name: str | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> ( Callable[[Callable[P, Awaitable[T]] | Callable[P, T]], _TaskFunction[P, T]] @@ -142,6 +159,14 @@ def task( name: An optional name for the task. If not provided, the function name will be used. retry_policy: An optional retry policy (or list of policies) to use for the task in case of a failure. cache_policy: An optional cache policy to use for the task. This allows caching of the task results. + timeout: Timeout for each task attempt. A number or `timedelta` is a hard + wall-clock cap and is not refreshed. Use `TimeoutPolicy` to configure + both a wall-clock `run_timeout` and an `idle_timeout` refreshed by + progress signals. For long-running work that doesn't naturally emit + progress, call `runtime.heartbeat()` from inside the task. When the + timeout fires, `NodeTimeoutError` is raised and the retry policy (if + any) decides whether to retry. Supported only for async tasks; sync + tasks cannot be safely cancelled in-process. Returns: A callable function when used as a decorator. @@ -196,6 +221,7 @@ def task( ) if retry_policy is None: retry_policy = retry # type: ignore[assignment] + timeout_policy = coerce_timeout_policy(timeout) retry_policies: Sequence[RetryPolicy] = ( () @@ -208,8 +234,15 @@ def task( def decorator( func: Callable[P, Awaitable[T]] | Callable[P, T], ) -> Callable[P, SyncAsyncFuture[T]]: + if timeout_policy is not None and not is_async_callable(func): + name_ = name or getattr(func, "__name__", func.__class__.__name__) + raise sync_timeout_unsupported(str(name_), kind="Task") return _TaskFunction( - func, retry_policy=retry_policies, cache_policy=cache_policy, name=name + func, + retry_policy=retry_policies, + cache_policy=cache_policy, + timeout=timeout_policy, + name=name, ) if __func_or_none__ is not None: @@ -268,6 +301,15 @@ class entrypoint(Generic[ContextT]): passed to the workflow. cache_policy: A cache policy to use for caching the results of the workflow. retry_policy: A retry policy (or list of policies) to use for the workflow in case of a failure. + timeout: Timeout for each workflow attempt. A number or `timedelta` is a + hard wall-clock cap and is not refreshed. Use `TimeoutPolicy` to + configure both a wall-clock `run_timeout` and an `idle_timeout` + refreshed by progress signals. For long-running work that doesn't + naturally emit progress, call `runtime.heartbeat()` from inside the + workflow. When the timeout fires, `NodeTimeoutError` is raised and + the retry policy (if any) decides whether to retry. Supported only + for async workflows; sync workflows cannot be safely cancelled + in-process. !!! warning "`config_schema` Deprecated" The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0. @@ -400,6 +442,7 @@ class entrypoint(Generic[ContextT]): context_schema: type[ContextT] | None = None, cache_policy: CachePolicy | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> None: """Initialize the entrypoint decorator.""" @@ -426,6 +469,7 @@ class entrypoint(Generic[ContextT]): self.cache = cache self.cache_policy = cache_policy self.retry_policy = retry_policy + self.timeout = coerce_timeout_policy(timeout) self.context_schema = context_schema @dataclass(**_DC_KWARGS) @@ -535,6 +579,7 @@ class entrypoint(Generic[ContextT]): bound=bound, triggers=[START], channels=START, + timeout=self.timeout, writers=[ ChannelWrite( [ diff --git a/libs/langgraph/langgraph/graph/_node.py b/libs/langgraph/langgraph/graph/_node.py index cadf097d9..d464f06f9 100644 --- a/libs/langgraph/langgraph/graph/_node.py +++ b/libs/langgraph/langgraph/graph/_node.py @@ -9,7 +9,7 @@ from langgraph.store.base import BaseStore from langgraph._internal._typing import EMPTY_SEQ from langgraph.runtime import Runtime -from langgraph.types import CachePolicy, RetryPolicy, StreamWriter +from langgraph.types import CachePolicy, RetryPolicy, StreamWriter, TimeoutPolicy from langgraph.typing import ContextT, NodeInputT, NodeInputT_contra @@ -90,3 +90,4 @@ class StateNodeSpec(Generic[NodeInputT, ContextT]): cache_policy: CachePolicy | 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 f545d4535..59704c607 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -7,6 +7,7 @@ import warnings from collections import defaultdict from collections.abc import Awaitable, Callable, Hashable, Sequence from dataclasses import is_dataclass +from datetime import timedelta from functools import partial from inspect import isclass, isfunction, ismethod, signature from types import FunctionType @@ -45,6 +46,7 @@ from langgraph._internal._fields import ( ) from langgraph._internal._pydantic import create_model from langgraph._internal._runnable import coerce_to_runnable +from langgraph._internal._timeout import coerce_timeout_policy from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -81,6 +83,7 @@ from langgraph.types import ( Command, RetryPolicy, Send, + TimeoutPolicy, ensure_valid_checkpointer, ) from langgraph.typing import ContextT, InputT, NodeInputT, OutputT, StateT @@ -300,6 +303,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Self: """Add a new node to the `StateGraph`, input schema is inferred as the state schema. @@ -367,6 +371,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Self: """Add a new node to the `StateGraph` where input schema is specified. @@ -439,6 +444,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Self: """Add a new node to the `StateGraph`, input schema is inferred as the state schema. @@ -506,6 +512,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Self: """Add a new node to the `StateGraph`, input schema is specified. @@ -580,6 +587,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Self: """Add a new node to the `StateGraph`. @@ -609,6 +617,14 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): !!! warning This is only used for graph rendering and doesn't have any effect on the graph execution. + timeout: Timeout for each node attempt. A number or `timedelta` is + a hard wall-clock cap and is not refreshed. Use `TimeoutPolicy` + to configure both a wall-clock `run_timeout` and an + `idle_timeout` refreshed by progress signals. When exceeded, a + [`NodeTimeoutError`][langgraph.errors.NodeTimeoutError] is raised + and the retry policy (if any) decides whether to retry. Timeouts + are supported only for async nodes; sync nodes cannot be safely + cancelled in-process. Example: ```python @@ -662,6 +678,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): ) if input_schema is None: input_schema = cast(type[NodeInputT] | None, input_) + timeout = coerce_timeout_policy(timeout) if not isinstance(node, str): action = node @@ -757,6 +774,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): cache_policy=cache_policy, ends=ends, defer=defer, + timeout=timeout, ) elif inferred_input_schema is not None: self.nodes[node] = StateNodeSpec( @@ -767,6 +785,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): cache_policy=cache_policy, ends=ends, defer=defer, + timeout=timeout, ) else: self.nodes[node] = StateNodeSpec[StateT, ContextT]( @@ -777,6 +796,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): cache_policy=cache_policy, ends=ends, defer=defer, + timeout=timeout, ) input_schema = input_schema or inferred_input_schema @@ -1341,6 +1361,7 @@ class CompiledStateGraph( retry_policy=node.retry_policy, cache_policy=node.cache_policy, bound=node.runnable, # type: ignore[arg-type] + timeout=node.timeout, ) else: raise RuntimeError diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index d7157e239..1e47cf67b 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -80,6 +80,7 @@ from langgraph.types import ( PregelTask, RetryPolicy, Send, + TimeoutPolicy, ) GetNextVersion = Callable[[V | None, None], V] @@ -114,13 +115,21 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry_policy", "cache_policy", "callbacks") + __slots__ = ( + "func", + "input", + "retry_policy", + "cache_policy", + "callbacks", + "timeout", + ) func: Callable input: tuple[tuple[Any, ...], dict[str, Any]] retry_policy: Sequence[RetryPolicy] | None cache_policy: CachePolicy | None callbacks: Callbacks + timeout: TimeoutPolicy | None def __init__( self, @@ -130,12 +139,14 @@ class Call: retry_policy: Sequence[RetryPolicy] | None, cache_policy: CachePolicy | None, callbacks: Callbacks, + timeout: TimeoutPolicy | None = None, ) -> None: self.func = func self.input = input self.retry_policy = retry_policy self.cache_policy = cache_policy self.callbacks = callbacks + self.timeout = timeout def should_interrupt( @@ -733,6 +744,7 @@ def prepare_single_task( task_path[:3], writers=proc.flat_writers, subgraphs=proc.subgraphs, + timeout=proc.timeout, ) else: return PregelTask(task_id, name, task_path[:3]) @@ -870,6 +882,7 @@ def prepare_push_task_functional( cache_key, task_id, in_progress_task_path, + timeout=call.timeout, ) else: return PregelTask(task_id, name, in_progress_task_path) @@ -1041,6 +1054,7 @@ def prepare_push_task_send( translated_task_path, writers=proc.flat_writers, subgraphs=proc.subgraphs, + timeout=proc.timeout, ) else: return PregelTask(task_id, packet.node, translated_task_path) diff --git a/libs/langgraph/langgraph/pregel/_call.py b/libs/langgraph/langgraph/pregel/_call.py index 0cd007042..6c3fb3856 100644 --- a/libs/langgraph/langgraph/pregel/_call.py +++ b/libs/langgraph/langgraph/pregel/_call.py @@ -8,6 +8,7 @@ import inspect import sys import types from collections.abc import Awaitable, Callable, Generator, Sequence +from datetime import timedelta from typing import Any, Generic, TypeVar, cast from langchain_core.runnables import Runnable @@ -20,9 +21,13 @@ from langgraph._internal._runnable import ( is_async_callable, run_in_executor, ) +from langgraph._internal._timeout import ( + coerce_timeout_policy, + sync_timeout_unsupported, +) from langgraph.config import get_config from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry -from langgraph.types import CachePolicy, RetryPolicy +from langgraph.types import CachePolicy, RetryPolicy, TimeoutPolicy ## # Utilities borrowed from cloudpickle. @@ -255,8 +260,31 @@ def call( *args: Any, retry_policy: Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, **kwargs: Any, ) -> SyncAsyncFuture[T]: + return _call_with_options( + func, + args, + kwargs, + retry_policy=retry_policy, + cache_policy=cache_policy, + timeout=coerce_timeout_policy(timeout), + ) + + +def _call_with_options( + func: Callable[P, Awaitable[T]] | Callable[P, T], + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + retry_policy: Sequence[RetryPolicy] | None = None, + cache_policy: CachePolicy | None = None, + timeout: TimeoutPolicy | None = None, +) -> SyncAsyncFuture[T]: + if timeout is not None and not is_async_callable(func): + name = getattr(func, "__name__", func.__class__.__name__) + raise sync_timeout_unsupported(name, kind="Task") config = get_config() impl = config[CONF][CONFIG_KEY_CALL] fut = impl( @@ -265,5 +293,6 @@ def call( retry_policy=retry_policy, cache_policy=cache_policy, callbacks=config["callbacks"], + timeout=timeout, ) return fut diff --git a/libs/langgraph/langgraph/pregel/_read.py b/libs/langgraph/langgraph/pregel/_read.py index 8d4c21135..d90a69483 100644 --- a/libs/langgraph/langgraph/pregel/_read.py +++ b/libs/langgraph/langgraph/pregel/_read.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from datetime import timedelta from functools import cached_property from typing import ( Any, @@ -11,10 +12,11 @@ from langchain_core.runnables import Runnable, RunnableConfig from langgraph._internal._config import merge_configs from langgraph._internal._constants import CONF, CONFIG_KEY_READ from langgraph._internal._runnable import RunnableCallable, RunnableSeq +from langgraph._internal._timeout import coerce_timeout_policy from langgraph.pregel._utils import find_subgraph_pregel from langgraph.pregel._write import ChannelWrite from langgraph.pregel.protocol import PregelProtocol -from langgraph.types import CachePolicy, RetryPolicy +from langgraph.types import CachePolicy, RetryPolicy, TimeoutPolicy READ_TYPE = Callable[[str | Sequence[str], bool], Any | dict[str, Any]] INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]] @@ -123,6 +125,13 @@ class PregelNode: cache_policy: CachePolicy | None """The cache policy to use when invoking the node.""" + timeout: TimeoutPolicy | None + """Timeout policy for a single invocation. + + If exceeded, `NodeTimeoutError` is raised and the retry policy (if any) + decides whether to retry. Supported only for async nodes. + """ + tags: Sequence[str] | None """Tags to attach to the node for tracing.""" @@ -145,6 +154,7 @@ class PregelNode: retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, subgraphs: Sequence[PregelProtocol] | None = None, + timeout: float | timedelta | TimeoutPolicy | None = None, ) -> None: self.channels = channels self.triggers = list(triggers) @@ -156,6 +166,7 @@ class PregelNode: self.retry_policy = (retry_policy,) else: self.retry_policy = retry_policy + self.timeout = coerce_timeout_policy(timeout) self.tags = tags self.metadata = metadata if subgraphs is not None: diff --git a/libs/langgraph/langgraph/pregel/_retry.py b/libs/langgraph/langgraph/pregel/_retry.py index 538d6b915..acade9643 100644 --- a/libs/langgraph/langgraph/pregel/_retry.py +++ b/libs/langgraph/langgraph/pregel/_retry.py @@ -4,32 +4,487 @@ import asyncio import logging import random import sys +import threading import time +import weakref from collections.abc import Awaitable, Callable, Sequence -from dataclasses import replace -from typing import Any +from contextlib import suppress +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from typing import Any, Literal, NamedTuple +from langchain_core.callbacks import BaseCallbackHandler from langchain_core.runnables import RunnableConfig -from langgraph._internal._config import patch_configurable, recast_checkpoint_ns +from langgraph._internal._config import ( + merge_configs, + patch_configurable, + recast_checkpoint_ns, +) from langgraph._internal._constants import ( CONF, + CONFIG_KEY_CALL, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING, CONFIG_KEY_RUNTIME, + CONFIG_KEY_SEND, + CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, CONFIG_KEY_THREAD_ID, + CONFIG_KEY_TIMED_ATTEMPT_OBSERVER, NS_SEP, ) -from langgraph.errors import GraphBubbleUp, ParentCommand +from langgraph._internal._runnable import create_task_in_config_context +from langgraph._internal._timeout import sync_timeout_unsupported +from langgraph.errors import GraphBubbleUp, NodeTimeoutError, ParentCommand +from langgraph.pregel.protocol import StreamProtocol from langgraph.runtime import ExecutionInfo, Runtime -from langgraph.types import Command, PregelExecutableTask, RetryPolicy +from langgraph.types import Command, PregelExecutableTask, RetryPolicy, TimeoutPolicy logger = logging.getLogger(__name__) SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) +def _timeout_secs(value: float | timedelta) -> float: + return value.total_seconds() if isinstance(value, timedelta) else value + + +@dataclass(frozen=True, slots=True) +class _ResolvedTimeout: + run_timeout_secs: float | None + idle_timeout_secs: float | None + refresh_on: Literal["auto", "heartbeat"] | None + + +def _resolve_timeout(timeout: TimeoutPolicy) -> _ResolvedTimeout: + idle_timeout_secs = ( + _timeout_secs(timeout.idle_timeout) + if timeout.idle_timeout is not None + else None + ) + return _ResolvedTimeout( + run_timeout_secs=( + _timeout_secs(timeout.run_timeout) + if timeout.run_timeout is not None + else None + ), + idle_timeout_secs=idle_timeout_secs, + refresh_on=timeout.refresh_on if idle_timeout_secs is not None else None, + ) + + +class _AttemptContext(NamedTuple): + """Immutable per-attempt metadata shared across start/progress/finish events. + + Built once at attempt start and referenced (not copied) by every emitted + `_AttemptEvent`, so per-event allocation is just the small event wrapper. + + Intentionally underscore-prefixed: this and `_AttemptEvent` are part of an + internal observer contract consumed by langgraph-server. Do not move to + `langgraph.types` — server imports them by this path. + """ + + task_id: str + task_name: str + attempt: int + run_id: str | None + thread_id: str | None + checkpoint_ns: str | None + started_at: datetime + run_timeout_secs: float | None + idle_timeout_secs: float | None + refresh_on: Literal["auto", "heartbeat"] | None + + +@dataclass(frozen=True, slots=True) +class _AttemptEvent: + """One lifecycle event for a timed attempt. + + Holds a reference to the shared `_AttemptContext` and the event-specific + fields. The observer must treat this and `context` as read-only — they + are reused across all events for the same attempt. + """ + + context: _AttemptContext + event: Literal["start", "progress", "finish"] + progress_at: datetime | None = None + finished_at: datetime | None = None + status: Literal["success", "error"] | None = None + error_type: str | None = None + error_message: str | None = None + + +class _TimedAttemptScope: + """Guarded-config window for timed attempts. + + The wrapped config marks writes, stream events, runtime stream writer calls, + child task scheduling, and any LangChain callback event emitted under the + node's run as observable progress when `refresh_on="auto"`. + `runtime.heartbeat()` exposes a manual progress signal for work that doesn't + otherwise emit any of these, and is the only progress signal when + `refresh_on="heartbeat"`. + Guarded writes are serialized with `close()` so cancelled background tasks + cannot persist writes past the timeout boundary. Stream/custom output is + best-effort: it is dropped after close is observed, but callbacks run outside + the lock because they may contain arbitrary user/runtime code. + """ + + __slots__ = ( + "__weakref__", + "_active", + "_last_progress", + "_last_progress_emit", + "_lock", + "_on_progress", + "_progress_min_interval", + "_refresh_on", + ) + + def __init__( + self, + on_progress: Callable[[], None] | None = None, + progress_min_interval: float = 0.0, + refresh_on: Literal["auto", "heartbeat"] | None = None, + ) -> None: + self._active = True + self._last_progress = time.monotonic() + self._lock = threading.Lock() + self._on_progress = on_progress + self._progress_min_interval = progress_min_interval + self._refresh_on = refresh_on + # `-inf` so the first touch always passes the rate-limit gate. + self._last_progress_emit: float = float("-inf") + + def wrap_config(self, config: RunnableConfig) -> RunnableConfig: + configurable = config.get(CONF, {}) + patch: dict[str, Any] = {} + if (send := configurable.get(CONFIG_KEY_SEND)) is not None: + patch[CONFIG_KEY_SEND] = self._guard_send(send) + if (stream := configurable.get(CONFIG_KEY_STREAM)) is not None: + patch[CONFIG_KEY_STREAM] = self._guard_stream(stream) + if (call := configurable.get(CONFIG_KEY_CALL)) is not None: + patch[CONFIG_KEY_CALL] = self._guard_call(call) + if isinstance(runtime := configurable.get(CONFIG_KEY_RUNTIME), Runtime): + if self._refresh_on is not None: + patch[CONFIG_KEY_RUNTIME] = runtime.override( + stream_writer=self._guard_stream_writer(runtime.stream_writer), + heartbeat=self.touch, + ) + else: + patch[CONFIG_KEY_RUNTIME] = runtime.override( + stream_writer=self._guard_stream_writer(runtime.stream_writer) + ) + new_config = patch_configurable(config, patch) if patch else config + if self._refresh_on == "auto": + return merge_configs( + new_config, {"callbacks": [_IdleProgressCallbackHandler(self)]} + ) + return new_config + + def touch(self) -> None: + # Avoid locking this hot progress path. We accept a small race window in + # timestamp ordering because idle_timeout is expected to be coarse compared + # with scheduler/thread timing. + now = time.monotonic() + self._last_progress = now + if self._on_progress is None: + return + # Best-effort rate limit: a benign race may emit a duplicate progress + # event under heavy concurrency, which observers must already tolerate + # (callbacks fire from arbitrary threads). + if now - self._last_progress_emit < self._progress_min_interval: + return + self._last_progress_emit = now + self._on_progress() + + def close(self) -> None: + with self._lock: + self._active = False + + async def wait_for_idle_timeout(self, idle_timeout_s: float) -> None: + while True: + with self._lock: + if not self._active: + return + remaining = self._last_progress + idle_timeout_s - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError + await asyncio.sleep(remaining) + + def _guard_send( + self, send: Callable[[Sequence[tuple[str, Any]]], None] + ) -> Callable[[Sequence[tuple[str, Any]]], None]: + def guarded_send(writes: Sequence[tuple[str, Any]]) -> None: + with self._lock: + if self._active: + if writes and self._refresh_on == "auto": + self._last_progress = time.monotonic() + send(writes) + + return guarded_send + + def _guard_stream(self, stream: StreamProtocol) -> StreamProtocol: + # No lock: stream callbacks fire from the event loop only, so the + # active-check + write happen atomically between awaits. + def guarded_stream(chunk: tuple[tuple[str, ...], str, Any]) -> None: + if not self._active: + return + if self._refresh_on == "auto": + self._last_progress = time.monotonic() + stream(chunk) + + return StreamProtocol(guarded_stream, stream.modes) + + def _guard_call(self, call: Callable[..., Any]) -> Callable[..., Any]: + # No lock: child-task scheduling happens from the event loop only. + def guarded_call(*args: Any, **kwargs: Any) -> Any: + if not self._active: + raise asyncio.CancelledError + if self._refresh_on == "auto": + self._last_progress = time.monotonic() + return call(*args, **kwargs) + + return guarded_call + + def _guard_stream_writer( + self, stream_writer: Callable[[Any], None] + ) -> Callable[[Any], None]: + def guarded_stream_writer(chunk: Any) -> None: + with self._lock: + if not self._active: + return + if self._refresh_on == "auto": + self._last_progress = time.monotonic() + stream_writer(chunk) + + return guarded_stream_writer + + +class _IdleProgressCallbackHandler(BaseCallbackHandler): + """Resets the idle timeout clock on any LangChain callback event. + + Inherits via `config["callbacks"]`, so it sees only events emitted by + runs descended from the node's attempt — sibling nodes do not bleed + through. Holds the scope by weakref so a child manager that outlives + the attempt cannot keep the scope alive. + """ + + # Run inline so progress is recorded in callback emission order; + # thread-pool dispatch would introduce extra reordering. + run_inline = True + + def __init__(self, scope: _TimedAttemptScope) -> None: + self._scope_ref = weakref.ref(scope) + + def _touch(self, *args: Any, **kwargs: Any) -> None: + if (scope := self._scope_ref()) is not None: + scope.touch() + + on_llm_start = _touch + on_chat_model_start = _touch + on_llm_new_token = _touch + on_llm_end = _touch + on_llm_error = _touch + on_chain_start = _touch + on_chain_end = _touch + on_chain_error = _touch + on_tool_start = _touch + on_tool_end = _touch + on_tool_error = _touch + on_retriever_start = _touch + on_retriever_end = _touch + on_retriever_error = _touch + on_agent_action = _touch + on_agent_finish = _touch + on_text = _touch + on_retry = _touch + on_custom_event = _touch + + +def _drain_cancelled(task: asyncio.Task[Any]) -> None: + # Mark the abandoned task's exception as retrieved so asyncio doesn't log it. + with suppress(asyncio.CancelledError): + task.exception() + + +def _start_timed_attempt( + task: PregelExecutableTask, config: RunnableConfig, timeout: _ResolvedTimeout +) -> _AttemptContext | None: + configurable = config.get(CONF, {}) + callback = configurable.get(CONFIG_KEY_TIMED_ATTEMPT_OBSERVER) + if callback is None: + return None + runtime = configurable.get(CONFIG_KEY_RUNTIME) + execution_info = runtime.execution_info if isinstance(runtime, Runtime) else None + context = _AttemptContext( + task_id=task.id, + task_name=task.name, + attempt=execution_info.node_attempt if execution_info is not None else 1, + run_id=execution_info.run_id if execution_info is not None else None, + thread_id=execution_info.thread_id if execution_info is not None else None, + checkpoint_ns=( + execution_info.checkpoint_ns if execution_info is not None else None + ), + started_at=datetime.now(timezone.utc), + run_timeout_secs=timeout.run_timeout_secs, + idle_timeout_secs=timeout.idle_timeout_secs, + refresh_on=timeout.refresh_on, + ) + _dispatch_observer(callback, _AttemptEvent(context=context, event="start")) + return context + + +def _finish_timed_attempt( + config: RunnableConfig, + context: _AttemptContext | None, + error: BaseException | None = None, +) -> None: + if context is None: + return + callback = config.get(CONF, {}).get(CONFIG_KEY_TIMED_ATTEMPT_OBSERVER) + if callback is None: + return + _dispatch_observer( + callback, + _AttemptEvent( + context=context, + event="finish", + finished_at=datetime.now(timezone.utc), + status="error" if error is not None else "success", + error_type=type(error).__name__ if error is not None else None, + error_message=str(error) if error is not None else None, + ), + ) + + +def _emit_progress( + callback: Callable[[_AttemptEvent], None], + context: _AttemptContext, +) -> None: + _dispatch_observer( + callback, + _AttemptEvent( + context=context, + event="progress", + progress_at=datetime.now(timezone.utc), + ), + ) + + +def _dispatch_observer( + callback: Callable[[_AttemptEvent], None], + event: _AttemptEvent, +) -> None: + try: + callback(event) + except Exception: + logger.warning("Timed attempt observer failed", exc_info=True) + + +async def _run_timeout_watchdog(run_timeout_s: float) -> None: + await asyncio.sleep(run_timeout_s) + raise asyncio.TimeoutError + + +async def _arun_with_timeout( + task: PregelExecutableTask, + config: RunnableConfig, + timeout: _ResolvedTimeout, + attempt_ctx: _AttemptContext | None, + *, + stream: bool, +) -> Any: + run_timeout_s = timeout.run_timeout_secs + idle_timeout_s = timeout.idle_timeout_secs + on_progress: Callable[[], None] | None = None + if attempt_ctx is not None: + callback = config.get(CONF, {}).get(CONFIG_KEY_TIMED_ATTEMPT_OBSERVER) + if callback is not None and idle_timeout_s is not None: + on_progress = lambda: _emit_progress(callback, attempt_ctx) # noqa: E731 + scope = _TimedAttemptScope( + on_progress=on_progress, + # Cap progress emission at ~4 events per idle window so token-rate + # callbacks don't flood the observer. + progress_min_interval=idle_timeout_s / 4 if idle_timeout_s is not None else 0.0, + refresh_on=timeout.refresh_on, + ) + scoped_config = scope.wrap_config(config) + start = time.monotonic() + if stream: + # Yielded chunks count as progress only under `refresh_on="auto"`. + # `refresh_on="heartbeat"` is the strict mode where only explicit + # `runtime.heartbeat()` calls reset the idle clock. + async def run() -> Any: + async for _ in task.proc.astream(task.input, scoped_config): + if timeout.refresh_on == "auto": + scope.touch() + + else: + + async def run() -> Any: + return await task.proc.ainvoke(task.input, scoped_config) + + bg = create_task_in_config_context(run, scoped_config) + watchdogs: dict[asyncio.Task[None], Literal["idle", "run"]] = {} + if idle_timeout_s is not None: + watchdogs[asyncio.create_task(scope.wait_for_idle_timeout(idle_timeout_s))] = ( + "idle" + ) + if run_timeout_s is not None: + watchdogs[asyncio.create_task(_run_timeout_watchdog(run_timeout_s))] = "run" + try: + done, _ = await asyncio.wait( + {bg, *watchdogs}, return_when=asyncio.FIRST_COMPLETED + ) + if bg in done: + # Task completed in time. + for watchdog in watchdogs: + watchdog.cancel() + # FIRST_COMPLETED can return both; a watchdog may have + # already raised TimeoutError before we cancelled it. + for watchdog in watchdogs: + with suppress(asyncio.CancelledError, asyncio.TimeoutError): + await watchdog + return await bg + # bg was not in `done`, so every member of `done` is one of our + # watchdogs. Only a watchdog's TimeoutError converts to + # NodeTimeoutError; any TimeoutError raised by the proc itself + # propagates unchanged. + for watchdog in done: + kind = watchdogs[watchdog] + try: + await watchdog + except asyncio.TimeoutError as exc: + elapsed = time.monotonic() - start + scope.close() + task.writes.clear() + bg.cancel() + bg.add_done_callback(_drain_cancelled) + raise NodeTimeoutError( + task.name, + elapsed, + kind=kind, + idle_timeout=idle_timeout_s, + run_timeout=run_timeout_s, + ) from exc + raise RuntimeError( + f"{kind} timeout watchdog completed without raising TimeoutError" + ) + raise RuntimeError("timeout wait completed without task or watchdog") + except asyncio.CancelledError: + scope.close() + bg.cancel() + for watchdog in watchdogs: + watchdog.cancel() + bg.add_done_callback(_drain_cancelled) + raise + finally: + scope.close() + for watchdog in watchdogs: + watchdog.cancel() + + def _ensure_execution_info( runtime: Runtime, config: RunnableConfig, task: PregelExecutableTask ) -> Runtime: @@ -90,6 +545,10 @@ def run_with_retry( ) -> None: """Run a task with retries.""" retry_policy = task.retry_policy or retry_policy + if task.timeout is not None: + # `validate_timeout_supported` catches sync nodes at compile time; + # this is a runtime safety net for paths that may bypass that validation. + raise sync_timeout_unsupported(task.name) attempts = 0 node_first_attempt_time = time.time() config = task.config @@ -195,6 +654,9 @@ async def arun_with_retry( ) -> None: """Run a task asynchronously with retries.""" retry_policy = task.retry_policy or retry_policy + resolved_timeout = ( + _resolve_timeout(task.timeout) if task.timeout is not None else None + ) attempts = 0 node_first_attempt_time = time.time() config = task.config @@ -229,35 +691,53 @@ async def arun_with_retry( ) }, ) + attempt_ctx = ( + _start_timed_attempt(task, config, resolved_timeout) + if resolved_timeout is not None + else None + ) try: - # clear any writes from previous attempts task.writes.clear() - # run the task + if resolved_timeout is None: + if stream: + async for _ in task.proc.astream(task.input, config): + pass + break + return await task.proc.ainvoke(task.input, config) + result = await _arun_with_timeout( + task, config, resolved_timeout, attempt_ctx, stream=stream + ) + _finish_timed_attempt(config, attempt_ctx) if stream: - async for _ in task.proc.astream(task.input, config): - pass # if successful, end break - else: - return await task.proc.ainvoke(task.input, config) + return result except ParentCommand as exc: ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] cmd = exc.args[0] # strip task_ids from namespace for comparison (ns format: "node1|node2:task_id") if cmd.graph in (ns, recast_checkpoint_ns(ns), task.name): - # this command is for the current graph, handle it - for w in task.writers: - w.invoke(cmd, config) + try: + # this command is for the current graph, handle it + for w in task.writers: + w.invoke(cmd, config) + except Exception as writer_exc: + _finish_timed_attempt(config, attempt_ctx, writer_exc) + raise + _finish_timed_attempt(config, attempt_ctx) break elif cmd.graph == Command.PARENT: # this command is for the parent graph, assign it to the parent. exc.args = (replace(cmd, graph=_checkpoint_ns_for_parent_command(ns)),) - # bubble up + _finish_timed_attempt(config, attempt_ctx) + # bubble up the exception to the parent graph raise except GraphBubbleUp: # if interrupted, end + _finish_timed_attempt(config, attempt_ctx) raise except Exception as exc: + _finish_timed_attempt(config, attempt_ctx, exc) if SUPPORTS_EXC_NOTES: exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") if not retry_policy: diff --git a/libs/langgraph/langgraph/pregel/_runner.py b/libs/langgraph/langgraph/pregel/_runner.py index fea4a7272..3945bbf01 100644 --- a/libs/langgraph/langgraph/pregel/_runner.py +++ b/libs/langgraph/langgraph/pregel/_runner.py @@ -46,6 +46,7 @@ from langgraph.types import ( CachePolicy, PregelExecutableTask, RetryPolicy, + TimeoutPolicy, ) F = TypeVar("F", concurrent.futures.Future, asyncio.Future) @@ -537,6 +538,7 @@ def _call( *, retry_policy: Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, + timeout: TimeoutPolicy | None = None, callbacks: Callbacks = None, futures: weakref.ref[FuturesDict], schedule_task: Callable[ @@ -560,6 +562,7 @@ def _call( retry_policy=retry_policy, cache_policy=cache_policy, callbacks=callbacks, + timeout=timeout, ), ): if fut := next( @@ -624,6 +627,7 @@ def _acall( *, retry_policy: Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, + timeout: TimeoutPolicy | None = None, callbacks: Callbacks = None, # injected dependencies futures: weakref.ref[FuturesDict], @@ -657,6 +661,7 @@ def _acall( input, retry_policy=retry_policy, cache_policy=cache_policy, + timeout=timeout, callbacks=callbacks, futures=futures, schedule_task=schedule_task, @@ -678,6 +683,7 @@ async def _acall_impl( *, retry_policy: Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, + timeout: TimeoutPolicy | None = None, callbacks: Callbacks = None, # injected dependencies futures: weakref.ref[FuturesDict[asyncio.Future, asyncio.Event]], @@ -703,6 +709,7 @@ async def _acall_impl( retry_policy=retry_policy, cache_policy=cache_policy, callbacks=callbacks, + timeout=timeout, ), ): if fut := next( diff --git a/libs/langgraph/langgraph/pregel/_utils.py b/libs/langgraph/langgraph/pregel/_utils.py index 0c8a14eec..f7f91e08b 100644 --- a/libs/langgraph/langgraph/pregel/_utils.py +++ b/libs/langgraph/langgraph/pregel/_utils.py @@ -4,16 +4,27 @@ import ast import inspect import re import textwrap -from collections.abc import Callable +from collections.abc import Callable, Sequence +from functools import partial from typing import Any -from langchain_core.runnables import Runnable, RunnableLambda, RunnableSequence +from langchain_core.runnables import ( + Runnable, + RunnableLambda, + RunnableParallel, + RunnableSequence, +) +from langchain_core.runnables.base import RunnableBindingBase +from langchain_core.runnables.config import run_in_executor from langgraph.checkpoint.base import ChannelVersions from typing_extensions import override from langgraph._internal._runnable import RunnableCallable, RunnableSeq +from langgraph._internal._timeout import sync_timeout_unsupported from langgraph.pregel.protocol import PregelProtocol +_SEQUENCE_TYPES = (RunnableSeq, RunnableSequence) + def get_new_channel_versions( previous_versions: ChannelVersions, current_versions: ChannelVersions @@ -64,6 +75,68 @@ def find_subgraph_pregel(candidate: Runnable) -> PregelProtocol | None: return None +def _sequence_steps(runnable: Runnable) -> Sequence[Runnable] | None: + if isinstance(runnable, _SEQUENCE_TYPES): + return runnable.steps + return None + + +def _parallel_steps(runnable: Runnable) -> Sequence[Runnable] | None: + if isinstance(runnable, RunnableParallel): + return tuple(runnable.steps__.values()) + return None + + +def _has_method_override(runnable: Runnable, method_name: str) -> bool: + method = getattr(type(runnable), method_name, None) + return method is not None and method is not getattr(Runnable, method_name) + + +def _is_executor_backed_afunc(afunc: Callable[..., Any] | None) -> bool: + return isinstance(afunc, partial) and afunc.func is run_in_executor + + +def _has_native_async(runnable: Runnable) -> bool: + if isinstance(runnable, RunnableCallable): + return runnable.afunc is not None and not _is_executor_backed_afunc( + runnable.afunc + ) + if isinstance(runnable, RunnableLambda): + return bool(getattr(runnable, "afunc", False)) + return _has_method_override(runnable, "ainvoke") + + +def _runnable_has_native_async(runnable: Runnable) -> bool: + """Return whether a runnable can be idle-timed without known sync code. + + For custom runnable subclasses, an `ainvoke` override is treated as the + async contract. We do not introspect whether that implementation delegates + to blocking work internally — e.g. a subclass whose `ainvoke` calls + `asyncio.to_thread(self.invoke, ...)` will pass this check but the wrapped + sync work is still uncancellable. Idle-timeout enforcement on such a + runnable will fire `NodeTimeoutError` correctly, but the background thread + will keep running until its sync work returns. + """ + + while isinstance(runnable, RunnableBindingBase): + runnable = runnable.bound + steps = _sequence_steps(runnable) + if steps is None: + steps = _parallel_steps(runnable) + if steps is not None: + return all(_runnable_has_native_async(step) for step in steps) + # Raw callables and the common composition wrappers created by graph + # builders fall through here. We do not exhaustively unwrap every Runnable + # wrapper — wrappers that provide `ainvoke` are treated as owning the async + # contract. + return _has_native_async(runnable) + + +def validate_timeout_supported(runnable: Runnable, *, name: str) -> None: + if not _runnable_has_native_async(runnable): + raise sync_timeout_unsupported(name) + + def get_function_nonlocals(func: Callable) -> list[Any]: """Get the nonlocal variables accessed by a function. diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 177fd1d8b..feedff3f0 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -17,6 +17,7 @@ from collections.abc import ( Sequence, ) from dataclasses import is_dataclass, replace +from datetime import timedelta from functools import partial from inspect import isclass from typing import ( @@ -96,6 +97,7 @@ from langgraph._internal._runnable import ( RunnableSeq, coerce_to_runnable, ) +from langgraph._internal._timeout import coerce_timeout_policy from langgraph._internal._typing import MISSING, DeprecatedKwargs from langgraph.callbacks import ( GraphInterruptEvent, @@ -142,7 +144,10 @@ from langgraph.pregel._read import DEFAULT_BOUND, PregelNode from langgraph.pregel._retry import RetryPolicy from langgraph.pregel._runner import PregelRunner from langgraph.pregel._tools import StreamToolCallHandler -from langgraph.pregel._utils import get_new_channel_versions +from langgraph.pregel._utils import ( + get_new_channel_versions, + validate_timeout_supported, +) from langgraph.pregel._validate import validate_graph, validate_keys from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes @@ -175,6 +180,7 @@ from langgraph.types import ( StateUpdate, StreamMode, StreamPart, + TimeoutPolicy, ensure_valid_checkpointer, ) from langgraph.typing import ContextT, InputT, OutputT, StateT @@ -200,6 +206,7 @@ class NodeBuilder: "_bound", "_retry_policy", "_cache_policy", + "_timeout", ) _channels: str | list[str] @@ -210,6 +217,7 @@ class NodeBuilder: _bound: Runnable _retry_policy: list[RetryPolicy] _cache_policy: CachePolicy | None + _timeout: TimeoutPolicy | None def __init__( self, @@ -222,6 +230,7 @@ class NodeBuilder: self._bound = DEFAULT_BOUND self._retry_policy = [] self._cache_policy = None + self._timeout = None def subscribe_only( self, @@ -340,6 +349,11 @@ class NodeBuilder: self._cache_policy = policy return self + def set_timeout(self, timeout: float | timedelta | TimeoutPolicy | None) -> Self: + """Set the per-attempt timeout policy for this node.""" + self._timeout = coerce_timeout_policy(timeout) + return self + def build(self) -> PregelNode: """Builds the node.""" return PregelNode( @@ -351,6 +365,7 @@ class NodeBuilder: bound=self._bound, retry_policy=self._retry_policy, cache_policy=self._cache_policy, + timeout=self._timeout, ) @@ -887,6 +902,9 @@ class Pregel( ) def validate(self) -> Self: + for name, node in self.nodes.items(): + if node.timeout is not None: + validate_timeout_supported(node.node or node.bound, name=name) validate_graph( self.nodes, {k: v for k, v in self.channels.items() if isinstance(v, BaseChannel)}, diff --git a/libs/langgraph/langgraph/runtime.py b/libs/langgraph/langgraph/runtime.py index 9de1b65bd..d1c94021d 100644 --- a/libs/langgraph/langgraph/runtime.py +++ b/libs/langgraph/langgraph/runtime.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass, field, replace from typing import Any, Generic, cast @@ -77,10 +78,14 @@ class ServerInfo: def _no_op_stream_writer(_: Any) -> None: ... +def _no_op_heartbeat() -> None: ... + + class _RuntimeOverrides(TypedDict, Generic[ContextT], total=False): context: ContextT store: BaseStore | None stream_writer: StreamWriter + heartbeat: Callable[[], None] previous: Any execution_info: ExecutionInfo server_info: ServerInfo | None @@ -171,6 +176,16 @@ class Runtime(Generic[ContextT]): stream_writer: StreamWriter = field(default=_no_op_stream_writer) """Function that writes to the custom stream.""" + heartbeat: Callable[[], None] = field(default=_no_op_heartbeat) + """Record progress for the current node's `idle_timeout`. + + Call this from inside long-running work that does not naturally emit + writes, stream chunks, child tasks, or LangChain callback events, to + prevent the node from being treated as idle. It is also the only + progress signal honored under `TimeoutPolicy(refresh_on="heartbeat")`. + Outside an idle-timed attempt this is a no-op. + """ + previous: Any = field(default=None) """The previous return value for the given thread. @@ -196,6 +211,9 @@ class Runtime(Generic[ContextT]): stream_writer=other.stream_writer if other.stream_writer is not _no_op_stream_writer else self.stream_writer, + heartbeat=other.heartbeat + if other.heartbeat is not _no_op_heartbeat + else self.heartbeat, previous=self.previous if other.previous is None else other.previous, execution_info=other.execution_info or self.execution_info, server_info=other.server_info or self.server_info, @@ -222,6 +240,7 @@ DEFAULT_RUNTIME = Runtime( context=None, store=None, stream_writer=_no_op_stream_writer, + heartbeat=_no_op_heartbeat, previous=None, execution_info=None, ) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d04d82da7..87ac33b6c 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -4,6 +4,7 @@ import sys from collections import deque from collections.abc import Callable, Hashable, Sequence from dataclasses import asdict, dataclass +from datetime import timedelta from typing import ( TYPE_CHECKING, Any, @@ -67,6 +68,7 @@ __all__ = ( "CheckpointPayload", "DebugPayload", "RetryPolicy", + "TimeoutPolicy", "CachePolicy", "Interrupt", "StateUpdate", @@ -423,6 +425,39 @@ class RetryPolicy(NamedTuple): """List of exception classes that should trigger a retry, or a callable that returns `True` for exceptions that should trigger a retry.""" +@dataclass(**_DC_KWARGS) +class TimeoutPolicy: + """Configuration for timing out node attempts. + + !!! note "Cooperative cancellation" + + Timeouts rely on asyncio cancellation. If your node uses synchronous + time.sleep() or other CPU-bound work that blocks the GIL, the timeout will not + be fired until after the event loop has been released. + + !!! note "Inline callback dispatch" + + Under `refresh_on="auto"`, an internal handler refreshes the timeout on any + callback event that occurs in the execution of the node or its nested descendants. + """ + + run_timeout: float | timedelta | None = None + """Hard wall-clock cap (in seconds) for a single node attempt. + + This timeout is never refreshed by progress signals or `runtime.heartbeat()`. + """ + + idle_timeout: float | timedelta | None = None + """Maximum time (in seconds) a single node attempt may go without observable progress.""" + + refresh_on: Literal["auto", "heartbeat"] = "auto" + """Which signals refresh `idle_timeout`. + + `"auto"` refreshes on standard graph progress signals and explicit heartbeats. + `"heartbeat"` refreshes only on explicit `runtime.heartbeat()` calls. + """ + + KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes]) @@ -548,6 +583,7 @@ class PregelExecutableTask: path: tuple[str | int | tuple, ...] writers: Sequence[Runnable] = () subgraphs: Sequence[PregelProtocol] = () + timeout: TimeoutPolicy | None = None class StateSnapshot(NamedTuple): diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index 3156f7599..4bb8755b5 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -1,7 +1,20 @@ +import asyncio +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 RunnableLambda, RunnableParallel from langgraph.checkpoint.memory import MemorySaver from typing_extensions import TypedDict @@ -10,18 +23,37 @@ from langgraph._internal._constants import ( 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.graph import START, StateGraph +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, 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 PregelExecutableTask, RetryPolicy +from langgraph.types import Command, PregelExecutableTask, RetryPolicy, TimeoutPolicy + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) def test_should_retry_on_single_exception(): @@ -567,3 +599,1102 @@ def test_run_with_retry_creates_execution_info_when_missing(): 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) + + 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_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) + + 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_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(0.5) + return {"x": state["x"] + 1} + + builder = StateGraph(_TimeoutState) + builder.add_node( + "flaky", + flaky, + timeout=TimeoutPolicy(idle_timeout=0.1), + 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 = [] + + class HeartbeatProc: + async def ainvoke(self, input, config): + runtime = config[CONF][CONFIG_KEY_RUNTIME] + for _ in range(8): + await asyncio.sleep(0.05) + runtime.heartbeat() + return "ok" + + task = _make_task(HeartbeatProc(), timeout=_idle_timeout(0.2), 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.05s; with 8 heartbeats spaced ~0.05s + # we should see at most ~one progress event per heartbeat (well below 8). + 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 == 0.2 + 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