fix(langgraph): harden graph callback dispatch semantics

Author: hntrl

libs/langgraph/langgraph/callbacks.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
 
 import asyncio
-import inspect
-from collections.abc import Coroutine, Sequence
+import atexit
+import functools
+import logging
+from collections.abc import Awaitable, Callable, Coroutine, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
+from functools import partial
 from typing import Any, TypeAlias, cast
 from uuid import UUID
 
@@ -21,6 +26,42 @@
 )
 
 GRAPH_CALLBACKS_KEY = "graph_callbacks"
+logger = logging.getLogger(__name__)
+
+
+@functools.lru_cache(maxsize=1)
+def _executor() -> ThreadPoolExecutor:
+    executor = ThreadPoolExecutor(max_workers=10)
+    atexit.register(executor.shutdown, wait=True)
+    return executor
+
+
+def _coerce_coroutine(awaitable: Awaitable[Any]) -> Coroutine[Any, Any, Any]:
+    if asyncio.iscoroutine(awaitable):
+        return awaitable
+
+    async def _await() -> Any:
+        return await awaitable
+
+    return _await()
+
+
+def _run_coroutine_with_new_loop(coro: Coroutine[Any, Any, Any]) -> None:
+    if hasattr(asyncio, "Runner"):
+        with asyncio.Runner() as runner:
+            runner.run(coro)
+            while pending := asyncio.all_tasks(runner.get_loop()):
+                runner.run(asyncio.wait(pending))
+    else:
+        loop = asyncio.new_event_loop()
+        try:
+            asyncio.set_event_loop(loop)
+            loop.run_until_complete(coro)
+            while pending := asyncio.all_tasks(loop):
+                loop.run_until_complete(asyncio.wait(pending))
+        finally:
+            asyncio.set_event_loop(None)
+            loop.close()
 
 
 class GraphCallbackHandler(BaseCallbackHandler):
@@ -131,6 +172,74 @@ def configure(
                 )
         return cls(handlers, run_id=run_id)
 
+    @staticmethod
+    def _maybe_raise_handler_error(
+        handler: GraphCallbackHandler,
+        event_name: str,
+        error: Exception,
+    ) -> None:
+        if getattr(handler, "raise_error", False):
+            raise error
+        logger.exception(
+            "Ignoring graph callback handler error in %s for %s",
+            event_name,
+            handler.__class__.__name__,
+            exc_info=error,
+        )
+
+    def _drain_awaitable(
+        self,
+        handler: GraphCallbackHandler,
+        event_name: str,
+        awaitable: Awaitable[Any],
+    ) -> None:
+        coro = _coerce_coroutine(awaitable)
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            try:
+                _run_coroutine_with_new_loop(coro)
+            except Exception as exc:
+                self._maybe_raise_handler_error(handler, event_name, exc)
+        else:
+            future = _executor().submit(
+                cast(Callable[..., Any], copy_context().run),
+                _run_coroutine_with_new_loop,
+                coro,
+            )
+            try:
+                future.result()
+            except Exception as exc:
+                self._maybe_raise_handler_error(handler, event_name, exc)
+
+    async def _ainvoke_handler(
+        self,
+        handler: GraphCallbackHandler,
+        callback: Callable[..., Any],
+        event_name: str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            if asyncio.iscoroutinefunction(callback):
+                await callback(*args, **kwargs)
+            elif handler.run_inline:
+                result = callback(*args, **kwargs)
+                if isinstance(result, Awaitable):
+                    await result
+            else:
+                result = await asyncio.get_running_loop().run_in_executor(
+                    None,
+                    cast(
+                        Callable[..., Any],
+                        partial(copy_context().run, callback, *args, **kwargs),
+                    ),
+                )
+                if isinstance(result, Awaitable):
+                    await result
+        except Exception as exc:
+            self._maybe_raise_handler_error(handler, event_name, exc)
+
     def on_interrupt(
         self,
         interrupts: Sequence[Interrupt],
@@ -142,16 +251,19 @@ def on_interrupt(
     ) -> None:
         for handler in self.handlers:
             graph_handler = cast(GraphCallbackHandler, handler)
-            result = graph_handler.on_interrupt(
-                interrupts,
-                run_id=self.run_id,
-                status=status,
-                checkpoint_id=checkpoint_id,
-                checkpoint_ns=checkpoint_ns,
-                is_nested=is_nested,
-            )
-            if inspect.iscoroutine(result):
-                self._drain_coroutine(result)
+            try:
+                result = graph_handler.on_interrupt(
+                    interrupts,
+                    run_id=self.run_id,
+                    status=status,
+                    checkpoint_id=checkpoint_id,
+                    checkpoint_ns=checkpoint_ns,
+                    is_nested=is_nested,
+                )
+                if isinstance(result, Awaitable):
+                    self._drain_awaitable(graph_handler, "on_interrupt", result)
+            except Exception as exc:
+                self._maybe_raise_handler_error(graph_handler, "on_interrupt", exc)
 
     def on_resume(
         self,
@@ -163,15 +275,18 @@ def on_resume(
     ) -> None:
         for handler in self.handlers:
             graph_handler = cast(GraphCallbackHandler, handler)
-            result = graph_handler.on_resume(
-                run_id=self.run_id,
-                status=status,
-                checkpoint_id=checkpoint_id,
-                checkpoint_ns=checkpoint_ns,
-                is_nested=is_nested,
-            )
-            if inspect.iscoroutine(result):
-                self._drain_coroutine(result)
+            try:
+                result = graph_handler.on_resume(
+                    run_id=self.run_id,
+                    status=status,
+                    checkpoint_id=checkpoint_id,
+                    checkpoint_ns=checkpoint_ns,
+                    is_nested=is_nested,
+                )
+                if isinstance(result, Awaitable):
+                    self._drain_awaitable(graph_handler, "on_resume", result)
+            except Exception as exc:
+                self._maybe_raise_handler_error(graph_handler, "on_resume", exc)
 
     async def aon_interrupt(
         self,
@@ -182,18 +297,40 @@ async def aon_interrupt(
         checkpoint_ns: tuple[str, ...],
         is_nested: bool,
     ) -> None:
-        for handler in self.handlers:
-            graph_handler = cast(GraphCallbackHandler, handler)
-            result = graph_handler.on_interrupt(
+        inline_handlers = [
+            cast(GraphCallbackHandler, handler)
+            for handler in self.handlers
+            if handler.run_inline
+        ]
+        for graph_handler in inline_handlers:
+            await self._ainvoke_handler(
+                graph_handler,
+                graph_handler.on_interrupt,
+                "on_interrupt",
                 interrupts,
                 run_id=self.run_id,
                 status=status,
                 checkpoint_id=checkpoint_id,
                 checkpoint_ns=checkpoint_ns,
                 is_nested=is_nested,
             )
-            if inspect.isawaitable(result):
-                await result
+        await asyncio.gather(
+            *(
+                self._ainvoke_handler(
+                    cast(GraphCallbackHandler, handler),
+                    cast(GraphCallbackHandler, handler).on_interrupt,
+                    "on_interrupt",
+                    interrupts,
+                    run_id=self.run_id,
+                    status=status,
+                    checkpoint_id=checkpoint_id,
+                    checkpoint_ns=checkpoint_ns,
+                    is_nested=is_nested,
+                )
+                for handler in self.handlers
+                if not handler.run_inline
+            )
+        )
 
     async def aon_resume(
         self,
@@ -203,26 +340,38 @@ async def aon_resume(
         checkpoint_ns: tuple[str, ...],
         is_nested: bool,
     ) -> None:
-        for handler in self.handlers:
-            graph_handler = cast(GraphCallbackHandler, handler)
-            result = graph_handler.on_resume(
+        inline_handlers = [
+            cast(GraphCallbackHandler, handler)
+            for handler in self.handlers
+            if handler.run_inline
+        ]
+        for graph_handler in inline_handlers:
+            await self._ainvoke_handler(
+                graph_handler,
+                graph_handler.on_resume,
+                "on_resume",
                 run_id=self.run_id,
                 status=status,
                 checkpoint_id=checkpoint_id,
                 checkpoint_ns=checkpoint_ns,
                 is_nested=is_nested,
             )
-            if inspect.isawaitable(result):
-                await result
-
-    @staticmethod
-    def _drain_coroutine(coro: Coroutine[Any, Any, Any]) -> None:
-        try:
-            loop = asyncio.get_running_loop()
-        except RuntimeError:
-            asyncio.run(coro)
-        else:
-            loop.create_task(coro)
+        await asyncio.gather(
+            *(
+                self._ainvoke_handler(
+                    cast(GraphCallbackHandler, handler),
+                    cast(GraphCallbackHandler, handler).on_resume,
+                    "on_resume",
+                    run_id=self.run_id,
+                    status=status,
+                    checkpoint_id=checkpoint_id,
+                    checkpoint_ns=checkpoint_ns,
+                    is_nested=is_nested,
+                )
+                for handler in self.handlers
+                if not handler.run_inline
+            )
+        )
 
 
 GraphCallbacks: TypeAlias = (

libs/langgraph/langgraph/pregel/_loop.py

@@ -215,7 +215,7 @@ class PregelLoop:
     tasks: dict[str, PregelExecutableTask]
     output: None | dict[str, Any] | Any = None
     updated_channels: set[str] | None = None
-    _graph_lifecycle_events: list[GraphLifecycleEvent]
+    _graph_lifecycle_events: deque[GraphLifecycleEvent]
 
     # public
 
@@ -265,7 +265,7 @@ def __init__(
         self.retry_policy = retry_policy
         self.cache_policy = cache_policy
         self.durability = durability
-        self._graph_lifecycle_events = []
+        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)
@@ -334,10 +334,10 @@ def _push_graph_lifecycle_event(
             )
         )
 
-    def shift_graph_lifecycle_event(self) -> GraphLifecycleEvent | None:
+    def _pop_lifecycle_event(self) -> GraphLifecycleEvent | None:
         if not self._graph_lifecycle_events:
             return None
-        return self._graph_lifecycle_events.pop(0)
+        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."""
@@ -923,9 +923,8 @@ def _suppress_interrupt(
             self._put_checkpoint(self.checkpoint_metadata)
             self._put_pending_writes()
         # suppress interrupt
-        suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
-        if suppress:
-            interrupt = cast(GraphInterrupt, exc_value)
+        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
@@ -954,12 +953,11 @@ def _suppress_interrupt(
                         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)
@@ -1177,7 +1175,7 @@ def put_writes(self, task_id: str, writes: WritesT) -> None:
     # context manager
 
     def __enter__(self) -> Self:
-        self._graph_lifecycle_events = []
+        self._graph_lifecycle_events = deque()
         if not self.checkpointer:
             saved = None
         elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
@@ -1377,7 +1375,7 @@ def put_writes(self, task_id: str, writes: WritesT) -> None:
     # context manager
 
     async def __aenter__(self) -> Self:
-        self._graph_lifecycle_events = []
+        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

@@ -2675,7 +2675,7 @@ def stream_writer(c: Any) -> None:
             _state_mapper = self._state_mapper if version == "v2" else None
 
             def emit_graph_lifecycle_events(loop: SyncPregelLoop) -> None:
-                while (event := loop.shift_graph_lifecycle_event()) is not None:
+                while (event := loop._pop_lifecycle_event()) is not None:
                     if event.kind == "resume":
                         graph_callback_manager.on_resume(
                             status=event.status,
@@ -3073,7 +3073,7 @@ def stream_writer(c: Any) -> None:
             _state_mapper = self._state_mapper if version == "v2" else None
 
             async def aemit_graph_lifecycle_events(loop: AsyncPregelLoop) -> None:
-                while (event := loop.shift_graph_lifecycle_event()) is not None:
+                while (event := loop._pop_lifecycle_event()) is not None:
                     if event.kind == "resume":
                         await graph_callback_manager.aon_resume(
                             status=event.status,