From aae1aa1b00a9972acea4a828f94e4343959f4eb8 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Fri, 24 Jul 2026 10:02:27 +0900 Subject: [PATCH] fix: improve provider compatibility and preserve session history on retries --- src/agents/extensions/models/any_llm_model.py | 46 ++++++- src/agents/run_internal/run_loop.py | 4 +- tests/models/test_any_llm_model.py | 129 ++++++++++++++++++ tests/test_agent_runner.py | 46 +++++++ 4 files changed, 222 insertions(+), 3 deletions(-) diff --git a/src/agents/extensions/models/any_llm_model.py b/src/agents/extensions/models/any_llm_model.py index 60d60da338..72f930dbab 100644 --- a/src/agents/extensions/models/any_llm_model.py +++ b/src/agents/extensions/models/any_llm_model.py @@ -731,6 +731,25 @@ async def _fetch_chat_response( extra_kwargs = self._build_chat_extra_kwargs(model_settings) extra_kwargs.pop("reasoning_effort", None) + headers = self._merge_headers(model_settings) + if self._provider_name in {"gemini", "vertexai"}: + http_options = extra_kwargs.get("http_options") + if isinstance(http_options, BaseModel): + existing_headers = getattr(http_options, "headers", None) or {} + extra_kwargs["http_options"] = http_options.model_copy( + update={"headers": {**existing_headers, **headers}} + ) + elif isinstance(http_options, dict): + existing_headers = http_options.get("headers") or {} + extra_kwargs["http_options"] = { + **http_options, + "headers": {**existing_headers, **headers}, + } + elif http_options is None: + extra_kwargs["http_options"] = {"headers": headers} + else: + extra_kwargs["extra_headers"] = headers + # The Chat Completions API requires logprobs=True whenever top_logprobs is set. Defer to a # caller-supplied logprobs (via extra_args, already merged into extra_kwargs) to avoid a # duplicate-key collision. @@ -753,7 +772,6 @@ async def _fetch_chat_response( stream_options=stream_options, reasoning_effort=reasoning_effort, top_logprobs=model_settings.top_logprobs, - extra_headers=self._merge_headers(model_settings), **extra_kwargs, ) @@ -959,6 +977,8 @@ def _get_provider(self) -> Any: api_key=self.api_key, api_base=self.base_url, ) + if self._provider_name in {"gemini", "vertexai"}: + self._normalize_google_tool_result_roles(base_provider) self._provider_cache[False] = base_provider if disable_provider_retries: @@ -968,6 +988,30 @@ def _get_provider(self) -> Any: return base_provider + @staticmethod + def _normalize_google_tool_result_roles(provider: Any) -> None: + convert_completion_params = getattr(provider, "_convert_completion_params", None) + if not callable(convert_completion_params): + return + + def convert_with_supported_tool_result_roles(*args: Any, **kwargs: Any) -> Any: + converted = convert_completion_params(*args, **kwargs) + contents = converted.get("contents") + if not isinstance(contents, list): + return converted + + converted["contents"] = [ + content.model_copy(update={"role": "user"}) + if isinstance(content, BaseModel) and getattr(content, "role", None) == "function" + else {**content, "role": "user"} + if isinstance(content, dict) and content.get("role") == "function" + else content + for content in contents + ] + return converted + + provider._convert_completion_params = convert_with_supported_tool_result_roles + def _clone_provider_without_retries(self, provider: Any) -> Any: client = getattr(provider, "client", None) with_options = getattr(client, "with_options", None) diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index aa68ad8f17..fa60d2299e 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -1955,9 +1955,9 @@ async def get_new_response( model_settings = model_settings_with_prompt_cache_key(model_settings, prompt_cache_key) async def rewind_model_request() -> None: - items_to_rewind = session_items_to_rewind if session_items_to_rewind is not None else [] - await rewind_session_items(session, items_to_rewind, server_conversation_tracker) if server_conversation_tracker is not None: + items_to_rewind = session_items_to_rewind if session_items_to_rewind is not None else [] + await rewind_session_items(session, items_to_rewind, server_conversation_tracker) server_conversation_tracker.rewind_input(filtered.input) with model_run_context(tool_use_tracker): diff --git a/tests/models/test_any_llm_model.py b/tests/models/test_any_llm_model.py index a8d155a389..c87477cd60 100644 --- a/tests/models/test_any_llm_model.py +++ b/tests/models/test_any_llm_model.py @@ -272,6 +272,135 @@ async def test_user_agent_header_any_llm_chat(override_ua: str | None, monkeypat assert provider.chat_calls[0]["extra_headers"]["User-Agent"] == expected_ua +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_name", ["gemini", "vertexai"]) +@pytest.mark.parametrize("options_type", ["unset", "dictionary", "model"]) +async def test_any_llm_google_chat_headers_use_http_options( + monkeypatch: pytest.MonkeyPatch, provider_name: str, options_type: str +) -> None: + class HttpOptions(BaseModel): + headers: dict[str, str] + timeout: int + + provider = FakeAnyLLMProvider(supports_responses=False, chat_response=_chat_completion("Hello")) + module, _create_calls = _import_any_llm_module(monkeypatch, provider) + model = module.AnyLLMModel(model=f"{provider_name}/gemini-2.5-flash") + + extra_args: dict[str, Any] = {} + configured_options: dict[str, Any] | HttpOptions | None = None + if options_type == "dictionary": + configured_options = {"headers": {"X-Existing": "existing"}, "timeout": 1000} + extra_args["http_options"] = configured_options + elif options_type == "model": + configured_options = HttpOptions(headers={"X-Existing": "existing"}, timeout=1000) + extra_args["http_options"] = configured_options + + await model.get_response( + system_instructions=None, + input="hi", + model_settings=ModelSettings( + extra_args=extra_args, + extra_headers={"X-Test-Header": "test"}, + ), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + prompt=None, + ) + + call = provider.chat_calls[0] + assert "extra_headers" not in call + http_options = call["http_options"] + if isinstance(http_options, BaseModel): + http_options = http_options.model_dump() + assert http_options["headers"]["User-Agent"] == f"Agents/Python {__version__}" + assert http_options["headers"]["X-Test-Header"] == "test" + if configured_options is not None: + assert http_options["headers"]["X-Existing"] == "existing" + assert http_options["timeout"] == 1000 + if isinstance(configured_options, BaseModel): + assert configured_options.headers == {"X-Existing": "existing"} + else: + assert configured_options["headers"] == {"X-Existing": "existing"} + + +@pytest.mark.parametrize("provider_name", ["gemini", "vertexai"]) +@pytest.mark.parametrize("content_type", ["model", "dictionary"]) +def test_any_llm_google_provider_normalizes_function_result_roles( + monkeypatch: pytest.MonkeyPatch, provider_name: str, content_type: str +) -> None: + class GoogleContent(BaseModel): + role: str + parts: list[dict[str, Any]] + + tool_result: dict[str, Any] = { + "role": "function", + "parts": [{"function_response": {"name": "get_weather", "response": {"result": "sunny"}}}], + } + original_tool_result: GoogleContent | dict[str, Any] + if content_type == "model": + original_tool_result = GoogleContent.model_validate(tool_result) + else: + original_tool_result = tool_result + + class GoogleProvider(FakeAnyLLMProvider): + @staticmethod + def _convert_completion_params(*args: Any, **kwargs: Any) -> dict[str, Any]: + return { + "model": "gemini-3.6-flash", + "contents": [ + GoogleContent(role="user", parts=[{"text": "Check the weather."}]), + original_tool_result, + GoogleContent(role="model", parts=[{"text": "Done."}]), + ], + } + + provider = GoogleProvider(supports_responses=False) + module, _create_calls = _import_any_llm_module(monkeypatch, provider) + model = module.AnyLLMModel(model=f"{provider_name}/gemini-3.6-flash") + + converted = model._get_provider()._convert_completion_params(object()) + contents = converted["contents"] + + assert [ + item.role if isinstance(item, GoogleContent) else item["role"] for item in contents + ] == [ + "user", + "user", + "model", + ] + normalized_tool_result = contents[1] + if isinstance(normalized_tool_result, BaseModel): + normalized_tool_result = normalized_tool_result.model_dump() + assert normalized_tool_result["parts"] == tool_result["parts"] + assert ( + original_tool_result.role + if isinstance(original_tool_result, GoogleContent) + else original_tool_result["role"] + ) == "function" + + +def test_any_llm_non_google_provider_does_not_normalize_function_result_roles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class NonGoogleProvider(FakeAnyLLMProvider): + @staticmethod + def _convert_completion_params(*args: Any, **kwargs: Any) -> dict[str, Any]: + return {"contents": [{"role": "function", "parts": [{"result": "ok"}]}]} + + provider = NonGoogleProvider(supports_responses=False) + module, _create_calls = _import_any_llm_module(monkeypatch, provider) + model = module.AnyLLMModel(model="openrouter/google/gemini-3.6-flash") + + converted = model._get_provider()._convert_completion_params(object()) + + assert converted["contents"][0]["role"] == "function" + + @pytest.mark.allow_call_model_methods @pytest.mark.asyncio async def test_any_llm_chat_path_is_used_when_responses_are_unsupported(monkeypatch) -> None: diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index e93e111464..12160d886e 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -2608,6 +2608,52 @@ async def test_conversation_lock_rewind_skips_when_no_snapshot() -> None: assert session.pop_calls == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("session_backend", ["memory", "sqlite"]) +async def test_non_streamed_model_retry_does_not_rewind_committed_session_input( + tmp_path: Path, session_backend: str +) -> None: + model = FakeModel() + model.add_multiple_turn_outputs( + [ + APIConnectionError( + message="connection error", + request=httpx.Request("POST", "https://example.com"), + ), + [get_text_message("done")], + ] + ) + agent = Agent( + name="test", + model=model, + model_settings=ModelSettings( + retry=ModelRetrySettings( + max_retries=1, + policy=retry_policies.network_error(), + ) + ), + ) + session: CountingSession | SQLiteSession + if session_backend == "sqlite": + session = SQLiteSession("retry-session", tmp_path / "retry.sqlite3") + await session.add_items([get_text_input_item("previous")]) + else: + session = CountingSession(history=[get_text_input_item("previous")]) + + try: + result = await Runner.run(agent, input="test", session=session) + saved_items = await session.get_items() + finally: + if isinstance(session, SQLiteSession): + session.close() + + assert result.final_output == "done" + assert [item.get("role") for item in saved_items] == ["user", "user", "assistant"] + assert [item.get("content") for item in saved_items[:2]] == ["previous", "test"] + if isinstance(session, CountingSession): + assert session.pop_calls == 0 + + @pytest.mark.asyncio async def test_get_new_response_uses_agent_retry_settings() -> None: model = FakeModel()