diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 3106e91ba40..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 @@ -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..771e94fb799 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,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}