WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-22 17:45:09 +02:00
Code Issues Packages Projects Releases Wiki Activity

chore: idle timeout (#7631)

Browse Source 
## 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 <will@langchain.dev>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
This commit is contained in:
William FH
2026-04-29 10:12:28 -07:00
committed by GitHub
co-authored by Will Fu-Hinthorn Claude Opus 4.7 copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
parent 3eb73e8ad2
commit 800071d0d4
17 changed files with 2038 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/_internal/_constants.py
+3
View File
@@ -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
libs/langgraph/langgraph/_internal/_runnable.py
+13
View File
@@ -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."""
libs/langgraph/langgraph/_internal/_timeout.py
+50
View File
@@ -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.")
libs/langgraph/langgraph/errors.py
+57 -1
View File
@@ -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
libs/langgraph/langgraph/func/__init__.py
+51 -6
View File
@@ -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(
[
libs/langgraph/langgraph/graph/_node.py
+2 -1
View File
@@ -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
libs/langgraph/langgraph/graph/state.py
+21
View File
@@ -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
libs/langgraph/langgraph/pregel/_algo.py
+15 -1
View File
@@ -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)
libs/langgraph/langgraph/pregel/_call.py
+30 -1
View File
@@ -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
libs/langgraph/langgraph/pregel/_read.py
+12 -1
View File
@@ -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:
libs/langgraph/langgraph/pregel/_retry.py
+495 -15
View File
@@ -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:
libs/langgraph/langgraph/pregel/_runner.py
+7
View File
@@ -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(
libs/langgraph/langgraph/pregel/_utils.py
+75 -2
View File
@@ -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.
libs/langgraph/langgraph/pregel/main.py
+19 -1
View File
@@ -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)},
libs/langgraph/langgraph/runtime.py
+19
View File
@@ -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,
)
libs/langgraph/langgraph/types.py
+36
View File
@@ -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):
libs/langgraph/tests/test_retry.py
+1133 -2
View File
File diff suppressed because it is too large Load Diff