From dfdb4d00ea3efb6d0ec3957aedad2100380246e0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 1 Jul 2024 18:12:37 +0100 Subject: [PATCH] Mark cancelled as seen --- libs/langgraph/langgraph/pregel/__init__.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 784af417f..a86674ca3 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1510,7 +1510,11 @@ class Pregel( except NameError: pass # wait for all background tasks to finish - await asyncio.shield(asyncio.gather(*bg)) + fut = asyncio.gather(*bg) + # mark the exception as retrieved + fut.add_done_callback(_mark_cancelled_as_seen) + # wait for the future to finish, shielded from cancellation + await asyncio.shield(fut) def invoke( self, @@ -2021,3 +2025,10 @@ def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any] yield (mode, chunk) else: yield from iter + + +def _mark_cancelled_as_seen(fut: concurrent.futures.Future) -> None: + try: + fut.exception() + except (asyncio.CancelledError, concurrent.futures.CancelledError): + pass