From 131501229a97ad9635a0dc6a363bb2558bc4e577 Mon Sep 17 00:00:00 2001 From: Adam Fortuna Date: Tue, 7 Jul 2026 09:21:48 +0200 Subject: [PATCH] fix(runtime): restore request_overrides after transport recovery Include request_overrides in primary runtime snapshots so transport recovery restores request-level model parameters. --- agent/agent_init.py | 1 + agent/agent_runtime_helpers.py | 2 + .../run_agent/test_primary_runtime_restore.py | 44 ++++++++++++++++++- 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/agent/agent_init.py b/agent/agent_init.py index 49a29f1e13..63e8f4c436 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -3105,6 +3105,7 @@ def init_agent( "base_url": agent.base_url, "api_mode": agent.api_mode, "api_key": getattr(agent, "api_key", ""), + "request_overrides": dict(getattr(agent, "request_overrides", {}) or {}), "client_kwargs": dict(agent._client_kwargs), "use_prompt_caching": agent._use_prompt_caching, "use_native_cache_layout": agent._use_native_cache_layout, diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 7b5b9b081a..c518935ce7 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -1486,6 +1486,7 @@ def try_recover_primary_transport( agent._transport_cache.clear() agent.api_key = rt["api_key"] agent._reasoning_echo_flag = rt.get("reasoning_echo_flag", False) + agent.request_overrides = dict(rt.get("request_overrides") or {}) if agent.api_mode == "anthropic_messages": from agent.anthropic_adapter import build_anthropic_client @@ -1748,6 +1749,7 @@ def restore_primary_runtime(agent) -> bool: agent._transport_cache.clear() agent.api_key = rt["api_key"] agent._reasoning_echo_flag = rt.get("reasoning_echo_flag", False) + agent.request_overrides = dict(rt.get("request_overrides") or {}) agent._client_kwargs = dict(rt["client_kwargs"]) agent._use_prompt_caching = rt["use_prompt_caching"] # Default to native layout when the restored snapshot predates the diff --git a/tests/run_agent/test_primary_runtime_restore.py b/tests/run_agent/test_primary_runtime_restore.py index a3344f4f65..cc5bcfb5b5 100644 --- a/tests/run_agent/test_primary_runtime_restore.py +++ b/tests/run_agent/test_primary_runtime_restore.py @@ -30,7 +30,12 @@ def _make_tool_defs(*names: str) -> list: ] -def _make_agent(fallback_model=None, provider="custom", base_url="https://my-llm.example.com/v1"): +def _make_agent( + fallback_model=None, + provider="custom", + base_url="https://my-llm.example.com/v1", + request_overrides=None, +): """Create a minimal AIAgent with optional fallback config.""" with ( patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")), @@ -45,6 +50,7 @@ def _make_agent(fallback_model=None, provider="custom", base_url="https://my-llm "agent.context_compressor.get_model_context_length", return_value=200_000, ), + patch("agent.anthropic_adapter.build_anthropic_client", return_value=MagicMock()), ): agent = AIAgent( api_key="test-key-12345678", @@ -54,6 +60,7 @@ def _make_agent(fallback_model=None, provider="custom", base_url="https://my-llm skip_context_files=True, skip_memory=True, fallback_model=fallback_model, + request_overrides=dict(request_overrides or {}), ) agent.client = MagicMock() return agent @@ -83,6 +90,13 @@ class TestPrimaryRuntimeSnapshot: assert "client_kwargs" in rt assert "compressor_context_length" in rt + def test_snapshot_includes_request_overrides(self): + overrides = {"extra_body": {"reasoning": {"effort": "medium"}}} + agent = _make_agent(request_overrides=overrides) + rt = agent._primary_runtime + assert rt["request_overrides"] == overrides + assert rt["request_overrides"] is not agent.request_overrides + def test_snapshot_includes_compressor_state(self): agent = _make_agent() rt = agent._primary_runtime @@ -280,6 +294,19 @@ class TestRestorePrimaryRuntime: assert agent._use_prompt_caching == original_caching + def test_restores_request_overrides(self): + original_overrides = {"extra_body": {"reasoning": {"effort": "medium"}}} + agent = _make_agent(request_overrides=original_overrides) + agent._fallback_activated = True + agent.request_overrides = {"extra_body": {"fallback_only": True}} + + with patch("run_agent.OpenAI", return_value=MagicMock()): + result = agent._restore_primary_runtime() + + assert result is True + assert agent.request_overrides == original_overrides + assert agent.request_overrides is not agent._primary_runtime["request_overrides"] + def test_restore_skips_cross_provider_pool_entry(self): """Restore must not swap in a fallback provider credential for the primary runtime.""" @@ -551,6 +578,21 @@ class TestTryRecoverPrimaryTransport: assert result is True + def test_recovery_restores_request_overrides(self): + original_overrides = {"extra_body": {"reasoning": {"effort": "medium"}}} + agent = _make_agent(provider="custom", request_overrides=original_overrides) + error = _make_transport_error("ReadTimeout") + agent.request_overrides = {"extra_body": {"fallback_only": True}} + + with patch("run_agent.OpenAI", return_value=MagicMock()), \ + patch("time.sleep"): + result = agent._try_recover_primary_transport( + error, retry_count=3, max_retries=3, + ) + + assert result is True + assert agent.request_overrides == original_overrides + assert agent.request_overrides is not agent._primary_runtime["request_overrides"]