diff --git a/README.md b/README.md index 64a89b3d7..9c44e66cb 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ PyPI - License PyPI - Downloads Version - Twitter / X + Twitter / X
diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index 558292e72..72044e179 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ [project.urls] Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/" diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index aa5877796..39bfc07aa 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -19,7 +19,7 @@ dependencies = [ [project.urls] Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-sqlite" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/" diff --git a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py index 35ac48665..1a5188f1f 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py @@ -73,6 +73,7 @@ SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset( ("langchain_core.documents.base", "Document"), # langgraph ("langgraph.types", "Send"), + ("langgraph.types", "TimeoutPolicy"), ("langgraph.types", "Interrupt"), ("langgraph.types", "Command"), ("langgraph.types", "StateSnapshot"), diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 5d9d7d5c2..281db5406 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -502,10 +502,13 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext: ), ) elif isinstance(obj, SendProtocol): + args: tuple[Any, ...] = (obj.node, obj.arg) + if (timeout := getattr(obj, "timeout", None)) is not None: + args = (obj.node, obj.arg, timeout) return ormsgpack.Ext( EXT_CONSTRUCTOR_POS_ARGS, _msgpack_enc( - (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)), + (obj.__class__.__module__, obj.__class__.__name__, args), ), ) elif dataclasses.is_dataclass(obj): @@ -556,6 +559,15 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext: raise TypeError(f"Object of type {obj.__class__.__name__} is not serializable") +def _send_from_args(args: Sequence[Any]) -> Any: + # ya we have a cyclic import here ¯\_(ツ)_/¯ + from langgraph.types import Send # type: ignore + + if len(args) == 2: + return Send(*args) + return Send(args[0], args[1], timeout=args[2]) + + def _create_msgpack_ext_hook( allowed_modules: set[tuple[str, ...]] | Literal[True] | None, ) -> Callable[[int, bytes], Any]: @@ -673,6 +685,8 @@ def _create_msgpack_ext_hook( ) if not _check_allowed(tup[0], tup[1]): return tup[2] + if tup[0] == "langgraph.types" and tup[1] == "Send": + return _send_from_args(tup[2]) # module, name, args return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2]) except Exception: @@ -786,9 +800,7 @@ def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any: option=ormsgpack.OPT_NON_STR_KEYS, ) if tup[0] == "langgraph.types" and tup[1] == "Send": - from langgraph.types import Send # type: ignore - - return Send(*tup[2]) + return _send_from_args(tup[2]) # module, name, args return tup[2] except Exception: diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index 8ed0bd3b5..6dd28f2e0 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ [project.urls] Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/" diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index 41e3ee547..424324315 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -29,7 +29,7 @@ inmem = [ [project.urls] Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/cli" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/" diff --git a/libs/langgraph/README.md b/libs/langgraph/README.md index deb7adb15..b02d64015 100644 --- a/libs/langgraph/README.md +++ b/libs/langgraph/README.md @@ -18,7 +18,7 @@ Version Open Issues Docs - Twitter / X + Twitter / X
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..11ff43527 --- /dev/null +++ b/libs/langgraph/langgraph/_internal/_timeout.py @@ -0,0 +1,25 @@ +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_policy( + value: float | timedelta | TimeoutPolicy | None, +) -> TimeoutPolicy | None: + """Normalize a timeout value to positive-second policy fields.""" + return TimeoutPolicy.coerce(value) + + +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 d2ee7527a..63f344c94 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 @@ -82,6 +84,7 @@ from langgraph.types import ( Command, RetryPolicy, Send, + TimeoutPolicy, ensure_valid_checkpointer, ) from langgraph.typing import ContextT, InputT, NodeInputT, OutputT, StateT @@ -301,6 +304,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. @@ -368,6 +372,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. @@ -440,6 +445,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. @@ -507,6 +513,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. @@ -581,6 +588,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`. @@ -610,6 +618,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 @@ -663,6 +679,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 @@ -758,6 +775,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( @@ -768,6 +786,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]( @@ -778,6 +797,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 @@ -1343,6 +1363,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..a77e31603 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=packet.timeout if packet.timeout is not None else proc.timeout, ) else: return PregelTask(task_id, packet.node, translated_task_path) @@ -1255,4 +1269,4 @@ def sanitize_untracked_values_in_send( for k, v in packet.arg.items() if not isinstance(channels.get(k), UntrackedValue) } - return Send(node=packet.node, arg=sanitized_arg) + return Send(node=packet.node, arg=sanitized_arg, timeout=packet.timeout) 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 eb9b46ead..d4a22ae05 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, @@ -143,7 +145,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 @@ -176,6 +181,7 @@ from langgraph.types import ( StateUpdate, StreamMode, StreamPart, + TimeoutPolicy, ensure_valid_checkpointer, ) from langgraph.typing import ContextT, InputT, OutputT, StateT @@ -201,6 +207,7 @@ class NodeBuilder: "_bound", "_retry_policy", "_cache_policy", + "_timeout", ) _channels: str | list[str] @@ -211,6 +218,7 @@ class NodeBuilder: _bound: Runnable _retry_policy: list[RetryPolicy] _cache_policy: CachePolicy | None + _timeout: TimeoutPolicy | None def __init__( self, @@ -223,6 +231,7 @@ class NodeBuilder: self._bound = DEFAULT_BOUND self._retry_policy = [] self._cache_policy = None + self._timeout = None def subscribe_only( self, @@ -341,6 +350,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( @@ -352,6 +366,7 @@ class NodeBuilder: bound=self._bound, retry_policy=self._retry_policy, cache_policy=self._cache_policy, + timeout=self._timeout, ) @@ -888,6 +903,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/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 1830597d0..294b760b2 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -14,15 +14,23 @@ from langgraph.stream.run_stream import ( ) from langgraph.stream.stream_channel import StreamChannel from langgraph.stream.transformers import ( + CheckpointsTransformer, + CustomTransformer, + DebugTransformer, LifecyclePayload, LifecycleTransformer, SubgraphStatus, SubgraphTransformer, + TasksTransformer, + UpdatesTransformer, ) __all__ = [ "AsyncGraphRunStream", "AsyncSubgraphRunStream", + "CheckpointsTransformer", + "CustomTransformer", + "DebugTransformer", "GraphRunStream", "LifecyclePayload", "LifecycleTransformer", @@ -32,4 +40,6 @@ __all__ = [ "SubgraphRunStream", "SubgraphStatus", "SubgraphTransformer", + "TasksTransformer", + "UpdatesTransformer", ] diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index fc2009b13..27ad12376 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -82,6 +82,76 @@ class ValuesTransformer(StreamTransformer): return True +class CustomTransformer(StreamTransformer): + """Capture custom events as a drainable stream of arbitrary payloads. + + Nodes emit custom data via `get_stream_writer()`. This transformer + surfaces those events on `run.custom` as a `StreamChannel[Any]`, + preserving payloads in arrival order. + + Only events at the run's own scope are captured; custom data from + deeper subgraphs is available on the respective subgraph handle's + `.custom` projection. + + Native transformer — `run.custom` is a direct attribute. + """ + + _native = True + required_stream_modes = ("custom",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: StreamChannel[Any] = StreamChannel() + self._scope_list: list[str] = list(scope) + + def init(self) -> dict[str, Any]: + return {"custom": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "custom": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + self._log.push(params["data"]) + return True + + +class UpdatesTransformer(StreamTransformer): + """Capture updates events as a drainable stream of node outputs. + + Surfaces `stream_mode="updates"` data on `run.updates` as a + `StreamChannel[dict[str, Any]]`. Each item is a dict mapping a node + (or task) name to the update it returned after a step. + + Only events at the run's own scope are captured; updates from deeper + subgraphs are available on the respective subgraph handle's + `.updates` projection. + + Native transformer — `run.updates` is a direct attribute. + """ + + _native = True + required_stream_modes = ("updates",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: StreamChannel[dict[str, Any]] = StreamChannel() + self._scope_list: list[str] = list(scope) + + def init(self) -> dict[str, Any]: + return {"updates": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "updates": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + self._log.push(params["data"]) + return True + + class MessagesTransformer(StreamTransformer): """Capture messages events as ChatModelStream objects. @@ -741,3 +811,118 @@ class SubgraphTransformer(_TasksLifecycleBase): handle.path, exc_info=True, ) + + +class CheckpointsTransformer(StreamTransformer): + """Capture checkpoint events as a drainable stream. + + Surfaces `stream_mode="checkpoints"` data on `run.checkpoints` as + a `StreamChannel[dict[str, Any]]`. Each item is in the same format + as returned by `get_state()`. + + Checkpoint events are only emitted when a checkpointer is configured + on the graph. When no checkpointer is present, the projection exists + but receives no events. + + Only events at the run's own scope are captured; checkpoint data from + deeper subgraphs is available on the respective subgraph handle's + `.checkpoints` projection. + + Native transformer — `run.checkpoints` is a direct attribute. + """ + + _native = True + required_stream_modes = ("checkpoints",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: StreamChannel[dict[str, Any]] = StreamChannel() + self._scope_list: list[str] = list(scope) + + def init(self) -> dict[str, Any]: + return {"checkpoints": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "checkpoints": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + self._log.push(params["data"]) + return True + + +class DebugTransformer(StreamTransformer): + """Capture debug events as a drainable stream. + + Surfaces `stream_mode="debug"` data on `run.debug` as a + `StreamChannel[dict[str, Any]]`. Each item is a debug event with + step-level detail (checkpoint snapshots, task payloads, and + task results wrapped with step number and timestamp). + + Only events at the run's own scope are captured; debug data from + deeper subgraphs is available on the respective subgraph handle's + `.debug` projection. + + Native transformer — `run.debug` is a direct attribute. + """ + + _native = True + required_stream_modes = ("debug",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: StreamChannel[dict[str, Any]] = StreamChannel() + self._scope_list: list[str] = list(scope) + + def init(self) -> dict[str, Any]: + return {"debug": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "debug": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + self._log.push(params["data"]) + return True + + +class TasksTransformer(StreamTransformer): + """Capture raw task events as a drainable stream. + + Surfaces `stream_mode="tasks"` data on `run.tasks` as a + `StreamChannel[dict[str, Any]]`. Each item is a task payload + (start or result). + + `LifecycleTransformer` and `SubgraphTransformer` also consume + `tasks` events for subgraph discovery and lifecycle tracking. + This transformer captures the raw payloads independently for + consumers who need task-level detail. + + Only events at the run's own scope are captured; task data from + deeper subgraphs is available on the respective subgraph handle's + `.tasks` projection. + + Native transformer — `run.tasks` is a direct attribute. + """ + + _native = True + required_stream_modes = ("tasks",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: StreamChannel[dict[str, Any]] = StreamChannel() + self._scope_list: list[str] = list(scope) + + def init(self) -> dict[str, Any]: + return {"tasks": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "tasks": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + self._log.push(params["data"]) + return True diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d04d82da7..fa0bdc685 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,83 @@ 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.""" +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 + + +@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. + """ + + @classmethod + def coerce( + cls, value: float | timedelta | TimeoutPolicy | None + ) -> TimeoutPolicy | None: + """Normalize a timeout value to positive-second policy fields.""" + if value is None: + return None + if isinstance(value, TimeoutPolicy): + # Fast path: a policy already produced by coerce() has float + # timeouts and a validated refresh_on, so we can return it as-is. + # `frozen=True` makes this safe to share. + rt, it = value.run_timeout, value.idle_timeout + if ( + value.refresh_on in ("auto", "heartbeat") + and (rt is None or (type(rt) is float and rt > 0)) + and (it is None or (type(it) is float and it > 0)) + and (rt is not None or it is not None) + ): + return value + else: + value = cls(run_timeout=value) + 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 = _coerce_timeout_seconds(value.idle_timeout, field="idle_timeout") + if run_timeout is None and idle_timeout is None: + return None + return cls( + run_timeout=run_timeout, + idle_timeout=idle_timeout, + refresh_on=value.refresh_on, + ) + + KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes]) @@ -548,6 +627,7 @@ class PregelExecutableTask: path: tuple[str | int | tuple, ...] writers: Sequence[Runnable] = () subgraphs: Sequence[PregelProtocol] = () + timeout: TimeoutPolicy | None = None class StateSnapshot(NamedTuple): @@ -587,6 +667,8 @@ class Send: Attributes: node (str): The name of the target node to send the message to. arg (Any): The state or message to send to the target node. + timeout (TimeoutPolicy | None): Optional timeout policy for this specific + pushed task. If omitted, the target node's timeout policy is used. !!! example @@ -616,33 +698,47 @@ class Send: ``` """ - __slots__ = ("node", "arg") + __slots__ = ("node", "arg", "timeout") node: str arg: Any + timeout: TimeoutPolicy | None - def __init__(self, /, node: str, arg: Any) -> None: + def __init__( + self, + /, + node: str, + arg: Any, + *, + timeout: float | timedelta | TimeoutPolicy | None = None, + ) -> None: """ Initialize a new instance of the `Send` class. Args: node: The name of the target node to send the message to. arg: The state or message to send to the target node. + timeout: Optional timeout policy for this specific pushed task. A + number or `timedelta` is treated as a hard `run_timeout`. """ self.node = node self.arg = arg + self.timeout = TimeoutPolicy.coerce(timeout) def __hash__(self) -> int: - return hash((self.node, self.arg)) + return hash((self.node, self.arg, self.timeout)) def __repr__(self) -> str: - return f"Send(node={self.node!r}, arg={self.arg!r})" + if self.timeout is None: + return f"Send(node={self.node!r}, arg={self.arg!r})" + return f"Send(node={self.node!r}, arg={self.arg!r}, timeout={self.timeout!r})" def __eq__(self, value: object) -> bool: return ( isinstance(value, Send) and self.node == value.node and self.arg == value.arg + and self.timeout == value.timeout ) diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 4613e48f9..d720c0b3b 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -38,7 +38,7 @@ Homepage = "https://docs.langchain.com/oss/python/langgraph/overview" Documentation = "https://reference.langchain.com/python/langgraph/" Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/langgraph" Changelog = "https://github.com/langchain-ai/langgraph/releases" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/" diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index 3156f7599..af538919b 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -1,8 +1,22 @@ +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 langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from typing_extensions import TypedDict from langgraph._internal._constants import ( @@ -10,18 +24,43 @@ 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, + Send, + 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 +606,1152 @@ 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) + assert Send("node", None, timeout=timedelta(milliseconds=250)).timeout == policy + + 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_coerce_timeout_policy_returns_same_instance_for_already_coerced(): + policy = coerce_timeout_policy(TimeoutPolicy(run_timeout=1.0, idle_timeout=2.0)) + assert coerce_timeout_policy(policy) is policy + assert TimeoutPolicy.coerce(policy) is policy + + +def test_send_timeout_round_trips_through_msgpack_serde(): + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + packet = Send( + "worker", + {"x": 1}, + timeout=TimeoutPolicy(run_timeout=1, idle_timeout=2), + ) + + assert serde.loads_typed(serde.dumps_typed(packet)) == packet + + +def test_send_without_timeout_round_trips_through_msgpack_serde(): + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + packet = Send("worker", {"x": 1}) + + assert serde.loads_typed(serde.dumps_typed(packet)) == packet + + +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) + + with pytest.raises(ValueError, match="greater than 0"): + Send("slow", {}, 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_send_timeout_overrides_target_node_timeout(): + async def slow(state: _TimeoutState) -> _TimeoutState: + await asyncio.sleep(0.2) + return {"x": state["x"] + 1} + + def route(state: _TimeoutState) -> list[Send]: + return [Send("slow", state, timeout=TimeoutPolicy(idle_timeout=0.05))] + + builder = StateGraph(_TimeoutState) + builder.add_node("slow", slow, timeout=TimeoutPolicy(idle_timeout=1.0)) + builder.add_conditional_edges(START, route) + builder.add_edge("slow", END) + graph = builder.compile() + + with pytest.raises(NodeTimeoutError) as excinfo: + await graph.ainvoke({"x": 1}) + assert excinfo.value.node == "slow" + assert excinfo.value.kind == "idle" + assert excinfo.value.idle_timeout == 0.05 + + +@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 diff --git a/libs/langgraph/tests/test_stream_data_transformers.py b/libs/langgraph/tests/test_stream_data_transformers.py new file mode 100644 index 000000000..61b105ccc --- /dev/null +++ b/libs/langgraph/tests/test_stream_data_transformers.py @@ -0,0 +1,694 @@ +"""Tests for CustomTransformer, UpdatesTransformer, CheckpointsTransformer, DebugTransformer, TasksTransformer. + +These transformers capture raw protocol events for their respective stream +modes and expose them as native projections on the run stream (run.custom, +run.updates, run.checkpoints, run.debug, run.tasks). Tests dispatch synthetic +protocol events through a StreamMux to isolate transformer logic; the final +group exercises real graphs through stream_v2. +""" + +from __future__ import annotations + +import operator +import time +from typing import Annotated, Any + +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.stream._mux import StreamMux +from langgraph.stream.stream_channel import StreamChannel +from langgraph.stream.transformers import ( + CheckpointsTransformer, + CustomTransformer, + DebugTransformer, + LifecycleTransformer, + TasksTransformer, + UpdatesTransformer, +) + +TS = int(time.time() * 1000) + + +def _custom_event(namespace: list[str], data: Any) -> dict[str, Any]: + return { + "type": "event", + "method": "custom", + "params": {"namespace": namespace, "timestamp": TS, "data": data}, + } + + +def _checkpoints_event(namespace: list[str], data: Any) -> dict[str, Any]: + return { + "type": "event", + "method": "checkpoints", + "params": {"namespace": namespace, "timestamp": TS, "data": data}, + } + + +def _debug_event(namespace: list[str], data: Any) -> dict[str, Any]: + return { + "type": "event", + "method": "debug", + "params": {"namespace": namespace, "timestamp": TS, "data": data}, + } + + +def _tasks_event(namespace: list[str], data: Any) -> dict[str, Any]: + return { + "type": "event", + "method": "tasks", + "params": {"namespace": namespace, "timestamp": TS, "data": data}, + } + + +def _updates_event(namespace: list[str], data: Any) -> dict[str, Any]: + return { + "type": "event", + "method": "updates", + "params": {"namespace": namespace, "timestamp": TS, "data": data}, + } + + +def _arm(mux: StreamMux, transformer: Any) -> None: + """Force projection logs to accept pushes (skip lazy-subscribe gate).""" + mux._events._subscribed = True + transformer._log._subscribed = True + + +def _drain(transformer: Any) -> list[Any]: + return list(transformer._log._items) + + +# --------------------------------------------------------------------------- +# CustomTransformer +# --------------------------------------------------------------------------- + + +def test_custom_captures_root_scope_events() -> None: + t = CustomTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_custom_event([], {"status": "processing"})) + mux.push(_custom_event([], {"status": "done"})) + + items = _drain(t) + assert items == [{"status": "processing"}, {"status": "done"}] + + +def test_custom_ignores_subgraph_scope_events() -> None: + t = CustomTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_custom_event(["subgraph:abc"], {"from": "child"})) + + assert _drain(t) == [] + + +def test_custom_scoped_transformer_captures_own_scope() -> None: + t = CustomTransformer(scope=("agent:abc",)) + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_custom_event([], {"from": "root"})) + mux.push(_custom_event(["agent:abc"], {"from": "self"})) + mux.push(_custom_event(["agent:abc", "deep:def"], {"from": "child"})) + + items = _drain(t) + assert items == [{"from": "self"}] + + +def test_custom_preserves_any_payload_type() -> None: + t = CustomTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_custom_event([], "string_payload")) + mux.push(_custom_event([], 42)) + mux.push(_custom_event([], [1, 2, 3])) + + assert _drain(t) == ["string_payload", 42, [1, 2, 3]] + + +def test_custom_does_not_suppress_from_main_log() -> None: + t = CustomTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_custom_event([], "data")) + + methods = [evt["method"] for evt in mux._events._items] + assert "custom" in methods + + +def test_custom_ignores_other_methods() -> None: + t = CustomTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "timestamp": TS, "data": {}}, + } + ) + assert _drain(t) == [] + + +def test_custom_required_stream_modes() -> None: + assert CustomTransformer.required_stream_modes == ("custom",) + + +def test_custom_is_native() -> None: + assert getattr(CustomTransformer, "_native", False) is True + + +def test_custom_init_returns_correct_key() -> None: + t = CustomTransformer() + projection = t.init() + assert "custom" in projection + assert isinstance(projection["custom"], StreamChannel) + + +# --------------------------------------------------------------------------- +# CheckpointsTransformer +# --------------------------------------------------------------------------- + + +def test_checkpoints_captures_root_scope_events() -> None: + t = CheckpointsTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + checkpoint_data = {"values": {"x": 1}, "next": ["node_b"]} + mux.push(_checkpoints_event([], checkpoint_data)) + + items = _drain(t) + assert items == [checkpoint_data] + + +def test_checkpoints_ignores_subgraph_events() -> None: + t = CheckpointsTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_checkpoints_event(["child:abc"], {"values": {"x": 1}})) + + assert _drain(t) == [] + + +def test_checkpoints_scoped_transformer() -> None: + t = CheckpointsTransformer(scope=("sub:abc",)) + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_checkpoints_event([], {"from": "root"})) + mux.push(_checkpoints_event(["sub:abc"], {"from": "self"})) + + assert _drain(t) == [{"from": "self"}] + + +def test_checkpoints_does_not_suppress_from_main_log() -> None: + t = CheckpointsTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_checkpoints_event([], {"values": {}})) + + methods = [evt["method"] for evt in mux._events._items] + assert "checkpoints" in methods + + +def test_checkpoints_required_stream_modes() -> None: + assert CheckpointsTransformer.required_stream_modes == ("checkpoints",) + + +def test_checkpoints_is_native() -> None: + assert getattr(CheckpointsTransformer, "_native", False) is True + + +# --------------------------------------------------------------------------- +# DebugTransformer +# --------------------------------------------------------------------------- + + +def test_debug_captures_root_scope_events() -> None: + t = DebugTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + debug_data = { + "step": 0, + "type": "checkpoint", + "timestamp": "2026-01-01T00:00:00Z", + "payload": {"values": {"x": 1}}, + } + mux.push(_debug_event([], debug_data)) + + items = _drain(t) + assert items == [debug_data] + + +def test_debug_ignores_subgraph_events() -> None: + t = DebugTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_debug_event(["child:abc"], {"step": 0, "type": "task"})) + + assert _drain(t) == [] + + +def test_debug_captures_multiple_event_types() -> None: + t = DebugTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_debug_event([], {"step": 0, "type": "checkpoint", "payload": {}})) + mux.push(_debug_event([], {"step": 1, "type": "task", "payload": {}})) + mux.push(_debug_event([], {"step": 1, "type": "task_result", "payload": {}})) + + items = _drain(t) + assert len(items) == 3 + assert [d["type"] for d in items] == ["checkpoint", "task", "task_result"] + + +def test_debug_does_not_suppress_from_main_log() -> None: + t = DebugTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_debug_event([], {"step": 0})) + + methods = [evt["method"] for evt in mux._events._items] + assert "debug" in methods + + +def test_debug_required_stream_modes() -> None: + assert DebugTransformer.required_stream_modes == ("debug",) + + +def test_debug_is_native() -> None: + assert getattr(DebugTransformer, "_native", False) is True + + +# --------------------------------------------------------------------------- +# TasksTransformer +# --------------------------------------------------------------------------- + + +def test_tasks_captures_root_scope_events() -> None: + t = TasksTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + task_start = {"id": "t1", "name": "my_node", "input": None, "triggers": []} + mux.push(_tasks_event([], task_start)) + + items = _drain(t) + assert items == [task_start] + + +def test_tasks_captures_start_and_result() -> None: + t = TasksTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + start = {"id": "t1", "name": "a", "input": None, "triggers": []} + result = {"id": "t1", "name": "a", "result": {"output": 42}, "error": None} + mux.push(_tasks_event([], start)) + mux.push(_tasks_event([], result)) + + items = _drain(t) + assert items == [start, result] + + +def test_tasks_ignores_subgraph_events() -> None: + t = TasksTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_tasks_event(["child:abc"], {"id": "t1", "name": "x"})) + + assert _drain(t) == [] + + +def test_tasks_scoped_transformer() -> None: + t = TasksTransformer(scope=("agent:abc",)) + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_tasks_event([], {"id": "t1"})) + mux.push(_tasks_event(["agent:abc"], {"id": "t2"})) + mux.push(_tasks_event(["agent:abc", "deep:def"], {"id": "t3"})) + + assert _drain(t) == [{"id": "t2"}] + + +def test_tasks_does_not_suppress_from_main_log() -> None: + """TasksTransformer returns True — it doesn't suppress tasks events. + + (LifecycleTransformer suppresses them, but that's independent.) + """ + t = TasksTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_tasks_event([], {"id": "t1"})) + + methods = [evt["method"] for evt in mux._events._items] + assert "tasks" in methods + + +def test_tasks_required_stream_modes() -> None: + assert TasksTransformer.required_stream_modes == ("tasks",) + + +def test_tasks_is_native() -> None: + assert getattr(TasksTransformer, "_native", False) is True + + +# --------------------------------------------------------------------------- +# UpdatesTransformer +# --------------------------------------------------------------------------- + + +def test_updates_captures_root_scope_events() -> None: + t = UpdatesTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + update = {"my_node": {"value": "hello!"}} + mux.push(_updates_event([], update)) + + items = _drain(t) + assert items == [update] + + +def test_updates_captures_multiple_steps() -> None: + t = UpdatesTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_updates_event([], {"node_a": {"x": 1}})) + mux.push(_updates_event([], {"node_b": {"x": 2}})) + + items = _drain(t) + assert items == [{"node_a": {"x": 1}}, {"node_b": {"x": 2}}] + + +def test_updates_ignores_subgraph_events() -> None: + t = UpdatesTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_updates_event(["child:abc"], {"inner_node": {"v": 1}})) + + assert _drain(t) == [] + + +def test_updates_scoped_transformer() -> None: + t = UpdatesTransformer(scope=("agent:abc",)) + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_updates_event([], {"from": "root"})) + mux.push(_updates_event(["agent:abc"], {"from": "self"})) + + assert _drain(t) == [{"from": "self"}] + + +def test_updates_does_not_suppress_from_main_log() -> None: + t = UpdatesTransformer() + mux = StreamMux([t], is_async=False) + _arm(mux, t) + + mux.push(_updates_event([], {"n": {}})) + + methods = [evt["method"] for evt in mux._events._items] + assert "updates" in methods + + +def test_updates_required_stream_modes() -> None: + assert UpdatesTransformer.required_stream_modes == ("updates",) + + +def test_updates_is_native() -> None: + assert getattr(UpdatesTransformer, "_native", False) is True + + +# --------------------------------------------------------------------------- +# Cross-transformer: unrelated events pass through +# --------------------------------------------------------------------------- + + +def test_unrelated_events_ignored_by_all() -> None: + """Non-matching method events don't land in any transformer's log.""" + transformers = [ + CustomTransformer(), + UpdatesTransformer(), + CheckpointsTransformer(), + DebugTransformer(), + TasksTransformer(), + ] + mux = StreamMux(transformers, is_async=False) + mux._events._subscribed = True + for t in transformers: + t._log._subscribed = True + + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "timestamp": TS, "data": {"x": 1}}, + } + ) + + for t in transformers: + assert list(t._log._items) == [] + + +# --------------------------------------------------------------------------- +# End-to-end: real graphs through stream_v2 +# --------------------------------------------------------------------------- + + +class _State(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _my_node(state: _State) -> dict[str, Any]: + from langgraph.config import get_stream_writer + + writer = get_stream_writer() + writer({"status": "working", "node": "my_node"}) + return {"value": state["value"] + "!", "items": ["done"]} + + +def _make_simple_graph() -> Any: + builder = StateGraph(_State, input_schema=_State) + builder.add_node("my_node", _my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + return builder.compile() + + +def test_stream_v2_custom_projection_opt_in() -> None: + """run.custom surfaces get_stream_writer() payloads when opted in.""" + graph = _make_simple_graph() + run = graph.stream_v2( + {"value": "hello", "items": []}, transformers=[CustomTransformer] + ) + + custom_events = list(run.custom) + assert len(custom_events) >= 1 + assert any(e.get("status") == "working" for e in custom_events) + + +def test_stream_v2_custom_and_values_coexist() -> None: + """Both run.custom and run.values work in the same run.""" + graph = _make_simple_graph() + run = graph.stream_v2( + {"value": "hello", "items": []}, transformers=[CustomTransformer] + ) + + custom_events = list(run.custom) + assert run.output is not None + assert run.output["value"] == "hello!" + assert len(custom_events) >= 1 + + +def test_stream_v2_tasks_projection_opt_in() -> None: + """run.tasks surfaces raw task events when opted in via transformers=.""" + graph = _make_simple_graph() + run = graph.stream_v2({"value": "x", "items": []}, transformers=[TasksTransformer]) + + tasks_events = list(run.tasks) + assert len(tasks_events) >= 1 + names = [t.get("name") for t in tasks_events if "name" in t] + assert "my_node" in names + + +def test_stream_v2_debug_projection_opt_in() -> None: + """run.debug surfaces debug events when opted in via transformers=.""" + graph = _make_simple_graph() + run = graph.stream_v2({"value": "x", "items": []}, transformers=[DebugTransformer]) + + debug_events = list(run.debug) + assert len(debug_events) >= 1 + types = {d.get("type") for d in debug_events} + assert types & {"checkpoint", "task", "task_result"} + + +def test_stream_v2_updates_projection_opt_in() -> None: + """run.updates surfaces node output dicts when opted in via transformers=.""" + graph = _make_simple_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, transformers=[UpdatesTransformer] + ) + + updates = list(run.updates) + assert len(updates) >= 1 + node_names = {k for u in updates for k in u if k != "__interrupt__"} + assert "my_node" in node_names + + +def test_stream_v2_all_transformers_interleaved() -> None: + """All five transformers registered together, consumed via interleave.""" + graph = _make_simple_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, + transformers=[ + CustomTransformer, + UpdatesTransformer, + CheckpointsTransformer, + DebugTransformer, + TasksTransformer, + ], + ) + + collected: dict[str, list[Any]] = { + "custom": [], + "updates": [], + "debug": [], + "tasks": [], + } + for name, item in run.interleave("custom", "updates", "debug", "tasks"): + collected[name].append(item) + + assert len(collected["custom"]) >= 1 + assert len(collected["updates"]) >= 1 + assert len(collected["tasks"]) >= 1 + assert len(collected["debug"]) >= 1 + types = {d.get("type") for d in collected["debug"]} + assert types & {"checkpoint", "task", "task_result"} + node_names = {k for u in collected["updates"] for k in u if k != "__interrupt__"} + assert "my_node" in node_names + + assert run.output is not None + assert run.output["value"] == "x!" + + +def test_stream_v2_all_transformers_with_checkpointer() -> None: + """All transformers with a checkpointer — run.checkpoints populated.""" + from langgraph.checkpoint.memory import InMemorySaver + + builder = StateGraph(_State, input_schema=_State) + builder.add_node("my_node", _my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + run = graph.stream_v2( + {"value": "x", "items": []}, + config={"configurable": {"thread_id": "test-all"}}, + transformers=[ + CustomTransformer, + UpdatesTransformer, + CheckpointsTransformer, + DebugTransformer, + TasksTransformer, + ], + ) + + collected: dict[str, list[Any]] = { + "custom": [], + "updates": [], + "checkpoints": [], + "debug": [], + "tasks": [], + } + for name, item in run.interleave( + "custom", "updates", "checkpoints", "debug", "tasks" + ): + collected[name].append(item) + + assert len(collected["checkpoints"]) >= 1 + assert len(collected["custom"]) >= 1 + + +def test_stream_v2_checkpoints_projection_opt_in() -> None: + """run.checkpoints surfaces checkpoint data when opted in with a checkpointer.""" + from langgraph.checkpoint.memory import InMemorySaver + + builder = StateGraph(_State, input_schema=_State) + builder.add_node("my_node", _my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + run = graph.stream_v2( + {"value": "x", "items": []}, + config={"configurable": {"thread_id": "test-ckpt-standalone"}}, + transformers=[CheckpointsTransformer], + ) + + checkpoints = list(run.checkpoints) + assert len(checkpoints) >= 1 + + +# --------------------------------------------------------------------------- +# TasksTransformer + LifecycleTransformer co-registration +# --------------------------------------------------------------------------- + + +def test_tasks_and_lifecycle_coregistration() -> None: + """When both are in the same StreamMux, LifecycleTransformer suppresses + tasks events from the main log (returns False) while TasksTransformer + still captures them into its own log. + """ + lifecycle = LifecycleTransformer() + tasks = TasksTransformer() + mux = StreamMux([lifecycle, tasks], is_async=False) + mux._events._subscribed = True + tasks._log._subscribed = True + lifecycle._channel._subscribed = True + + task_data = {"id": "t1", "name": "my_node", "input": None, "triggers": []} + mux.push(_tasks_event([], task_data)) + + assert _drain(tasks) == [task_data] + + methods = [evt["method"] for evt in mux._events._items] + assert "tasks" not in methods + + +def test_tasks_and_lifecycle_coregistration_e2e() -> None: + """E2e: TasksTransformer captures task events even when LifecycleTransformer + is present and suppressing them from the main log. + """ + graph = _make_simple_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, + transformers=[TasksTransformer], + ) + + tasks_events = list(run.tasks) + assert len(tasks_events) >= 1 + names = [t.get("name") for t in tasks_events if "name" in t] + assert "my_node" in names diff --git a/libs/prebuilt/pyproject.toml b/libs/prebuilt/pyproject.toml index ad4f5c3c1..0c658f7c5 100644 --- a/libs/prebuilt/pyproject.toml +++ b/libs/prebuilt/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ [project.urls] Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/prebuilt" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/" diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml index c856b2757..22f398e37 100644 --- a/libs/sdk-py/pyproject.toml +++ b/libs/sdk-py/pyproject.toml @@ -18,7 +18,7 @@ path = "langgraph_sdk/__init__.py" [project.urls] Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/sdk-py" -Twitter = "https://x.com/LangChain" +Twitter = "https://x.com/langchain_oss" Slack = "https://www.langchain.com/join-community" Reddit = "https://www.reddit.com/r/LangChain/"