From 40ab009c62fe0a4b2e99d3119a72c0c6b73edd32 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Wed, 29 Apr 2026 15:23:31 -0700 Subject: [PATCH] feat: allow graph to graceful shutdown/drain by request (#7274) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds cooperative drain support for Pregel runs so a graph can be asked to stop at the next superstep boundary, persist its checkpoint, and surface a resumable terminal exception. - New `RunControl` (in `langgraph.runtime`) — a thread-safe handle whose `request_drain(reason="shutdown")` sets a single flag. - New `GraphDrained(GraphBubbleUp)` exception (in `langgraph.errors`) raised when a run exits early due to drain. Carries the `reason` string. - New `control: RunControl | None` kwarg on `invoke` / `ainvoke` / `stream` / `astream` / `stream_v2` / `astream_v2`. Wired through to `Runtime.control`, so nodes can read `runtime.control.drain_requested` / `drain_reason` and even call `request_drain()` from inside a node. - Stream transformers learn `"drained"` as a terminal `SubgraphStatus`. The intended use is hooking SIGTERM (or any external supervisor signal) to `control.request_drain("sigterm")` so an in-flight graph run can stop cleanly and be resumed later from the saved checkpoint. ## Semantics: cooperative, between-superstep `request_drain()` flips a flag. The Pregel loop checks it at the top of each `tick()`, **after** the previous superstep's writes have been applied and checkpointed. It never preempts work that is already running. | Scenario | Behavior | |---|---| | Node mid-execution (blocking I/O, sleeps, etc.) | Runs to completion. Drain takes effect on the next superstep. | | Node with a retry policy currently retrying | Retry loop runs to exhaustion or success (drain is not checked between retries). Drain takes effect on the next superstep. | | Functional API: `@entrypoint` with pending `@task` futures | Entrypoint and all dispatched tasks complete; drain takes effect after the entrypoint returns. | | Graph naturally finishes on the same tick where drain was requested (no more tasks) | Treated as `done`; returns normally. **No `GraphDrained` is raised.** The caller can inspect `control.drain_requested` afterwards to distinguish a drained-but-completed run from a normal one. | | More tasks remain | Raises `GraphDrained(reason)`. The checkpoint of the last completed superstep is saved (also under `durability="exit"`). Resume with `invoke(None, config)` / `ainvoke(None, config)`. | | Subgraph requests drain | `GraphDrained` bubbles up through the parent loop and stops it at its own next superstep boundary; the parent's checkpoint is saved and resumable. | Drain does **not** cancel asyncio tasks or kill threads. Pair it with a graceful timeout + `task.cancel()` (or process exit) if you need a hard upper bound — see `test_drain_then_cancel_after_graceful_timeout` for the recommended pattern. ## Usage ```python from langgraph.runtime import RunControl from langgraph.errors import GraphDrained control = RunControl() # In a signal handler, supervisor, etc.: # control.request_drain("sigterm") try: result = graph.invoke(input, config, control=control) if control.drain_requested: # finished naturally on the same tick where drain was requested ... except GraphDrained as e: # checkpoint saved; resume later with the same config log.info("graph drained: %s", e.reason) ``` ## Test plan - [x] Sync + async drain stops the next superstep (`test_run_control_request_drain_stops_future_steps[_async]`) - [x] Drain on the terminal step finishes normally (`test_drain_requested_in_terminal_step_finishes_normally[_async]`) - [x] `durability=\"exit\"` persists a resumable checkpoint on drain (`test_drain_with_exit_durability_persists_resume_checkpoint`) - [x] Subgraph drain bubbles up and parent resumes correctly (`test_drain_from_subgraph_can_resume_parent`) - [x] External thread / task triggering drain mid-run (`test_external_drain_concurrent_sync` / `_async`) - [x] Drain + hard cancel after graceful timeout (`test_drain_then_cancel_after_graceful_timeout`) - [x] Functional API: in-flight `@task` futures still resolve after `request_drain()` (`test_request_drain_allows_inflight_[a]call_scheduling`) - [x] `control` kwarg wired through `stream_v2` (`test_stream_v2_accepts_control_for_drain`) - [x] `Runtime.merge` preserves `control` (`test_merge_runtime_preserves_run_control`) --------- Co-authored-by: Quanzheng Long Co-authored-by: Will Fu-Hinthorn Co-authored-by: Claude Opus 4.7 (1M context) --- libs/langgraph/langgraph/errors.py | 22 +- libs/langgraph/langgraph/pregel/_loop.py | 19 +- libs/langgraph/langgraph/pregel/main.py | 40 ++ libs/langgraph/langgraph/runtime.py | 51 +- .../langgraph/stream/transformers.py | 12 +- libs/langgraph/tests/test_pregel.py | 23 + libs/langgraph/tests/test_pregel_async.py | 24 + libs/langgraph/tests/test_runtime.py | 517 +++++++++++++++++- libs/langgraph/tests/test_stream_v2.py | 28 + 9 files changed, 721 insertions(+), 15 deletions(-) diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index aef1bf92f..a99546f10 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -15,6 +15,7 @@ from langgraph.warnings import LangGraphDeprecatedSinceV10 __all__ = ( "EmptyChannelError", "ErrorCode", + "GraphDrained", "GraphRecursionError", "InvalidUpdateError", "GraphBubbleUp", @@ -43,6 +44,23 @@ def create_error_message(*, message: str, error_code: ErrorCode) -> str: ) +class GraphBubbleUp(Exception): + pass + + +class GraphDrained(GraphBubbleUp): + """Raised when a graph run exits early due to a drain request. + + This indicates the graph stopped cooperatively at a superstep boundary + because `RunControl.request_drain()` was called (e.g., in response to + SIGTERM). The checkpoint is saved and the run can be resumed later. + """ + + def __init__(self, reason: str = "shutdown") -> None: + self.reason = reason + super().__init__(f"Graph drained: {reason}") + + class GraphRecursionError(RecursionError): """Raised when the graph has exhausted the maximum number of steps. @@ -78,10 +96,6 @@ class InvalidUpdateError(Exception): pass -class GraphBubbleUp(Exception): - pass - - class GraphInterrupt(GraphBubbleUp): """Raised when a subgraph is interrupted, suppressed by the root graph. Never raised directly, or surfaced to the user.""" diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index e6aea17d8..80b9e7b42 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -45,6 +45,7 @@ from langgraph._internal._constants import ( CONFIG_KEY_REPLAY_STATE, CONFIG_KEY_RESUME_MAP, CONFIG_KEY_RESUMING, + CONFIG_KEY_RUNTIME, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, @@ -119,6 +120,7 @@ from langgraph.pregel.debug import ( map_debug_tasks, ) from langgraph.pregel.protocol import StreamChunk, StreamProtocol +from langgraph.runtime import RunControl, Runtime from langgraph.types import ( All, CachePolicy, @@ -206,10 +208,12 @@ class PregelLoop: "input", "pending", "done", + "draining", "interrupt_before", "interrupt_after", "out_of_steps", ] + control: RunControl | None tasks: dict[str, PregelExecutableTask] output: None | dict[str, Any] | Any = None updated_channels: set[str] | None = None @@ -317,6 +321,8 @@ class PregelLoop: else () ) self.prev_checkpoint_config = None + runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME) + self.control = runtime.control if isinstance(runtime, Runtime) else None def _push_graph_lifecycle_event( self, @@ -324,11 +330,16 @@ class PregelLoop: *, interrupts: tuple[Interrupt, ...] = (), ) -> None: + # drain status never reaches lifecycle events: tick() returns False + # before pushing, and interrupts are raised through GraphInterrupt + if self.status == "draining": + raise RuntimeError("Draining status cannot emit lifecycle events") + status = self.status if kind == "resume": self._graph_lifecycle_events.append( GraphResumeEvent( run_id=None, - status=self.status, + status=status, checkpoint_id=self.checkpoint["id"], checkpoint_ns=self.checkpoint_ns, ) @@ -337,7 +348,7 @@ class PregelLoop: self._graph_lifecycle_events.append( GraphInterruptEvent( run_id=None, - status=self.status, + status=status, checkpoint_id=self.checkpoint["id"], checkpoint_ns=self.checkpoint_ns, interrupts=interrupts, @@ -569,6 +580,10 @@ class PregelLoop: self.status = "done" return False + if self.control is not None and self.control.drain_requested: + self.status = "draining" + return False + # if there are pending writes from a previous loop, apply them if not self.is_replaying and self.checkpoint_pending_writes: self._match_writes(self.tasks) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index d4a22ae05..fc9428694 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -111,6 +111,7 @@ from langgraph.config import get_config from langgraph.constants import END from langgraph.errors import ( ErrorCode, + GraphDrained, GraphRecursionError, InvalidUpdateError, create_error_message, @@ -156,6 +157,7 @@ from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtoco from langgraph.runtime import ( DEFAULT_RUNTIME, BaseUser, + RunControl, Runtime, ServerInfo, ) @@ -2570,6 +2572,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, subgraphs: bool = False, debug: bool | None = None, version: Literal["v2"], @@ -2589,6 +2592,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, subgraphs: bool = False, debug: bool | None = None, version: Literal["v1"] = ..., @@ -2607,6 +2611,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, subgraphs: bool = False, debug: bool | None = None, version: Literal["v1", "v2"] = "v1", @@ -2651,6 +2656,7 @@ class Pregel( - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. + control: Optional run control used to request cooperative drain. subgraphs: Whether to stream events from inside subgraphs, defaults to `False`. If `True`, the events will be emitted as tuples `(namespace, data)`, @@ -2815,6 +2821,7 @@ class Pregel( previous=None, execution_info=None, server_info=server_info, + control=control or parent_runtime.control or RunControl(), ) runtime = parent_runtime.merge(runtime) config[CONF][CONFIG_KEY_RUNTIME] = runtime @@ -2945,6 +2952,10 @@ class Pregel( error_code=ErrorCode.GRAPH_RECURSION_LIMIT, ) raise GraphRecursionError(msg) + elif loop.status == "draining": + if loop.control is None: + raise RuntimeError("Draining status requires run control") + raise GraphDrained(loop.control.drain_reason or "shutdown") # set final channel values as run output run_manager.on_chain_end(loop.output) except BaseException as e: @@ -2965,6 +2976,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, subgraphs: bool = False, debug: bool | None = None, version: Literal["v2"], @@ -2984,6 +2996,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, subgraphs: bool = False, debug: bool | None = None, version: Literal["v1"] = ..., @@ -3002,6 +3015,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, subgraphs: bool = False, debug: bool | None = None, version: Literal["v1", "v2"] = "v1", @@ -3046,6 +3060,7 @@ class Pregel( - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. + control: Optional run control used to request cooperative drain. subgraphs: Whether to stream events from inside subgraphs, defaults to `False`. If `True`, the events will be emitted as tuples `(namespace, data)`, @@ -3245,6 +3260,7 @@ class Pregel( previous=None, execution_info=None, server_info=server_info, + control=control or parent_runtime.control or RunControl(), ) runtime = parent_runtime.merge(runtime) config[CONF][CONFIG_KEY_RUNTIME] = runtime @@ -3413,6 +3429,10 @@ class Pregel( error_code=ErrorCode.GRAPH_RECURSION_LIMIT, ) raise GraphRecursionError(msg) + elif loop.status == "draining": + if loop.control is None: + raise RuntimeError("Draining status requires run control") + raise GraphDrained(loop.control.drain_reason or "shutdown") # set final channel values as run output await run_manager.on_chain_end(loop.output) except BaseException as e: @@ -3427,6 +3447,7 @@ class Pregel( *, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, + control: RunControl | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, ) -> Any: """Start a sync v2 streaming run driven by transformer projections. @@ -3456,6 +3477,7 @@ class Pregel( config: Optional runnable config forwarded to the graph. interrupt_before: Nodes to interrupt before, if any. interrupt_after: Nodes to interrupt after, if any. + control: Optional run control used to request cooperative drain. transformers: Extra transformer classes or configured factories appended after compile-time `stream_transformers`. Factories are called as `factory(scope)` so they can propagate to @@ -3490,6 +3512,7 @@ class Pregel( version="v2", interrupt_before=interrupt_before, interrupt_after=interrupt_after, + control=control, ) ) return GraphRunStream(graph_iter, mux) @@ -3501,6 +3524,7 @@ class Pregel( *, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, + control: RunControl | None = None, transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, ) -> Any: """Async counterpart to `stream_v2`. @@ -3523,6 +3547,7 @@ class Pregel( config: Optional runnable config forwarded to the graph. interrupt_before: Nodes to interrupt before, if any. interrupt_after: Nodes to interrupt after, if any. + control: Optional run control used to request cooperative drain. transformers: Extra transformer classes or configured factories appended after compile-time `stream_transformers`. Factories are called as `factory(scope)` so they can propagate to @@ -3553,6 +3578,7 @@ class Pregel( version="v2", interrupt_before=interrupt_before, interrupt_after=interrupt_after, + control=control, ).__aiter__() return AsyncGraphRunStream(graph_aiter, mux) @@ -3569,6 +3595,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v2"], **kwargs: Any, ) -> GraphOutput[OutputT]: ... @@ -3586,6 +3613,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v2"], **kwargs: Any, ) -> list[StreamPart[StateT, OutputT]]: ... @@ -3603,6 +3631,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v1"] = ..., **kwargs: Any, ) -> dict[str, Any] | Any: ... @@ -3619,6 +3648,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: @@ -3643,6 +3673,7 @@ class Pregel( - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. + control: Optional run control used to request cooperative drain. version: The streaming format version. `"v1"` (default) returns the traditional format, `"v2"` returns `StreamPart` typed dicts when `stream_mode` is not `"values"`. @@ -3670,6 +3701,7 @@ class Pregel( interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, + control=control, version=version, **kwargs, ): @@ -3693,6 +3725,7 @@ class Pregel( interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, + control=control, **kwargs, ): if stream_mode == "values": @@ -3739,6 +3772,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v2"], **kwargs: Any, ) -> GraphOutput[OutputT]: ... @@ -3756,6 +3790,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v2"], **kwargs: Any, ) -> list[StreamPart[StateT, OutputT]]: ... @@ -3773,6 +3808,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v1"] = ..., **kwargs: Any, ) -> dict[str, Any] | Any: ... @@ -3789,6 +3825,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + control: RunControl | None = None, version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: @@ -3813,6 +3850,7 @@ class Pregel( - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. + control: Optional run control used to request cooperative drain. version: The streaming format version. `"v1"` (default) returns the traditional format, `"v2"` returns `StreamPart` typed dicts when `stream_mode` is not `"values"`. @@ -3840,6 +3878,7 @@ class Pregel( interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, + control=control, version=version, **kwargs, ): @@ -3863,6 +3902,7 @@ class Pregel( interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, + control=control, **kwargs, ): if stream_mode == "values": diff --git a/libs/langgraph/langgraph/runtime.py b/libs/langgraph/langgraph/runtime.py index d1c94021d..9e6d32bf8 100644 --- a/libs/langgraph/langgraph/runtime.py +++ b/libs/langgraph/langgraph/runtime.py @@ -16,6 +16,7 @@ from langgraph.typing import ContextT __all__ = ( "BaseUser", "ExecutionInfo", + "RunControl", "Runtime", "ServerInfo", "get_runtime", @@ -75,6 +76,34 @@ class ServerInfo: """ +class RunControl: + """Run-scoped control surface for cooperative draining. + + Intended for a single graph run. Create a fresh `RunControl` per run; + reusing a control after `request_drain()` leaves it drained. + + Safe to call from any thread: the drain request is represented by a + single attribute write, so no lock is needed for this signal. + If more mutable state is added here, add synchronization. + """ + + __slots__ = ("_drain_reason",) + + def __init__(self) -> None: + self._drain_reason: str | None = None + + def request_drain(self, reason: str = "shutdown") -> None: + self._drain_reason = reason + + @property + def drain_requested(self) -> bool: + return self._drain_reason is not None + + @property + def drain_reason(self) -> str | None: + return self._drain_reason + + def _no_op_stream_writer(_: Any) -> None: ... @@ -89,6 +118,7 @@ class _RuntimeOverrides(TypedDict, Generic[ContextT], total=False): previous: Any execution_info: ExecutionInfo server_info: ServerInfo | None + control: RunControl | None @dataclass(**_DC_KWARGS) @@ -167,7 +197,7 @@ class Runtime(Generic[ContextT]): context: ContextT = field(default=None) # type: ignore[assignment] """Static context for the graph run, like `user_id`, `db_conn`, etc. - + Can also be thought of as 'run dependencies'.""" store: BaseStore | None = field(default=None) @@ -188,7 +218,7 @@ class Runtime(Generic[ContextT]): previous: Any = field(default=None) """The previous return value for the given thread. - + Only available with the functional API when a checkpointer is provided. """ @@ -200,6 +230,13 @@ class Runtime(Generic[ContextT]): server_info: ServerInfo | None = field(default=None) """Metadata injected by LangGraph Server. None when running open-source LangGraph without LangSmith deployments.""" + control: RunControl | None = field(default=None) + """Run-scoped control plane for cooperative draining. + + Populated automatically during graph runs. None outside an active + graph runtime. + """ + def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]: """Merge two runtimes together. @@ -217,6 +254,7 @@ class Runtime(Generic[ContextT]): 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, + control=other.control or self.control, ) def override( @@ -235,6 +273,14 @@ class Runtime(Generic[ContextT]): execution_info=self.execution_info.patch(**overrides), ) + @property + def drain_requested(self) -> bool: + return self.control.drain_requested if self.control is not None else False + + @property + def drain_reason(self) -> str | None: + return self.control.drain_reason if self.control is not None else None + DEFAULT_RUNTIME = Runtime( context=None, @@ -243,6 +289,7 @@ DEFAULT_RUNTIME = Runtime( heartbeat=_no_op_heartbeat, previous=None, execution_info=None, + control=None, ) diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 27ad12376..ef0ef4bf9 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -12,7 +12,7 @@ from langchain_core.messages import AIMessageChunk, BaseMessage from langchain_protocol.protocol import MessagesData from typing_extensions import NotRequired, TypedDict -from langgraph.errors import GraphInterrupt +from langgraph.errors import GraphDrained, GraphInterrupt from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import AsyncSubgraphRunStream, SubgraphRunStream from langgraph.stream.stream_channel import StreamChannel @@ -327,7 +327,7 @@ class MessagesTransformer(StreamTransformer): self._by_run.clear() -SubgraphStatus = Literal["started", "completed", "failed", "interrupted"] +SubgraphStatus = Literal["started", "completed", "failed", "interrupted", "drained"] def _parse_ns_segment(segment: str) -> tuple[str, str | None]: @@ -472,10 +472,8 @@ class _TasksLifecycleBase(StreamTransformer): self._open.clear() def fail(self, err: BaseException) -> None: - """Emit `failed` / `interrupted` for any tracked namespace still open.""" - is_interrupt = isinstance(err, GraphInterrupt) - status: SubgraphStatus = "interrupted" if is_interrupt else "failed" - error_str = None if is_interrupt else str(err) + """Emit terminal status for any tracked namespace still open.""" + status, error_str = _status_from_exception(err) for ns in list(self._open): self._on_terminal(ns, status, error_str) self._open.clear() @@ -483,6 +481,8 @@ class _TasksLifecycleBase(StreamTransformer): def _status_from_exception(err: BaseException) -> tuple[SubgraphStatus, str | None]: """Map a run exception to a subgraph terminal status and error string.""" + if isinstance(err, GraphDrained): + return "drained", None if isinstance(err, GraphInterrupt): return "interrupted", None return "failed", str(err) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d010eecbe..d4cb624e9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -122,6 +122,29 @@ def test_graph_validation() -> None: graph.invoke({"hello": "there"}) +def test_request_drain_allows_inflight_call_scheduling( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from langgraph.runtime import RunControl + + @task + def child(x: int) -> int: + return x + 1 + + control = RunControl() + + @entrypoint(checkpointer=sync_checkpointer) + def graph(x: int) -> int: + control.request_drain() + fut = child(x) + return fut.result() + + config = {"configurable": {"thread_id": "drain-call-sync"}} + + assert graph.invoke(1, config=config, control=control) == 2 + assert control.drain_requested + + def test_invalid_checkpointer_type() -> None: class State(TypedDict): foo: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 19f2ce88c..e61a40ecd 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -215,6 +215,30 @@ async def test_checkpoint_errors() -> None: pass +@NEEDS_CONTEXTVARS +async def test_request_drain_allows_inflight_acall_scheduling( + async_checkpointer: BaseCheckpointSaver, +) -> None: + from langgraph.runtime import RunControl + + @task + async def child(x: int) -> int: + return x + 1 + + control = RunControl() + + @entrypoint(checkpointer=async_checkpointer) + async def graph(x: int) -> int: + control.request_drain() + fut = child(x) + return await fut + + config = {"configurable": {"thread_id": "drain-call-async"}} + + assert await graph.ainvoke(1, config=config, control=control) == 2 + assert control.drain_requested + + async def test_py_async_with_cancel_behavior() -> None: """This test confirms that in all versions of Python we support, __aexit__ is not cancelled when the coroutine containing the async with block is cancelled.""" diff --git a/libs/langgraph/tests/test_runtime.py b/libs/langgraph/tests/test_runtime.py index a85abe45f..f1de3ba03 100644 --- a/libs/langgraph/tests/test_runtime.py +++ b/libs/langgraph/tests/test_runtime.py @@ -1,3 +1,6 @@ +import asyncio +import threading +import time from dataclasses import dataclass from typing import Any @@ -6,8 +9,15 @@ from langgraph.checkpoint.memory import MemorySaver from pydantic import BaseModel, ValidationError from typing_extensions import TypedDict +from langgraph.errors import GraphDrained from langgraph.graph import END, START, StateGraph -from langgraph.runtime import ExecutionInfo, Runtime, ServerInfo, get_runtime +from langgraph.runtime import ( + ExecutionInfo, + RunControl, + Runtime, + ServerInfo, + get_runtime, +) def test_injected_runtime() -> None: @@ -79,6 +89,183 @@ def test_merge_runtime() -> None: assert runtime1.merge(runtime3).context.api_key == "abc" # type: ignore +def test_merge_runtime_preserves_run_control() -> None: + control = RunControl() + runtime1 = Runtime(control=control) + runtime2 = Runtime(context=None) + + assert runtime1.merge(runtime2).control is control + + +def test_run_control_request_drain_stops_future_steps() -> None: + class State(TypedDict, total=False): + first: str + second: str + + control = RunControl() + + def first_node(state: State) -> dict[str, str]: + control.request_drain() + return {"first": "done"} + + def second_node(state: State) -> dict[str, str]: + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", first_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + with pytest.raises(GraphDrained, match="shutdown"): + graph.compile().invoke({}, control=control) + + +@pytest.mark.anyio +async def test_run_control_request_drain_stops_future_steps_async() -> None: + class State(TypedDict, total=False): + first: str + second: str + + control = RunControl() + + async def first_node(state: State) -> dict[str, str]: + control.request_drain() + return {"first": "done"} + + async def second_node(state: State) -> dict[str, str]: + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", first_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + with pytest.raises(GraphDrained, match="shutdown"): + await graph.compile().ainvoke({}, control=control) + + +def test_drain_requested_in_terminal_step_finishes_normally() -> None: + class State(TypedDict, total=False): + value: str + + control = RunControl() + + def node(state: State) -> dict[str, str]: + control.request_drain() + return {"value": "done"} + + graph = StateGraph(State) + graph.add_node("node", node) + graph.add_edge(START, "node") + graph.add_edge("node", END) + + assert graph.compile().invoke({}, control=control) == {"value": "done"} + assert control.drain_requested + + +def test_drain_with_exit_durability_persists_resume_checkpoint() -> None: + class State(TypedDict, total=False): + first: str + second: str + + control = RunControl() + + def first_node(state: State) -> dict[str, str]: + control.request_drain("sigterm") + return {"first": "done"} + + def second_node(state: State) -> dict[str, str]: + return {"second": "done"} + + graph = StateGraph(State) + graph.add_node("first", first_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + compiled = graph.compile(checkpointer=MemorySaver()) + config = {"configurable": {"thread_id": "drain-exit"}} + + with pytest.raises(GraphDrained, match="sigterm"): + compiled.invoke({}, config, durability="exit", control=control) + + assert compiled.invoke(None, config, durability="exit") == { + "first": "done", + "second": "done", + } + + +def test_drain_from_subgraph_can_resume_parent() -> None: + class State(TypedDict, total=False): + child_first: str + child_second: str + parent_second: str + + control = RunControl() + + def child_first(state: State) -> dict[str, str]: + control.request_drain("sigterm") + return {"child_first": "done"} + + def child_second(state: State) -> dict[str, str]: + return {"child_second": "done"} + + child_builder = StateGraph(State) + child_builder.add_node("child_first", child_first) + child_builder.add_node("child_second", child_second) + child_builder.add_edge(START, "child_first") + child_builder.add_edge("child_first", "child_second") + child_builder.add_edge("child_second", END) + child_graph = child_builder.compile(checkpointer=True) + + def parent_second(state: State) -> dict[str, str]: + return {"parent_second": "done"} + + parent_builder = StateGraph(State) + parent_builder.add_node("child", child_graph) + parent_builder.add_node("parent_second", parent_second) + parent_builder.add_edge(START, "child") + parent_builder.add_edge("child", "parent_second") + parent_builder.add_edge("parent_second", END) + + compiled = parent_builder.compile(checkpointer=MemorySaver()) + config = {"configurable": {"thread_id": "drain-subgraph"}} + + with pytest.raises(GraphDrained, match="sigterm"): + compiled.invoke({}, config, control=control) + + assert compiled.invoke(None, config) == { + "child_first": "done", + "child_second": "done", + "parent_second": "done", + } + + +@pytest.mark.anyio +async def test_drain_requested_in_terminal_step_finishes_normally_async() -> None: + class State(TypedDict, total=False): + value: str + + control = RunControl() + + async def node(state: State) -> dict[str, str]: + control.request_drain() + return {"value": "done"} + + graph = StateGraph(State) + graph.add_node("node", node) + graph.add_edge(START, "node") + graph.add_edge("node", END) + + assert await graph.compile().ainvoke({}, control=control) == {"value": "done"} + assert control.drain_requested + + def test_runtime_propogated_to_subgraph() -> None: @dataclass class Context: @@ -392,6 +579,334 @@ def test_context_coercion_pydantic_validation_errors() -> None: ) +def test_external_drain_concurrent_sync() -> None: + """External thread calls request_drain() while graph is mid-execution.""" + + class State(TypedDict, total=False): + first: str + second: str + + started = threading.Event() + + def first_node(state: State) -> dict[str, str]: + started.set() + time.sleep(0.05) + return {"first": "done"} + + def second_node(state: State) -> dict[str, str]: + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", first_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + control = RunControl() + compiled = graph.compile() + + exc_holder: list[BaseException | None] = [None] + + def run_graph() -> None: + try: + compiled.invoke({}, control=control) + except GraphDrained as e: + exc_holder[0] = e + + t = threading.Thread(target=run_graph) + t.start() + + started.wait(timeout=5) + control.request_drain("sigterm") + + t.join(timeout=10) + + exc = exc_holder[0] + assert isinstance(exc, GraphDrained) + assert exc.reason == "sigterm" + + +@pytest.mark.anyio +async def test_external_drain_concurrent_async() -> None: + """External task calls request_drain() while graph is mid-execution.""" + + class State(TypedDict, total=False): + first: str + second: str + + started = asyncio.Event() + + async def first_node(state: State) -> dict[str, str]: + started.set() + await asyncio.sleep(0.05) + return {"first": "done"} + + async def second_node(state: State) -> dict[str, str]: + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", first_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + control = RunControl() + compiled = graph.compile() + + async def drain_after_start() -> None: + await started.wait() + control.request_drain("sigterm") + + drain_task = asyncio.create_task(drain_after_start()) + + with pytest.raises(GraphDrained, match="sigterm"): + await compiled.ainvoke({}, control=control) + + await drain_task + + +@pytest.mark.anyio +async def test_drain_then_cancel_after_graceful_timeout() -> None: + """Simulate: drain requested -> node still running -> graceful timeout -> cancel. + + This shows what happens when a long-running node doesn't finish within + the graceful period after drain is requested. + """ + + class State(TypedDict, total=False): + first: str + second: str + + node_started = asyncio.Event() + node_cancelled = asyncio.Event() + node_finished = asyncio.Event() + + async def slow_node(state: State) -> dict[str, str]: + node_started.set() + try: + await asyncio.sleep(30) # very long operation + except asyncio.CancelledError: + node_cancelled.set() + raise + node_finished.set() + return {"first": "done"} + + async def second_node(state: State) -> dict[str, str]: + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", slow_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + control = RunControl() + compiled = graph.compile() + + # Phase 1: start graph + graph_task = asyncio.create_task(compiled.ainvoke({}, control=control)) + + # Phase 2: wait for node to start, then request drain + await node_started.wait() + control.request_drain("sigterm") + + # Phase 3: graceful timeout — node is still running, cancel after 1s + graceful_timeout = 1.0 + await asyncio.sleep(graceful_timeout) + + assert not node_finished.is_set(), "node should still be running" + assert not node_cancelled.is_set(), "node should not be cancelled yet" + + # Phase 4: force cancel + graph_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await graph_task + + # The node received CancelledError at the await point + assert node_cancelled.is_set(), "node should have received CancelledError" + assert not node_finished.is_set(), "node should NOT have finished normally" + + +@pytest.mark.anyio +async def test_cancel_ainvoke_with_async_node() -> None: + """Cancel ainvoke running an async node: CancelledError is delivered + at the await point and the node stops immediately.""" + + class State(TypedDict, total=False): + first: str + second: str + + timeline: list[str] = [] + node_started = asyncio.Event() + + async def slow_async_node(state: State) -> dict[str, str]: + timeline.append(f"async_node:start thread={threading.current_thread().name}") + node_started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + timeline.append("async_node:cancelled") + raise + timeline.append("async_node:finished") + return {"first": "done"} + + async def second_node(state: State) -> dict[str, str]: + timeline.append("second_node:run") + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", slow_async_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + compiled = graph.compile() + graph_task = asyncio.create_task(compiled.ainvoke({})) + + await node_started.wait() + timeline.append("test:cancel") + graph_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await graph_task + timeline.append("test:done") + + # async node runs on the event loop thread (MainThread) + assert any("MainThread" in e for e in timeline if "async_node:start" in e) + # CancelledError was delivered at the await point — node stopped + assert "async_node:cancelled" in timeline + # Node did NOT run to completion + assert "async_node:finished" not in timeline + # Second node never ran + assert "second_node:run" not in timeline + + +@pytest.mark.anyio +async def test_cancel_ainvoke_with_sync_node() -> None: + """Cancel ainvoke running a sync node. + + Sync nodes in ainvoke run on a separate thread (via run_in_executor), + NOT on the event loop thread. Cancelling the asyncio task disconnects + from the thread future, but the thread keeps running as an orphan and + completes on its own. + + Key difference from async nodes: + - async node: CancelledError stops the coroutine at an await point + - sync node: cancel only disconnects asyncio; the thread runs to completion + + In shutdown case, we will ignore this because the instance will be destroyed soon. + """ + + class State(TypedDict, total=False): + first: str + second: str + + timeline: list[str] = [] + node_started = threading.Event() + node_finished = threading.Event() + + def slow_sync_node(state: State) -> dict[str, str]: + timeline.append(f"sync_node:start thread={threading.current_thread().name}") + node_started.set() + time.sleep(1) + timeline.append("sync_node:after_sleep") + node_finished.set() + return {"first": "done"} + + def second_node(state: State) -> dict[str, str]: + timeline.append("second_node:run") + return {"second": "should-not-run"} + + graph = StateGraph(State) + graph.add_node("first", slow_sync_node) + graph.add_node("second", second_node) + graph.add_edge(START, "first") + graph.add_edge("first", "second") + graph.add_edge("second", END) + + control = RunControl() + compiled = graph.compile() + + timeline.append(f"test:main thread={threading.current_thread().name}") + graph_task = asyncio.create_task(compiled.ainvoke({}, control=control)) + + loop = asyncio.get_event_loop() + await loop.run_in_executor(None, node_started.wait, 5) + + timeline.append("test:cancel+drain") + graph_task.cancel() + control.request_drain("sigterm") + + with pytest.raises(asyncio.CancelledError): + await graph_task + timeline.append("test:exc=CancelledError") + + # Sync node runs on a background thread (asyncio_*), NOT MainThread + sync_start = next(e for e in timeline if "sync_node:start" in e) + assert "MainThread" not in sync_start, ( + "sync node should run on a background thread, not the event loop thread" + ) + + # At this point, the asyncio task is done but the thread is orphaned. + # The sync node has NOT finished yet — cancel only disconnected asyncio. + assert not node_finished.is_set(), ( + "sync node should still be running in its background thread" + ) + + # Wait for the orphaned thread to complete on its own. + await loop.run_in_executor(None, node_finished.wait, 5) + assert node_finished.is_set() + + # After the orphaned thread finishes, the full timeline looks like: + # test:main thread=MainThread + # sync_node:start thread=asyncio_N <- background thread + # test:cancel+drain <- cancel + drain fired + # test:exc=CancelledError <- asyncio disconnected + # sync_node:after_sleep <- thread ran to completion anyway + assert "sync_node:after_sleep" in timeline + # Second node never ran + assert "second_node:run" not in timeline + + # Verify timeline ordering: cancel happened before node finished + cancel_idx = timeline.index("test:cancel+drain") + sleep_idx = timeline.index("sync_node:after_sleep") + assert cancel_idx < sleep_idx, ( + "cancel was issued while the sync node was still sleeping" + ) + + +def test_drain_with_control_parameter_sync() -> None: + """Control parameter is wired through invoke -> stream.""" + + class State(TypedDict, total=False): + value: str + + ran = False + + def node(state: State) -> dict[str, str]: + nonlocal ran + ran = True + return {"value": "done"} + + graph = StateGraph(State) + graph.add_node("node", node) + graph.add_edge(START, "node") + graph.add_edge("node", END) + + # Pre-drained control stops before executing the first pending task. + control = RunControl() + control.request_drain("pre-drained") + + with pytest.raises(GraphDrained, match="pre-drained"): + graph.compile().invoke({}, control=control) + assert not ran + + # --- ExecutionInfo unit tests --- diff --git a/libs/langgraph/tests/test_stream_v2.py b/libs/langgraph/tests/test_stream_v2.py index 29d34e678..b56310af0 100644 --- a/libs/langgraph/tests/test_stream_v2.py +++ b/libs/langgraph/tests/test_stream_v2.py @@ -19,9 +19,11 @@ from typing_extensions import TypedDict, assert_type from langgraph._internal._constants import INTERRUPT from langgraph.constants import END, START +from langgraph.errors import GraphDrained from langgraph.func import entrypoint from langgraph.graph import StateGraph from langgraph.graph.message import MessagesState +from langgraph.runtime import RunControl from langgraph.types import ( CheckpointPayload, CheckpointStreamPart, @@ -229,6 +231,32 @@ class TestV2Stream: for c in chunks: _assert_stream_part_shape(c) + def test_stream_v2_accepts_control_for_drain(self) -> None: + class DrainState(TypedDict, total=False): + value: str + skipped: str + + control = RunControl() + + def first_node(state: DrainState) -> dict[str, str]: + control.request_drain("sigterm") + return {"value": "done"} + + def second_node(state: DrainState) -> dict[str, str]: + return {"skipped": "nope"} + + builder = StateGraph(DrainState) + builder.add_node("first", first_node) + builder.add_node("second", second_node) + builder.add_edge(START, "first") + builder.add_edge("first", "second") + builder.add_edge("second", END) + graph = builder.compile() + + run = graph.stream_v2({}, control=control) + with pytest.raises(GraphDrained, match="sigterm"): + list(run.values) + def test_subgraphs_ns(self) -> None: outer = _make_subgraph() chunks = list(