From 57a72d4610f33a8d726126b923eac5ded09b2b67 Mon Sep 17 00:00:00 2001 From: quifox <289420841+quifox@users.noreply.github.com> Date: Tue, 6 Oct 2026 01:58:53 +0000 Subject: [PATCH 1/3] Python: Close delegated provider streams on public stream release --- .../ag-ui/agent_framework_ag_ui/_client.py | 60 +- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 156 ++++++ .../packages/core/agent_framework/_tools.py | 128 ++--- .../core/test_function_invocation_logic.py | 521 +++++++++++++++++- 4 files changed, 774 insertions(+), 91 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index f47bb020333..3b466b58cd4 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -503,7 +503,7 @@ async def _streaming_impl( available_interrupts = options.get("available_interrupts", options.get("availableInterrupts")) - async for event in self._http_service.post_run( + events = self._http_service.post_run( thread_id=thread_id, run_id=run_id, messages=agui_messages, @@ -511,29 +511,35 @@ async def _streaming_impl( tools=agui_tools, available_interrupts=_serialize_available_interrupts(cast(Sequence[Any] | None, available_interrupts)), resume=_serialize_resume(options.get("resume")), - ): - logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}") - update = converter.convert_event(event) - if update is not None: - logger.debug( - "[AGUIChatClient] Converted update", - extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]}, - ) - # Distinguish client vs server tools - for i, content in enumerate(update.contents): - if content.type == "function_call": - logger.debug( - f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}" - ) - if content.name in client_tool_set: - # Client tool - let function invocation execute it - if not content.additional_properties: - content.additional_properties = {} - content.additional_properties["agui_thread_id"] = thread_id - else: - # Server tool - wrap so function invocation ignores it - logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}") - self._register_server_tool_placeholder(content.name) # type: ignore[arg-type] - update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore - - yield update + ) + try: + async for event in events: + logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}") + update = converter.convert_event(event) + if update is not None: + logger.debug( + "[AGUIChatClient] Converted update", + extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]}, + ) + # Distinguish client vs server tools + for i, content in enumerate(update.contents): + if content.type == "function_call": + logger.debug( + f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}" + ) + if content.name in client_tool_set: + # Client tool - let function invocation execute it + if not content.additional_properties: + content.additional_properties = {} + content.additional_properties["agui_thread_id"] = thread_id + else: + # Server tool - wrap so function invocation ignores it + logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}") + self._register_server_tool_placeholder(content.name) # type: ignore[arg-type] + update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore + + yield update + finally: + close = getattr(events, "aclose", None) + if close is not None: + await close() diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 679cfb11eec..b9d1d578208 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -12,6 +12,8 @@ import pytest from ag_ui.core import Interrupt, ResumeEntry from agent_framework import ( + Agent, + AgentResponseUpdate, ChatOptions, ChatResponse, ChatResponseUpdate, @@ -57,6 +59,51 @@ def inner_get_response( return self._inner_get_response(messages=messages, options=options, stream=stream) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] +class _GatedSSEStream(httpx.AsyncByteStream): + """Yield one SSE chunk and keep the next event gated until another pull.""" + + def __init__(self) -> None: + self._tail_release = asyncio.Event() + self._iterator: AsyncGenerator[bytes, None] | None = None + self.chunk_indexes_yielded: list[int] = [] + self.tail_requested = False + self.tail_yielded = False + self.exhausted = False + self.aclose_calls = 0 + self.closed = False + self.iterator_finally = False + events = [ + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "first"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "later"}, + ] + self._chunks = [f"data: {json.dumps(event)}\n\n".encode() for event in events] + + def __aiter__(self) -> AsyncGenerator[bytes, None]: + if self._iterator is None: + self._iterator = self._iterate() + return self._iterator + + async def _iterate(self) -> AsyncGenerator[bytes, None]: + try: + self.chunk_indexes_yielded.append(0) + yield self._chunks[0] + self.tail_requested = True + await self._tail_release.wait() + self.chunk_indexes_yielded.append(1) + self.tail_yielded = True + yield self._chunks[1] + self.exhausted = True + finally: + self.iterator_finally = True + + async def aclose(self) -> None: + self.aclose_calls += 1 + self.closed = True + self._tail_release.set() + if self._iterator is not None: + await self._iterator.aclose() + + class TestAGUIChatClient: """Test suite for AGUIChatClient.""" @@ -703,6 +750,115 @@ async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str assert first_content.text == "Hello" assert second_content.text == " world" + @pytest.mark.parametrize( + ("caller", "function_invocation_enabled", "close_mode"), + [ + pytest.param("client", True, "explicit_close", id="client-fi_default-explicit_close"), + pytest.param("client", True, "async_with_break", id="client-fi_default-async_with_break"), + pytest.param("client", False, "explicit_close", id="client-fi_disabled-explicit_close"), + pytest.param("client", False, "async_with_break", id="client-fi_disabled-async_with_break"), + pytest.param("agent", True, "explicit_close", id="agent-fi_default-explicit_close"), + pytest.param("agent", True, "async_with_break", id="agent-fi_default-async_with_break"), + pytest.param("agent", False, "explicit_close", id="agent-fi_disabled-explicit_close"), + pytest.param("agent", False, "async_with_break", id="agent-fi_disabled-async_with_break"), + ], + ) + async def test_public_stream_close_releases_http_response( + self, + caller: str, + function_invocation_enabled: bool, + close_mode: str, + ) -> None: + """Closing a public stream releases its HTTP response and preserves supplied clients.""" + body = _GatedSSEStream() + response: httpx.Response | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal response + if request.method == "GET": + return httpx.Response(204, request=request) + await request.aread() + response = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=body, + request=request, + ) + return response + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client_options: dict[str, Any] = {} + if not function_invocation_enabled: + client_options["function_invocation_configuration"] = {"enabled": False} + client = AGUIChatClient( + endpoint="https://agui.example.test/run", + http_client=http_client, + **client_options, + ) + assert client.function_invocation_configuration["enabled"] is function_invocation_enabled + + agent = Agent(client=client) if caller == "agent" else None + messages = [Message(role="user", contents=["Question"])] + stream = ( + agent.run(messages, stream=True) if agent is not None else client.get_response(messages, stream=True) + ) + + def assert_first_chunk_only() -> None: + assert response is not None + assert not response.is_closed + assert body.aclose_calls == 0 + assert body.chunk_indexes_yielded == [0] + assert not body.tail_requested + assert not body.tail_yielded + assert not body.exhausted + + try: + if close_mode == "explicit_close": + first = await asyncio.wait_for(anext(stream), timeout=5.0) + assert isinstance(first, (ChatResponseUpdate, AgentResponseUpdate)) + assert first.text == "first" + assert_first_chunk_only() + await stream.close() + else: + async with stream: + async for update in stream: + assert isinstance(update, (ChatResponseUpdate, AgentResponseUpdate)) + assert update.text == "first" + assert_first_chunk_only() + break + + assert response is not None + assert response.is_closed + assert body.aclose_calls == 1 + assert body.closed + assert body.iterator_finally + assert body.chunk_indexes_yielded == [0] + assert not body.tail_yielded + assert not body.exhausted + + await stream.close() + assert body.aclose_calls == 1 + + if agent is not None: + await agent.close() + else: + await client.close() + assert not http_client.is_closed + health_response = await http_client.get("https://agui.example.test/health") + assert health_response.status_code == 204 + finally: + try: + await stream.close() + finally: + try: + if response is not None and not response.is_closed: + await response.aclose() + finally: + if agent is not None: + await agent.close() + else: + await client.close() + async def test_get_response_non_streaming(self, monkeypatch: MonkeyPatch) -> None: """Test non-streaming response method.""" mock_events = [ diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 602da9dc4cb..a34b5523d64 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -5358,64 +5358,65 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non streamed_names_by_call_id: dict[str, str] = {} last_streamed_identity: tuple[str, str] | None = None warned_empty_call_ids: set[str] = set() - async for update in _iterate_provider_stream(inner_stream, on_invalidated): - for content in update.contents: - if content.type != "function_call": - continue - if not _is_actionable_function_call(content): - continue - had_occurrence_id = content.id is not None - provider_call_id = content.call_id - identity = streamed_identities_by_call_id.get(provider_call_id) if provider_call_id else None - if ( - identity is not None - and provider_call_id is not None - and content.id is None - and content.name - and ( - streamed_names_by_call_id.get(provider_call_id) != content.name - or isinstance(content.arguments, Mapping) - ) - ): - identity = None - if identity is None and not provider_call_id and not content.name: - identity = last_streamed_identity - - if identity is None: - occurrence_id = content.id or _generate_function_call_occurrence_id() - effective_call_id = provider_call_id or ("" if had_occurrence_id else occurrence_id) - else: - occurrence_id, effective_call_id = identity - if content.id is not None: - occurrence_id = content.id - if provider_call_id: - effective_call_id = provider_call_id - - content.id = occurrence_id - if not content.call_id and not had_occurrence_id: - content.call_id = effective_call_id - if identity is None and occurrence_id not in warned_empty_call_ids: - warnings.warn( - "An actionable function_call had an empty call_id. Agent Framework used its generated " - "Content.id for local correlation. Providers should supply and preserve their service " - "call_id; this fallback will be removed in a future release.", - FutureWarning, - stacklevel=3, + async with inner_stream: + async for update in _iterate_provider_stream(inner_stream, on_invalidated): + for content in update.contents: + if content.type != "function_call": + continue + if not _is_actionable_function_call(content): + continue + had_occurrence_id = content.id is not None + provider_call_id = content.call_id + identity = streamed_identities_by_call_id.get(provider_call_id) if provider_call_id else None + if ( + identity is not None + and provider_call_id is not None + and content.id is None + and content.name + and ( + streamed_names_by_call_id.get(provider_call_id) != content.name + or isinstance(content.arguments, Mapping) ) - warned_empty_call_ids.add(occurrence_id) - identity = (occurrence_id, effective_call_id) - if effective_call_id: - streamed_identities_by_call_id[effective_call_id] = identity - if content.name: - streamed_names_by_call_id[effective_call_id] = content.name - last_streamed_identity = identity - if drop_unexecutable_calls: - update = _drop_unexecutable_tool_contents_from_update(update) - if update is None: - continue - yield update - - response = await _await_provider_call(inner_stream.get_final_response, on_invalidated) + ): + identity = None + if identity is None and not provider_call_id and not content.name: + identity = last_streamed_identity + + if identity is None: + occurrence_id = content.id or _generate_function_call_occurrence_id() + effective_call_id = provider_call_id or ("" if had_occurrence_id else occurrence_id) + else: + occurrence_id, effective_call_id = identity + if content.id is not None: + occurrence_id = content.id + if provider_call_id: + effective_call_id = provider_call_id + + content.id = occurrence_id + if not content.call_id and not had_occurrence_id: + content.call_id = effective_call_id + if identity is None and occurrence_id not in warned_empty_call_ids: + warnings.warn( + "An actionable function_call had an empty call_id. Agent Framework used its " + "generated Content.id for local correlation. Providers should supply and preserve " + "their service call_id; this fallback will be removed in a future release.", + FutureWarning, + stacklevel=3, + ) + warned_empty_call_ids.add(occurrence_id) + identity = (occurrence_id, effective_call_id) + if effective_call_id: + streamed_identities_by_call_id[effective_call_id] = identity + if content.name: + streamed_names_by_call_id[effective_call_id] = content.name + last_streamed_identity = identity + if drop_unexecutable_calls: + update = _drop_unexecutable_tool_contents_from_update(update) + if update is None: + continue + yield update + + response = await _await_provider_call(inner_stream.get_final_response, on_invalidated) fallback_added = False if options.get("tool_choice") == "none" and budget_state.get("truncated"): fallback_added = _ensure_function_invocation_limit_fallback_response(response) @@ -5510,12 +5511,13 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non client_kwargs=request_kwargs, ), ) - async for update in _iterate_provider_stream(final_inner_stream, on_invalidated): - update = _drop_unexecutable_tool_contents_from_update(update) - if update is None: - continue - yield update - final_response = await _await_provider_call(final_inner_stream.get_final_response, on_invalidated) + async with final_inner_stream: + async for update in _iterate_provider_stream(final_inner_stream, on_invalidated): + update = _drop_unexecutable_tool_contents_from_update(update) + if update is None: + continue + yield update + final_response = await _await_provider_call(final_inner_stream.get_final_response, on_invalidated) fallback_added = _ensure_function_invocation_limit_fallback_response(final_response) self._update_function_invocation_continuation_state( request_kwargs, diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 243dd0dcff0..cde80da7735 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -6,7 +6,7 @@ import math import threading import warnings -from collections.abc import AsyncIterable, Awaitable, Callable, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Sequence from typing import Any, Literal from unittest.mock import Mock @@ -11194,3 +11194,522 @@ async def test_unbound_local_approval_response_is_filtered_before_mixed_batch_va assert executed == [] # The warning proves the response was filtered before the completeness check reported it missing. assert any("tool-approval response" in record.message for record in caplog.records) + + +def _make_lifecycle_noop_tool() -> FunctionTool: + @tool(name="lifecycle_noop", approval_mode="never_require") + def lifecycle_noop() -> str: + return "done" + + return lifecycle_noop + + +def _lifecycle_public_stream(client: Any, surface: Literal["client", "agent"], tools: list[FunctionTool]) -> Any: + options: ChatOptions[Any] = {"tool_choice": "auto"} + if surface == "client": + options["tools"] = tools + return client.get_response([Message(role="user", contents=["run"])], stream=True, options=options) + + agent = Agent(client=client, tools=tools) + return agent.run("run", stream=True, options=options) + + +def _install_lifecycle_provider_streams( + client: Any, + monkeypatch: pytest.MonkeyPatch, + retained_provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]], + retained_provider_sources: list[AsyncGenerator[ChatResponseUpdate, None]], + turns: Sequence[Sequence[ChatResponseUpdate]], +) -> tuple[ + list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]], + list[AsyncGenerator[ChatResponseUpdate, None]], + list[int], + list[int], + list[int], + list[str | None], +]: + provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]] = [] + provider_sources: list[AsyncGenerator[ChatResponseUpdate, None]] = [] + started_turns: list[int] = [] + generator_finally_turns: list[int] = [] + cleanup_hook_turns: list[int] = [] + requested_tool_choices: list[str | None] = [] + + def get_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: + turn_index = len(provider_streams) + requested_tool_choices.append(options.get("tool_choice")) + + async def updates() -> AsyncGenerator[ChatResponseUpdate, None]: + client.call_count += 1 + started_turns.append(turn_index) + try: + for update in turns[turn_index]: + yield update + finally: + generator_finally_turns.append(turn_index) + + source = updates() + provider_sources.append(source) + retained_provider_sources.append(source) + + async def cleanup() -> None: + cleanup_hook_turns.append(turn_index) + + def finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: + return ChatResponse.from_updates(updates, output_format_type=options.get("response_format")) + + stream = ResponseStream(source, finalizer=finalize, cleanup_hooks=[cleanup]) + provider_streams.append(stream) + retained_provider_streams.append(stream) + return stream + + monkeypatch.setattr(client, "_get_streaming_response", get_streaming_response) + return ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) + + +class TestPublicFunctionInvocationStreamLifecycle: + _retained_provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]] + _retained_provider_sources: list[AsyncGenerator[ChatResponseUpdate, None]] + + @pytest.fixture(autouse=True) + async def _close_retained_provider_streams(self) -> AsyncGenerator[None, None]: + self._retained_provider_streams = [] + self._retained_provider_sources = [] + try: + yield + finally: + for provider_stream in reversed(self._retained_provider_streams): + await provider_stream.close() + for provider_source in reversed(self._retained_provider_sources): + await provider_source.aclose() + + @pytest.mark.parametrize("surface", ["client", "agent"], ids=["client", "agent"]) + async def test_close_releases_active_provider_turn_and_is_idempotent( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + surface: Literal["client", "agent"], + ) -> None: + """Explicit close releases the suspended provider generator and awaits its cleanup hook.""" + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [ + [ + ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant"), + ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant"), + ] + ], + ) + stream = _lifecycle_public_stream(chat_client_base, surface, [_make_lifecycle_noop_tool()]) + + first_update = await stream.__anext__() + assert first_update.text == "first" + assert requested_tool_choices == ["auto"] + assert started_turns == [0] + assert len(provider_streams) == len(provider_sources) == 1 + + await stream.close() + assert generator_finally_turns == [0] + assert cleanup_hook_turns == [0] + + await stream.close() + assert generator_finally_turns == [0] + assert cleanup_hook_turns == [0] + + @pytest.mark.parametrize("surface", ["client", "agent"], ids=["client", "agent"]) + async def test_async_context_closes_stream_after_break( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + surface: Literal["client", "agent"], + ) -> None: + """Leaving an async context after async-for break closes the active provider turn.""" + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [ + [ + ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant"), + ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant"), + ] + ], + ) + stream = _lifecycle_public_stream(chat_client_base, surface, [_make_lifecycle_noop_tool()]) + + async with stream: + async for update in stream: + assert update.text == "first" + break + + assert requested_tool_choices == ["auto"] + assert started_turns == [0] + assert len(provider_streams) == len(provider_sources) == 1 + assert generator_finally_turns == [0] + assert cleanup_hook_turns == [0] + + async def test_close_before_start_does_not_call_provider( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Closing an unconsumed public client stream does not start a provider request.""" + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [], + ) + stream = _lifecycle_public_stream(chat_client_base, "client", [_make_lifecycle_noop_tool()]) + + assert requested_tool_choices == [] + assert provider_streams == [] + assert provider_sources == [] + assert started_turns == [] + assert generator_finally_turns == [] + assert cleanup_hook_turns == [] + + await stream.close() + + assert requested_tool_choices == [] + assert provider_streams == [] + assert provider_sources == [] + assert started_turns == [] + assert generator_finally_turns == [] + assert cleanup_hook_turns == [] + + async def test_zero_iteration_final_no_tools_stream_closes_on_early_stop( + self, chat_client_base: SupportsChatGetResponse, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The final no-tools provider stream is released when a zero-turn loop is stopped early.""" + chat_client_base.function_invocation_configuration["max_iterations"] = 0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [[ChatResponseUpdate(contents=[Content.from_text("final")], role="assistant")]], + ) + stream = chat_client_base.get_response( + [Message(role="user", contents=["run"])], + stream=True, + options={"tool_choice": "auto", "tools": [_make_lifecycle_noop_tool()]}, + ) + + first_update = await stream.__anext__() + assert first_update.text == "final" + assert requested_tool_choices == ["none"] + assert started_turns == [0] + assert len(provider_streams) == len(provider_sources) == 1 + + await stream.close() + + assert generator_finally_turns == [0] + assert cleanup_hook_turns == [0] + + @pytest.mark.parametrize( + "stop_on_second_turn", + [False, True], + ids=["complete", "close-during-second-turn"], + ) + async def test_normal_local_tool_roundtrip_cleans_each_turn_and_keeps_final_response( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + stop_on_second_turn: bool, + ) -> None: + """Normal streaming executes once and cleans both turns, including a close mid-final-response.""" + side_effects: list[int] = [] + + @tool(name="lifecycle_increment", approval_mode="never_require") + def lifecycle_increment(value: int) -> int: + side_effects.append(value) + return value + 1 + + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [ + [ + ChatResponseUpdate( + contents=[ + Content.from_function_call( + call_id="call-1", + name="lifecycle_increment", + arguments='{"value": 1}', + ) + ], + role="assistant", + finish_reason="tool_calls", + ) + ], + [ + ChatResponseUpdate(contents=[Content.from_text("completed")], role="assistant"), + ChatResponseUpdate(contents=[Content.from_text("!")], role="assistant", finish_reason="stop"), + ], + ], + ) + stream = chat_client_base.get_response( + [Message(role="user", contents=["run"])], + stream=True, + options={"tool_choice": "auto", "tools": [lifecycle_increment]}, + ) + + if stop_on_second_turn: + updates: list[ChatResponseUpdate] = [] + async with stream: + async for update in stream: + updates.append(update) + if update.text == "completed": + break + response = None + else: + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert len(provider_streams) == len(provider_sources) == 2 + assert requested_tool_choices == ["auto", "auto"] + assert started_turns == [0, 1] + assert generator_finally_turns == [0, 1] + assert cleanup_hook_turns == [0, 1] + assert side_effects == [1] + assert any( + content.type == "function_call" and content.call_id == "call-1" + for update in updates + for content in update.contents + ) + assert any( + content.type == "function_result" and content.result == "2" + for update in updates + for content in update.contents + ) + if stop_on_second_turn: + assert updates[-1].text == "completed" + else: + assert len(updates) == 4 + assert response is not None + assert response.text == "completed!" + assert any( + content.type == "function_call" and content.call_id == "call-1" + for message in response.messages + for content in message.contents + ) + assert any( + content.type == "function_result" and content.result == "2" + for message in response.messages + for content in message.contents + ) + + @pytest.mark.parametrize("failure_kind", ["runtime", "cancelled"], ids=["runtime-error", "cancelled-error"]) + async def test_outer_update_transform_error_closes_inner_provider_stream( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + failure_kind: Literal["runtime", "cancelled"], + ) -> None: + """An outer transform error closes the suspended provider stream and preserves its identity.""" + failure: BaseException + if failure_kind == "runtime": + failure = RuntimeError("transform failed") + else: + failure = asyncio.CancelledError("transform cancelled") + + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [ + [ + ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant"), + ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant"), + ] + ], + ) + stream = _lifecycle_public_stream(chat_client_base, "client", [_make_lifecycle_noop_tool()]) + + def fail_transform(update: ChatResponseUpdate) -> ChatResponseUpdate: + raise failure + + stream.with_transform_hook(fail_transform) + with pytest.raises(type(failure)) as error: + await stream.__anext__() + + assert error.value is failure + assert requested_tool_choices == ["auto"] + assert started_turns == [0] + assert len(provider_streams) == len(provider_sources) == 1 + assert generator_finally_turns == [0] + assert cleanup_hook_turns == [0] + + @pytest.mark.parametrize("failure_kind", ["runtime", "cancelled"], ids=["runtime-error", "cancelled-error"]) + async def test_provider_error_identity_survives_successful_inner_cleanup( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + failure_kind: Literal["runtime", "cancelled"], + ) -> None: + """A provider failure remains the same exception after its generator and hook are released.""" + failure: BaseException + if failure_kind == "runtime": + failure = RuntimeError("provider failed") + else: + failure = asyncio.CancelledError("provider cancelled") + + generator_finally: list[bool] = [] + cleanup_hook_calls: list[bool] = [] + provider_sources: list[AsyncGenerator[ChatResponseUpdate, None]] = [] + provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]] = [] + + def get_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: + async def updates() -> AsyncGenerator[ChatResponseUpdate, None]: + try: + yield ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant") + raise failure + finally: + generator_finally.append(True) + + async def cleanup() -> None: + cleanup_hook_calls.append(True) + + source = updates() + provider_sources.append(source) + self._retained_provider_sources.append(source) + stream: ResponseStream[ChatResponseUpdate, ChatResponse[Any]] = ResponseStream( + source, finalizer=ChatResponse.from_updates, cleanup_hooks=[cleanup] + ) + provider_streams.append(stream) + self._retained_provider_streams.append(stream) + return stream + + monkeypatch.setattr(chat_client_base, "_get_streaming_response", get_streaming_response) + stream = _lifecycle_public_stream(chat_client_base, "client", [_make_lifecycle_noop_tool()]) + + with pytest.raises(type(failure)) as error: + await stream.__anext__() + await stream.__anext__() + + assert error.value is failure + assert len(provider_streams) == len(provider_sources) == 1 + assert generator_finally == [True] + assert cleanup_hook_calls == [True] + + async def test_cancelling_consumer_task_closes_suspended_provider_stream( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Cancelling a blocked consumer awaits cleanup of the active provider turn.""" + ( + provider_streams, + provider_sources, + started_turns, + generator_finally_turns, + cleanup_hook_turns, + requested_tool_choices, + ) = _install_lifecycle_provider_streams( + chat_client_base, + monkeypatch, + self._retained_provider_streams, + self._retained_provider_sources, + [ + [ + ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant"), + ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant"), + ] + ], + ) + stream = _lifecycle_public_stream(chat_client_base, "client", [_make_lifecycle_noop_tool()]) + transform_entered = asyncio.Event() + release_transform = asyncio.Event() + + async def wait_in_transform(update: ChatResponseUpdate) -> ChatResponseUpdate: + transform_entered.set() + await release_transform.wait() + return update + + stream.with_transform_hook(wait_in_transform) + consumer: asyncio.Task[ChatResponseUpdate] = asyncio.create_task(stream.__anext__()) + try: + await asyncio.wait_for(transform_entered.wait(), timeout=5) + assert requested_tool_choices == ["auto"] + assert started_turns == [0] + assert generator_finally_turns == [] + assert cleanup_hook_turns == [] + assert len(provider_streams) == len(provider_sources) == 1 + + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + + assert generator_finally_turns == [0] + assert cleanup_hook_turns == [0] + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) From fc985f5822fa4ed55a3dcd5506368701322ddd7c Mon Sep 17 00:00:00 2001 From: quifox <289420841+quifox@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:14:03 +0800 Subject: [PATCH 2/3] Python: Keep stream-close test consumption in one task --- python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index b9d1d578208..7b0200f279a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -750,6 +750,7 @@ async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str assert first_content.text == "Hello" assert second_content.text == " world" + @pytest.mark.timeout(10) @pytest.mark.parametrize( ("caller", "function_invocation_enabled", "close_mode"), [ @@ -814,7 +815,8 @@ def assert_first_chunk_only() -> None: try: if close_mode == "explicit_close": - first = await asyncio.wait_for(anext(stream), timeout=5.0) + # Keep pull and close in one task; wait_for creates a child task on Python 3.10/3.11. + first = await anext(stream) assert isinstance(first, (ChatResponseUpdate, AgentResponseUpdate)) assert first.text == "first" assert_first_chunk_only() From f96aabc848ac839aabc824cd632f7d10523ada66 Mon Sep 17 00:00:00 2001 From: quifox <289420841+quifox@users.noreply.github.com> Date: Wed, 7 Oct 2026 23:05:20 +0800 Subject: [PATCH 3/3] Python: preserve stream failures during delegated cleanup --- .../ag-ui/agent_framework_ag_ui/_client.py | 14 +- .../agent_framework_ag_ui/_http_service.py | 29 +- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 314 +++++++++++++++++- .../packages/core/agent_framework/_types.py | 57 +++- .../core/test_function_invocation_logic.py | 194 +++++++++++ python/packages/core/tests/core/test_types.py | 241 +++++++++++++- 6 files changed, 822 insertions(+), 27 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index 3b466b58cd4..4c3e7ae625b 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -512,6 +512,7 @@ async def _streaming_impl( available_interrupts=_serialize_available_interrupts(cast(Sequence[Any] | None, available_interrupts)), resume=_serialize_resume(options.get("resume")), ) + body_failed = False try: async for event in events: logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}") @@ -539,7 +540,18 @@ async def _streaming_impl( update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore yield update + except GeneratorExit: + # An explicit close() has no body failure to take precedence over cleanup. + raise + except BaseException: + body_failed = True + raise finally: close = getattr(events, "aclose", None) if close is not None: - await close() + try: + await close() + except Exception: + if not body_failed: + raise + logger.warning("AG-UI event cleanup failed while handling an existing exception.", exc_info=True) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py b/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py index 43461e5cc5e..f13d749080e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py @@ -7,7 +7,9 @@ import json import logging from collections.abc import AsyncIterable, Mapping, Sequence +from contextlib import AsyncExitStack from http.cookiejar import CookieJar, DefaultCookiePolicy +from types import TracebackType from typing import Any, cast import httpx @@ -238,12 +240,27 @@ async def post_run( ) # Stream the response using SSE - async with self.http_client.stream( - "POST", - self.endpoint, - json=request_data, - headers={"Accept": "text/event-stream"}, - ) as response: + async with AsyncExitStack() as stack: + stream_context = self.http_client.stream( + "POST", + self.endpoint, + json=request_data, + headers={"Accept": "text/event-stream"}, + ) + response = await stream_context.__aenter__() + + async def exit_response( + exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> bool | None: + try: + return await stream_context.__aexit__(exc_type, exc, tb) + except Exception: + if exc is None or isinstance(exc, GeneratorExit): + raise + logger.warning("HTTP response cleanup failed while handling an existing exception.", exc_info=True) + return False + + stack.push_async_exit(exit_response) try: response.raise_for_status() except httpx.HTTPStatusError as e: diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 7b0200f279a..fc2276d1bd7 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -25,6 +25,7 @@ from pytest import MonkeyPatch from agent_framework_ag_ui._client import AGUIChatClient +from agent_framework_ag_ui._event_converters import AGUIEventConverter from agent_framework_ag_ui._http_service import AGUIHttpService @@ -62,9 +63,17 @@ def inner_get_response( class _GatedSSEStream(httpx.AsyncByteStream): """Yield one SSE chunk and keep the next event gated until another pull.""" - def __init__(self) -> None: + def __init__( + self, + *, + aclose_error: BaseException | None = None, + iteration_error: BaseException | None = None, + ) -> None: self._tail_release = asyncio.Event() self._iterator: AsyncGenerator[bytes, None] | None = None + self.aclose_error = aclose_error + self.iteration_error = iteration_error + self.tail_waiting = asyncio.Event() self.chunk_indexes_yielded: list[int] = [] self.tail_requested = False self.tail_yielded = False @@ -88,7 +97,10 @@ async def _iterate(self) -> AsyncGenerator[bytes, None]: self.chunk_indexes_yielded.append(0) yield self._chunks[0] self.tail_requested = True + self.tail_waiting.set() await self._tail_release.wait() + if self.iteration_error is not None: + raise self.iteration_error self.chunk_indexes_yielded.append(1) self.tail_yielded = True yield self._chunks[1] @@ -96,12 +108,61 @@ async def _iterate(self) -> AsyncGenerator[bytes, None]: finally: self.iterator_finally = True + def release_tail(self) -> None: + self._tail_release.set() + async def aclose(self) -> None: self.aclose_calls += 1 self.closed = True self._tail_release.set() if self._iterator is not None: await self._iterator.aclose() + if self.aclose_error is not None and self.aclose_calls == 1: + raise self.aclose_error + + +class _FailingAGUIEventStream: + """Yield a public AG-UI event, then fail or wait while its close fails.""" + + def __init__( + self, + event: dict[str, Any], + *, + close_error: OSError, + iteration_error: BaseException | None = None, + wait_after_event: bool = False, + ) -> None: + self._event = event + self._close_error = close_error + self._iteration_error = iteration_error + self._wait_after_event = wait_after_event + self._index = 0 + self._release = asyncio.Event() + self.waiting = asyncio.Event() + self.aclose_calls = 0 + self.aclose_completed = 0 + + def __aiter__(self) -> "_FailingAGUIEventStream": + return self + + async def __anext__(self) -> dict[str, Any]: + if self._index == 0: + self._index += 1 + return self._event + if self._wait_after_event: + self.waiting.set() + await self._release.wait() + if self._iteration_error is not None: + raise self._iteration_error + raise StopAsyncIteration + + async def aclose(self) -> None: + self.aclose_calls += 1 + self._release.set() + await asyncio.sleep(0) + self.aclose_completed += 1 + if self.aclose_calls == 1: + raise self._close_error class TestAGUIChatClient: @@ -861,6 +922,257 @@ def assert_first_chunk_only() -> None: else: await client.close() + @pytest.mark.timeout(10) + @pytest.mark.parametrize("caller", ["client", "agent"], ids=["client", "agent"]) + @pytest.mark.parametrize("failure_kind", ["iteration", "conversion", "cancel"]) + async def test_primary_failure_survives_event_or_http_body_aclose_error( + self, + monkeypatch: MonkeyPatch, + caller: str, + failure_kind: str, + ) -> None: + """AG-UI keeps event, conversion, and cancellation failures over close errors.""" + primary = RuntimeError("AG-UI primary failure") + close_failure = OSError("stream close failed") + body = _GatedSSEStream(aclose_error=close_failure) + response: httpx.Response | None = None + event_source = ( + _FailingAGUIEventStream( + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "first"}, + close_error=close_failure, + iteration_error=primary if failure_kind == "iteration" else None, + wait_after_event=failure_kind == "cancel", + ) + if failure_kind != "conversion" + else None + ) + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal response + if request.method == "GET": + return httpx.Response(204, request=request) + await request.aread() + response = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=body, + request=request, + ) + return response + + if failure_kind == "conversion": + + def fail_convert_event(_converter: AGUIEventConverter, _event: dict[str, Any]) -> ChatResponseUpdate | None: + raise primary + + monkeypatch.setattr(AGUIEventConverter, "convert_event", fail_convert_event) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = AGUIChatClient( + endpoint="https://agui.example.test/run", + http_client=http_client, + function_invocation_configuration={"enabled": False}, + ) + agent = Agent(client=client) if caller == "agent" else None + messages = [Message(role="user", contents=["Question"])] + stream = ( + agent.run(messages, stream=True) if agent is not None else client.get_response(messages, stream=True) + ) + + if event_source is not None: + monkeypatch.setattr(client._http_service, "post_run", lambda *args, **kwargs: event_source) + + if failure_kind == "iteration": + first = await stream.__anext__() + assert first.text == "first" + with pytest.raises(RuntimeError) as error: + await stream.__anext__() + assert error.value is primary + elif failure_kind == "conversion": + with pytest.raises(RuntimeError) as error: + await stream.__anext__() + assert error.value is primary + else: + first = await stream.__anext__() + assert first.text == "first" + consumer = asyncio.create_task(stream.__anext__()) + try: + assert event_source is not None + await asyncio.wait_for(event_source.waiting.wait(), timeout=5) + consumer.cancel("consumer cancelled") + with pytest.raises(asyncio.CancelledError) as cancel_error: + await consumer + assert isinstance(cancel_error.value, asyncio.CancelledError) + assert consumer.cancelled() + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + if failure_kind == "conversion": + assert response is not None + assert body.aclose_calls >= 1 + assert body.closed + assert body.iterator_finally + else: + assert event_source is not None + assert event_source.aclose_calls >= 1 + assert event_source.aclose_completed == event_source.aclose_calls + + @pytest.mark.timeout(10) + @pytest.mark.parametrize("caller", ["client", "agent"], ids=["client", "agent"]) + async def test_http_body_task_cancel_survives_aclose_error(self, caller: str) -> None: + """A task canceled during HTTP body iteration remains canceled when body close fails.""" + close_failure = OSError("HTTP response body close failed") + body = _GatedSSEStream(aclose_error=close_failure) + response: httpx.Response | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal response + if request.method == "GET": + return httpx.Response(204, request=request) + await request.aread() + response = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=body, + request=request, + ) + return response + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = AGUIChatClient( + endpoint="https://agui.example.test/run", + http_client=http_client, + function_invocation_configuration={"enabled": False}, + ) + agent = Agent(client=client) if caller == "agent" else None + messages = [Message(role="user", contents=["Question"])] + stream = ( + agent.run(messages, stream=True) if agent is not None else client.get_response(messages, stream=True) + ) + first = await stream.__anext__() + assert first.text == "first" + consumer = asyncio.create_task(stream.__anext__()) + try: + await asyncio.wait_for(body.tail_waiting.wait(), timeout=5) + consumer.cancel("consumer cancelled") + with pytest.raises(asyncio.CancelledError): + await consumer + assert consumer.cancelled() + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + assert response is not None + assert body.aclose_calls >= 1 + assert body.closed + assert body.iterator_finally + + @pytest.mark.timeout(10) + @pytest.mark.parametrize("caller", ["client", "agent"], ids=["client", "agent"]) + async def test_http_body_iteration_error_survives_aclose_error( + self, + caller: str, + ) -> None: + """An HTTP body iteration failure should remain primary when body close fails.""" + primary = RuntimeError("HTTP event iteration failed") + close_failure = OSError("HTTP response body close failed") + body = _GatedSSEStream(aclose_error=close_failure, iteration_error=primary) + response: httpx.Response | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal response + if request.method == "GET": + return httpx.Response(204, request=request) + await request.aread() + response = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=body, + request=request, + ) + return response + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = AGUIChatClient( + endpoint="https://agui.example.test/run", + http_client=http_client, + function_invocation_configuration={"enabled": False}, + ) + agent = Agent(client=client) if caller == "agent" else None + messages = [Message(role="user", contents=["Question"])] + stream = ( + agent.run(messages, stream=True) if agent is not None else client.get_response(messages, stream=True) + ) + + first = await stream.__anext__() + assert first.text == "first" + body.release_tail() + with pytest.raises(RuntimeError) as error: + await stream.__anext__() + + assert error.value is primary + assert response is not None + assert body.aclose_calls >= 1 + assert body.closed + assert body.iterator_finally + + @pytest.mark.timeout(10) + @pytest.mark.parametrize("caller", ["client", "agent"], ids=["client", "agent"]) + @pytest.mark.parametrize("close_mode", ["explicit", "complete"], ids=["public-close", "normal-completion"]) + async def test_http_body_aclose_error_propagates_without_primary_failure( + self, + caller: str, + close_mode: str, + ) -> None: + """A body close failure remains visible when no operation failure takes precedence.""" + close_failure = OSError("HTTP response body close failed") + body = _GatedSSEStream(aclose_error=close_failure) + response: httpx.Response | None = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal response + if request.method == "GET": + return httpx.Response(204, request=request) + await request.aread() + response = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=body, + request=request, + ) + return response + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = AGUIChatClient( + endpoint="https://agui.example.test/run", + http_client=http_client, + function_invocation_configuration={"enabled": False}, + ) + agent = Agent(client=client) if caller == "agent" else None + messages = [Message(role="user", contents=["Question"])] + stream = ( + agent.run(messages, stream=True) if agent is not None else client.get_response(messages, stream=True) + ) + + if close_mode == "explicit": + first = await stream.__anext__() + assert first.text == "first" + with pytest.raises(OSError) as error: + await stream.close() + else: + body.release_tail() + with pytest.raises(OSError) as error: + async for _ in stream: + pass + + assert error.value is close_failure + + assert response is not None + assert body.aclose_calls >= 1 + assert body.closed + assert body.iterator_finally + async def test_get_response_non_streaming(self, monkeypatch: MonkeyPatch) -> None: """Test non-streaming response method.""" mock_events = [ diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index f98743794a4..a576aca36de 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -3841,7 +3841,13 @@ async def __aenter__(self) -> ResponseStream[UpdateT, FinalT]: async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: """Close the stream on block exit, including an early consumer break.""" - await self.close() + try: + await self.close() + except Exception: + # GeneratorExit also represents an explicit close(), whose failure must be reported. + if exc is None or isinstance(exc, GeneratorExit): + raise + logger.warning("Stream cleanup failed while handling an existing exception.", exc_info=True) def _start_content_pipeline(self) -> None: if self._content_pipeline_started: @@ -3988,16 +3994,8 @@ async def _run_release_error_hooks(self, exc: BaseException) -> None: async def _abort_buffered_materialization(self, exc: BaseException) -> None: self._stream_error = exc try: - iterator = self._iterator - if iterator is not None: - if isinstance(iterator, ResponseStream): - await cast(ResponseStream[UpdateT, Any], iterator).close() - else: - close = getattr(iterator, "aclose", None) - if close is not None: - await close() - self._consumed = True - await self._run_cleanup_hooks() + # close() also runs the hooks if releasing the iterator fails. + await self.close() finally: self._stream_error = None @@ -4131,9 +4129,18 @@ async def _ensure_buffered_materialized(self) -> None: self._buffered_materialized = True except BaseException as exc: try: - if release_phase_started: - await self._run_release_error_hooks(exc) - await self._abort_buffered_materialization(exc) + try: + if release_phase_started: + await self._run_release_error_hooks(exc) + except Exception: + logger.warning("Stream error hook failed while handling an existing exception.", exc_info=True) + finally: + try: + await self._abort_buffered_materialization(exc) + except Exception: + logger.warning( + "Buffered stream cleanup failed while handling an existing exception.", exc_info=True + ) except BaseException as cleanup_exc: self._buffered_materialization_error = cleanup_exc raise @@ -4166,12 +4173,17 @@ async def __anext__(self) -> UpdateT: # provider stream suspended until GC. Hooks run first because they # read self._stream_error, and close() would consume the one-shot # cleanup run without it. close() stays in finally so the provider - # stream is released even when a hook raises; the original - # exception always re-raises. + # stream is released even when a hook raises. Ordinary cleanup + # errors must not replace this failure; a new cancellation still propagates. try: await self._handle_stream_error(exc) + except Exception: + logger.warning("Stream cleanup hook failed while handling an existing exception.", exc_info=True) finally: - await self.close() + try: + await self.close() + except Exception: + logger.warning("Stream cleanup failed while handling an existing exception.", exc_info=True) raise async def close(self) -> None: @@ -4179,6 +4191,7 @@ async def close(self) -> None: This method is idempotent and also closes nested ``ResponseStream`` wrappers. """ + iterator_failed = False try: iterator: AsyncIterator[UpdateT] | None = self._iterator if iterator is not None: @@ -4188,9 +4201,17 @@ async def close(self) -> None: close = getattr(iterator, "aclose", None) if close is not None: await close() + except BaseException: + iterator_failed = True + raise finally: self._consumed = True - await self._run_cleanup_hooks() + try: + await self._run_cleanup_hooks() + except Exception: + if not iterator_failed: + raise + logger.warning("Stream cleanup hook failed after iterator close failed.", exc_info=True) async def _resolve_stream_with_pull_contexts(self) -> AsyncIterable[UpdateT]: """Resolve the underlying stream while activating any registered pull context managers. diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index cde80da7735..ec262f94820 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -11279,6 +11279,53 @@ def finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: ) +class _LifecycleCloseErrorIterator: + """Provider iterator whose awaited close can fail after an operation has failed.""" + + def __init__( + self, + updates: Sequence[ChatResponseUpdate], + *, + close_error: OSError, + iteration_error: BaseException | None = None, + wait_after_updates: bool = False, + ) -> None: + self._updates = list(updates) + self._close_error = close_error + self._iteration_error = iteration_error + self._wait_after_updates = wait_after_updates + self._index = 0 + self._release = asyncio.Event() + self.waiting = asyncio.Event() + self.aclose_calls = 0 + self.aclose_completed = 0 + self.closed = False + + def __aiter__(self) -> "_LifecycleCloseErrorIterator": + return self + + async def __anext__(self) -> ChatResponseUpdate: + if self._index < len(self._updates): + update = self._updates[self._index] + self._index += 1 + return update + if self._wait_after_updates: + self.waiting.set() + await self._release.wait() + if self._iteration_error is not None: + raise self._iteration_error + raise StopAsyncIteration + + async def aclose(self) -> None: + self.aclose_calls += 1 + self._release.set() + await asyncio.sleep(0) + self.aclose_completed += 1 + self.closed = True + if self.aclose_calls == 1: + raise self._close_error + + class TestPublicFunctionInvocationStreamLifecycle: _retained_provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]] _retained_provider_sources: list[AsyncGenerator[ChatResponseUpdate, None]] @@ -11713,3 +11760,150 @@ async def wait_in_transform(update: ChatResponseUpdate) -> ChatResponseUpdate: finally: consumer.cancel() await asyncio.gather(consumer, return_exceptions=True) + + @pytest.mark.timeout(10) + @pytest.mark.parametrize( + ("surface", "iteration_limit"), + [ + pytest.param("client", 2, id="client-ordinary-turn"), + pytest.param("agent", 2, id="agent-ordinary-turn"), + pytest.param("client", 0, id="client-final-no-tools"), + pytest.param("agent", 0, id="agent-final-no-tools"), + ], + ) + @pytest.mark.parametrize("failure_kind", ["iteration", "finalizer", "transform", "cancel"]) + async def test_public_function_loop_preserves_primary_failure_when_provider_close_fails( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + surface: Literal["client", "agent"], + iteration_limit: int, + failure_kind: Literal["iteration", "finalizer", "transform", "cancel"], + ) -> None: + """Provider cleanup errors do not replace public function-loop failures or cancellation.""" + primary = RuntimeError("primary function-loop failure") + close_failure = OSError("provider iterator close failed") + source = _LifecycleCloseErrorIterator( + [ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant")], + close_error=close_failure, + iteration_error=primary if failure_kind == "iteration" else None, + wait_after_updates=failure_kind == "cancel", + ) + provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]] = [] + requested_tool_choices: list[str | None] = [] + chat_client_base.function_invocation_configuration["max_iterations"] = iteration_limit # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + def get_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: + requested_tool_choices.append(options.get("tool_choice")) + + def finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: + if failure_kind == "finalizer": + raise primary + return ChatResponse.from_updates(updates, output_format_type=options.get("response_format")) + + provider = ResponseStream(source, finalizer=finalize) + provider_streams.append(provider) + self._retained_provider_streams.append(provider) + return provider + + monkeypatch.setattr(chat_client_base, "_get_streaming_response", get_streaming_response) + stream = _lifecycle_public_stream(chat_client_base, surface, [_make_lifecycle_noop_tool()]) + + if failure_kind == "transform": + + def fail_transform(_update: Any) -> Any: + raise primary + + stream.with_transform_hook(fail_transform) + + async def consume_public_stream() -> None: + async for _ in stream: + pass + + if failure_kind == "cancel": + consumer: asyncio.Task[None] = asyncio.create_task(consume_public_stream()) + try: + await asyncio.wait_for(source.waiting.wait(), timeout=5) + consumer.cancel("consumer cancelled") + with pytest.raises(asyncio.CancelledError): + await consumer + assert consumer.cancelled() + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + else: + with pytest.raises(RuntimeError) as error: + await consume_public_stream() + assert error.value is primary + + assert len(provider_streams) == 1 + assert requested_tool_choices == (["auto"] if iteration_limit == 2 else ["none"]) + assert source.aclose_calls >= 1 + assert source.aclose_completed == source.aclose_calls + assert source.closed + + @pytest.mark.timeout(10) + @pytest.mark.parametrize( + ("surface", "iteration_limit"), + [ + pytest.param("client", 2, id="client-ordinary-turn"), + pytest.param("agent", 2, id="agent-ordinary-turn"), + pytest.param("client", 0, id="client-final-no-tools"), + pytest.param("agent", 0, id="agent-final-no-tools"), + ], + ) + @pytest.mark.parametrize("close_mode", ["explicit", "complete"], ids=["public-close", "normal-completion"]) + async def test_public_function_loop_propagates_provider_close_error_without_primary_failure( + self, + chat_client_base: SupportsChatGetResponse, + monkeypatch: pytest.MonkeyPatch, + surface: Literal["client", "agent"], + iteration_limit: int, + close_mode: Literal["explicit", "complete"], + ) -> None: + """Provider close errors still propagate on public close or normal completion.""" + close_failure = OSError("provider iterator close failed") + source = _LifecycleCloseErrorIterator( + [ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant")], + close_error=close_failure, + ) + provider_streams: list[ResponseStream[ChatResponseUpdate, ChatResponse[Any]]] = [] + requested_tool_choices: list[str | None] = [] + chat_client_base.function_invocation_configuration["max_iterations"] = iteration_limit # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + def get_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: + requested_tool_choices.append(options.get("tool_choice")) + provider = ResponseStream(source, finalizer=ChatResponse.from_updates) + provider_streams.append(provider) + self._retained_provider_streams.append(provider) + return provider + + monkeypatch.setattr(chat_client_base, "_get_streaming_response", get_streaming_response) + stream = _lifecycle_public_stream(chat_client_base, surface, [_make_lifecycle_noop_tool()]) + + if close_mode == "explicit": + first = await stream.__anext__() + assert first.text == "first" + with pytest.raises(OSError) as error: + await stream.close() + else: + with pytest.raises(OSError) as error: + async for _ in stream: + pass + + assert error.value is close_failure + assert len(provider_streams) == 1 + assert requested_tool_choices == (["auto"] if iteration_limit == 2 else ["none"]) + assert source.aclose_calls >= 1 + assert source.aclose_completed == source.aclose_calls + assert source.closed diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 05d474687f4..31d69ad3b50 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -5,7 +5,7 @@ import contextlib import json import warnings -from collections.abc import AsyncIterable, Awaitable, Callable, Sequence +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Sequence from dataclasses import dataclass from datetime import datetime, timezone from functools import partial @@ -4590,6 +4590,41 @@ def uppercase_transform(update: ChatResponseUpdate) -> ChatResponseUpdate: assert (await pending_update).text == "UPDATE_0" +class _FailingCloseIterator: + def __init__( + self, + failure: BaseException | None, + close_error: BaseException | None, + *, + close_waiting: asyncio.Event | None = None, + ) -> None: + self._failure = failure + self._close_error = close_error + self._close_waiting = close_waiting + self._yielded = False + self.close_calls = 0 + + def __aiter__(self) -> AsyncIterator[str]: + return self + + async def __anext__(self) -> str: + if not self._yielded: + self._yielded = True + return "first" + if self._failure is not None: + raise self._failure + raise StopAsyncIteration + + async def aclose(self) -> None: + self.close_calls += 1 + if self._close_waiting is not None and self.close_calls == 1: + self._close_waiting.set() + await asyncio.Event().wait() + await asyncio.sleep(0) + if self._close_error is not None: + raise self._close_error + + class TestResponseStreamCleanupHooks: """Tests for cleanup hooks (after stream consumption, before finalizer).""" @@ -4893,6 +4928,210 @@ def failing_mapper(update: ChatResponseUpdate) -> ChatResponseUpdate: assert cleanup_called["value"] is True + @pytest.mark.parametrize("stream_updates", [True, False], ids=["live", "buffered"]) + @pytest.mark.parametrize("cleanup_kind", ["iterator", "hook"]) + @pytest.mark.parametrize("failure_kind", ["runtime", "cancelled"]) + async def test_primary_error_survives_cleanup_failure( + self, + stream_updates: bool, + cleanup_kind: Literal["iterator", "hook"], + failure_kind: Literal["runtime", "cancelled"], + caplog: pytest.LogCaptureFixture, + ) -> None: + failure = RuntimeError("operation failed") if failure_kind == "runtime" else asyncio.CancelledError("cancelled") + cleanup_error = OSError("cleanup failed") + source = _FailingCloseIterator(failure, cleanup_error if cleanup_kind == "iterator" else None) + cleanup_calls = 0 + + async def cleanup() -> None: + nonlocal cleanup_calls + cleanup_calls += 1 + await asyncio.sleep(0) + if cleanup_kind == "hook": + raise cleanup_error + + stream: ResponseStream[str, str] = ResponseStream( + source, + finalizer=lambda values: "".join(values), + cleanup_hooks=[cleanup], + stream_updates=stream_updates, + ) + with pytest.raises(type(failure)) as error: + async for _ in stream: + pass + + assert error.value is failure + assert source.close_calls >= 1 + assert cleanup_calls == 1 + assert any(record.exc_info is not None for record in caplog.records) + if not stream_updates: + close_calls = source.close_calls + with pytest.raises(type(failure)) as repeated_error: + await anext(stream) + assert repeated_error.value is failure + assert source.close_calls == close_calls + assert cleanup_calls == 1 + + @pytest.mark.parametrize("failure_kind", ["runtime", "cancelled"]) + async def test_context_body_error_survives_close_failure( + self, failure_kind: Literal["runtime", "cancelled"] + ) -> None: + failure = RuntimeError("body failed") if failure_kind == "runtime" else asyncio.CancelledError("body cancelled") + source = _FailingCloseIterator(None, OSError("close failed")) + stream: ResponseStream[str, str] = ResponseStream(source, finalizer=lambda values: "".join(values)) + + with pytest.raises(type(failure)) as error: + async with stream: + assert await anext(stream) == "first" + raise failure + + assert error.value is failure + assert source.close_calls == 1 + + @pytest.mark.parametrize("release_kind", ["close", "context", "complete", "buffered_complete"]) + @pytest.mark.parametrize("cleanup_kind", ["iterator", "hook"]) + async def test_cleanup_failure_without_body_error_is_reported( + self, + release_kind: Literal["close", "context", "complete", "buffered_complete"], + cleanup_kind: Literal["iterator", "hook"], + ) -> None: + cleanup_error = OSError("cleanup failed") + source = _FailingCloseIterator(None, cleanup_error if cleanup_kind == "iterator" else None) + cleanup_calls = 0 + + async def cleanup() -> None: + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_kind == "hook": + raise cleanup_error + + stream: ResponseStream[str, str] = ResponseStream( + source, + finalizer=lambda values: "".join(values), + cleanup_hooks=[cleanup], + stream_updates=release_kind != "buffered_complete", + ) + with pytest.raises(OSError) as error: + if release_kind == "close": + await anext(stream) + await stream.close() + else: + async with stream: + if release_kind == "context": + await anext(stream) + else: + async for _ in stream: + pass + + assert error.value is cleanup_error + assert source.close_calls >= 1 + assert cleanup_calls == 1 + + @pytest.mark.parametrize("stream_updates", [True, False], ids=["live", "buffered"]) + async def test_finalizer_error_survives_close_failure(self, stream_updates: bool) -> None: + failure = ValueError("finalizer failed") + source = _FailingCloseIterator(None, OSError("close failed")) + + def finalize(values: Sequence[str]) -> str: + raise failure + + stream: ResponseStream[str, str] = ResponseStream(source, finalizer=finalize, stream_updates=stream_updates) + with pytest.raises(ValueError) as error: + async with stream: + async for _ in stream: + pass + + assert error.value is failure + assert source.close_calls >= 1 + + @pytest.mark.parametrize("stream_updates", [True, False], ids=["live", "buffered"]) + @pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-success", "cleanup-failure"]) + async def test_new_cancellation_during_error_cleanup_propagates( + self, stream_updates: bool, cleanup_fails: bool + ) -> None: + cleanup_waiting = asyncio.Event() + source = _FailingCloseIterator(RuntimeError("operation failed"), None, close_waiting=cleanup_waiting) + cleanup_calls = 0 + + def cleanup() -> None: + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_fails: + raise OSError("cleanup hook failed") + + stream: ResponseStream[str, str] = ResponseStream( + source, + finalizer=lambda values: "".join(values), + cleanup_hooks=[cleanup], + stream_updates=stream_updates, + ) + if stream_updates: + await anext(stream) + task = asyncio.create_task(anext(stream)) + try: + await asyncio.wait_for(cleanup_waiting.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert task.cancelled() + assert cleanup_calls == 1 + if not stream_updates: + close_calls = source.close_calls + with pytest.raises(asyncio.CancelledError): + await anext(stream) + assert source.close_calls == close_calls + assert cleanup_calls == 1 + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.parametrize("failure_kind", ["runtime", "cancelled"]) + async def test_iterator_close_error_survives_cleanup_hook_failure( + self, failure_kind: Literal["runtime", "cancelled"] + ) -> None: + failure = OSError("iterator close failed") if failure_kind == "runtime" else asyncio.CancelledError("cancelled") + source = _FailingCloseIterator(None, failure) + cleanup_calls = 0 + + def cleanup() -> None: + nonlocal cleanup_calls + cleanup_calls += 1 + raise ValueError("cleanup hook failed") + + stream: ResponseStream[str, str] = ResponseStream( + source, finalizer=lambda values: "".join(values), cleanup_hooks=[cleanup] + ) + await anext(stream) + with pytest.raises(type(failure)) as error: + await stream.close() + + assert error.value is failure + assert source.close_calls == 1 + assert cleanup_calls == 1 + + async def test_buffered_error_hook_failure_does_not_skip_cleanup(self) -> None: + failure = ValueError("release failed") + source = _FailingCloseIterator(None, OSError("close failed")) + + def fail_release() -> None: + raise failure + + def fail_error_hook(error: BaseException) -> None: + assert error is failure + raise RuntimeError("error hook failed") + + stream: ResponseStream[str, str] = ResponseStream( + source, finalizer=lambda values: "".join(values), stream_updates=False + ) + stream._with_release_hook(fail_release)._with_release_error_hook(fail_error_hook) + + with pytest.raises(ValueError) as error: + await anext(stream) + + assert error.value is failure + assert source.close_calls == 1 + assert stream.updates == [] + class TestResponseStreamResultHooks: """Tests for result hooks (after finalizer)."""