Fix exception handling for imperative tasks (#4802)

libs/langgraph/langgraph/pregel/loop.py

@@ -135,7 +135,6 @@ P = ParamSpec("P")
 INPUT_DONE = object()
 INPUT_RESUMING = object()
 INPUT_SHOULD_VALIDATE = object()
-SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED)
 WritesT = Sequence[tuple[str, Any]]
 
 

libs/langgraph/langgraph/pregel/runner.py

@@ -56,6 +56,10 @@ EXCLUDED_FRAME_FNAMES = (
     "concurrent/futures/_base.py",
 )
 
+SKIP_RERAISE_SET: weakref.WeakSet[Union[concurrent.futures.Future, asyncio.Future]] = (
+    weakref.WeakSet()
+)
+
 
 class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
     event: E
@@ -165,7 +169,6 @@ class PregelRunner:
                             futures=weakref.ref(futures),
                             schedule_task=schedule_task,
                             submit=self.submit,
-                            reraise=reraise,
                         ),
                     },
                 )
@@ -207,7 +210,6 @@ class PregelRunner:
                         futures=weakref.ref(futures),
                         schedule_task=schedule_task,
                         submit=self.submit,
-                        reraise=reraise,
                     ),
                 },
                 __reraise_on_exit__=reraise,
@@ -302,7 +304,6 @@ class PregelRunner:
                             futures=weakref.ref(futures),
                             schedule_task=schedule_task,
                             submit=self.submit,
-                            reraise=reraise,
                             loop=loop,
                         ),
                     },
@@ -349,7 +350,6 @@ class PregelRunner:
                             futures=weakref.ref(futures),
                             schedule_task=schedule_task,
                             submit=self.submit,
-                            reraise=reraise,
                             loop=loop,
                         ),
                     },
@@ -434,7 +434,8 @@ class PregelRunner:
                 raise exception
             else:
                 # save error to checkpointer
-                self.put_writes()(task.id, [(ERROR, exception)])  # type: ignore[misc]
+                task.writes.append((ERROR, exception))
+                self.put_writes()(task.id, task.writes)  # type: ignore[misc]
         else:
             if self.node_finished and (
                 task.config is None or TAG_HIDDEN not in task.config.get("tags", [])
@@ -494,7 +495,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 +505,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 +532,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 +583,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 +616,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 +645,6 @@ def _acall(
             schedule_task=schedule_task,
             submit=submit,
             loop=loop,
-            reraise=reraise,
             stream=stream,
         ),
         loop,
@@ -669,7 +670,6 @@ async def _acall_impl(
     ],
     submit: weakref.ref[Submit],
     loop: asyncio.AbstractEventLoop,
-    reraise: bool = False,
     stream: bool = False,
 ) -> None:
     try:
@@ -726,17 +726,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.py

@@ -8769,3 +8769,30 @@ def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None:
 
     assert json.dumps(graph.get_graph().to_json(), indent=2) == snapshot
     assert graph.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+
+def test_imp_exception(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    @task()
+    def my_task(number: int):
+        time.sleep(0.1)
+        return number * 2
+
+    @task()
+    def task_with_exception(number: int):
+        time.sleep(0.1)
+        raise Exception("This is a test exception")
+
+    @entrypoint(checkpointer=sync_checkpointer)
+    def my_workflow(number: int):
+        my_task(number).result()
+        try:
+            task_with_exception(number).result()
+        except Exception as e:
+            print(f"Exception caught: {e}")
+        my_task(number).result()
+        return "done"
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    assert my_workflow.invoke(1, thread1) == "done"

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(0.1)
+        return number * 2
+
+    @task()
+    async def task_with_exception(number: int):
+        await asyncio.sleep(0.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"