From 6bc8d68ad7d752f0e29fa9e24a3c4847127a3b4a Mon Sep 17 00:00:00 2001 From: Israel Lot Date: Fri, 24 Jul 2026 09:53:31 +0000 Subject: [PATCH] fix(codex): scope 24h retention to Bedrock Mantle --- agent/auxiliary_client.py | 18 ++++------- agent/chat_completion_helpers.py | 1 + agent/transports/codex.py | 31 +++++++++++++------ tests/agent/test_auxiliary_client.py | 15 ++++++++- .../agent/transports/test_codex_transport.py | 30 ++++++++++++------ .../test_run_agent_codex_responses.py | 19 ++++++++++++ 6 files changed, 82 insertions(+), 32 deletions(-) diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 1b3658677c0f..151ed79512d8 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -1060,7 +1060,7 @@ class _CodexCompletionsAdapter: try: from agent.transports.codex import ( _content_cache_key, - _prompt_cache_retention_for_model, + _default_prompt_cache_retention_for_request, ) from utils import base_url_host_matches @@ -1070,21 +1070,15 @@ class _CodexCompletionsAdapter: base_url_host_matches(_host_src, "githubcopilot.com") or base_url_host_matches(_host_src, "models.github.ai") ) - _is_codex_backend = ( - base_url_host_matches(_host_src, "chatgpt.com") - and "/backend-api/codex" in _host_src.lower() - ) if not _is_xai and not _is_github and "prompt_cache_key" not in resp_kwargs: _cache_key = _content_cache_key(instructions, resp_kwargs.get("tools")) if _cache_key: resp_kwargs["prompt_cache_key"] = _cache_key - if ( - not _is_xai - and not _is_github - and not _is_codex_backend - and "prompt_cache_retention" not in resp_kwargs - ): - _cache_retention = _prompt_cache_retention_for_model(model) + if "prompt_cache_retention" not in resp_kwargs: + _cache_retention = _default_prompt_cache_retention_for_request( + model, + _host_src, + ) if _cache_retention: resp_kwargs["prompt_cache_retention"] = _cache_retention except Exception: diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 70e4ea002a51..3edcddec8279 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -1076,6 +1076,7 @@ def build_api_kwargs(agent, api_messages: list) -> dict: tools=tools_for_api, reasoning_config=agent.reasoning_config, session_id=getattr(agent, "session_id", None), + base_url=agent.base_url, max_tokens=agent.max_tokens, timeout=agent._resolved_api_call_timeout(), request_overrides=agent.request_overrides, diff --git a/agent/transports/codex.py b/agent/transports/codex.py index a8502027259d..15dd3409e3fc 100644 --- a/agent/transports/codex.py +++ b/agent/transports/codex.py @@ -48,8 +48,23 @@ _EXTENDED_PROMPT_CACHE_MODEL_RE = re.compile( ) -def _prompt_cache_retention_for_model(model: str) -> Optional[str]: - """Return ``24h`` for models documented to support extended retention.""" +def _default_prompt_cache_retention_for_request( + model: str, + base_url: Any, +) -> Optional[str]: + """Return ``24h`` for supported models on Amazon Bedrock Mantle.""" + from utils import base_url_hostname + + hostname_parts = base_url_hostname(str(base_url or "")).split(".") + is_bedrock_mantle = ( + len(hostname_parts) == 4 + and hostname_parts[0] == "bedrock-mantle" + and bool(hostname_parts[1]) + and hostname_parts[2:] == ["api", "aws"] + ) + if not is_bedrock_mantle: + return None + normalized = str(model or "").strip().lower().replace("_", "-") if _EXTENDED_PROMPT_CACHE_MODEL_RE.search(normalized): return "24h" @@ -313,13 +328,11 @@ class ResponsesApiTransport(ProviderTransport): if not is_github_responses and not is_xai_responses and cache_key: kwargs["prompt_cache_key"] = cache_key - cache_retention = _prompt_cache_retention_for_model(model) - if ( - cache_retention - and not is_github_responses - and not is_xai_responses - and not is_codex_backend - ): + cache_retention = _default_prompt_cache_retention_for_request( + model, + params.get("base_url"), + ) + if cache_retention: kwargs.setdefault("prompt_cache_retention", cache_retention) if reasoning_enabled and is_xai_responses: diff --git a/tests/agent/test_auxiliary_client.py b/tests/agent/test_auxiliary_client.py index a24577da0bdc..fae90d2a89e4 100644 --- a/tests/agent/test_auxiliary_client.py +++ b/tests/agent/test_auxiliary_client.py @@ -4965,7 +4965,7 @@ class TestCodexAdapterPromptCacheKey: ]) def test_extended_cache_models_set_prompt_cache_retention(self, model): adapter, captured = self._build_adapter( - base_url="https://responses.example.com/v1", + base_url="https://bedrock-mantle.us-west-2.api.aws/v1", model=model, ) adapter.create(messages=[ @@ -4982,6 +4982,19 @@ class TestCodexAdapterPromptCacheKey: ]) assert "prompt_cache_retention" not in captured + @pytest.mark.parametrize("base_url", [ + "https://api.openai.com/v1", + "https://example.services.ai.azure.com/openai/v1", + "https://responses.example.com/v1", + ]) + def test_prompt_cache_retention_skipped_for_other_compatible_endpoints(self, base_url): + adapter, captured = self._build_adapter(base_url=base_url) + adapter.create(messages=[ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "hi"}, + ]) + assert "prompt_cache_retention" not in captured + def test_prompt_cache_retention_skipped_for_xai_and_github_hosts(self): adapter, captured = self._build_adapter(base_url="https://api.x.ai/v1") adapter.create(messages=[ diff --git a/tests/agent/transports/test_codex_transport.py b/tests/agent/transports/test_codex_transport.py index 040acf6ade16..8563ca7a607c 100644 --- a/tests/agent/transports/test_codex_transport.py +++ b/tests/agent/transports/test_codex_transport.py @@ -264,6 +264,7 @@ class TestCodexBuildKwargs: kw = transport.build_kwargs( model=model, messages=messages, tools=[], session_id="test-session", + base_url="https://bedrock-mantle.us-west-2.api.aws/v1", ) assert kw["prompt_cache_retention"] == "24h" @@ -274,20 +275,29 @@ class TestCodexBuildKwargs: messages=[{"role": "user", "content": "Hi"}], tools=[], session_id="test-session", + base_url="https://bedrock-mantle.us-west-2.api.aws/v1", ) assert "prompt_cache_retention" not in kw - def test_prompt_cache_retention_skipped_for_known_incompatible_backends(self, transport): - messages = [{"role": "user", "content": "Hi"}] - assert "prompt_cache_retention" not in transport.build_kwargs( - model="gpt-5.4", messages=messages, tools=[], is_xai_responses=True - ) - assert "prompt_cache_retention" not in transport.build_kwargs( - model="gpt-5.4", messages=messages, tools=[], is_github_responses=True - ) - assert "prompt_cache_retention" not in transport.build_kwargs( - model="gpt-5.4", messages=messages, tools=[], is_codex_backend=True + @pytest.mark.parametrize("base_url", [ + "https://api.openai.com/v1", + "https://example.openai.azure.com/openai/v1", + "https://api.x.ai/v1", + "https://models.github.ai/inference", + "https://api.githubcopilot.com", + "https://chatgpt.com/backend-api/codex", + "https://responses.example.com/v1", + "https://bedrock-mantle.us-west-2.api.aws.example/v1", + "https://example.com/bedrock-mantle.us-west-2.api.aws/v1", + ]) + def test_prompt_cache_retention_omitted_for_non_mantle_endpoints(self, transport, base_url): + kw = transport.build_kwargs( + model="gpt-5.4", + messages=[{"role": "user", "content": "Hi"}], + tools=[], + base_url=base_url, ) + assert "prompt_cache_retention" not in kw def test_xai_responses_sends_cache_key_via_extra_body(self, transport): """xAI's Responses API documents ``prompt_cache_key`` as the diff --git a/tests/run_agent/test_run_agent_codex_responses.py b/tests/run_agent/test_run_agent_codex_responses.py index 071596e20465..6c425e72da0c 100644 --- a/tests/run_agent/test_run_agent_codex_responses.py +++ b/tests/run_agent/test_run_agent_codex_responses.py @@ -396,6 +396,25 @@ def test_build_api_kwargs_codex(monkeypatch): assert "extra_body" not in kwargs +def test_build_api_kwargs_mantle_sets_extended_prompt_cache_retention(monkeypatch): + _patch_agent_bootstrap(monkeypatch) + agent = run_agent.AIAgent( + model="openai.gpt-5.5", + provider="custom", + api_mode="codex_responses", + base_url="https://bedrock-mantle.us-west-2.api.aws/v1", + api_key="test-token", + quiet_mode=True, + max_iterations=1, + skip_context_files=True, + skip_memory=True, + ) + + kwargs = agent._build_api_kwargs([{"role": "user", "content": "Ping"}]) + + assert kwargs["prompt_cache_retention"] == "24h" + + def test_build_api_kwargs_codex_clamps_minimal_effort(monkeypatch): """'minimal' reasoning effort is clamped to 'low' on the Responses API.