feat: allow graph to graceful shutdown/drain by request (#7274)

## 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 <long@langchain.dev>
Co-authored-by: Will Fu-Hinthorn <will@langchain.dev>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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

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)

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

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

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)

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

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

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

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(