chore: dynamic push-task timeouts (#7646)

You can Send(..., timeout=...) now. Much fun.


This would allow us to do something like adding support for an
annotation to let the LLM to pick a timeout for a given tool call, etc.

---------

Co-authored-by: Will Fu-Hinthorn <will@langchain.dev>

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"),

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:

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(

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)

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
         )
 
 

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."""