Python: Pass stream errors to nested streams when closing them

python/packages/core/agent_framework/_types.py

@@ -4594,7 +4594,7 @@ async def _abort_buffered_materialization(self, exc: BaseException) -> None:
         self._stream_error = exc
         try:
             # close() also runs the hooks if releasing the iterator fails.
-            await self.close()
+            await self.close(exc)
         finally:
             self._stream_error = None
 
@@ -4780,22 +4780,25 @@ async def __anext__(self) -> UpdateT:
                 logger.warning("Stream cleanup hook failed while handling an existing exception.", exc_info=True)
             finally:
                 try:
-                    await self.close()
+                    await self.close(exc)
                 except Exception:
                     logger.warning("Stream cleanup failed while handling an existing exception.", exc_info=True)
             raise
 
-    async def close(self) -> None:
+    async def close(self, exc: BaseException | None = None) -> None:
         """Close the active iterator and run cleanup hooks.
 
         This method is idempotent and also closes nested ``ResponseStream`` wrappers.
+
+        Args:
+            exc: The error that ended the stream, recorded on this and nested streams while cleanup hooks run.
         """
         iterator_failed = False
         try:
             iterator: AsyncIterator[UpdateT] | None = self._iterator
             if iterator is not None:
                 if isinstance(iterator, ResponseStream):
-                    await cast(ResponseStream[UpdateT, Any], iterator).close()
+                    await cast(ResponseStream[UpdateT, Any], iterator).close(exc)
                 else:
                     close = getattr(iterator, "aclose", None)
                     if close is not None:
@@ -4806,7 +4809,10 @@ async def close(self) -> None:
         finally:
             self._consumed = True
             try:
-                await self._run_cleanup_hooks()
+                if exc is None:
+                    await self._run_cleanup_hooks()
+                else:
+                    await self._handle_stream_error(exc)
             except Exception:
                 if not iterator_failed:
                     raise

python/packages/core/tests/core/test_types.py

@@ -18,20 +18,26 @@
 from pytest import fixture, mark, raises
 
 from agent_framework import (
+    Agent,
     AgentResponse,
     AgentResponseUpdate,
     Annotation,
+    BaseChatClient,
     ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
     ComputerSafetyCheck,
     Content,
+    ContextProvider,
+    FunctionInvocationLayer,
     FunctionTool,
+    InMemoryHistoryProvider,
     Message,
     ResponseStream,
     TextSpanRegion,
     ToolMode,
     UsageDetails,
+    agent_middleware,
     detect_media_type_from_base64,
     get_namespaced_value,
     merge_additional_properties,
@@ -5616,6 +5622,51 @@ async def cancelling_transform(update: ChatResponseUpdate) -> ChatResponseUpdate
 
         assert events == ["iterator"]
 
+    @pytest.mark.parametrize("buffered", [False, True])
+    @pytest.mark.parametrize("failure", [RuntimeError("hook failed"), asyncio.CancelledError()])
+    async def test_outer_stream_error_skips_after_run_with_agent_middleware(
+        self, span_exporter: Any, failure: BaseException, buffered: bool
+    ) -> None:
+        """An error in a middleware-wrapped agent stream reaches the inner stream, so after_run does not fire."""
+
+        class StreamingClient(FunctionInvocationLayer, BaseChatClient):
+            def _inner_get_response(self, *, messages: Any, stream: bool, options: Any, **kwargs: Any) -> Any:
+                async def updates() -> AsyncIterable[ChatResponseUpdate]:
+                    for i in range(3):
+                        yield ChatResponseUpdate(contents=[Content.from_text(f"part{i} ")], role="assistant")
+
+                return ResponseStream(updates(), finalizer=ChatResponse.from_updates)
+
+        class RecordingProvider(ContextProvider):
+            def __init__(self) -> None:
+                super().__init__("recorder")
+                self.after_run_calls = 0
+
+            async def after_run(self, *, agent: Any, session: Any, context: Any, state: Any) -> None:
+                self.after_run_calls += 1
+
+        @agent_middleware
+        async def passthrough(context: Any, call_next: Any) -> None:
+            await call_next()
+
+        def failing_hook(update: AgentResponseUpdate) -> AgentResponseUpdate:
+            raise failure
+
+        recorder = RecordingProvider()
+        history = InMemoryHistoryProvider()
+        agent = Agent(client=StreamingClient(), context_providers=[history, recorder], middleware=[passthrough])
+        session = agent.create_session()
+
+        stream = agent.run("hi", stream=True, session=session).with_transform_hook(failing_hook)
+        if buffered:
+            stream = stream.buffer_updates()
+        with pytest.raises(type(failure)):
+            async for _ in stream:
+                pass
+
+        assert recorder.after_run_calls == 0
+        assert await history.get_messages(session.session_id, state=session.state.get(history.source_id, {})) == []
+
     async def test_cleanup_hook_called_after_iteration(self) -> None:
         """Cleanup hook is called after iteration completes."""
         cleanup_called = {"value": False}