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 98902d502..86b72759c 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -492,10 +492,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): @@ -546,6 +549,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]: @@ -655,6 +667,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: @@ -768,9 +782,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/langgraph/langgraph/_internal/_timeout.py b/libs/langgraph/langgraph/_internal/_timeout.py index f100e36cb..11ff43527 100644 --- a/libs/langgraph/langgraph/_internal/_timeout.py +++ b/libs/langgraph/langgraph/_internal/_timeout.py @@ -11,36 +11,11 @@ _SYNC_TIMEOUT_PREFIX = ( ) -def _coerce_timeout_seconds( - value: float | timedelta | None, *, field: str -) -> float | None: - if value is None: - return None - seconds = value.total_seconds() if isinstance(value, timedelta) else float(value) - if seconds <= 0: - raise ValueError(f"{field} must be greater than 0") - return seconds - - def coerce_timeout_policy( value: float | timedelta | TimeoutPolicy | None, ) -> TimeoutPolicy | None: """Normalize a timeout value to positive-second policy fields.""" - if value is not None and not isinstance(value, TimeoutPolicy): - value = TimeoutPolicy(run_timeout=value) - if value is None: - return None - if value.refresh_on not in ("auto", "heartbeat"): - raise ValueError("refresh_on must be 'auto' or 'heartbeat'") - run_timeout = _coerce_timeout_seconds(value.run_timeout, field="run_timeout") - idle_timeout_s = _coerce_timeout_seconds(value.idle_timeout, field="idle_timeout") - if run_timeout is None and idle_timeout_s is None: - return None - return TimeoutPolicy( - run_timeout=run_timeout, - idle_timeout=idle_timeout_s, - refresh_on=value.refresh_on, - ) + return TimeoutPolicy.coerce(value) def sync_timeout_unsupported( diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 1e47cf67b..a77e31603 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -1054,7 +1054,7 @@ def prepare_push_task_send( translated_task_path, writers=proc.flat_writers, subgraphs=proc.subgraphs, - timeout=proc.timeout, + timeout=packet.timeout if packet.timeout is not None else proc.timeout, ) else: return PregelTask(task_id, packet.node, translated_task_path) @@ -1269,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/types.py b/libs/langgraph/langgraph/types.py index 87ac33b6c..fa0bdc685 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -425,6 +425,17 @@ 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. @@ -457,6 +468,39 @@ class TimeoutPolicy: `"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]) @@ -623,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 @@ -652,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/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index 4bb8755b5..af538919b 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -16,6 +16,7 @@ from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, Huma 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 ( @@ -48,7 +49,13 @@ from langgraph.pregel._retry import ( ) from langgraph.pregel.protocol import StreamProtocol from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime -from langgraph.types import Command, PregelExecutableTask, RetryPolicy, TimeoutPolicy +from langgraph.types import ( + Command, + PregelExecutableTask, + RetryPolicy, + Send, + TimeoutPolicy, +) NEEDS_CONTEXTVARS = pytest.mark.skipif( sys.version_info < (3, 11), @@ -647,6 +654,7 @@ 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) @@ -655,6 +663,30 @@ def test_coerce_timeout_policy_scalar_is_run_timeout(): 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 @@ -1155,6 +1187,9 @@ def test_timeout_validation_is_eager_across_apis(): 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) @@ -1389,6 +1424,28 @@ async def test_state_graph_add_node_timeout_e2e(): 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."""