WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(langgraph): add graph lifecycle callback handlers (#7429)

Browse Source 
## 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 <eyurtsev@gmail.com>
This commit is contained in:
Hunter Lovell
2026-04-10 20:43:48 +00:00
committed by GitHub
co-authored by Eugene Yurtsev
parent 1142ebf921
commit 3a5b5c9821
4 changed files with 806 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/callbacks.py
+412
View File
@@ -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,
)
libs/langgraph/langgraph/pregel/_loop.py
+60 -6
View File
@@ -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):
libs/langgraph/langgraph/pregel/main.py
+57 -1
View File
@@ -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,
libs/langgraph/tests/test_graph_callbacks.py
+277
View File
@@ -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