From 55c0e3874f6889c7db92e161f8fc146cf193eb5c Mon Sep 17 00:00:00 2001 From: "zhao.wang" <57819425+Excelius-Wang@users.noreply.github.com> Date: Mon, 31 Aug 2026 00:54:30 +0800 Subject: [PATCH 1/2] fix: close model providers created by Runner --- src/agents/result.py | 62 +++- src/agents/run.py | 81 ++++- tests/test_cancel_streaming.py | 20 +- tests/test_runner_model_provider_lifecycle.py | 320 ++++++++++++++++++ 4 files changed, 469 insertions(+), 14 deletions(-) create mode 100644 tests/test_runner_model_provider_lifecycle.py diff --git a/src/agents/result.py b/src/agents/result.py index 0ceb0d7187..e5e7a8ed67 100644 --- a/src/agents/result.py +++ b/src/agents/result.py @@ -693,6 +693,21 @@ class RunResultStreaming(RunResultBase): ) _sandbox_cleanup_task: asyncio.Task[None] | None = field(default=None, init=False, repr=False) _sandbox_cleanup_callback_registered: bool = field(default=False, init=False, repr=False) + _sandbox_wrapped_run_loop_task: asyncio.Task[Any] | None = field( + default=None, + init=False, + repr=False, + ) + _model_provider_cleanup: Callable[[], Awaitable[None]] | None = field( + default=None, + init=False, + repr=False, + ) + _model_provider_cleanup_task: asyncio.Task[None] | None = field( + default=None, + init=False, + repr=False, + ) def __post_init__(self, _run_impl_task: asyncio.Task[Any] | None) -> None: self._current_agent_ref = weakref.ref(self.current_agent) @@ -757,6 +772,7 @@ def ensure_sandbox_cleanup_on_completion(self) -> None: return original_task = self.run_loop_task + self._sandbox_wrapped_run_loop_task = original_task self._sandbox_cleanup_callback_registered = True original_task.add_done_callback( lambda _task: asyncio.create_task(self._run_sandbox_cleanup()) @@ -784,10 +800,50 @@ async def _await_run_and_cleanup() -> Any: await self._run_sandbox_cleanup() return result - self.run_loop_task = asyncio.create_task( + cleanup_wrapper_task = asyncio.create_task( _await_data_redacted_error_boundary(_await_run_and_cleanup) ) + def cancel_original_if_wrapper_cancelled(task: asyncio.Task[Any]) -> None: + # A task cancelled before its first event-loop step cannot enter its coroutine body. + if task.cancelled() and not original_task.done(): + original_task.cancel() + + cleanup_wrapper_task.add_done_callback(cancel_original_if_wrapper_cancelled) + self.run_loop_task = cleanup_wrapper_task + + def _ensure_model_provider_cleanup_on_completion( + self, + cleanup: Callable[[], Awaitable[None]], + ) -> None: + """Register one cleanup task that also starts if the run task never enters its body.""" + self._model_provider_cleanup = cleanup + if self.run_loop_task is not None: + self.run_loop_task.add_done_callback(lambda _task: self._start_model_provider_cleanup()) + + def _start_model_provider_cleanup(self) -> asyncio.Task[None] | None: + task = self._model_provider_cleanup_task + if task is not None: + return task + + cleanup = self._model_provider_cleanup + if cleanup is None: + return None + + self._model_provider_cleanup = None + + async def run_cleanup() -> None: + await cleanup() + + task = asyncio.create_task(run_cleanup()) + self._model_provider_cleanup_task = task + return task + + async def _await_model_provider_cleanup(self) -> None: + task = self._start_model_provider_cleanup() + if task is not None: + await asyncio.shield(task) + @property def run_loop_exception(self) -> BaseException | None: """The exception raised by the background run loop, if any. @@ -1003,6 +1059,7 @@ def register_current_consumer() -> None: self._cleanup_tasks() if not cancelled: + await self._await_model_provider_cleanup() await self._run_sandbox_cleanup() finally: # Allow any pending callbacks (e.g., cancellation handlers) to enqueue their @@ -1093,6 +1150,9 @@ def _cleanup_tasks(self): if self.run_loop_task and not self.run_loop_task.done(): self.run_loop_task.cancel() + if self._sandbox_wrapped_run_loop_task and not self._sandbox_wrapped_run_loop_task.done(): + self._sandbox_wrapped_run_loop_task.cancel() + if self._input_guardrails_task and not self._input_guardrails_task.done(): self._input_guardrails_task.cancel() diff --git a/src/agents/run.py b/src/agents/run.py index 37dd66582e..afb350859c 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -37,8 +37,14 @@ TResponseInputItem, ) from .lifecycle import RunHooks -from .logger import log_model_and_tool_action_warning, log_tool_action_warning, logger +from .logger import ( + log_model_action_warning, + log_model_and_tool_action_warning, + log_tool_action_warning, + logger, +) from .memory import Session +from .models.interface import ModelProvider from .result import RunResult, RunResultStreaming from .run_config import ( DEFAULT_MAX_TURNS, @@ -541,6 +547,44 @@ def run_streamed( ) +def _normalize_run_config_for_runner( + value: RunConfig | dict[str, Any] | None, +) -> tuple[RunConfig, bool]: + """Normalize a run config and report whether Runner created its model provider.""" + owns_model_provider = value is None or ( + isinstance(value, dict) and "model_provider" not in value + ) + run_config = RunConfig() if value is None else _coerce_run_config(value) + return run_config, owns_model_provider + + +async def _close_runner_owned_model_provider(model_provider: ModelProvider) -> None: + """Finish provider cleanup despite repeated cancellation, then restore cancellation.""" + + async def close() -> None: + try: + await model_provider.aclose() + except Exception as error: + log_model_action_warning( + logger, + "Failed to close model provider created for run", + error, + ) + + close_task = asyncio.create_task(close()) + try: + await asyncio.shield(close_task) + except asyncio.CancelledError: + while not close_task.done(): + try: + await asyncio.shield(close_task) + except asyncio.CancelledError: + continue + if not close_task.cancelled(): + close_task.result() + raise + + class AgentRunner: """ WARNING: this class is experimental and not part of the public API @@ -553,18 +597,25 @@ async def run( input: str | list[TResponseInputItem] | RunState[TContext], **kwargs: Unpack[RunOptions[TContext]], ) -> RunResult: + run_config, owns_model_provider = _normalize_run_config_for_runner(kwargs.get("run_config")) + cast(dict[str, Any], kwargs)["run_config"] = run_config redacted_error: BaseException | None = None try: - return await self._run_impl(starting_agent, input, **kwargs) - except BaseException as error: - if not _is_error_data_redacted(error): - raise - _detach_data_redacted_error_traceback(error) - redacted_error = error + try: + return await self._run_impl(starting_agent, input, **kwargs) + except BaseException as error: + if not _is_error_data_redacted(error): + raise + _detach_data_redacted_error_traceback(error) + redacted_error = error + finally: + if owns_model_provider: + await _close_runner_owned_model_provider(run_config.model_provider) self = cast(Any, None) starting_agent = cast(Any, None) input = cast(Any, None) + run_config = cast(Any, None) cast(dict[str, Any], kwargs).clear() assert redacted_error is not None _detach_data_redacted_error_traceback(redacted_error) @@ -586,7 +637,7 @@ async def _run_impl( conversation_id = kwargs.get("conversation_id") session = kwargs.get("session") - run_config = RunConfig() if run_config is None else _coerce_run_config(run_config) + run_config = cast(RunConfig, run_config) is_resumed_state = isinstance(input, RunState) run_state: RunState[TContext] | None = ( @@ -2347,7 +2398,7 @@ def run_streamed( conversation_id = kwargs.get("conversation_id") session = kwargs.get("session") - run_config = RunConfig() if run_config is None else _coerce_run_config(run_config) + run_config, owns_model_provider = _normalize_run_config_for_runner(run_config) # Handle RunState input is_resumed_state = isinstance(input, RunState) @@ -2565,8 +2616,8 @@ def run_streamed( sandbox_runtime.apply_result_metadata(streamed_result) # Kick off the actual agent loop in the background and return the streamed result object. - streamed_result.run_loop_task = asyncio.create_task( - _await_data_redacted_error_boundary( + async def run_loop() -> None: + await _await_data_redacted_error_boundary( lambda: start_streaming( starting_input=input_for_result, streamed_result=streamed_result, @@ -2586,7 +2637,13 @@ def run_streamed( sandbox_runtime=sandbox_runtime, ) ) - ) + + streamed_result.run_loop_task = asyncio.create_task(run_loop()) + if owns_model_provider: + model_provider = run_config.model_provider + streamed_result._ensure_model_provider_cleanup_on_completion( + lambda: _close_runner_owned_model_provider(model_provider) + ) if sandbox_runtime.enabled: streamed_result.ensure_sandbox_cleanup_on_completion() return streamed_result diff --git a/tests/test_cancel_streaming.py b/tests/test_cancel_streaming.py index 3fbf2571ac..ea7cb85635 100644 --- a/tests/test_cancel_streaming.py +++ b/tests/test_cancel_streaming.py @@ -1,4 +1,5 @@ import asyncio +import gc import json import time @@ -7,6 +8,7 @@ from agents import Agent, Runner from agents.guardrail import input_guardrail +from agents.models.multi_provider import MultiProvider from agents.stream_events import RawResponsesStreamEvent from agents.testing import ScriptedModel @@ -109,13 +111,29 @@ async def test_cancel_is_idempotent(): @pytest.mark.asyncio -async def test_cancel_before_streaming(): +async def test_cancel_before_streaming( + monkeypatch: pytest.MonkeyPatch, + recwarn: pytest.WarningsRecorder, +) -> None: + closed: list[MultiProvider] = [] + + async def record_close(provider: MultiProvider) -> None: + closed.append(provider) + + monkeypatch.setattr(MultiProvider, "aclose", record_close) model = ScriptedModel() agent = Agent(name="Joker", model=model) result = Runner.run_streamed(agent, input="Please tell me 5 jokes.") result.cancel() # Cancel before streaming events = [e async for e in result.stream_events()] + gc.collect() + assert events == [], "No events should be yielded if cancel() is called before streaming." + assert len(closed) == 1 + assert not any( + warning.category is RuntimeWarning and "was never awaited" in str(warning.message) + for warning in recwarn + ) @pytest.mark.asyncio diff --git a/tests/test_runner_model_provider_lifecycle.py b/tests/test_runner_model_provider_lifecycle.py new file mode 100644 index 0000000000..fdfe8fcb21 --- /dev/null +++ b/tests/test_runner_model_provider_lifecycle.py @@ -0,0 +1,320 @@ +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import AsyncIterator +from typing import Any + +import pytest +from openai import AsyncOpenAI +from openai.types.responses import ResponseCompletedEvent, ResponseOutputItemDoneEvent +from websockets.asyncio.server import ServerConnection, serve + +from agents import ( + Agent, + RunConfig, + Runner, + set_default_openai_client, + set_default_openai_responses_transport, +) +from agents.models.interface import Model, ModelProvider +from agents.models.multi_provider import MultiProvider +from agents.testing import ModelStep, ScriptedModel +from tests.model_test_helpers import get_response_obj + +from .test_responses import get_text_message + + +def _scripted_agent(*steps: Any) -> Agent[None]: + return Agent(name="test", model=ScriptedModel(steps)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "run_config", + [None, {"tracing_disabled": True}], + ids=["omitted", "dictionary-with-default-provider"], +) +async def test_run_closes_implicitly_created_model_provider( + monkeypatch: pytest.MonkeyPatch, + run_config: dict[str, Any] | None, +) -> None: + closed: list[MultiProvider] = [] + + async def record_close(provider: MultiProvider) -> None: + closed.append(provider) + + monkeypatch.setattr(MultiProvider, "aclose", record_close) + agent = _scripted_agent([get_text_message("done")]) + + result = ( + await Runner.run(agent, "hello") + if run_config is None + else await Runner.run(agent, "hello", run_config=run_config) + ) + + assert result.final_output == "done" + assert len(closed) == 1 + + +def test_run_sync_closes_implicitly_created_model_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + closed: list[MultiProvider] = [] + + async def record_close(provider: MultiProvider) -> None: + closed.append(provider) + + monkeypatch.setattr(MultiProvider, "aclose", record_close) + agent = _scripted_agent([get_text_message("done")]) + + result = Runner.run_sync(agent, "hello") + + assert result.final_output == "done" + assert len(closed) == 1 + + +@pytest.mark.asyncio +async def test_run_preserves_primary_error_when_provider_cleanup_fails( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + run_error = RuntimeError("run failed") + close_error = RuntimeError("close failed") + + async def fail_close(_provider: MultiProvider) -> None: + raise close_error + + monkeypatch.setattr(MultiProvider, "aclose", fail_close) + agent = _scripted_agent(ModelStep.raise_error(run_error)) + + with caplog.at_level(logging.WARNING, logger="openai.agents"): + with pytest.raises(RuntimeError) as exc_info: + await Runner.run(agent, "hello") + + assert exc_info.value is run_error + assert "Failed to close model provider created for run" in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dictionary_config", [False, True], ids=["RunConfig", "dictionary"]) +async def test_run_keeps_explicit_model_provider_open_for_reuse(dictionary_config: bool) -> None: + class ReusableProvider(ModelProvider): + def __init__(self, model: Model) -> None: + self.model = model + self.close_calls = 0 + + def get_model(self, model_name: str | None) -> Model: + return self.model + + async def aclose(self) -> None: + self.close_calls += 1 + + model = ScriptedModel( + [ + [get_text_message("first")], + [get_text_message("second")], + ] + ) + provider = ReusableProvider(model) + run_config: RunConfig | dict[str, Any] = ( + {"model_provider": provider} if dictionary_config else RunConfig(model_provider=provider) + ) + agent = Agent(name="test", model="test-model") + + first = await Runner.run(agent, "first", run_config=run_config) + second = await Runner.run(agent, "second", run_config=run_config) + + assert first.final_output == "first" + assert second.final_output == "second" + assert provider.close_calls == 0 + + +@pytest.mark.asyncio +async def test_run_streamed_closes_provider_only_after_run_settles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + finish = asyncio.Event() + closed: list[MultiProvider] = [] + output = get_text_message("done") + + async def events(_call: object) -> AsyncIterator[Any]: + started.set() + await finish.wait() + yield ResponseOutputItemDoneEvent( + type="response.output_item.done", + item=output, + output_index=0, + sequence_number=0, + ) + yield ResponseCompletedEvent( + type="response.completed", + response=get_response_obj([output]), + sequence_number=1, + ) + + async def record_close(provider: MultiProvider) -> None: + closed.append(provider) + + monkeypatch.setattr(MultiProvider, "aclose", record_close) + agent = _scripted_agent(ModelStep.stream(events)) + result = Runner.run_streamed(agent, "hello") + + await asyncio.wait_for(started.wait(), timeout=1) + assert closed == [] + + finish.set() + async for _event in result.stream_events(): + pass + + assert result.final_output == "done" + assert len(closed) == 1 + + +@pytest.mark.asyncio +async def test_run_streamed_closes_provider_after_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + stream_closed = asyncio.Event() + blocked = asyncio.Event() + closed: list[MultiProvider] = [] + + async def events(_call: object) -> AsyncIterator[Any]: + started.set() + try: + await blocked.wait() + finally: + stream_closed.set() + if False: # pragma: no cover - makes this an async generator + yield None + + async def record_close(provider: MultiProvider) -> None: + closed.append(provider) + + monkeypatch.setattr(MultiProvider, "aclose", record_close) + agent = _scripted_agent(ModelStep.stream(events)) + result = Runner.run_streamed(agent, "hello") + await asyncio.wait_for(started.wait(), timeout=1) + + result.cancel() + async for _event in result.stream_events(): + pass + + assert stream_closed.is_set() + assert len(closed) == 1 + + +@pytest.mark.asyncio +async def test_run_streamed_repeated_cancellation_waits_for_provider_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stream_started = asyncio.Event() + blocked = asyncio.Event() + close_started = asyncio.Event() + close_release = asyncio.Event() + close_completed = asyncio.Event() + + async def events(_call: object) -> AsyncIterator[Any]: + stream_started.set() + await blocked.wait() + if False: # pragma: no cover - makes this an async generator + yield None + + async def slow_close(_provider: MultiProvider) -> None: + close_started.set() + await close_release.wait() + close_completed.set() + + monkeypatch.setattr(MultiProvider, "aclose", slow_close) + agent = _scripted_agent(ModelStep.stream(events)) + result = Runner.run_streamed(agent, "hello") + await asyncio.wait_for(stream_started.wait(), timeout=1) + + result.cancel() + await asyncio.wait_for(close_started.wait(), timeout=1) + result.cancel() + close_release.set() + async for _event in result.stream_events(): + pass + + assert close_completed.is_set() + + +@pytest.mark.asyncio +async def test_run_streamed_cancel_before_start_propagates_through_cleanup_wrapper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sandbox_cleanup_completed = asyncio.Event() + + async def record_provider_close(_provider: MultiProvider) -> None: + return None + + async def record_sandbox_cleanup() -> None: + sandbox_cleanup_completed.set() + + monkeypatch.setattr(MultiProvider, "aclose", record_provider_close) + result = Runner.run_streamed(Agent(name="test", model=ScriptedModel()), "hello") + original_task = result.run_loop_task + assert original_task is not None + result._sandbox_cleanup = record_sandbox_cleanup + result.ensure_sandbox_cleanup_on_completion() + + result.cancel() + async for _event in result.stream_events(): + pass + + assert original_task.cancelled() + assert sandbox_cleanup_completed.is_set() + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_run_streamed_closes_implicit_responses_websocket_connection() -> None: + connection_closed = asyncio.Event() + + async def handle(connection: ServerConnection) -> None: + try: + async for request_json in connection: + request = json.loads(request_json) + assert request["type"] == "response.create" + response = get_response_obj( + [get_text_message("done")], + response_id="resp-runner-provider-cleanup", + ) + await connection.send( + json.dumps( + { + "type": "response.completed", + "response": response.model_dump(), + "sequence_number": 1, + } + ) + ) + finally: + connection_closed.set() + + async with serve(handle, "127.0.0.1", 0) as server: + server_socket = next(iter(server.sockets)) + host, port = server_socket.getsockname()[:2] + client = AsyncOpenAI( + api_key="test-key", + base_url=f"http://{host}:{port}/v1", + websocket_base_url=f"ws://{host}:{port}/v1", + max_retries=0, + ) + set_default_openai_client(client, use_for_tracing=False) + set_default_openai_responses_transport("websocket") + agent = Agent(name="test", model="gpt-4.1-mini") + + try: + result = Runner.run_streamed(agent, "hello") + async for _event in result.stream_events(): + pass + + assert result.final_output == "done" + await asyncio.wait_for(connection_closed.wait(), timeout=1) + finally: + await client.close() From 99853c0045031b13408cbb3963d128e1e791db85 Mon Sep 17 00:00:00 2001 From: "zhao.wang" <57819425+Excelius-Wang@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:19:25 +0800 Subject: [PATCH 2/2] refactor: move provider lifecycle helpers into run_internal --- src/agents/run.py | 51 ++----------------- .../run_internal/model_provider_lifecycle.py | 46 +++++++++++++++++ 2 files changed, 51 insertions(+), 46 deletions(-) create mode 100644 src/agents/run_internal/model_provider_lifecycle.py diff --git a/src/agents/run.py b/src/agents/run.py index afb350859c..4a636b7784 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -37,14 +37,8 @@ TResponseInputItem, ) from .lifecycle import RunHooks -from .logger import ( - log_model_action_warning, - log_model_and_tool_action_warning, - log_tool_action_warning, - logger, -) +from .logger import log_model_and_tool_action_warning, log_tool_action_warning, logger from .memory import Session -from .models.interface import ModelProvider from .result import RunResult, RunResultStreaming from .run_config import ( DEFAULT_MAX_TURNS, @@ -61,7 +55,6 @@ ToolExecutionConfig, ToolNameCollisionPolicy as ToolNameCollisionPolicy, ToolNotFoundBehavior, - _coerce_run_config, ) from .run_context import RunContextWrapper, TContext from .run_error_handlers import RunErrorHandlers @@ -114,6 +107,10 @@ normalize_resumed_input, reconcile_nested_history_owned_input_after_rewrite, ) +from .run_internal.model_provider_lifecycle import ( + _close_runner_owned_model_provider, + _normalize_run_config_for_runner, +) from .run_internal.oai_conversation import OpenAIServerConversationTracker from .run_internal.prompt_cache_key import PromptCacheKeyResolver from .run_internal.run_grouping import resolve_run_grouping_id @@ -547,44 +544,6 @@ def run_streamed( ) -def _normalize_run_config_for_runner( - value: RunConfig | dict[str, Any] | None, -) -> tuple[RunConfig, bool]: - """Normalize a run config and report whether Runner created its model provider.""" - owns_model_provider = value is None or ( - isinstance(value, dict) and "model_provider" not in value - ) - run_config = RunConfig() if value is None else _coerce_run_config(value) - return run_config, owns_model_provider - - -async def _close_runner_owned_model_provider(model_provider: ModelProvider) -> None: - """Finish provider cleanup despite repeated cancellation, then restore cancellation.""" - - async def close() -> None: - try: - await model_provider.aclose() - except Exception as error: - log_model_action_warning( - logger, - "Failed to close model provider created for run", - error, - ) - - close_task = asyncio.create_task(close()) - try: - await asyncio.shield(close_task) - except asyncio.CancelledError: - while not close_task.done(): - try: - await asyncio.shield(close_task) - except asyncio.CancelledError: - continue - if not close_task.cancelled(): - close_task.result() - raise - - class AgentRunner: """ WARNING: this class is experimental and not part of the public API diff --git a/src/agents/run_internal/model_provider_lifecycle.py b/src/agents/run_internal/model_provider_lifecycle.py new file mode 100644 index 0000000000..a3c22e8d53 --- /dev/null +++ b/src/agents/run_internal/model_provider_lifecycle.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +from ..logger import log_model_action_warning, logger +from ..models.interface import ModelProvider +from ..run_config import RunConfig, _coerce_run_config + + +def _normalize_run_config_for_runner( + value: RunConfig | dict[str, Any] | None, +) -> tuple[RunConfig, bool]: + """Normalize a run config and report whether Runner created its model provider.""" + owns_model_provider = value is None or ( + isinstance(value, dict) and "model_provider" not in value + ) + run_config = RunConfig() if value is None else _coerce_run_config(value) + return run_config, owns_model_provider + + +async def _close_runner_owned_model_provider(model_provider: ModelProvider) -> None: + """Finish provider cleanup despite repeated cancellation, then restore cancellation.""" + + async def close() -> None: + try: + await model_provider.aclose() + except Exception as error: + log_model_action_warning( + logger, + "Failed to close model provider created for run", + error, + ) + + close_task = asyncio.create_task(close()) + try: + await asyncio.shield(close_task) + except asyncio.CancelledError: + while not close_task.done(): + try: + await asyncio.shield(close_task) + except asyncio.CancelledError: + continue + if not close_task.cancelled(): + close_task.result() + raise