From 3a5b5c9821f23a1be4bd562c66d7f736e264a67c Mon Sep 17 00:00:00 2001 From: Hunter Lovell <40191806+hntrl@users.noreply.github.com> Date: Fri, 10 Apr 2026 13:43:48 -0700 Subject: [PATCH] feat(langgraph): add graph lifecycle callback handlers (#7429) ## Summary This change adds first-class graph lifecycle callbacks to LangGraph so interrupt and resume transitions can be observed without overloading the existing LangChain custom event system. It introduces a dedicated graph callback manager and wires lifecycle emission through Pregel execution in both sync and async paths. ## Changes - **`libs/langgraph/langgraph/callbacks.py`**: Adds `GraphCallbackHandler` and `GraphCallbackManager` (built on LangChain base callback classes), plus config plumbing via `graph_callbacks` and `get_graph_callback_manager_for_config`. - **`libs/langgraph/langgraph/pregel/_loop.py`**: Introduces `GraphLifecycleEvent` and records lifecycle transitions (`resume`, `interrupt`) into an internal FIFO queue with `shift_graph_lifecycle_event()`. - **`libs/langgraph/langgraph/pregel/main.py`**: Resolves graph callback manager from config and drains lifecycle events while loop execution progresses, dispatching `on_resume` and `on_interrupt` consistently in sync and async runtimes. - **`libs/langgraph/tests/test_graph_callbacks.py`**: Adds sync and async coverage verifying lifecycle callbacks fire correctly and remain distinct from LangChain `on_custom_event` handlers. --------- Co-authored-by: Eugene Yurtsev --- libs/langgraph/langgraph/callbacks.py | 412 +++++++++++++++++++ libs/langgraph/langgraph/pregel/_loop.py | 66 ++- libs/langgraph/langgraph/pregel/main.py | 58 ++- libs/langgraph/tests/test_graph_callbacks.py | 277 +++++++++++++ 4 files changed, 806 insertions(+), 7 deletions(-) create mode 100644 libs/langgraph/langgraph/callbacks.py create mode 100644 libs/langgraph/tests/test_graph_callbacks.py diff --git a/libs/langgraph/langgraph/callbacks.py b/libs/langgraph/langgraph/callbacks.py new file mode 100644 index 000000000..d5933b1b6 --- /dev/null +++ b/libs/langgraph/langgraph/callbacks.py @@ -0,0 +1,412 @@ +"""Graph lifecycle callback interfaces and event payloads. + +This module defines the public callback surface for observing LangGraph-specific +lifecycle transitions such as interrupt and resume. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Literal, TypeAlias, TypeVar +from uuid import UUID + +from langchain_core.callbacks import BaseCallbackHandler, BaseCallbackManager +from langchain_core.callbacks.manager import ahandle_event, handle_event +from langchain_core.runnables import RunnableConfig + +from langgraph.types import Interrupt + +__all__ = ( + "GraphCallbackHandler", + "GraphInterruptEvent", + "GraphLifecycleEvent", + "GraphLifecycleStatus", + "GraphResumeEvent", + "get_async_graph_callback_manager_for_config", + "get_sync_graph_callback_manager_for_config", +) + + +GraphLifecycleStatus: TypeAlias = Literal[ + "input", + "pending", + "done", + "interrupt_before", + "interrupt_after", + "out_of_steps", +] +"""Allowed lifecycle statuses reported in graph lifecycle callback events.""" + + +@dataclass(frozen=True) +class GraphInterruptEvent: + """Graph lifecycle event emitted when execution pauses for interrupts.""" + + run_id: UUID | None + """Run id for the current graph execution, if available.""" + + status: GraphLifecycleStatus + """Loop status when the interrupt was captured.""" + + checkpoint_id: str + """Checkpoint id associated with the interrupted execution.""" + + checkpoint_ns: tuple[str, ...] + """Checkpoint namespace path for the current graph or subgraph.""" + + interrupts: tuple[Interrupt, ...] + """Interrupt payloads that caused the graph to pause.""" + + +@dataclass(frozen=True) +class GraphResumeEvent: + """Graph lifecycle event emitted when execution resumes from a checkpoint.""" + + run_id: UUID | None + """Run id for the current graph execution, if available.""" + + status: GraphLifecycleStatus + """Loop status when the resume was captured.""" + + checkpoint_id: str + """Checkpoint id the graph resumed from.""" + + checkpoint_ns: tuple[str, ...] + """Checkpoint namespace path for the current graph or subgraph.""" + + +GraphLifecycleEvent: TypeAlias = GraphInterruptEvent | GraphResumeEvent +"""Union of all public graph lifecycle callback event payloads. + +Use this alias when a callback or helper can receive either interrupt or resume +lifecycle events. +""" + + +class GraphCallbackHandler(BaseCallbackHandler): + """Base class for graph-level lifecycle callbacks. + + Subclass this handler to observe graph lifecycle transitions that are + specific to LangGraph execution, rather than generic LangChain runnable + callbacks. + + Instances can be passed through `config["callbacks"]` when invoking a + graph. Only handlers that inherit from `GraphCallbackHandler` receive these + lifecycle events. + """ + + def on_interrupt(self, event: GraphInterruptEvent) -> Any: + """Run when graph execution pauses due to one or more interrupts. + + Args: + event: Interrupt lifecycle event payload. + """ + + def on_resume(self, event: GraphResumeEvent) -> Any: + """Run when graph execution resumes from a persisted checkpoint. + + Args: + event: Resume lifecycle event payload. + """ + + +_MISSING = object() + + +def _filter_graph_handlers( + handlers: list[BaseCallbackHandler], +) -> list[GraphCallbackHandler]: + return [h for h in handlers if isinstance(h, GraphCallbackHandler)] + + +def _init_base_manager( + manager: BaseCallbackManager, + handlers: Sequence[GraphCallbackHandler] | None, + inheritable_handlers: Sequence[GraphCallbackHandler] | None, + parent_run_id: UUID | None, + *, + tags: list[str] | None, + inheritable_tags: list[str] | None, + metadata: dict[str, Any] | None, + inheritable_metadata: dict[str, Any] | None, + run_id: UUID | None, +) -> None: + base_handlers: list[BaseCallbackHandler] = [] + base_inheritable_handlers: list[BaseCallbackHandler] = [] + if handlers is not None: + base_handlers.extend(handlers) + if inheritable_handlers is not None: + base_inheritable_handlers.extend(inheritable_handlers) + BaseCallbackManager.__init__( + manager, + handlers=base_handlers, + inheritable_handlers=base_inheritable_handlers, + parent_run_id=parent_run_id, + tags=tags, + inheritable_tags=inheritable_tags, + metadata=metadata, + inheritable_metadata=inheritable_metadata, + ) + manager.run_id = run_id # type: ignore[attr-defined] + + +def _configure_graph_callbacks( + cls: type[_GraphManagerT], + callbacks: object | None, + *, + run_id: UUID | None, +) -> _GraphManagerT: + if callbacks is None: + return cls(run_id=run_id) + if isinstance(callbacks, cls): + return callbacks.copy(run_id=run_id) + if isinstance(callbacks, (_GraphCallbackManager, _AsyncGraphCallbackManager)): + # Cross-type: extract handlers into the requested cls. + return cls( + handlers=_filter_graph_handlers(callbacks.handlers), + inheritable_handlers=_filter_graph_handlers(callbacks.inheritable_handlers), + parent_run_id=callbacks.parent_run_id, + tags=callbacks.tags.copy(), + inheritable_tags=callbacks.inheritable_tags.copy(), + metadata=callbacks.metadata.copy(), + inheritable_metadata=callbacks.inheritable_metadata.copy(), + run_id=run_id, + ) + if isinstance(callbacks, BaseCallbackManager): + return cls( + handlers=_filter_graph_handlers(callbacks.handlers), + inheritable_handlers=_filter_graph_handlers(callbacks.inheritable_handlers), + parent_run_id=callbacks.parent_run_id, + tags=callbacks.tags.copy(), + inheritable_tags=callbacks.inheritable_tags.copy(), + metadata=callbacks.metadata.copy(), + inheritable_metadata=callbacks.inheritable_metadata.copy(), + run_id=run_id, + ) + if isinstance(callbacks, GraphCallbackHandler): + return cls((callbacks,), run_id=run_id) + if isinstance(callbacks, (str, bytes)) or not isinstance(callbacks, Sequence): + raise TypeError("callbacks must be a handler, sequence, or manager") + return cls(_filter_graph_handlers(list(callbacks)), run_id=run_id) + + +def _copy_graph_manager( + manager: _GraphCallbackManager | _AsyncGraphCallbackManager, + cls: type[_GraphManagerT], + run_id: UUID | None | object, +) -> _GraphManagerT: + resolved_run_id: UUID | None + if run_id is _MISSING: + resolved_run_id = manager.run_id + else: + if run_id is not None and not isinstance(run_id, UUID): + raise TypeError("run_id must be a UUID or None") + resolved_run_id = run_id + + return cls( + handlers=_filter_graph_handlers(manager.handlers), + inheritable_handlers=_filter_graph_handlers(manager.inheritable_handlers), + parent_run_id=manager.parent_run_id, + tags=manager.tags.copy(), + inheritable_tags=manager.inheritable_tags.copy(), + metadata=manager.metadata.copy(), + inheritable_metadata=manager.inheritable_metadata.copy(), + run_id=resolved_run_id, + ) + + +class _GraphCallbackManager(BaseCallbackManager): + """Sync dispatcher for graph lifecycle events.""" + + run_id: UUID | None + + def __init__( + self, + handlers: Sequence[GraphCallbackHandler] | None = None, + inheritable_handlers: Sequence[GraphCallbackHandler] | None = None, + parent_run_id: UUID | None = None, + *, + tags: list[str] | None = None, + inheritable_tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + inheritable_metadata: dict[str, Any] | None = None, + run_id: UUID | None = None, + ) -> None: + _init_base_manager( + self, + handlers, + inheritable_handlers, + parent_run_id, + tags=tags, + inheritable_tags=inheritable_tags, + metadata=metadata, + inheritable_metadata=inheritable_metadata, + run_id=run_id, + ) + + def add_handler( + self, + handler: BaseCallbackHandler, + inherit: bool = True, # noqa: FBT001,FBT002 + ) -> None: + if not isinstance(handler, GraphCallbackHandler): + raise TypeError("handlers must inherit GraphCallbackHandler") + super().add_handler(handler, inherit=inherit) + + def copy( + self, + *, + run_id: UUID | None | object = _MISSING, + ) -> _GraphCallbackManager: + return _copy_graph_manager(self, _GraphCallbackManager, run_id) + + @classmethod + def configure( + cls, + callbacks: object | None = None, + *, + run_id: UUID | None = None, + ) -> _GraphCallbackManager: + return _configure_graph_callbacks(cls, callbacks, run_id=run_id) + + def on_interrupt(self, event: GraphInterruptEvent) -> None: + handle_event( + self.handlers, + "on_interrupt", + None, + event, + ) + + def on_resume(self, event: GraphResumeEvent) -> None: + handle_event( + self.handlers, + "on_resume", + None, + event, + ) + + +class _AsyncGraphCallbackManager(BaseCallbackManager): + """Async dispatcher for graph lifecycle events.""" + + run_id: UUID | None + + @property + def is_async(self) -> bool: + """Return whether the manager is async.""" + return True + + def __init__( + self, + handlers: Sequence[GraphCallbackHandler] | None = None, + inheritable_handlers: Sequence[GraphCallbackHandler] | None = None, + parent_run_id: UUID | None = None, + *, + tags: list[str] | None = None, + inheritable_tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + inheritable_metadata: dict[str, Any] | None = None, + run_id: UUID | None = None, + ) -> None: + _init_base_manager( + self, + handlers, + inheritable_handlers, + parent_run_id, + tags=tags, + inheritable_tags=inheritable_tags, + metadata=metadata, + inheritable_metadata=inheritable_metadata, + run_id=run_id, + ) + + def add_handler( + self, + handler: BaseCallbackHandler, + inherit: bool = True, # noqa: FBT001,FBT002 + ) -> None: + if not isinstance(handler, GraphCallbackHandler): + raise TypeError("handlers must inherit GraphCallbackHandler") + super().add_handler(handler, inherit=inherit) + + def copy( + self, + *, + run_id: UUID | None | object = _MISSING, + ) -> _AsyncGraphCallbackManager: + return _copy_graph_manager(self, _AsyncGraphCallbackManager, run_id) + + @classmethod + def configure( + cls, + callbacks: object | None = None, + *, + run_id: UUID | None = None, + ) -> _AsyncGraphCallbackManager: + return _configure_graph_callbacks(cls, callbacks, run_id=run_id) + + async def on_interrupt(self, event: GraphInterruptEvent) -> None: + await ahandle_event( + self.handlers, + "on_interrupt", + None, + event, + ) + + async def on_resume(self, event: GraphResumeEvent) -> None: + await ahandle_event( + self.handlers, + "on_resume", + None, + event, + ) + + +_GraphManagerT = TypeVar( + "_GraphManagerT", _GraphCallbackManager, _AsyncGraphCallbackManager +) + +GraphCallbacks: TypeAlias = ( + _GraphCallbackManager + | _AsyncGraphCallbackManager + | BaseCallbackManager + | GraphCallbackHandler + | Sequence[BaseCallbackHandler] + | Sequence[GraphCallbackHandler] + | None +) + + +def get_sync_graph_callback_manager_for_config( + config: RunnableConfig, + *, + run_id: UUID | None = None, +) -> _GraphCallbackManager: + """Build a sync graph lifecycle callback manager from a runnable config. + + This helper filters `config["callbacks"]` down to handlers that inherit + from `GraphCallbackHandler` and binds the provided `run_id` onto the + returned manager. + """ + return _GraphCallbackManager.configure( + config.get("callbacks"), + run_id=run_id, + ) + + +def get_async_graph_callback_manager_for_config( + config: RunnableConfig, + *, + run_id: UUID | None = None, +) -> _AsyncGraphCallbackManager: + """Build an async graph lifecycle callback manager from a runnable config. + + This helper filters `config["callbacks"]` down to handlers that inherit + from `GraphCallbackHandler` and binds the provided `run_id` onto the + returned manager. + """ + return _AsyncGraphCallbackManager.configure( + config.get("callbacks"), + run_id=run_id, + ) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 4f9c55d2d..aec4b74cd 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -62,6 +62,11 @@ from langgraph._internal._constants import ( from langgraph._internal._replay import ReplayState from langgraph._internal._scratchpad import PregelScratchpad from langgraph._internal._typing import EMPTY_SEQ, MISSING +from langgraph.callbacks import ( + GraphInterruptEvent, + GraphLifecycleEvent, + GraphResumeEvent, +) from langgraph.channels.base import BaseChannel from langgraph.channels.untracked_value import UntrackedValue from langgraph.constants import TAG_HIDDEN @@ -117,6 +122,7 @@ from langgraph.types import ( CachePolicy, Command, Durability, + Interrupt, PregelExecutableTask, RetryPolicy, Send, @@ -203,6 +209,8 @@ class PregelLoop: tasks: dict[str, PregelExecutableTask] output: None | dict[str, Any] | Any = None updated_channels: set[str] | None = None + _graph_lifecycle_events: deque[GraphLifecycleEvent] + _has_graph_lifecycle_callbacks: bool # public @@ -228,6 +236,7 @@ class PregelLoop: migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + has_graph_lifecycle_callbacks: bool = False, ) -> None: self.stream = stream self.config = config @@ -252,6 +261,8 @@ class PregelLoop: self.retry_policy = retry_policy self.cache_policy = cache_policy self.durability = durability + self._has_graph_lifecycle_callbacks = has_graph_lifecycle_callbacks + self._graph_lifecycle_events = deque() if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD) @@ -303,6 +314,40 @@ class PregelLoop: ) self.prev_checkpoint_config = None + def _push_graph_lifecycle_event( + self, + kind: Literal["resume", "interrupt"], + *, + interrupts: tuple[Interrupt, ...] = (), + ) -> None: + if kind == "resume": + self._graph_lifecycle_events.append( + GraphResumeEvent( + run_id=None, + status=self.status, + checkpoint_id=self.checkpoint["id"], + checkpoint_ns=self.checkpoint_ns, + ) + ) + elif kind == "interrupt": + self._graph_lifecycle_events.append( + GraphInterruptEvent( + run_id=None, + status=self.status, + checkpoint_id=self.checkpoint["id"], + checkpoint_ns=self.checkpoint_ns, + interrupts=interrupts, + ) + ) + else: + msg = f"Unknown graph lifecycle event type: {kind}" + raise AssertionError(msg) + + def _pop_lifecycle_event(self) -> GraphLifecycleEvent | None: + if not self._graph_lifecycle_events: + return None + return self._graph_lifecycle_events.popleft() + def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" if not writes: @@ -785,6 +830,8 @@ class PregelLoop: ) # set flag self.status = "pending" + if is_resuming: + self._push_graph_lifecycle_event("resume") return updated_channels def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: @@ -885,8 +932,10 @@ class PregelLoop: self._put_checkpoint(self.checkpoint_metadata) self._put_pending_writes() # suppress interrupt - suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested - if suppress: + if isinstance(exc_value, GraphInterrupt) and not self.is_nested: + interrupt = exc_value + interrupts = tuple(interrupt.args[0]) if interrupt.args else () + self._push_graph_lifecycle_event("interrupt", interrupts=interrupts) # emit one last "values" event, with pending writes applied if ( hasattr(self, "tasks") @@ -913,12 +962,11 @@ class PregelLoop: self.channels, ) # emit INTERRUPT if exception is empty (otherwise emitted by put_writes) - if exc_value is not None and (not exc_value.args or not exc_value.args[0]): + if not interrupt.args or not interrupt.args[0]: + interrupt_payload = interrupt.args[0] if interrupt.args else () self._emit( "updates", - lambda: iter( - [{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}] - ), + lambda: iter([{INTERRUPT: interrupt_payload}]), ) # save final output self.output = read_channels(self.channels, self.output_keys) @@ -1040,6 +1088,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + has_graph_lifecycle_callbacks: bool = False, ) -> None: super().__init__( input, @@ -1061,6 +1110,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): retry_policy=retry_policy, cache_policy=cache_policy, durability=durability, + has_graph_lifecycle_callbacks=has_graph_lifecycle_callbacks, ) self.stack = ExitStack() if checkpointer: @@ -1136,6 +1186,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): # context manager def __enter__(self) -> Self: + self._graph_lifecycle_events = deque() if not self.checkpointer: saved = None elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID): @@ -1236,6 +1287,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + has_graph_lifecycle_callbacks: bool = False, ) -> None: super().__init__( input, @@ -1257,6 +1309,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): retry_policy=retry_policy, cache_policy=cache_policy, durability=durability, + has_graph_lifecycle_callbacks=has_graph_lifecycle_callbacks, ) self.stack = AsyncExitStack() if checkpointer: @@ -1335,6 +1388,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): # context manager async def __aenter__(self) -> Self: + self._graph_lifecycle_events = deque() if not self.checkpointer: saved = None elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID): diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index d754db391..02d0ee11e 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -16,7 +16,7 @@ from collections.abc import ( Mapping, Sequence, ) -from dataclasses import is_dataclass +from dataclasses import is_dataclass, replace from functools import partial from inspect import isclass from typing import ( @@ -96,6 +96,12 @@ from langgraph._internal._runnable import ( coerce_to_runnable, ) from langgraph._internal._typing import MISSING, DeprecatedKwargs +from langgraph.callbacks import ( + GraphInterruptEvent, + GraphResumeEvent, + get_async_graph_callback_manager_for_config, + get_sync_graph_callback_manager_for_config, +) from langgraph.channels.base import BaseChannel from langgraph.channels.topic import Topic from langgraph.config import get_config @@ -2585,6 +2591,10 @@ class Pregel( name=config.get("run_name", self.get_name()), run_id=config.get("run_id"), ) + graph_callback_manager = get_sync_graph_callback_manager_for_config( + config, + run_id=run_manager.run_id, + ) try: # assign defaults ( @@ -2669,6 +2679,17 @@ class Pregel( _output_mapper = self._output_mapper if version == "v2" else None _state_mapper = self._state_mapper if version == "v2" else None + def emit_graph_lifecycle_events(loop: SyncPregelLoop) -> None: + while (event := loop._pop_lifecycle_event()) is not None: + if isinstance(event, GraphResumeEvent): + graph_callback_manager.on_resume( + replace(event, run_id=graph_callback_manager.run_id) + ) + else: + graph_callback_manager.on_interrupt( + replace(event, run_id=graph_callback_manager.run_id) + ) + with SyncPregelLoop( input, stream=StreamProtocol(stream.put, stream_modes), @@ -2689,7 +2710,9 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers), ) as loop: + emit_graph_lifecycle_events(loop) # create runner runner = PregelRunner( submit=config[CONF].get( @@ -2751,9 +2774,11 @@ class Pregel( _state_mapper, ) loop.after_tick() + emit_graph_lifecycle_events(loop) # wait for checkpoint if durability_ == "sync": loop._put_checkpoint_fut.result() + emit_graph_lifecycle_events(loop) # emit output yield from _output( stream_mode, @@ -2928,6 +2953,10 @@ class Pregel( name=config.get("run_name", self.get_name()), run_id=config.get("run_id"), ) + graph_callback_manager = get_async_graph_callback_manager_for_config( + config, + run_id=run_manager.run_id, + ) # if running from astream_log() run each proc with streaming do_stream = ( next( @@ -3042,6 +3071,28 @@ class Pregel( _output_mapper = self._output_mapper if version == "v2" else None _state_mapper = self._state_mapper if version == "v2" else None + async def aemit_graph_lifecycle_events(loop: AsyncPregelLoop) -> None: + while (event := loop._pop_lifecycle_event()) is not None: + if isinstance(event, GraphResumeEvent): + await graph_callback_manager.on_resume( + GraphResumeEvent( + run_id=graph_callback_manager.run_id, + status=event.status, + checkpoint_id=event.checkpoint_id, + checkpoint_ns=event.checkpoint_ns, + ) + ) + else: + await graph_callback_manager.on_interrupt( + GraphInterruptEvent( + run_id=graph_callback_manager.run_id, + status=event.status, + checkpoint_id=event.checkpoint_id, + checkpoint_ns=event.checkpoint_ns, + interrupts=event.interrupts, + ) + ) + async with AsyncPregelLoop( input, stream=StreamProtocol(stream.put_nowait, stream_modes), @@ -3062,7 +3113,9 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers), ) as loop: + await aemit_graph_lifecycle_events(loop) # create runner runner = PregelRunner( submit=config[CONF].get( @@ -3144,6 +3197,7 @@ class Pregel( ): yield o loop.after_tick() + await aemit_graph_lifecycle_events(loop) # wait for checkpoint if durability_ == "sync": await cast(asyncio.Future, loop._put_checkpoint_fut) @@ -3152,6 +3206,8 @@ class Pregel( if _cleanup_waiter is not None: await _cleanup_waiter() + await aemit_graph_lifecycle_events(loop) + # emit output for o in _output( stream_mode, diff --git a/libs/langgraph/tests/test_graph_callbacks.py b/libs/langgraph/tests/test_graph_callbacks.py new file mode 100644 index 000000000..ca825f904 --- /dev/null +++ b/libs/langgraph/tests/test_graph_callbacks.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +import sys +from typing import Any + +import pytest +from langchain_core.callbacks.base import BaseCallbackHandler +from langchain_core.callbacks.manager import CallbackManager +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import TypedDict + +from langgraph.callbacks import ( + GraphCallbackHandler, + GraphInterruptEvent, + GraphResumeEvent, +) +from langgraph.graph import START, StateGraph +from langgraph.types import Command, Interrupt, interrupt + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +class _GraphEventHandler(GraphCallbackHandler): + def __init__(self) -> None: + self.interrupt_events: list[GraphInterruptEvent] = [] + self.resume_events: list[GraphResumeEvent] = [] + + def on_interrupt(self, event: GraphInterruptEvent) -> Any: + self.interrupt_events.append(event) + + def on_resume(self, event: GraphResumeEvent) -> Any: + self.resume_events.append(event) + + +class _LangChainCustomEventHandler(BaseCallbackHandler): + run_inline = True + + def __init__(self) -> None: + self.events: list[str] = [] + + def on_custom_event(self, name: str, data: Any, **kwargs: Any) -> Any: + self.events.append(name) + + +class _RaisingGraphEventHandler(GraphCallbackHandler): + def __init__( + self, + *, + raise_on_interrupt: bool = False, + raise_on_resume: bool = False, + raise_error: bool = False, + ) -> None: + self.raise_on_interrupt = raise_on_interrupt + self.raise_on_resume = raise_on_resume + self.raise_error = raise_error + + def on_interrupt(self, event: GraphInterruptEvent) -> Any: + if self.raise_on_interrupt: + raise ValueError("boom-interrupt") + + def on_resume(self, event: GraphResumeEvent) -> Any: + if self.raise_on_resume: + raise ValueError("boom-resume") + + +class _AsyncRaisingGraphEventHandler(GraphCallbackHandler): + def __init__( + self, + *, + raise_on_interrupt: bool = False, + raise_on_resume: bool = False, + raise_error: bool = False, + ) -> None: + self.raise_on_interrupt = raise_on_interrupt + self.raise_on_resume = raise_on_resume + self.raise_error = raise_error + + async def on_interrupt(self, event: GraphInterruptEvent) -> Any: + if self.raise_on_interrupt: + raise ValueError("boom-interrupt") + + async def on_resume(self, event: GraphResumeEvent) -> Any: + if self.raise_on_resume: + raise ValueError("boom-resume") + + +class _State(TypedDict): + answer: str | None + + +def _build_interrupt_graph() -> Any: + def ask(state: _State) -> _State: + answer = interrupt("Provide value") + return {"answer": answer} + + builder = StateGraph(_State) + builder.add_node("ask", ask) + builder.add_edge(START, "ask") + return builder.compile(checkpointer=InMemorySaver()) + + +def test_graph_callbacks_interrupt_and_resume_sync() -> None: + graph = _build_interrupt_graph() + handler = _GraphEventHandler() + langchain_handler = _LangChainCustomEventHandler() + config = { + "configurable": {"thread_id": "graph-callback-sync"}, + "callbacks": [langchain_handler, handler], + } + + first = graph.invoke({"answer": None}, config) + assert "__interrupt__" in first + + assert len(handler.interrupt_events) == 1 + assert handler.interrupt_events[0].interrupts + assert isinstance(handler.interrupt_events[0].interrupts[0], Interrupt) + assert handler.interrupt_events[0].checkpoint_ns == () + assert langchain_handler.events == [] + + handler.resume_events.clear() + resumed = graph.invoke(Command(resume="done"), config) + assert resumed == {"answer": "done"} + + assert len(handler.resume_events) == 1 + assert handler.resume_events[0].checkpoint_ns == () + assert langchain_handler.events == [] + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +async def test_graph_callbacks_interrupt_and_resume_async() -> None: + graph = _build_interrupt_graph() + handler = _GraphEventHandler() + langchain_handler = _LangChainCustomEventHandler() + config = { + "configurable": {"thread_id": "graph-callback-async"}, + "callbacks": [langchain_handler, handler], + } + + first = await graph.ainvoke({"answer": None}, config) + assert "__interrupt__" in first + + assert len(handler.interrupt_events) == 1 + assert handler.interrupt_events[0].interrupts + assert isinstance(handler.interrupt_events[0].interrupts[0], Interrupt) + assert handler.interrupt_events[0].checkpoint_ns == () + assert langchain_handler.events == [] + + handler.resume_events.clear() + resumed = await graph.ainvoke(Command(resume="done"), config) + assert resumed == {"answer": "done"} + + assert len(handler.resume_events) == 1 + assert handler.resume_events[0].checkpoint_ns == () + assert langchain_handler.events == [] + + +def test_graph_callbacks_continue_when_interrupt_handler_raises_sync() -> None: + graph = _build_interrupt_graph() + raising_handler = _RaisingGraphEventHandler(raise_on_interrupt=True) + recording_handler = _GraphEventHandler() + + first = graph.invoke( + {"answer": None}, + { + "configurable": {"thread_id": "graph-callback-sync-raises"}, + "callbacks": [raising_handler, recording_handler], + }, + ) + + assert "__interrupt__" in first + assert len(recording_handler.interrupt_events) == 1 + + +def test_graph_callbacks_continue_when_resume_handler_raises_sync() -> None: + graph = _build_interrupt_graph() + raising_handler = _RaisingGraphEventHandler(raise_on_resume=True) + recording_handler = _GraphEventHandler() + config = { + "configurable": {"thread_id": "graph-callback-sync-raises-resume"}, + "callbacks": [raising_handler, recording_handler], + } + + first = graph.invoke({"answer": None}, config) + assert "__interrupt__" in first + + resumed = graph.invoke(Command(resume="done"), config) + assert resumed == {"answer": "done"} + assert len(recording_handler.resume_events) == 1 + + +def test_graph_callbacks_raise_error_propagates_sync() -> None: + graph = _build_interrupt_graph() + raising_handler = _RaisingGraphEventHandler( + raise_on_interrupt=True, + raise_error=True, + ) + + with pytest.raises(ValueError, match="boom-interrupt"): + graph.invoke( + {"answer": None}, + { + "configurable": {"thread_id": "graph-callback-sync-raise-error"}, + "callbacks": [raising_handler], + }, + ) + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +async def test_graph_callbacks_continue_when_handler_raises_async() -> None: + graph = _build_interrupt_graph() + raising_interrupt_handler = _AsyncRaisingGraphEventHandler(raise_on_interrupt=True) + recording_handler = _GraphEventHandler() + config = { + "configurable": {"thread_id": "graph-callback-async-raises-interrupt"}, + "callbacks": [raising_interrupt_handler, recording_handler], + } + + first = await graph.ainvoke({"answer": None}, config) + assert "__interrupt__" in first + assert len(recording_handler.interrupt_events) == 1 + + graph = _build_interrupt_graph() + raising_resume_handler = _AsyncRaisingGraphEventHandler(raise_on_resume=True) + recording_handler = _GraphEventHandler() + config = { + "configurable": {"thread_id": "graph-callback-async-raises-resume"}, + "callbacks": [raising_resume_handler, recording_handler], + } + + first = await graph.ainvoke({"answer": None}, config) + assert "__interrupt__" in first + resumed = await graph.ainvoke(Command(resume="done"), config) + assert resumed == {"answer": "done"} + assert len(recording_handler.resume_events) == 1 + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +async def test_graph_callbacks_raise_error_propagates_async() -> None: + graph = _build_interrupt_graph() + raising_handler = _AsyncRaisingGraphEventHandler( + raise_on_interrupt=True, + raise_error=True, + ) + + with pytest.raises(ValueError, match="boom-interrupt"): + await graph.ainvoke( + {"answer": None}, + { + "configurable": {"thread_id": "graph-callback-async-raise-error"}, + "callbacks": [raising_handler], + }, + ) + + +def test_graph_callbacks_accept_base_callback_manager() -> None: + graph = _build_interrupt_graph() + graph_handler = _GraphEventHandler() + custom_handler = _LangChainCustomEventHandler() + manager = CallbackManager.configure(inheritable_callbacks=[custom_handler]) + manager.add_handler(graph_handler) + + first = graph.invoke( + {"answer": None}, + { + "configurable": {"thread_id": "graph-callback-base-manager"}, + "callbacks": manager, + }, + ) + + assert "__interrupt__" in first + assert len(graph_handler.interrupt_events) == 1