From fb377fc60ece2841ec18578b7bbedb38917896bb Mon Sep 17 00:00:00 2001 From: digit50 <93092234+digit50@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:11:12 +0300 Subject: [PATCH 1/2] Python: Pass stream errors to nested streams when closing them When a ResponseStream failed or was cancelled, close() shut nested ResponseStreams without the error, so their cleanup hooks treated the run as successful. With agent middleware this made the agent telemetry stream call get_final_response(), which ran every ContextProvider.after_run and saved the partial turn to history. close() now takes the terminal error and records it on itself and on nested streams while their cleanup hooks run. --- .../packages/core/agent_framework/_types.py | 14 ++++-- python/packages/core/tests/core/test_types.py | 48 +++++++++++++++++++ 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 3106e91ba40..1ebbc119be7 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -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 diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index fe4f5dc11c5..21e966cf6d0 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/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,48 @@ async def cancelling_transform(update: ChatResponseUpdate) -> ChatResponseUpdate assert events == ["iterator"] + @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 + ) -> 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) + 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} From 87944c25464ee06005133eefc24683284624c0dc Mon Sep 17 00:00:00 2001 From: digit50 <93092234+digit50@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:43:39 +0300 Subject: [PATCH 2/2] Python: Pass stream errors to nested streams in buffered mode too A transform failure or cancellation in a buffered outer stream went through _abort_buffered_materialization, which closed the stream without the error, so the nested stream still ran its cleanup as a normal completion. It now passes the recorded error to close(), and the regression test gains a buffered variant. --- python/packages/core/agent_framework/_types.py | 2 +- python/packages/core/tests/core/test_types.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 1ebbc119be7..7d45fa2803d 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/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 diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 21e966cf6d0..771e94fb799 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -5622,9 +5622,10 @@ 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 + 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.""" @@ -5657,6 +5658,8 @@ def failing_hook(update: AgentResponseUpdate) -> AgentResponseUpdate: 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