diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 412266ad3..d4c075a99 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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]] diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 96a256ad9..bbdb90bae 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/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) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 5558daa68..81b5eadf8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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" diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3c95a2cec..f77612241 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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"