From 972cdde8da3d081cdb1f455f7ef4e1cdaaf85025 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Wed, 2 Sep 2026 18:20:33 +0800 Subject: [PATCH] fix(provider): preserve Codex prompt cache affinity --- nanobot/providers/openai_codex_provider.py | 26 +++++++++++++++---- tests/providers/test_openai_codex_provider.py | 23 +++++++++++++--- 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/nanobot/providers/openai_codex_provider.py b/nanobot/providers/openai_codex_provider.py index f2cfb0689..edbd90b17 100644 --- a/nanobot/providers/openai_codex_provider.py +++ b/nanobot/providers/openai_codex_provider.py @@ -111,6 +111,7 @@ class OpenAICodexProvider(LLMProvider): model=_strip_model_prefix(model), ) session_id = provider_context.session_id if provider_context is not None else None + session_routing_key = _prompt_cache_key(session_id) if session_id else None body: dict[str, Any] = { "model": _strip_model_prefix(model), @@ -122,8 +123,8 @@ class OpenAICodexProvider(LLMProvider): "tool_choice": tool_choice or "auto", "parallel_tool_calls": True, } - if session_id: - body["prompt_cache_key"] = _prompt_cache_key(session_id) + if session_routing_key: + body["prompt_cache_key"] = session_routing_key body["include"] = ["reasoning.encrypted_content"] reasoning_options = _build_reasoning_options(reasoning_effort) if replayed and "gpt-5.6" in _strip_model_prefix(model).lower(): @@ -136,13 +137,20 @@ class OpenAICodexProvider(LLMProvider): if self._extra_body: # Apply explicit provider overrides last, matching other provider backends. body.update(self._extra_body) + effective_cache_key = body.get("prompt_cache_key") stage = "oauth_token" native_compaction_applied = False native_compaction_state: ProviderConversationState | None = None try: token = await asyncio.to_thread(get_codex_token, proxy=self.proxy) - headers = _build_headers(cast(str, token.account_id), token.access) + headers = _build_headers( + cast(str, token.account_id), + token.access, + session_routing_key=( + effective_cache_key if isinstance(effective_cache_key, str) else None + ), + ) async def _send( request_body: dict[str, Any], @@ -416,8 +424,13 @@ def _build_reasoning_options(reasoning_effort: str | None) -> dict[str, str] | N return options -def _build_headers(account_id: str, token: str) -> dict[str, str]: - return { +def _build_headers( + account_id: str, + token: str, + *, + session_routing_key: str | None = None, +) -> dict[str, str]: + headers = { "Authorization": f"Bearer {token}", "chatgpt-account-id": account_id, "OpenAI-Beta": "responses=experimental", @@ -426,6 +439,9 @@ def _build_headers(account_id: str, token: str) -> dict[str, str]: "accept": "text/event-stream", "content-type": "application/json", } + if session_routing_key: + headers["session-id"] = session_routing_key + return headers class _CodexHTTPError(RuntimeError): diff --git a/tests/providers/test_openai_codex_provider.py b/tests/providers/test_openai_codex_provider.py index 687ecc884..c5391a652 100644 --- a/tests/providers/test_openai_codex_provider.py +++ b/tests/providers/test_openai_codex_provider.py @@ -265,6 +265,7 @@ async def test_codex_request_uses_configured_proxy(monkeypatch) -> None: @pytest.mark.asyncio async def test_codex_omits_prompt_cache_key_without_session_id(monkeypatch) -> None: bodies: list[dict[str, Any]] = [] + headers_seen: list[dict[str, str]] = [] _mock_codex_token(monkeypatch) @@ -280,6 +281,7 @@ async def test_codex_omits_prompt_cache_key_without_session_id(monkeypatch) -> N ): _ = proxy, on_thinking_delta, on_tool_call_delta bodies.append(body) + headers_seen.append(headers) return provider_base.LLMResponse(content="ok") monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request) @@ -294,16 +296,19 @@ async def test_codex_omits_prompt_cache_key_without_session_id(monkeypatch) -> N ) assert "prompt_cache_key" not in bodies[0] + assert "session-id" not in headers_seen[0] assert "service_tier" not in bodies[0] @pytest.mark.asyncio async def test_codex_prompt_cache_key_prefers_stable_session_id(monkeypatch) -> None: bodies: list[dict[str, Any]] = [] + headers_seen: list[dict[str, str]] = [] _mock_codex_token(monkeypatch) - async def fake_request(_url, _headers, body, **_kwargs): + async def fake_request(_url, headers, body, **_kwargs): bodies.append(body) + headers_seen.append(headers) return provider_base.LLMResponse(content="ok") monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request) @@ -326,15 +331,22 @@ async def test_codex_prompt_cache_key_prefers_stable_session_id(monkeypatch) -> assert bodies[0]["prompt_cache_key"] == bodies[1]["prompt_cache_key"] assert bodies[0]["prompt_cache_key"] != bodies[2]["prompt_cache_key"] + assert headers_seen[0]["session-id"] != "session-a" + assert headers_seen[2]["session-id"] != "session-b" + assert headers_seen[0]["session-id"] == bodies[0]["prompt_cache_key"] + assert headers_seen[1]["session-id"] == bodies[1]["prompt_cache_key"] + assert headers_seen[2]["session-id"] == bodies[2]["prompt_cache_key"] @pytest.mark.asyncio async def test_codex_provider_applies_extra_body_from_config(monkeypatch) -> None: bodies: list[dict[str, Any]] = [] + headers_seen: list[dict[str, str]] = [] _mock_codex_token(monkeypatch) - async def fake_request(_url, _headers, body, **_kwargs): + async def fake_request(_url, headers, body, **_kwargs): bodies.append(body) + headers_seen.append(headers) return provider_base.LLMResponse(content="ok") monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request) @@ -347,7 +359,10 @@ async def test_codex_provider_applies_extra_body_from_config(monkeypatch) -> Non }, "providers": { "openaiCodex": { - "extraBody": {"service_tier": "priority"}, + "extraBody": { + "service_tier": "priority", + "prompt_cache_key": "explicit-cache-key", + }, }, }, }) @@ -357,6 +372,8 @@ async def test_codex_provider_applies_extra_body_from_config(monkeypatch) -> Non assert response.content == "ok" assert bodies[0]["service_tier"] == "priority" + assert bodies[0]["prompt_cache_key"] == "explicit-cache-key" + assert headers_seen[0]["session-id"] == "explicit-cache-key" @pytest.mark.asyncio