feat(langgraph): add Runtime.on_exit for cleanup that runs however a run ends

Bumps langgraph from 1.2.18 to 1.2.19.

Includes:

- List a `ToolNode`'s tool names in its `get_graph()` node metadata
(#9449)
- Apply a node's `trace_policy` to tracers only, so it no longer changes
what streams see (#9486)
- Give `messages` and `tools` events the namespace of the graph they ran
in (#9458)
- Stream messages a node returns from a graph streamed inside another
node (#9463)
- Respect `subgraphs=False` in `RemoteGraph.stream` when the parent
streams subgraphs (#9453)
- Stop reporting a tool paused on `interrupt()` as failed on the `tools`
stream (#9406)
- Return an awaited `@task`'s value while streaming instead of `None`
(#9407)

libs/checkpoint-sqlite/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-checkpoint-sqlite"
-version = "3.1.2"
+version = "3.1.3"
 description = "Library with a SQLite implementation of LangGraph checkpoint saver."
 authors = []
 requires-python = ">=3.10"
@@ -12,7 +12,7 @@ readme = "README.md"
 license = "MIT"
 license-files = ['LICENSE']
 dependencies = [
-    "langgraph-checkpoint>=4.3.0,<5.0.0",
+    "langgraph-checkpoint>=4.3.1,<5.0.0",
     "aiosqlite>=0.20",
     "sqlite-vec>=0.1.6",
 ]

libs/checkpoint-sqlite/uv.lock

@@ -360,7 +360,7 @@ test = [
 
 [[package]]
 name = "langgraph-checkpoint-sqlite"
-version = "3.1.2"
+version = "3.1.3"
 source = { editable = "." }
 dependencies = [
     { name = "aiosqlite" },

libs/langgraph/langgraph/_internal/_runnable.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 import enum
 import inspect
 import logging
@@ -23,9 +24,16 @@ from typing import (
     Optional,
     Protocol,
     TypeGuard,
+    TypeVar,
     cast,
 )
 
+from langchain_core.callbacks.base import BaseCallbackHandler, BaseCallbackManager
+from langchain_core.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    BaseRunManager,
+    CallbackManagerForChainRun,
+)
 from langchain_core.runnables.base import (
     Runnable,
     RunnableConfig,
@@ -41,6 +49,7 @@ from langchain_core.runnables.config import (
     var_child_runnable_config,
 )
 from langchain_core.runnables.utils import Input, Output
+from langchain_core.tracers.base import BaseTracer
 from langchain_core.tracers.langchain import LangChainTracer
 from langgraph.store.base import BaseStore
 
@@ -57,7 +66,7 @@ from langgraph._internal._constants import (
 )
 from langgraph._internal._typing import MISSING
 from langgraph.errors import NodeError
-from langgraph.types import StreamWriter
+from langgraph.types import StreamWriter, TracePolicy
 
 try:
     from langchain_core.tracers._streaming import _StreamingCallbackHandler
@@ -85,6 +94,29 @@ def _trace_payload(value: Any, transform: Callable[[Any], Any] | None) -> Any:
         return value
 
 
+def _is_stream_handler(handler: BaseCallbackHandler) -> bool:
+    # Inheritance, not isinstance on the runtime-checkable protocol, which also
+    # matches any exporter that merely has the tap methods: stream handlers get a
+    # node's real payload and its streamed output.
+    return _StreamingCallbackHandler in type(handler).__mro__
+
+
+_M = TypeVar("_M", bound=BaseCallbackManager | BaseRunManager)
+
+
+def _split_handlers(manager: _M) -> tuple[_M, _M]:
+    # The recorded payload exists to keep tracer ingest small, so only tracers get it.
+    # `astream_log`'s handler is a tracer too, but it feeds the caller's own stream.
+    real, tracers = copy.copy(manager), copy.copy(manager)
+    tracers.handlers = [
+        h
+        for h in manager.handlers
+        if isinstance(h, BaseTracer) and not _is_stream_handler(h)
+    ]
+    real.handlers = [h for h in manager.handlers if h not in tracers.handlers]
+    return real, tracers
+
+
 def _set_config_context(
     config: RunnableConfig, run: Any = None
 ) -> Token[RunnableConfig | None]:
@@ -593,13 +625,15 @@ class RunnableSeq(Runnable):
         *steps: RunnableLike,
         name: str | None = None,
         trace_inputs: Callable[[Any], Any] | None = None,
-        trace_outputs: Callable[[Any], Any] | None = None,
+        trace_policy: TracePolicy | None = None,
     ) -> None:
         """Create a new RunnableSeq.
 
         Args:
             steps: The steps to include in the sequence.
             name: The name of the `Runnable`.
+            trace_inputs: Transforms the input every callback handler receives.
+            trace_policy: Processes what tracers record, after `trace_inputs`.
 
         Raises:
             ValueError: If the sequence has less than 2 steps.
@@ -619,7 +653,67 @@ class RunnableSeq(Runnable):
         self.steps = steps_flat
         self.name = name
         self.trace_inputs = trace_inputs
-        self.trace_outputs = trace_outputs
+        self.trace_policy = trace_policy or TracePolicy()
+
+    def _start(self, input: Any, config: RunnableConfig) -> CallbackManagerForChainRun:
+        callback_manager = get_callback_manager_for_config(config)
+        payload = _trace_payload(input, self.trace_inputs)
+        recorded = _trace_payload(payload, self.trace_policy.process_inputs)
+        name = config.get("run_name") or self.get_name()
+        run_id = config.pop("run_id", None)
+        if recorded is payload:
+            return callback_manager.on_chain_start(
+                None, payload, name=name, run_id=run_id
+            )
+        real, tracers = _split_handlers(callback_manager)
+        run_manager = real.on_chain_start(None, payload, name=name, run_id=run_id)
+        tracers.on_chain_start(None, recorded, name=name, run_id=run_manager.run_id)
+        run_manager.handlers = callback_manager.handlers
+        return run_manager
+
+    async def _astart(
+        self, input: Any, config: RunnableConfig
+    ) -> AsyncCallbackManagerForChainRun:
+        callback_manager = get_async_callback_manager_for_config(config)
+        payload = _trace_payload(input, self.trace_inputs)
+        recorded = _trace_payload(payload, self.trace_policy.process_inputs)
+        name = config.get("run_name") or self.get_name()
+        run_id = config.pop("run_id", None)
+        if recorded is payload:
+            return await callback_manager.on_chain_start(
+                None, payload, name=name, run_id=run_id
+            )
+        real, tracers = _split_handlers(callback_manager)
+        # Tracers start last, so a cancellation before then leaves no tracer run open.
+        run_manager = await real.on_chain_start(None, payload, name=name, run_id=run_id)
+        await tracers.on_chain_start(
+            None, recorded, name=name, run_id=run_manager.run_id
+        )
+        run_manager.handlers = callback_manager.handlers
+        return run_manager
+
+    def _end(self, run_manager: CallbackManagerForChainRun, output: Any) -> None:
+        recorded = _trace_payload(output, self.trace_policy.process_outputs)
+        if recorded is output:
+            run_manager.on_chain_end(output)
+        else:
+            real, tracers = _split_handlers(run_manager)
+            tracers.on_chain_end(recorded)
+            real.on_chain_end(output)
+
+    async def _aend(
+        self, run_manager: AsyncCallbackManagerForChainRun, output: Any
+    ) -> None:
+        recorded = _trace_payload(output, self.trace_policy.process_outputs)
+        if recorded is output:
+            await run_manager.on_chain_end(output)
+        else:
+            real, tracers = _split_handlers(run_manager)
+            # Tracers end first: on_chain_end is shielded once it starts, so a
+            # cancellation can only skip the other handlers' end, never leave a
+            # tracer run open.
+            await tracers.on_chain_end(recorded)
+            await real.on_chain_end(output)
 
     def __or__(
         self,
@@ -676,15 +770,8 @@ class RunnableSeq(Runnable):
     ) -> Any:
         if config is None:
             config = ensure_config()
-        # setup callbacks and context
-        callback_manager = get_callback_manager_for_config(config)
         # start the root run
-        run_manager = callback_manager.on_chain_start(
-            None,
-            _trace_payload(input, self.trace_inputs),
-            name=config.get("run_name") or self.get_name(),
-            run_id=config.pop("run_id", None),
-        )
+        run_manager = self._start(input, config)
         # invoke all steps in sequence
         try:
             for i, step in enumerate(self.steps):
@@ -712,7 +799,7 @@ class RunnableSeq(Runnable):
             run_manager.on_chain_error(e)
             raise
         else:
-            run_manager.on_chain_end(_trace_payload(input, self.trace_outputs))
+            self._end(run_manager, input)
             return input
 
     async def ainvoke(
@@ -723,15 +810,8 @@ class RunnableSeq(Runnable):
     ) -> Any:
         if config is None:
             config = ensure_config()
-        # setup callbacks
-        callback_manager = get_async_callback_manager_for_config(config)
         # start the root run
-        run_manager = await callback_manager.on_chain_start(
-            None,
-            _trace_payload(input, self.trace_inputs),
-            name=config.get("run_name") or self.get_name(),
-            run_id=config.pop("run_id", None),
-        )
+        run_manager = await self._astart(input, config)
 
         # invoke all steps in sequence
         try:
@@ -765,7 +845,7 @@ class RunnableSeq(Runnable):
             await run_manager.on_chain_error(e)
             raise
         else:
-            await run_manager.on_chain_end(_trace_payload(input, self.trace_outputs))
+            await self._aend(run_manager, input)
             return input
 
     def stream(
@@ -776,15 +856,8 @@ class RunnableSeq(Runnable):
     ) -> Iterator[Any]:
         if config is None:
             config = ensure_config()
-        # setup callbacks
-        callback_manager = get_callback_manager_for_config(config)
         # start the root run
-        run_manager = callback_manager.on_chain_start(
-            None,
-            _trace_payload(input, self.trace_inputs),
-            name=config.get("run_name") or self.get_name(),
-            run_id=config.pop("run_id", None),
-        )
+        run_manager = self._start(input, config)
         # get the run object
         for h in run_manager.handlers:
             if isinstance(h, LangChainTracer):
@@ -814,10 +887,9 @@ class RunnableSeq(Runnable):
                         )
                         iterator = step.transform(iterator, config)
                 # populates streamed_output in astream_log() output if needed
-                if _StreamingCallbackHandler is not None:
-                    for h in run_manager.handlers:
-                        if isinstance(h, _StreamingCallbackHandler):
-                            iterator = h.tap_output_iter(run_manager.run_id, iterator)
+                for h in run_manager.handlers:
+                    if _is_stream_handler(h):
+                        iterator = h.tap_output_iter(run_manager.run_id, iterator)
                 # consume into final output
                 output = context.run(_consume_iter, iterator)
                 # sequence doesn't emit output, yield to mark as generator
@@ -826,7 +898,7 @@ class RunnableSeq(Runnable):
                 run_manager.on_chain_error(e)
                 raise
             else:
-                run_manager.on_chain_end(_trace_payload(output, self.trace_outputs))
+                self._end(run_manager, output)
 
     async def astream(
         self,
@@ -836,15 +908,8 @@ class RunnableSeq(Runnable):
     ) -> AsyncIterator[Any]:
         if config is None:
             config = ensure_config()
-        # setup callbacks
-        callback_manager = get_async_callback_manager_for_config(config)
         # start the root run
-        run_manager = await callback_manager.on_chain_start(
-            None,
-            _trace_payload(input, self.trace_inputs),
-            name=config.get("run_name") or self.get_name(),
-            run_id=config.pop("run_id", None),
-        )
+        run_manager = await self._astart(input, config)
         # stream the last steps
         # transform the input stream of each step with the next
         # steps that don't natively support transforming an input stream will
@@ -880,12 +945,11 @@ class RunnableSeq(Runnable):
                             if hasattr(aiterator, "aclose"):
                                 stack.push_async_callback(aiterator.aclose)
                         # populates streamed_output in astream_log() output if needed
-                        if _StreamingCallbackHandler is not None:
-                            for h in run_manager.handlers:
-                                if isinstance(h, _StreamingCallbackHandler):
-                                    aiterator = h.tap_output_aiter(
-                                        run_manager.run_id, aiterator
-                                    )
+                        for h in run_manager.handlers:
+                            if _is_stream_handler(h):
+                                aiterator = h.tap_output_aiter(
+                                    run_manager.run_id, aiterator
+                                )
                         # consume into final output
                         output = await asyncio.create_task(
                             _consume_aiter(aiterator), context=context
@@ -896,9 +960,7 @@ class RunnableSeq(Runnable):
                     await run_manager.on_chain_error(e)
                     raise
                 else:
-                    await run_manager.on_chain_end(
-                        _trace_payload(output, self.trace_outputs)
-                    )
+                    await self._aend(run_manager, output)
         else:
             try:
                 async with AsyncExitStack() as stack:
@@ -914,12 +976,11 @@ class RunnableSeq(Runnable):
                         if hasattr(aiterator, "aclose"):
                             stack.push_async_callback(aiterator.aclose)
                     # populates streamed_output in astream_log() output if needed
-                    if _StreamingCallbackHandler is not None:
-                        for h in run_manager.handlers:
-                            if isinstance(h, _StreamingCallbackHandler):
-                                aiterator = h.tap_output_aiter(
-                                    run_manager.run_id, aiterator
-                                )
+                    for h in run_manager.handlers:
+                        if _is_stream_handler(h):
+                            aiterator = h.tap_output_aiter(
+                                run_manager.run_id, aiterator
+                            )
                     # consume into final output
                     output = await _consume_aiter(aiterator)
                     # sequence doesn't emit output, yield to mark as generator
@@ -928,9 +989,7 @@ class RunnableSeq(Runnable):
                 await run_manager.on_chain_error(e)
                 raise
             else:
-                await run_manager.on_chain_end(
-                    _trace_payload(output, self.trace_outputs)
-                )
+                await self._aend(run_manager, output)
 
 
 def _consume_iter(it: Iterator[Any]) -> Any:

libs/langgraph/langgraph/graph/message.py

@@ -22,7 +22,7 @@ from langchain_core.messages import (
 )
 from typing_extensions import TypedDict, deprecated
 
-from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
+from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
 from langgraph.graph.state import StateGraph
 from langgraph.pregel._messages import (
     _messages_delta_reducer as _messages_delta_reducer,
@@ -343,7 +343,7 @@ def push_message(
     )
 
     from langgraph.config import get_config
-    from langgraph.pregel._messages import StreamMessagesHandler
+    from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
 
     config = get_config()
     message = next(x for x in convert_to_messages([message]))
@@ -359,15 +359,12 @@ def push_message(
     ):
         handlers = config["callbacks"]
 
-    if stream_handler := next(
-        (x for x in handlers if isinstance(x, StreamMessagesHandler)), None
-    ):
-        metadata = config["metadata"]
-        message_meta = (
-            tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
-            metadata,
-        )
-        stream_handler._emit(message_meta, message, dedupe=False)
+    metadata = config["metadata"]
+    for handler in handlers:
+        if isinstance(handler, StreamMessagesHandler):
+            ns = graph_namespace(metadata)
+            if handler.subgraphs or not ns or ns == handler.parent_ns:
+                handler._emit((ns, metadata), message, dedupe=False)
 
     if state_key:
         config[CONF][CONFIG_KEY_SEND]([(state_key, message)])

libs/langgraph/langgraph/pregel/_loop.py

@@ -132,7 +132,7 @@ from langgraph.pregel.debug import (
     map_debug_tasks,
 )
 from langgraph.pregel.protocol import StreamChunk, StreamProtocol
-from langgraph.runtime import RunControl, Runtime
+from langgraph.runtime import RunControl, Runtime, _RunExits
 from langgraph.types import (
     All,
     CachePolicy,
@@ -1497,6 +1497,10 @@ class PregelLoop:
             if self._delta_write_futs is not None:
                 self._delta_write_futs.append(fut)
 
+    def _run_exits(self) -> _RunExits | None:
+        runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
+        return runtime._run_exits if isinstance(runtime, Runtime) else None
+
     def _suppress_interrupt(
         self,
         exc_type: type[BaseException] | None,
@@ -1874,6 +1878,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         self._exit_delta_writes = (
             [] if self.durability == "exit" and self.checkpointer is not None else None
         )
+        # Bottom of the stack, so run-exit callbacks run last: after the
+        # executor has stopped every task and the final checkpoint is saved.
+        if (exits := self._run_exits()) is not None and exits.bind(is_async=False):
+            self.stack.push(exits.exit)
         self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
         self.channels, self.managed = channels_from_checkpoint(
             self.specs,
@@ -2130,6 +2138,10 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         self._exit_delta_writes = (
             [] if self.durability == "exit" and self.checkpointer is not None else None
         )
+        # Bottom of the stack, so run-exit callbacks run last: after the
+        # executor has stopped every task and the final checkpoint is saved.
+        if (exits := self._run_exits()) is not None and exits.bind(is_async=True):
+            self.stack.push_async_exit(exits.aexit)
         self.submit = await self.stack.enter_async_context(
             AsyncBackgroundExecutor(self.config)
         )

libs/langgraph/langgraph/pregel/_messages.py

@@ -23,7 +23,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResul
 from pydantic import BaseModel
 
 from langgraph._internal._config import filter_to_user_tags
-from langgraph._internal._constants import NS_SEP
+from langgraph._internal._constants import NS_SEP, PUSH
 from langgraph.channels.binop import _get_overwrite
 from langgraph.channels.delta import DeltaChannel
 from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
@@ -55,6 +55,23 @@ def _state_values(obj: Any) -> Sequence[Any]:
     return ()
 
 
+def graph_namespace(metadata: Mapping[str, Any]) -> tuple[str, ...]:
+    """Namespace of the graph a task runs in, from the task's run metadata.
+
+    Drops the task's own segment, plus one per enclosing `@task` call: a call
+    runs in its caller's graph but nests its segment under the caller's. A
+    call's `langgraph_path` is `(PUSH, <caller's path>, <index>, True)`, while
+    a `Send` task's ends in `False`.
+    """
+    ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
+    path = metadata.get("langgraph_path", ())
+    drop = 1
+    while len(path) == 4 and path[0] == PUSH and path[-1] is True:
+        drop += 1
+        path = path[1]
+    return ns[:-drop]
+
+
 class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
     """A callback handler that implements stream_mode=messages.
 
@@ -148,9 +165,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
         **kwargs: Any,
     ) -> Any:
         if metadata and (not tags or (TAG_NOSTREAM not in tags)):
-            ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
-                :-1
-            ]
+            ns = graph_namespace(metadata)
             if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
                 return
             if (filtered_tags := filter_to_user_tags(tags)) is not None:
@@ -213,9 +228,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
             and kwargs.get("name") == metadata.get("langgraph_node")
             and (not tags or TAG_HIDDEN not in tags)
         ):
-            ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
-                :-1
-            ]
+            ns = graph_namespace(metadata)
             if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
                 return
             self.metadata[run_id] = (ns, metadata)

libs/langgraph/langgraph/pregel/_read.py

@@ -227,23 +227,14 @@ class PregelNode:
     def node(self) -> Runnable[Any, Any] | None:
         """Get a runnable that combines `bound` and `writers`."""
         writers = self.flat_writers
-        trace_inputs = self.trace_policy.process_inputs if self.trace_policy else None
-        trace_outputs = self.trace_policy.process_outputs if self.trace_policy else None
         if self.bound is DEFAULT_BOUND and not writers:
             return None
         elif self.bound is DEFAULT_BOUND and len(writers) == 1:
             return writers[0]
         elif self.bound is DEFAULT_BOUND:
-            return RunnableSeq(
-                *writers, trace_inputs=trace_inputs, trace_outputs=trace_outputs
-            )
+            return RunnableSeq(*writers, trace_policy=self.trace_policy)
         elif writers:
-            return RunnableSeq(
-                self.bound,
-                *writers,
-                trace_inputs=trace_inputs,
-                trace_outputs=trace_outputs,
-            )
+            return RunnableSeq(self.bound, *writers, trace_policy=self.trace_policy)
         else:
             return self.bound
 

libs/langgraph/langgraph/pregel/_tools.py

@@ -7,9 +7,9 @@ from uuid import UUID
 
 from langchain_core.callbacks import BaseCallbackHandler
 
-from langgraph._internal._constants import NS_SEP
 from langgraph.constants import TAG_NOSTREAM
 from langgraph.errors import GraphDrained, GraphInterrupt
+from langgraph.pregel._messages import graph_namespace
 from langgraph.pregel.protocol import StreamChunk
 
 try:
@@ -92,12 +92,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
     ) -> tuple[str, ...] | None:
         """Resolve the namespace this tool call should emit at, or `None` to skip.
 
-        Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
-        derivation: parses `langgraph_checkpoint_ns` (which ends with
-        the `node_name:task_id` of the calling node), drops that
-        trailing segment, and returns the containing subgraph's own
-        namespace. Returns `None` when the call should be silently
-        suppressed:
+        Mirrors `StreamMessagesHandler.on_chat_model_start`: the namespace
+        is that of the graph the tool runs in (`graph_namespace`). Returns
+        `None` when the call should be silently suppressed:
 
         - `metadata` is missing — handler is attached to a context
           without Pregel routing info.
@@ -110,11 +107,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
             return None
         if tags and TAG_NOSTREAM in tags:
             return None
-        nskey = metadata.get("langgraph_checkpoint_ns")
-        if not nskey:
-            ns: tuple[str, ...] = ()
-        else:
-            ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
+        ns = (
+            graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
+        )
         if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
             return None
         return ns

libs/langgraph/langgraph/pregel/main.py

@@ -97,6 +97,7 @@ from langgraph._internal._runnable import (
     Runnable,
     RunnableLike,
     RunnableSeq,
+    _is_stream_handler,
     coerce_to_runnable,
 )
 from langgraph._internal._timeout import coerce_timeout_policy
@@ -174,6 +175,7 @@ from langgraph.runtime import (
     RunControl,
     Runtime,
     ServerInfo,
+    _RunExits,
 )
 from langgraph.stream._mux import StreamMux
 from langgraph.stream._types import StreamTransformer
@@ -203,11 +205,6 @@ from langgraph.types import (
 from langgraph.typing import ContextT, InputT, OutputT, StateT
 from langgraph.warnings import LangGraphDeprecatedSinceV10
 
-try:
-    from langchain_core.tracers._streaming import _StreamingCallbackHandler
-except ImportError:
-    _StreamingCallbackHandler = None  # type: ignore
-
 __all__ = ("NodeBuilder", "Pregel")
 
 _WriteValue = Callable[[Input], Output] | Any
@@ -2880,6 +2877,20 @@ class Pregel(
                     for h in callback_manager.inheritable_handlers
                     if not isinstance(h, StreamMessagesHandlerV2)
                 ]
+            if CONFIG_KEY_TASK_ID not in config[CONF]:
+                # Not run as a task (a child invoked with its own thread, say),
+                # this graph's namespace starts at the root, so a caller's
+                # handler would take this graph's events for its own graph's.
+                callback_manager.handlers = [
+                    h
+                    for h in callback_manager.handlers
+                    if not _is_stream_mode_handler(h)
+                ]
+                callback_manager.inheritable_handlers = [
+                    h
+                    for h in callback_manager.inheritable_handlers
+                    if not _is_stream_mode_handler(h)
+                ]
             if "ls_integration" not in callback_manager.metadata:
                 callback_manager.add_metadata({"ls_integration": "langgraph"})
             run_manager = callback_manager.on_chain_start(
@@ -2898,9 +2909,9 @@ class Pregel(
                 )
             # set up subgraph checkpointing
             config = self._own_checkpoint_config(config)
+            stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
             # set up messages stream mode
             if "messages" in stream_modes:
-                ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
                 use_stream_messages_v2 = bool(
                     version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
                 )
@@ -2909,24 +2920,12 @@ class Pregel(
                     if use_stream_messages_v2
                     else StreamMessagesHandler
                 )
-                run_manager.inheritable_handlers.append(
-                    messages_handler_cls(
-                        stream.put,
-                        subgraphs,
-                        parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
-                    )
-                )
+                stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
 
             # set up tools stream mode
             if "tools" in stream_modes:
-                ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
-                run_manager.inheritable_handlers.append(
-                    StreamToolCallHandler(
-                        stream.put,
-                        subgraphs,
-                        parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
-                    )
-                )
+                stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
+            run_manager.inheritable_handlers.extend(stream_handlers)
 
             # set up custom stream mode
             if "custom" in stream_modes:
@@ -2968,6 +2967,8 @@ class Pregel(
                 execution_info=None,
                 server_info=server_info,
                 control=control or parent_runtime.control or RunControl(),
+                # a subgraph shares the root run's exit callbacks
+                _run_exits=parent_runtime._run_exits or _RunExits(),
             )
             runtime = parent_runtime.merge(runtime)
             config[CONF][CONFIG_KEY_RUNTIME] = runtime
@@ -3009,6 +3010,10 @@ class Pregel(
                 cache_policy=self.cache_policy,
                 has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
             ) as loop:
+                # Known only now: the loop numbers the namespace of every
+                # subgraph a task runs after its first.
+                for handler in stream_handlers:
+                    handler.parent_ns = loop.checkpoint_ns
                 emit_graph_lifecycle_events(loop)
                 # create runner
                 runner = PregelRunner(
@@ -3291,6 +3296,20 @@ class Pregel(
                     for h in callback_manager.inheritable_handlers
                     if not isinstance(h, StreamMessagesHandlerV2)
                 ]
+            if CONFIG_KEY_TASK_ID not in config[CONF]:
+                # Not run as a task (a child invoked with its own thread, say),
+                # this graph's namespace starts at the root, so a caller's
+                # handler would take this graph's events for its own graph's.
+                callback_manager.handlers = [
+                    h
+                    for h in callback_manager.handlers
+                    if not _is_stream_mode_handler(h)
+                ]
+                callback_manager.inheritable_handlers = [
+                    h
+                    for h in callback_manager.inheritable_handlers
+                    if not _is_stream_mode_handler(h)
+                ]
             if "ls_integration" not in callback_manager.metadata:
                 callback_manager.add_metadata({"ls_integration": "langgraph"})
             run_manager = await callback_manager.on_chain_start(
@@ -3304,20 +3323,10 @@ class Pregel(
                 run_id=run_manager.run_id,
             )
             # if running from astream_log() run each proc with streaming
-            do_stream = (
-                next(
-                    (
-                        True
-                        for h in run_manager.handlers
-                        if isinstance(h, _StreamingCallbackHandler)
-                        and not isinstance(
-                            h, (StreamMessagesHandler, StreamToolCallHandler)
-                        )
-                    ),
-                    False,
-                )
-                if _StreamingCallbackHandler is not None
-                else False
+            do_stream = any(
+                _is_stream_handler(h)
+                and not isinstance(h, (StreamMessagesHandler, StreamToolCallHandler))
+                for h in run_manager.handlers
             )
             if checkpointer is None and durability is not None:
                 warnings.warn(
@@ -3325,10 +3334,9 @@ class Pregel(
                 )
             # set up subgraph checkpointing
             config = self._own_checkpoint_config(config)
+            stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
             # set up messages stream mode
             if "messages" in stream_modes:
-                # namespace can be None in a root level graph?
-                ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
                 use_stream_messages_v2 = bool(
                     version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
                 )
@@ -3337,24 +3345,12 @@ class Pregel(
                     if use_stream_messages_v2
                     else StreamMessagesHandler
                 )
-                run_manager.inheritable_handlers.append(
-                    messages_handler_cls(
-                        stream_put,
-                        subgraphs,
-                        parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
-                    )
-                )
+                stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
 
             # set up tools stream mode
             if "tools" in stream_modes:
-                ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
-                run_manager.inheritable_handlers.append(
-                    StreamToolCallHandler(
-                        stream_put,
-                        subgraphs,
-                        parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
-                    )
-                )
+                stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
+            run_manager.inheritable_handlers.extend(stream_handlers)
 
             # set up custom stream mode
             def stream_writer(c: Any) -> None:
@@ -3411,6 +3407,8 @@ class Pregel(
                 execution_info=None,
                 server_info=server_info,
                 control=control or parent_runtime.control or RunControl(),
+                # a subgraph shares the root run's exit callbacks
+                _run_exits=parent_runtime._run_exits or _RunExits(),
             )
             runtime = parent_runtime.merge(runtime)
             config[CONF][CONFIG_KEY_RUNTIME] = runtime
@@ -3463,6 +3461,10 @@ class Pregel(
                 cache_policy=self.cache_policy,
                 has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
             ) as loop:
+                # Known only now: the loop numbers the namespace of every
+                # subgraph a task runs after its first.
+                for handler in stream_handlers:
+                    handler.parent_ns = loop.checkpoint_ns
                 await aemit_graph_lifecycle_events(loop)
                 # create runner
                 runner = PregelRunner(
@@ -4460,6 +4462,10 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
         payload["values"] = mapper(payload["values"])
 
 
+def _is_stream_mode_handler(handler: object) -> bool:
+    return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
+
+
 def _resolve_parent_ns(
     graph_config: RunnableConfig | None, call_config: RunnableConfig | None
 ) -> tuple[str, ...]:

libs/langgraph/langgraph/runtime.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Callable
+import inspect
+import logging
+import threading
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field, replace
+from types import TracebackType
 from typing import Any, Generic, cast
 
 from langgraph.store.base import BaseStore
@@ -104,6 +108,129 @@ class RunControl:
         return self._drain_reason
 
 
+logger = logging.getLogger(__name__)
+
+ExitCallback = Callable[[], Awaitable[None] | None]
+"""A cleanup callback registered with `Runtime.on_exit`."""
+
+
+class _RunExits:
+    """Cleanup callbacks for one root graph invocation.
+
+    Created by the root `stream`/`astream` call and shared by every subgraph
+    it runs, through `Runtime`. The root Pregel loop binds it at the bottom of
+    its exit stack, so callbacks run after the loop has stopped every task,
+    persisted its final checkpoint and handled any interrupt, whatever the
+    reason the run ended.
+    """
+
+    __slots__ = ("_callbacks", "_is_async", "_lock", "_state")
+
+    def __init__(self) -> None:
+        self._callbacks: list[ExitCallback] = []
+        self._lock = threading.Lock()
+        self._is_async = False
+        # "new" -> "bound" (owned by a loop) -> "closed" (callbacks ran)
+        self._state = "new"
+
+    def bind(self, *, is_async: bool) -> bool:
+        """Claim these exits for a loop. Only the first (root) loop wins."""
+        with self._lock:
+            if self._state != "new":
+                return False
+            self._state = "bound"
+            self._is_async = is_async
+            return True
+
+    def register(self, callback: ExitCallback) -> None:
+        if not callable(callback):
+            msg = "on_exit() expects a callable"
+            raise TypeError(msg)
+        with self._lock:
+            if self._state == "closed":
+                msg = "on_exit() called after the graph run has already exited"
+                raise RuntimeError(msg)
+            if not self._is_async and _is_async_callable(callback):
+                msg = (
+                    "on_exit() got an async callback in a synchronous graph run "
+                    "(invoke/stream). Register a sync callback, or run the graph "
+                    "with ainvoke/astream."
+                )
+                raise TypeError(msg)
+            self._callbacks.append(callback)
+
+    def _drain(self) -> list[ExitCallback]:
+        with self._lock:
+            self._state = "closed"
+            callbacks = self._callbacks[::-1]  # LIFO
+            self._callbacks = []
+        return callbacks
+
+    def exit(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        errors: list[Exception] = []
+        for callback in self._drain():
+            try:
+                result = callback()
+                if inspect.isawaitable(result):
+                    if inspect.iscoroutine(result):
+                        result.close()
+                    msg = (
+                        f"on_exit() callback {callback!r} returned an awaitable "
+                        "in a synchronous graph run; it was not awaited."
+                    )
+                    raise TypeError(msg)
+            except Exception as e:
+                errors.append(e)
+        return self._finish(errors, exc_value)
+
+    async def aexit(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        errors: list[Exception] = []
+        for callback in self._drain():
+            try:
+                result = callback()
+                if inspect.isawaitable(result):
+                    await result
+            except Exception as e:
+                errors.append(e)
+        return self._finish(errors, exc_value)
+
+    @staticmethod
+    def _finish(errors: list[Exception], exc_value: BaseException | None) -> bool:
+        # Every callback runs even if an earlier one failed. A failing callback
+        # never replaces the run's own exception (error, cancellation, stream
+        # close): it is logged instead. If the run otherwise ended cleanly, the
+        # first callback error is raised so cleanup bugs aren't silent.
+        if not errors:
+            return False
+        if exc_value is None:
+            for e in errors[1:]:
+                logger.error("on_exit() callback failed", exc_info=e)
+            raise errors[0]
+        for e in errors:
+            logger.error(
+                "on_exit() callback failed while the graph run was exiting with %r",
+                exc_value,
+                exc_info=e,
+            )
+        return False
+
+
+def _is_async_callable(callback: Any) -> bool:
+    return inspect.iscoroutinefunction(callback) or inspect.iscoroutinefunction(
+        getattr(callback, "__call__", None)
+    )
+
+
 def _no_op_stream_writer(_: Any) -> None: ...
 
 
@@ -237,6 +364,40 @@ class Runtime(Generic[ContextT]):
     graph runtime.
     """
 
+    _run_exits: _RunExits | None = field(default=None, repr=False, compare=False)
+    """Exit callbacks of the root graph run. Use `on_exit` to register one."""
+
+    def on_exit(self, callback: ExitCallback) -> None:
+        """Register cleanup to run when the current graph run ends.
+
+        The callback runs once the root invocation (the `invoke`/`stream`/
+        `ainvoke`/`astream` call the caller made) exits for any reason: it
+        completed, raised, was cancelled, paused on an interrupt, or the caller
+        closed the stream or broke out of it. Calls from subgraphs register on
+        that same root run.
+
+        Callbacks take no arguments and run in reverse registration order,
+        after every task of the run has stopped and the final checkpoint has
+        been written. In `ainvoke`/`astream` a callback may be sync or async; in
+        `invoke`/`stream` it must be sync. A failing callback never hides the
+        run's own exception: it is logged, and raised only if the run otherwise
+        ended without one.
+
+        A process that is killed does not run its callbacks, so use them to
+        free resources early, not as the only guarantee they are freed.
+
+        Args:
+            callback: A zero-argument sync or async callable.
+
+        Raises:
+            RuntimeError: When called outside a graph run, or after it exited.
+            TypeError: When given an async callback in a synchronous run.
+        """
+        if self._run_exits is None:
+            msg = "on_exit() can only be called during a graph run"
+            raise RuntimeError(msg)
+        self._run_exits.register(callback)
+
     def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
         """Merge two runtimes together.
 
@@ -255,6 +416,7 @@ class Runtime(Generic[ContextT]):
             execution_info=other.execution_info or self.execution_info,
             server_info=other.server_info or self.server_info,
             control=other.control or self.control,
+            _run_exits=other._run_exits or self._run_exits,
         )
 
     def override(

libs/langgraph/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph"
-version = "1.2.18"
+version = "1.2.19"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 requires-python = ">=3.10"

libs/langgraph/tests/test_run_exit.py

@@ -0,0 +1,339 @@
+"""Tests for `Runtime.on_exit`: cleanup that runs however a graph run ends."""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+
+import pytest
+from langgraph.checkpoint.memory import InMemorySaver
+from typing_extensions import TypedDict
+
+from langgraph.graph import END, START, StateGraph
+from langgraph.runtime import Runtime, get_runtime
+from langgraph.types import Command, interrupt
+
+pytestmark = pytest.mark.anyio
+
+NEEDS_CONTEXTVARS = pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+
+
+class State(TypedDict, total=False):
+    value: str
+
+
+def _two_node_graph(first, second=None, *, checkpointer=None):
+    builder = StateGraph(State)
+    builder.add_node("first", first)
+    builder.add_edge(START, "first")
+    if second is None:
+        builder.add_edge("first", END)
+    else:
+        builder.add_node("second", second)
+        builder.add_edge("first", "second")
+        builder.add_edge("second", END)
+    return builder.compile(checkpointer=checkpointer)
+
+
+# sync
+
+
+def test_runs_on_success_in_lifo_order() -> None:
+    events: list[str] = []
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup-a"))
+        runtime.on_exit(lambda: events.append("cleanup-b"))
+        events.append("first")
+        return {"value": "x"}
+
+    def second(state: State) -> State:
+        events.append("second")
+        return {}
+
+    graph = _two_node_graph(first, second)
+    assert graph.invoke({}) == {"value": "x"}
+    assert events == ["first", "second", "cleanup-b", "cleanup-a"]
+
+
+def test_runs_on_error_and_keeps_the_run_exception() -> None:
+    events: list[str] = []
+
+    def failing_cleanup() -> None:
+        events.append("failing")
+        raise KeyError("cleanup failed")
+
+    def first(state: State) -> State:
+        runtime = get_runtime()
+        runtime.on_exit(lambda: events.append("cleanup"))
+        runtime.on_exit(failing_cleanup)
+        raise ValueError("node failed")
+
+    graph = _two_node_graph(first)
+    with pytest.raises(ValueError, match="node failed"):
+        graph.invoke({})
+    # the failing callback did not stop the earlier one from running
+    assert events == ["failing", "cleanup"]
+
+
+def test_callback_error_is_raised_when_the_run_succeeded() -> None:
+    def failing_cleanup() -> None:
+        raise KeyError("cleanup failed")
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(failing_cleanup)
+        return {}
+
+    graph = _two_node_graph(first)
+    with pytest.raises(KeyError, match="cleanup failed"):
+        graph.invoke({})
+
+
+def test_runs_on_interrupt() -> None:
+    events: list[str] = []
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup"))
+        return {"value": interrupt("approve?")}
+
+    graph = _two_node_graph(first, checkpointer=InMemorySaver())
+    config = {"configurable": {"thread_id": "1"}}
+    result = graph.invoke({}, config)
+    assert "__interrupt__" in result
+    assert events == ["cleanup"]
+
+    # each resumed run is its own invocation with its own callbacks
+    graph.invoke(Command(resume="ok"), config)
+    assert events == ["cleanup", "cleanup"]
+
+
+def test_runs_when_the_stream_is_closed_early() -> None:
+    events: list[str] = []
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup"))
+        return {"value": "x"}
+
+    def second(state: State) -> State:
+        events.append("second")
+        return {}
+
+    graph = _two_node_graph(first, second)
+    stream = graph.stream({}, stream_mode="updates")
+    assert next(stream) == {"first": {"value": "x"}}
+    assert events == []
+    stream.close()
+    assert events == ["cleanup"]
+
+
+def test_runs_when_breaking_out_of_stream() -> None:
+    events: list[str] = []
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup"))
+        return {"value": "x"}
+
+    graph = _two_node_graph(first, lambda s: {})
+    for _ in graph.stream({}, stream_mode="updates"):
+        break
+    # the dropped generator is closed by refcounting on CPython
+    assert events == ["cleanup"]
+
+
+def test_subgraph_registers_on_the_root_run() -> None:
+    events: list[str] = []
+
+    def inner(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup"))
+        events.append("inner")
+        return {}
+
+    sub = StateGraph(State)
+    sub.add_node("inner", inner)
+    sub.add_edge(START, "inner")
+    subgraph = sub.compile()
+
+    def call_subgraph(state: State) -> State:
+        subgraph.invoke(state)
+        events.append("subgraph returned")
+        return {}
+
+    def after(state: State) -> State:
+        events.append("after")
+        return {}
+
+    builder = StateGraph(State)
+    builder.add_node("as_node", subgraph)
+    builder.add_node("call_subgraph", call_subgraph)
+    builder.add_node("after", after)
+    builder.add_edge(START, "as_node")
+    builder.add_edge("as_node", "call_subgraph")
+    builder.add_edge("call_subgraph", "after")
+    graph = builder.compile()
+
+    graph.invoke({})
+    # nothing runs when either subgraph ends; both run when the root ends
+    assert events == [
+        "inner",
+        "inner",
+        "subgraph returned",
+        "after",
+        "cleanup",
+        "cleanup",
+    ]
+
+
+def test_async_callback_rejected_in_sync_run() -> None:
+    async def acleanup() -> None:
+        pass
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(acleanup)
+        return {}
+
+    graph = _two_node_graph(first)
+    with pytest.raises(TypeError, match="async callback in a synchronous"):
+        graph.invoke({})
+
+
+def test_outside_a_run_raises() -> None:
+    with pytest.raises(RuntimeError, match="during a graph run"):
+        Runtime().on_exit(lambda: None)
+
+
+def test_each_invocation_gets_its_own_callbacks() -> None:
+    calls: list[int] = []
+
+    def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: calls.append(1))
+        return {}
+
+    graph = _two_node_graph(first)
+    graph.invoke({})
+    graph.invoke({})
+    assert calls == [1, 1]
+
+
+# async
+
+
+async def test_async_runs_on_success_sync_and_async_callbacks() -> None:
+    events: list[str] = []
+
+    async def acleanup() -> None:
+        await asyncio.sleep(0)
+        events.append("async-cleanup")
+
+    async def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("sync-cleanup"))
+        runtime.on_exit(acleanup)
+        events.append("first")
+        return {"value": "x"}
+
+    graph = _two_node_graph(first)
+    assert await graph.ainvoke({}) == {"value": "x"}
+    assert events == ["first", "async-cleanup", "sync-cleanup"]
+
+
+async def test_async_runs_on_error_and_keeps_the_run_exception() -> None:
+    events: list[str] = []
+
+    async def failing_cleanup() -> None:
+        raise KeyError("cleanup failed")
+
+    async def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup"))
+        runtime.on_exit(failing_cleanup)
+        raise ValueError("node failed")
+
+    graph = _two_node_graph(first)
+    with pytest.raises(ValueError, match="node failed"):
+        await graph.ainvoke({})
+    assert events == ["cleanup"]
+
+
+async def test_async_runs_on_interrupt() -> None:
+    events: list[str] = []
+
+    async def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(lambda: events.append("cleanup"))
+        return {"value": interrupt("approve?")}
+
+    graph = _two_node_graph(first, checkpointer=InMemorySaver())
+    result = await graph.ainvoke({}, {"configurable": {"thread_id": "1"}})
+    assert "__interrupt__" in result
+    assert events == ["cleanup"]
+
+
+async def test_async_runs_on_cancel_after_the_node_unwinds() -> None:
+    events: list[str] = []
+    started = asyncio.Event()
+
+    async def acleanup() -> None:
+        events.append("cleanup")
+
+    async def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(acleanup)
+        started.set()
+        try:
+            await asyncio.sleep(60)
+        finally:
+            events.append("node finally")
+        return {}
+
+    graph = _two_node_graph(first)
+    task = asyncio.create_task(graph.ainvoke({}))
+    await started.wait()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    # cleanup ran after the node had stopped
+    assert events == ["node finally", "cleanup"]
+
+
+async def test_async_runs_when_the_stream_is_closed_early() -> None:
+    events: list[str] = []
+
+    async def acleanup() -> None:
+        events.append("cleanup")
+
+    async def first(state: State, runtime: Runtime) -> State:
+        runtime.on_exit(acleanup)
+        return {"value": "x"}
+
+    async def second(state: State) -> State:
+        events.append("second")
+        return {}
+
+    graph = _two_node_graph(first, second)
+    stream = graph.astream({}, stream_mode="updates")
+    assert await stream.__anext__() == {"first": {"value": "x"}}
+    assert events == []
+    await stream.aclose()
+    assert events == ["cleanup"]
+
+
+@NEEDS_CONTEXTVARS
+async def test_async_subgraph_registers_on_the_root_run() -> None:
+    events: list[str] = []
+
+    async def inner(state: State) -> State:
+        get_runtime().on_exit(lambda: events.append("cleanup"))
+        return {}
+
+    sub = StateGraph(State)
+    sub.add_node("inner", inner)
+    sub.add_edge(START, "inner")
+    subgraph = sub.compile()
+
+    async def call_subgraph(state: State) -> State:
+        await subgraph.ainvoke(state)
+        events.append("subgraph returned")
+        return {}
+
+    graph = _two_node_graph(call_subgraph, lambda s: events.append("after") or {})
+    await graph.ainvoke({})
+    assert events == ["subgraph returned", "after", "cleanup"]

libs/langgraph/tests/test_stream_namespaces.py

@@ -0,0 +1,308 @@
+import sys
+from typing import Any
+
+import pytest
+from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+from langchain_core.messages import AIMessage
+from langchain_core.tools import tool
+
+from langgraph.func import entrypoint, task
+from langgraph.graph import START, MessagesState, StateGraph
+from langgraph.graph.message import push_message
+from langgraph.graph.state import CompiledStateGraph
+
+pytestmark = pytest.mark.anyio
+
+NEEDS_CONTEXTVARS = pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+
+
+@tool
+def echo(text: str) -> str:
+    """Echo the text."""
+    return text
+
+
+def _one_node_graph(node: Any) -> CompiledStateGraph:
+    return (
+        StateGraph(MessagesState)
+        .add_node(node)
+        .add_edge(START, node.__name__)
+        .compile()
+    )
+
+
+def _chat_and_echo(state: MessagesState) -> dict:
+    echo.invoke({"text": "done"})
+    return {"messages": [GenericFakeChatModel(messages=iter(["one two"])).invoke("hi")]}
+
+
+async def _achat_and_echo(state: MessagesState) -> dict:
+    await echo.ainvoke({"text": "done"})
+    model = GenericFakeChatModel(messages=iter(["one two"]))
+    return {"messages": [await model.ainvoke("hi")]}
+
+
+def _segment_names(ns: tuple[str, ...]) -> tuple[str, ...]:
+    return tuple(segment.split(":")[0] for segment in ns)
+
+
+def _summary(part: dict) -> tuple[tuple[str, ...], str]:
+    if part["type"] == "messages":
+        return _segment_names(part["ns"]), part["data"][0].content
+    return _segment_names(part["ns"]), part["data"]["event"]
+
+
+_CHAT_AND_ECHO_EVENTS = [
+    ((), "tool-started"),
+    ((), "tool-finished"),
+    ((), "one"),
+    ((), " "),
+    ((), "two"),
+]
+
+_CHILD_CONFIGS = pytest.mark.parametrize(
+    "child_config",
+    [
+        None,
+        {"configurable": {"thread_id": "child"}},
+        {"configurable": {"thread_id": "parent"}},
+    ],
+    ids=["no_config", "own_thread", "caller_thread"],
+)
+
+
+_TASK_CALLERS = pytest.mark.parametrize(
+    "nested", [False, True], ids=["from_entrypoint", "from_task_in_task"]
+)
+
+
+@_TASK_CALLERS
+@pytest.mark.parametrize("subgraphs", [False, True])
+def test_task_events_stream_under_the_graph_namespace(
+    nested: bool, subgraphs: bool
+) -> None:
+    @task
+    def chat_and_echo() -> None:
+        _chat_and_echo({"messages": []})
+
+    @task
+    def outer() -> None:
+        chat_and_echo().result()
+
+    @entrypoint()
+    def workflow(inputs: dict) -> None:
+        (outer() if nested else chat_and_echo()).result()
+
+    parts = workflow.stream(
+        {}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
+    )
+
+    assert [_summary(part) for part in parts] == _CHAT_AND_ECHO_EVENTS
+
+
+@NEEDS_CONTEXTVARS
+@_TASK_CALLERS
+@pytest.mark.parametrize("subgraphs", [False, True])
+async def test_task_events_stream_under_the_graph_namespace_async(
+    nested: bool, subgraphs: bool
+) -> None:
+    @task
+    async def chat_and_echo() -> None:
+        await _achat_and_echo({"messages": []})
+
+    @task
+    async def outer() -> None:
+        await chat_and_echo()
+
+    @entrypoint()
+    async def workflow(inputs: dict) -> None:
+        await (outer() if nested else chat_and_echo())
+
+    parts = workflow.astream(
+        {}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
+    )
+
+    assert [_summary(part) async for part in parts] == _CHAT_AND_ECHO_EVENTS
+
+
+def test_task_events_stream_under_the_namespace_of_their_subgraph() -> None:
+    @task
+    def chat_and_echo() -> None:
+        _chat_and_echo({"messages": []})
+
+    def node(state: MessagesState) -> dict:
+        chat_and_echo().result()
+        return {}
+
+    sub = _one_node_graph(node)
+    parent = StateGraph(MessagesState).add_node("sub", sub).add_edge(START, "sub")
+
+    parts = parent.compile().stream(
+        {"messages": []},
+        stream_mode=["messages", "tools"],
+        subgraphs=True,
+        version="v2",
+    )
+
+    assert [_summary(part) for part in parts] == [
+        (("sub",), event) for _, event in _CHAT_AND_ECHO_EVENTS
+    ]
+
+
+@_CHILD_CONFIGS
+def test_child_graph_streams_only_to_the_node_streaming_it(
+    child_config: dict | None,
+) -> None:
+    child = _one_node_graph(_chat_and_echo)
+    received = []
+
+    def node(state: MessagesState) -> dict:
+        received.extend(
+            child.stream(
+                {"messages": []}, child_config, stream_mode=["messages", "tools"]
+            )
+        )
+        return {}
+
+    parts = _one_node_graph(node).stream(
+        {"messages": []},
+        {"configurable": {"thread_id": "parent"}},
+        stream_mode=["messages", "tools"],
+    )
+
+    assert list(parts) == []
+    assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
+
+
+@NEEDS_CONTEXTVARS
+@_CHILD_CONFIGS
+async def test_child_graph_streams_only_to_the_node_streaming_it_async(
+    child_config: dict | None,
+) -> None:
+    child = _one_node_graph(_achat_and_echo)
+    received = []
+
+    async def node(state: MessagesState) -> dict:
+        stream = child.astream(
+            {"messages": []}, child_config, stream_mode=["messages", "tools"]
+        )
+        received.extend([part async for part in stream])
+        return {}
+
+    parts = _one_node_graph(node).astream(
+        {"messages": []},
+        {"configurable": {"thread_id": "parent"}},
+        stream_mode=["messages", "tools"],
+    )
+
+    assert [part async for part in parts] == []
+    assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
+
+
+@pytest.mark.parametrize(
+    "child_config",
+    [
+        {"configurable": {"thread_id": "child"}},
+        {"configurable": {"thread_id": "parent"}},
+    ],
+    ids=["own_thread", "caller_thread"],
+)
+def test_child_graph_on_its_own_thread_skips_a_caller_streaming_subgraphs(
+    child_config: dict,
+) -> None:
+    child = _one_node_graph(_chat_and_echo)
+
+    def node(state: MessagesState) -> dict:
+        child.invoke({"messages": []}, child_config)
+        return {}
+
+    parts = _one_node_graph(node).stream(
+        {"messages": []},
+        {"configurable": {"thread_id": "parent"}},
+        stream_mode=["messages", "tools"],
+        subgraphs=True,
+    )
+
+    assert list(parts) == []
+
+
+def test_node_gets_the_events_of_every_subgraph_it_streams() -> None:
+    sub = _one_node_graph(_chat_and_echo)
+    counts = []
+
+    def node(state: MessagesState) -> dict:
+        for _ in range(2):
+            stream = sub.stream({"messages": []}, stream_mode=["messages", "tools"])
+            counts.append(len(list(stream)))
+        return {}
+
+    _one_node_graph(node).invoke({"messages": []})
+
+    assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
+
+
+@NEEDS_CONTEXTVARS
+async def test_node_gets_the_events_of_every_subgraph_it_streams_async() -> None:
+    sub = _one_node_graph(_achat_and_echo)
+    counts = []
+
+    async def node(state: MessagesState) -> dict:
+        for _ in range(2):
+            stream = sub.astream({"messages": []}, stream_mode=["messages", "tools"])
+            counts.append(len([part async for part in stream]))
+        return {}
+
+    await _one_node_graph(node).ainvoke({"messages": []})
+
+    assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
+
+
+def _push(state: MessagesState) -> dict:
+    push_message(AIMessage(id="pushed", content="hi"), state_key=None)
+    return {}
+
+
+def _namespaces(graph: CompiledStateGraph, subgraphs: bool) -> list[tuple[str, ...]]:
+    parts = graph.stream(
+        {"messages": []}, stream_mode="messages", subgraphs=subgraphs, version="v2"
+    )
+    return [_segment_names(part["ns"]) for part in parts]
+
+
+def _streams_pushing_graph(received: list) -> Any:
+    pushing = _one_node_graph(_push)
+
+    def node(state: MessagesState) -> dict:
+        for message, _ in pushing.stream({"messages": []}, stream_mode="messages"):
+            received.append(message.id)
+        return {}
+
+    return node
+
+
+def test_pushed_message_streams_under_the_namespace_of_its_graph() -> None:
+    pushing = _one_node_graph(_push)
+    parent = StateGraph(MessagesState).add_node("sub", pushing).add_edge(START, "sub")
+
+    assert _namespaces(pushing, subgraphs=True) == [()]
+    assert _namespaces(parent.compile(), subgraphs=False) == []
+    assert _namespaces(parent.compile(), subgraphs=True) == [("sub",)]
+
+
+def test_pushed_message_streams_to_the_node_that_streams_its_graph() -> None:
+    received: list[str] = []
+    graph = _one_node_graph(_streams_pushing_graph(received))
+
+    assert _namespaces(graph, subgraphs=False) == []
+    assert received == ["pushed"]
+
+
+def test_pushed_message_streams_to_every_stream_that_takes_it() -> None:
+    received: list[str] = []
+    graph = _one_node_graph(_streams_pushing_graph(received))
+
+    assert _namespaces(graph, subgraphs=True) == [("node",)]
+    assert received == ["pushed"]

libs/langgraph/tests/test_trace_policy.py

@@ -1,11 +1,24 @@
 """End-to-end tests for `TracePolicy` input processing on node trace runs."""
 
+import asyncio
+from collections.abc import AsyncIterator, Iterator
 from typing import Any
+from uuid import UUID
 
 import pytest
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.runnables import RunnableConfig
 from typing_extensions import TypedDict
 
-from langgraph.graph import END, START, StateGraph
+from langgraph._internal._runnable import (
+    ASYNCIO_ACCEPTS_CONTEXT,
+    RunnableCallable,
+    RunnableSeq,
+)
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.graph.state import CompiledStateGraph
 from langgraph.types import TracePolicy
 from tests.fake_tracer import FakeTracer, Run
 
@@ -147,3 +160,179 @@ async def test_trace_policy_transforms_recorded_inputs_async() -> None:
     run = _node_run(tracer, "n")
     assert run.inputs == {"scrubbed_in": True}
     assert run.outputs == {"value": 6}
+
+
+_RECORDING_POLICY = TracePolicy(
+    process_inputs=lambda _: {"recorded": "inputs"},
+    process_outputs=lambda _: {"recorded": "outputs"},
+)
+
+
+def _chat_graph() -> CompiledStateGraph:
+    model = FakeMessagesListChatModel(
+        responses=[AIMessage("model reply", id="ai-model")]
+    )
+
+    def chat(state: MessagesState) -> dict[str, Any]:
+        reply = model.invoke(state["messages"])
+        return {
+            "messages": [
+                *state["messages"],
+                reply,
+                AIMessage("from node", id="ai-node"),
+            ]
+        }
+
+    return (
+        StateGraph(MessagesState)
+        .add_node("chat", chat, trace_policy=_RECORDING_POLICY)
+        .add_edge(START, "chat")
+        .compile()
+    )
+
+
+def _assert_scrubbed_node_run(tracer: FakeTracer, *, model_is_child: bool) -> None:
+    run = _node_run(tracer, "chat")
+    assert run.inputs == {"recorded": "inputs"}
+    assert run.outputs == {"recorded": "outputs"}
+    if model_is_child:
+        assert [c.name for c in run.child_runs] == ["FakeMessagesListChatModel"]
+
+
+def test_trace_policy_scrubs_tracer_not_messages_stream() -> None:
+    tracer = FakeTracer()
+    streamed = [
+        message.content
+        for message, _ in _chat_graph().stream(
+            {"messages": [HumanMessage("hello", id="h-1")]},
+            {"callbacks": [tracer]},
+            stream_mode="messages",
+        )
+    ]
+
+    assert streamed == ["model reply", "from node"]
+    _assert_scrubbed_node_run(tracer, model_is_child=True)
+
+
+@pytest.mark.anyio
+async def test_trace_policy_scrubs_tracer_not_messages_stream_async() -> None:
+    tracer = FakeTracer()
+    streamed = [
+        message.content
+        async for message, _ in _chat_graph().astream(
+            {"messages": [HumanMessage("hello", id="h-1")]},
+            {"callbacks": [tracer]},
+            stream_mode="messages",
+        )
+    ]
+
+    assert streamed == ["model reply", "from node"]
+    _assert_scrubbed_node_run(tracer, model_is_child=ASYNCIO_ACCEPTS_CONTEXT)
+
+
+@pytest.mark.anyio
+async def test_trace_policy_scrubs_tracer_not_astream_events() -> None:
+    tracer = FakeTracer()
+    chat_events = [
+        event
+        async for event in _chat_graph().astream_events(
+            {"messages": [HumanMessage("hello", id="h-1")]},
+            {"callbacks": [tracer]},
+            version="v2",
+        )
+        if event["name"] == "chat"
+        and event["event"] in ("on_chain_start", "on_chain_end")
+    ]
+
+    assert [e["event"] for e in chat_events] == ["on_chain_start", "on_chain_end"]
+    start, end = (e["data"] for e in chat_events)
+    assert [m.content for m in start["input"]["messages"]] == ["hello"]
+    assert [m.content for m in end["output"]["messages"]] == [
+        "hello",
+        "model reply",
+        "from node",
+    ]
+    _assert_scrubbed_node_run(tracer, model_is_child=ASYNCIO_ACCEPTS_CONTEXT)
+
+
+class _HandlerWithTapMethods(BaseCallbackHandler):
+    def __init__(self) -> None:
+        self.node_payloads: list[Any] = []
+        self.node_run_ids: set[UUID] = set()
+
+    def _record(self, run_id: UUID, payload: Any) -> None:
+        if run_id in self.node_run_ids:
+            self.node_payloads.append(payload)
+
+    def tap_output_iter(self, run_id: UUID, output: Iterator[Any]) -> Iterator[Any]:
+        for chunk in output:
+            self._record(run_id, chunk)
+            yield chunk
+
+    async def tap_output_aiter(
+        self, run_id: UUID, output: AsyncIterator[Any]
+    ) -> AsyncIterator[Any]:
+        async for chunk in output:
+            self._record(run_id, chunk)
+            yield chunk
+
+    def on_chain_start(
+        self, serialized: Any, inputs: Any, *, run_id: UUID, **kwargs: Any
+    ) -> None:
+        if kwargs.get("name") == "n":
+            self.node_run_ids.add(run_id)
+            self._record(run_id, inputs)
+
+    def on_chain_end(self, outputs: Any, *, run_id: UUID, **kwargs: Any) -> None:
+        self._record(run_id, outputs)
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("via_events", [False, True], ids=["ainvoke", "astream_events"])
+async def test_trace_policy_reaches_tracers_but_not_custom_handlers(
+    via_events: bool,
+) -> None:
+    graph = (
+        StateGraph(State)
+        .add_node("n", _incr, trace_policy=_RECORDING_POLICY)
+        .add_edge(START, "n")
+        .add_edge("n", END)
+        .compile()
+    )
+    handler = _HandlerWithTapMethods()
+    tracer = FakeTracer()
+    config: RunnableConfig = {"callbacks": [handler, tracer]}
+
+    if via_events:
+        async for _ in graph.astream_events({"value": 1}, config, version="v2"):
+            pass
+    else:
+        assert await graph.ainvoke({"value": 1}, config) == {"value": 2}
+    assert handler.node_payloads == [{"value": 1}, {"value": 2}]
+    run = _node_run(tracer, "n")
+    assert (run.inputs, run.outputs) == (
+        {"recorded": "inputs"},
+        {"recorded": "outputs"},
+    )
+
+
+@pytest.mark.anyio
+async def test_trace_policy_cancel_at_node_end_still_closes_tracer_run() -> None:
+    def cancel_own_task(value: Any) -> Any:
+        asyncio.current_task().cancel()
+        return value
+
+    tracer = FakeTracer()
+    seq = RunnableSeq(
+        RunnableCallable(lambda x: x, trace=False),
+        RunnableCallable(cancel_own_task, trace=False),
+        trace_policy=_RECORDING_POLICY,
+    )
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.create_task(seq.ainvoke({}, {"callbacks": [tracer]}))
+    for _ in range(100):
+        if tracer.runs:
+            break
+        await asyncio.sleep(0.01)
+
+    assert [r.outputs for r in tracer.runs] == [{"recorded": "outputs"}]

libs/langgraph/uv.lock

@@ -1439,7 +1439,7 @@ wheels = [
 
 [[package]]
 name = "langgraph"
-version = "1.2.18"
+version = "1.2.19"
 source = { editable = "." }
 dependencies = [
     { name = "langchain-core" },
@@ -1722,7 +1722,7 @@ test = [
 
 [[package]]
 name = "langgraph-checkpoint-sqlite"
-version = "3.1.2"
+version = "3.1.3"
 source = { editable = "../checkpoint-sqlite" }
 dependencies = [
     { name = "aiosqlite" },

libs/prebuilt/uv.lock

@@ -286,7 +286,7 @@ wheels = [
 
 [[package]]
 name = "langgraph"
-version = "1.2.18"
+version = "1.2.19"
 source = { editable = "../langgraph" }
 dependencies = [
     { name = "langchain-core" },
@@ -467,7 +467,7 @@ test = [
 
 [[package]]
 name = "langgraph-checkpoint-sqlite"
-version = "3.1.2"
+version = "3.1.3"
 source = { editable = "../checkpoint-sqlite" }
 dependencies = [
     { name = "aiosqlite" },

libs/sdk-py/langgraph_sdk/__init__.py

@@ -5,7 +5,7 @@ from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
 from langgraph_sdk.priority import Priority
 from langgraph_sdk.priority.types import PriorityContext, PrioritySettings
 
-__version__ = "0.4.10"
+__version__ = "0.4.11"
 
 __all__ = [
     "Auth",

libs/sdk-py/langgraph_sdk/_async/http.py

@@ -151,7 +151,8 @@ class HttpClient:
         on_response: Callable[[httpx.Response], None] | None = None,
         reconnect_limit: int = 5,
     ) -> Any:
-        """Send a request that automatically reconnects to Location header."""
+        """Send a request that reconnects to its Location header if the response is
+        cut off, with backoff between attempts."""
         request_headers, content = await _aencode_json(json)
         if headers:
             request_headers.update(headers)
@@ -181,13 +182,27 @@ class HttpClient:
                     stacklevel=2,
                 )
                 await r.aclose()
-                return await self.request_reconnect(
-                    loc,
-                    "GET",
-                    headers=headers,
-                    # don't pass on_response so it's only called once
-                    reconnect_limit=reconnect_limit - 1,
-                )
+                for attempt in range(reconnect_limit):
+                    await asyncio.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
+                    # A restart or a deploy looks like a refused connection, a
+                    # dropped body, a 429 or a 5xx; other errors won't go away.
+                    try:
+                        return await self.request_reconnect(
+                            loc,
+                            "GET",
+                            headers=headers,
+                            # don't pass on_response so it's only called once
+                            reconnect_limit=0,
+                        )
+                    except httpx.HTTPStatusError as err:
+                        status = err.response.status_code
+                        if attempt == reconnect_limit - 1 or (
+                            status != 429 and status < 500
+                        ):
+                            raise
+                    except httpx.TransportError:
+                        if attempt == reconnect_limit - 1:
+                            raise
 
     async def stream(
         self,
@@ -230,6 +245,17 @@ class HttpClient:
                 }
             if last_event_id:
                 current_headers["Last-Event-ID"] = last_event_id
+            if reconnect_path is not None and not any(
+                key.lower() == "last-event-id" and value
+                for key, value in current_headers.items()
+            ):
+                warnings.warn(
+                    "Reconnecting without a Last-Event-ID, so events sent while "
+                    "disconnected are not replayed. Only runs created with "
+                    "stream_resumable=True send the event IDs a reconnect resumes "
+                    "from.",
+                    stacklevel=2,
+                )
 
             current_method = method if reconnect_path is None else "GET"
             current_content = content if reconnect_path is None else None

libs/sdk-py/langgraph_sdk/_sync/http.py

@@ -152,7 +152,8 @@ class SyncHttpClient:
         on_response: Callable[[httpx.Response], None] | None = None,
         reconnect_limit: int = 5,
     ) -> Any:
-        """Send a request that automatically reconnects to Location header."""
+        """Send a request that reconnects to its Location header if the response is
+        cut off, with backoff between attempts."""
         request_headers, content = _encode_json(json)
         if headers:
             request_headers.update(headers)
@@ -182,13 +183,27 @@ class SyncHttpClient:
                     stacklevel=2,
                 )
                 r.close()
-                return self.request_reconnect(
-                    loc,
-                    "GET",
-                    headers=headers,
-                    # don't pass on_response so it's only called once
-                    reconnect_limit=reconnect_limit - 1,
-                )
+                for attempt in range(reconnect_limit):
+                    time.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
+                    # A restart or a deploy looks like a refused connection, a
+                    # dropped body, a 429 or a 5xx; other errors won't go away.
+                    try:
+                        return self.request_reconnect(
+                            loc,
+                            "GET",
+                            headers=headers,
+                            # don't pass on_response so it's only called once
+                            reconnect_limit=0,
+                        )
+                    except httpx.HTTPStatusError as err:
+                        status = err.response.status_code
+                        if attempt == reconnect_limit - 1 or (
+                            status != 429 and status < 500
+                        ):
+                            raise
+                    except httpx.TransportError:
+                        if attempt == reconnect_limit - 1:
+                            raise
 
     def stream(
         self,
@@ -233,6 +248,17 @@ class SyncHttpClient:
                 }
             if last_event_id:
                 current_headers["Last-Event-ID"] = last_event_id
+            if reconnect_path is not None and not any(
+                key.lower() == "last-event-id" and value
+                for key, value in current_headers.items()
+            ):
+                warnings.warn(
+                    "Reconnecting without a Last-Event-ID, so events sent while "
+                    "disconnected are not replayed. Only runs created with "
+                    "stream_resumable=True send the event IDs a reconnect resumes "
+                    "from.",
+                    stacklevel=2,
+                )
 
             current_method = method if reconnect_path is None else "GET"
             current_content = content if reconnect_path is None else None

libs/sdk-py/tests/conftest.py

@@ -0,0 +1,17 @@
+import asyncio
+import time
+
+import pytest
+
+
+@pytest.fixture
+def no_sleep(monkeypatch):
+    """Record backoff sleeps without actually waiting."""
+    sleeps: list[float] = []
+
+    async def fake_async_sleep(d):
+        sleeps.append(d)
+
+    monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
+    monkeypatch.setattr(time, "sleep", lambda d: sleeps.append(d))
+    return sleeps

libs/sdk-py/tests/test_client_stream.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import asyncio
-import time
 from collections.abc import Callable, Iterator, Sequence
 from pathlib import Path
 from typing import Any
@@ -70,19 +68,6 @@ class ListByteStream(httpx.ByteStream):
         return None
 
 
-@pytest.fixture
-def no_sleep(monkeypatch):
-    """Record backoff sleeps without actually waiting."""
-    sleeps: list[float] = []
-
-    async def fake_async_sleep(d):
-        sleeps.append(d)
-
-    monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
-    monkeypatch.setattr(time, "sleep", lambda d: sleeps.append(d))
-    return sleeps
-
-
 def iter_lines_raw(payload: list[bytes]) -> Iterator[BytesLike]:
     decoder = BytesLineDecoder()
     for part in payload:
@@ -501,9 +486,10 @@ _CALLER_CURSORS = pytest.mark.parametrize(
 def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
     empty_id: bytes, headers: dict[str, str] | None
 ):
-    reconnect = _sync_reconnect(
-        [b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
-    )
+    with pytest.warns(UserWarning, match="Last-Event-ID"):
+        reconnect = _sync_reconnect(
+            [b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
+        )
 
     assert "last-event-id" not in reconnect.headers
 
@@ -514,9 +500,10 @@ def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
 async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
     empty_id: bytes, headers: dict[str, str] | None
 ):
-    reconnect = await _async_reconnect(
-        [b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
-    )
+    with pytest.warns(UserWarning, match="Last-Event-ID"):
+        reconnect = await _async_reconnect(
+            [b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
+        )
 
     assert "last-event-id" not in reconnect.headers
 
@@ -555,6 +542,45 @@ async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id
     assert reconnect.headers.get_list("last-event-id") == ["5"]
 
 
+# A reconnect that sends no Last-Event-ID can't replay what was sent while it was
+# disconnected, so it warns.
+_RECONNECT_CURSORS = pytest.mark.parametrize(
+    ("chunks", "headers", "warns"),
+    [
+        ([b"event: values\ndata: 1\n\n"], None, True),
+        ([b"id: 1\nevent: values\ndata: 1\n\n"], None, False),
+        ([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}, False),
+    ],
+    ids=["no-cursor", "stream-cursor", "caller-cursor"],
+)
+
+
+@_RECONNECT_CURSORS
+@pytest.mark.usefixtures("no_sleep")
+def test_sync_reconnect_warns_only_without_a_cursor(
+    chunks: list[bytes],
+    headers: dict[str, str] | None,
+    warns: bool,
+    recwarn: pytest.WarningsRecorder,
+):
+    _sync_reconnect(chunks, headers)
+
+    assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
+
+
+@_RECONNECT_CURSORS
+@pytest.mark.usefixtures("no_sleep")
+async def test_async_reconnect_warns_only_without_a_cursor(
+    chunks: list[bytes],
+    headers: dict[str, str] | None,
+    warns: bool,
+    recwarn: pytest.WarningsRecorder,
+):
+    await _async_reconnect(chunks, headers)
+
+    assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
+
+
 # --- reconnects survive a server restart or a deploy ---
 
 SSE = {"Content-Type": "text/event-stream"}
@@ -772,6 +798,39 @@ async def test_async_long_stream_survives_many_spaced_out_drops():
     assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
 
 
+# Each reconnect resumes after the newest event the stream got, so a reconnect
+# loop on a resumable run replays only what it missed. A failed reconnect keeps
+# the cursor for the next one.
+RECONNECT_LOOPS = {
+    "progress_between_drops": (
+        lambda stream_cls: _long_stream_steps(stream_cls, drops=3),
+        [None, "0", "1", "2"],
+    ),
+    "failed_reconnect": (
+        lambda stream_cls: _transient_steps(stream_cls, _status(502)),
+        [None, "1", "1"],
+    ),
+}
+
+
+@pytest.mark.usefixtures("no_sleep")
+@pytest.mark.parametrize(
+    ("steps", "cursors"), RECONNECT_LOOPS.values(), ids=RECONNECT_LOOPS.keys()
+)
+def test_sync_reconnect_loop_resumes_after_the_newest_event(steps, cursors):
+    _, seen = _run_sync(steps(ListByteStream))
+    assert [r.headers.get("last-event-id") for r in seen] == cursors
+
+
+@pytest.mark.usefixtures("no_sleep")
+@pytest.mark.parametrize(
+    ("steps", "cursors"), RECONNECT_LOOPS.values(), ids=RECONNECT_LOOPS.keys()
+)
+async def test_async_reconnect_loop_resumes_after_the_newest_event(steps, cursors):
+    _, seen = await _run_async(steps(AsyncListByteStream))
+    assert [r.headers.get("last-event-id") for r in seen] == cursors
+
+
 # Reconnects should back off rather than hammer the server.
 
 

libs/sdk-py/tests/test_runs_client.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import json
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Callable, Iterator, Sequence
+from typing import Any
 
 import httpx
 import pytest
@@ -24,7 +25,16 @@ class _BodyDropsMidRead(httpx.SyncByteStream, httpx.AsyncByteStream):
         raise httpx.ReadError("connection dropped")
 
 
-def _wait_handler(requests: list[httpx.Request]):
+def _done(_request: httpx.Request) -> httpx.Response:
+    return httpx.Response(200, json={"done": True})
+
+
+def _wait_handler(
+    requests: list[httpx.Request],
+    reconnects: Sequence[Callable[[httpx.Request], httpx.Response]] = (_done,),
+):
+    """POST drops mid-body; the reconnects answer from `reconnects` in order."""
+
     def handler(request: httpx.Request) -> httpx.Response:
         requests.append(request)
         if request.method == "POST":
@@ -33,7 +43,7 @@ def _wait_handler(requests: list[httpx.Request]):
                 headers={"Location": "/threads/t1/runs/r1/join"},
                 stream=_BodyDropsMidRead(),
             )
-        return httpx.Response(200, json={"done": True})
+        return reconnects[len(requests) - 2](request)
 
     return handler
 
@@ -46,6 +56,7 @@ def _assert_bodyless_get(request: httpx.Request) -> None:
     assert request.read() == b""
 
 
+@pytest.mark.usefixtures("no_sleep")
 def test_sync_wait_reconnects_with_a_bodyless_get():
     requests: list[httpx.Request] = []
     transport = httpx.MockTransport(_wait_handler(requests))
@@ -58,6 +69,7 @@ def test_sync_wait_reconnects_with_a_bodyless_get():
     _assert_bodyless_get(requests[1])
 
 
+@pytest.mark.usefixtures("no_sleep")
 async def test_async_wait_reconnects_with_a_bodyless_get():
     requests: list[httpx.Request] = []
     transport = httpx.MockTransport(_wait_handler(requests))
@@ -72,6 +84,89 @@ async def test_async_wait_reconnects_with_a_bodyless_get():
     _assert_bodyless_get(requests[1])
 
 
+def _connect_error(request: httpx.Request) -> httpx.Response:
+    raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
+
+
+async def _wait(
+    handler: Callable[[httpx.Request], httpx.Response], *, asynchronous: bool
+) -> Any:
+    transport = httpx.MockTransport(handler)
+    if asynchronous:
+        async with httpx.AsyncClient(
+            transport=transport, base_url="https://example.com"
+        ) as client:
+            return await RunsClient(HttpClient(client)).wait("t1", "agent", input={})
+    with httpx.Client(transport=transport, base_url="https://example.com") as client:
+        return SyncRunsClient(SyncHttpClient(client)).wait("t1", "agent", input={})
+
+
+# What a server restart or a deploy looks like to the reconnect GET: each one is
+# retried after a backoff.
+RECONNECT_FAILURES = {
+    "429": lambda request: httpx.Response(429, request=request),
+    "503": lambda request: httpx.Response(503, request=request),
+    "connect_error": _connect_error,
+    "body_drops": lambda _request: httpx.Response(
+        200,
+        headers={"Location": "/threads/t1/runs/r1/join"},
+        stream=_BodyDropsMidRead(),
+    ),
+}
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize(
+    "failure", RECONNECT_FAILURES.values(), ids=RECONNECT_FAILURES.keys()
+)
+async def test_wait_retries_a_failed_reconnect_after_a_backoff(
+    failure, asynchronous, no_sleep
+):
+    requests: list[httpx.Request] = []
+    with pytest.warns(UserWarning, match="attempting reconnect"):
+        result = await _wait(
+            _wait_handler(requests, [failure, _done]), asynchronous=asynchronous
+        )
+
+    assert result == {"done": True}
+    assert len(requests) == 3
+    assert len(no_sleep) == 2 and min(no_sleep) >= 1, no_sleep
+
+
+@pytest.mark.usefixtures("no_sleep")
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_wait_raises_after_five_failed_reconnects(asynchronous):
+    requests: list[httpx.Request] = []
+    with (
+        pytest.warns(UserWarning, match="attempting reconnect"),
+        pytest.raises(httpx.HTTPStatusError, match="503"),
+    ):
+        await _wait(
+            _wait_handler(requests, [RECONNECT_FAILURES["503"]] * 5),
+            asynchronous=asynchronous,
+        )
+
+    assert len(requests) == 6
+
+
+@pytest.mark.usefixtures("no_sleep")
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_wait_does_not_retry_a_client_error(asynchronous):
+    requests: list[httpx.Request] = []
+    with (
+        pytest.warns(UserWarning, match="attempting reconnect"),
+        pytest.raises(httpx.HTTPStatusError, match="404"),
+    ):
+        await _wait(
+            _wait_handler(
+                requests, [lambda request: httpx.Response(404, request=request)]
+            ),
+            asynchronous=asynchronous,
+        )
+
+    assert len(requests) == 2
+
+
 def test_sync_wait_raises_the_run_error():
     requests: list[httpx.Request] = []
 

libs/sdk-py/uv.lock

@@ -301,7 +301,7 @@ wheels = [
 
 [[package]]
 name = "langgraph"
-version = "1.2.18"
+version = "1.2.19"
 source = { editable = "../langgraph" }
 dependencies = [
     { name = "langchain-core" },