Fix exception handling for imperative tasks

- These exceptions should not be re-raised at end of tick, given they're handled explicitly by the developer in their entrypoint

libs/langgraph/langgraph/pregel/runner.py

@@ -56,6 +56,8 @@ EXCLUDED_FRAME_FNAMES = (
     "concurrent/futures/_base.py",
 )
 
+SKIP_RERAISE_SET = weakref.WeakSet()
+
 
 class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
     event: E
@@ -165,7 +167,6 @@ class PregelRunner:
                             futures=weakref.ref(futures),
                             schedule_task=schedule_task,
                             submit=self.submit,
-                            reraise=reraise,
                         ),
                     },
                 )
@@ -207,7 +208,6 @@ class PregelRunner:
                         futures=weakref.ref(futures),
                         schedule_task=schedule_task,
                         submit=self.submit,
-                        reraise=reraise,
                     ),
                 },
                 __reraise_on_exit__=reraise,
@@ -302,7 +302,6 @@ class PregelRunner:
                             futures=weakref.ref(futures),
                             schedule_task=schedule_task,
                             submit=self.submit,
-                            reraise=reraise,
                             loop=loop,
                         ),
                     },
@@ -349,7 +348,6 @@ class PregelRunner:
                             futures=weakref.ref(futures),
                             schedule_task=schedule_task,
                             submit=self.submit,
-                            reraise=reraise,
                             loop=loop,
                         ),
                     },
@@ -494,7 +492,8 @@ def _panic_or_proceed(
     interrupts: list[GraphInterrupt] = []
     while done:
         # if any task failed
-        if exc := _exception(done.pop()):
+        fut = done.pop()
+        if exc := _exception(fut):
             # cancel all pending tasks
             while inflight:
                 inflight.pop().cancel()
@@ -503,7 +502,7 @@ def _panic_or_proceed(
                 if isinstance(exc, GraphInterrupt):
                     # collect interrupts
                     interrupts.append(exc)
-                else:
+                elif fut not in SKIP_RERAISE_SET:
                     raise exc
     # raise combined interrupts
     if interrupts:
@@ -530,7 +529,6 @@ def _call(
         [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
     ],
     submit: weakref.ref[Submit],
-    reraise: bool,
 ) -> concurrent.futures.Future[Any]:
     if asyncio.iscoroutinefunction(func):
         raise RuntimeError("In an sync context async tasks cannot be called")
@@ -582,14 +580,16 @@ def _call(
                         callbacks=callbacks,
                         schedule_task=schedule_task,
                         submit=submit,
-                        reraise=reraise,
                     ),
                 },
-                __reraise_on_exit__=reraise,
+                __reraise_on_exit__=False,
                 # starting a new task in the next tick ensures
                 # updates from this tick are committed/streamed first
                 __next_tick__=True,
             )
+            # exceptions for call() tasks are raised into the parent task
+            # so we should not re-raise at the end of the tick
+            SKIP_RERAISE_SET.add(fut)
             futures()[fut] = next_task  # type: ignore[index]
     fut = cast(Union[asyncio.Future, concurrent.futures.Future], fut)
     # return a chained future to ensure commit() callback is called
@@ -613,7 +613,6 @@ def _acall(
     ],
     submit: weakref.ref[Submit],
     loop: asyncio.AbstractEventLoop,
-    reraise: bool = False,
     stream: bool = False,
 ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
     # return a chained future to ensure commit() callback is called
@@ -643,7 +642,6 @@ def _acall(
             schedule_task=schedule_task,
             submit=submit,
             loop=loop,
-            reraise=reraise,
             stream=stream,
         ),
         loop,
@@ -669,7 +667,6 @@ async def _acall_impl(
     ],
     submit: weakref.ref[Submit],
     loop: asyncio.AbstractEventLoop,
-    reraise: bool = False,
     stream: bool = False,
 ) -> None:
     try:
@@ -726,17 +723,19 @@ async def _acall_impl(
                                 schedule_task=schedule_task,
                                 submit=submit,
                                 loop=loop,
-                                reraise=reraise,
                             ),
                         },
-                        __name__=task().name,  # type: ignore[union-attr]
+                        __name__=next_task.name,
                         __cancel_on_exit__=True,
-                        __reraise_on_exit__=reraise,
+                        __reraise_on_exit__=False,
                         # starting a new task in the next tick ensures
                         # updates from this tick are committed/streamed first
                         __next_tick__=True,
                     ),
                 )
+                # exceptions for call() tasks are raised into the parent task
+                # so we should not re-raise at the end of the tick
+                SKIP_RERAISE_SET.add(fut)
                 futures()[fut] = next_task  # type: ignore[index]
         if fut is not None:
             chain_future(fut, destination)

libs/langgraph/tests/test_pregel_async.py

@@ -9148,3 +9148,31 @@ async def test_draw_invalid():
             {"source": "nothing", "target": "__end__"},
         ],
     }
+
+
+@NEEDS_CONTEXTVARS
+async def test_imp_exception(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    @task()
+    async def my_task(number: int):
+        await asyncio.sleep(1)
+        return number * 2
+
+    @task()
+    async def task_with_exception(number: int):
+        await asyncio.sleep(1)
+        raise Exception("This is a test exception")
+
+    @entrypoint(checkpointer=async_checkpointer)
+    async def my_workflow(number: int):
+        await my_task(number)
+        try:
+            await task_with_exception(number)
+        except Exception as e:
+            print(f"Exception caught: {e}")
+        await my_task(number)
+        return "done"
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    assert await my_workflow.ainvoke(1, thread1) == "done"