Fix pydantic_graph.beta.GraphRun GeneratorExit handling (#3525)

Co-authored-by: Douwe Maan <[email protected]>

Author: michgur

pydantic_graph/pydantic_graph/beta/graph.py

@@ -43,9 +43,9 @@
 from pydantic_graph.nodes import BaseNode, End
 
 if sys.version_info < (3, 11):
-    from exceptiongroup import ExceptionGroup as ExceptionGroup  # pragma: lax no cover
+    from exceptiongroup import BaseExceptionGroup as BaseExceptionGroup  # pragma: lax no cover
 else:
-    ExceptionGroup = ExceptionGroup  # pragma: lax no cover
+    BaseExceptionGroup = BaseExceptionGroup  # pragma: lax no cover
 
 if TYPE_CHECKING:
     from pydantic_graph.beta.mermaid import StateDiagramDirection
@@ -970,7 +970,7 @@ def _unwrap_exception_groups():
     else:
         try:
             yield
-        except ExceptionGroup as e:
+        except BaseExceptionGroup as e:
             exception = e.exceptions[0]
             if exception.__cause__ is None:
                 # bizarrely, this prevents recursion errors when formatting the exception for logfire

tests/graph/beta/test_graph_execution.py

@@ -380,3 +380,20 @@ async def add_three(ctx: StepContext[None, None, int]) -> list[int]:
     graph = g.build()
     result = await graph.run()
     assert sorted(result) == [11, 12, 13]
+
+
+async def test_early_termination_from_nested_generator():
+    """Test that a generator wrapping an iteration can be terminated early."""
+    g = GraphBuilder()
+    g.add_edge(g.start_node, g.end_node)
+    graph = g.build()
+
+    async def stream_graph():
+        async with graph.iter() as run:
+            async for node in run:  # pragma: no branch
+                yield node
+
+    gen = stream_graph()
+    async for _ in gen:  # pragma: no branch
+        break
+    await gen.aclose()