From 42f37dc4c0e409eac7818a82d165c5f2757bc075 Mon Sep 17 00:00:00 2001 From: chengyongru <61816729+chengyongru@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:36:00 +0800 Subject: [PATCH] fix(provider): stabilize Codex prompt cache routing (#5540) --- nanobot/agent/runner.py | 1 + nanobot/providers/base.py | 4 ++ nanobot/providers/conversation_state.py | 10 +++- nanobot/providers/fallback_provider.py | 2 + nanobot/providers/openai_codex_provider.py | 9 ++-- tests/agent/test_runner_fallback.py | 6 ++- tests/agent/test_runner_runtime_identity.py | 7 +++ tests/providers/test_conversation_state.py | 21 ++++++++ tests/providers/test_openai_codex_provider.py | 51 ++++++++++++------- tests/providers/test_provider_retry.py | 6 ++- 10 files changed, 91 insertions(+), 26 deletions(-) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index a777bb6c6..94955bfaf 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -495,6 +495,7 @@ class AgentRunner: model=spec.runtime.model, messages=messages, state=spec.provider_state, + session_id=spec.session_key, ) governance_config = ContextGovernanceConfig( provider=spec.runtime.provider, diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index 23a43b32d..eb8c6c545 100644 --- a/nanobot/providers/base.py +++ b/nanobot/providers/base.py @@ -252,10 +252,13 @@ class ProviderCallContext: The regular ``chat`` contract stays provider-agnostic. Responses-capable providers consume this context through the opt-in ``chat_with_context`` hooks, while every other provider inherits the context-free delegation. + ``session_id`` gives providers a stable conversation-scoped routing key + without exposing that identity in the public message transcript. """ conversation_state: ProviderConversationState | None = field(default=None, repr=False) context_window_tokens: int | None = None + session_id: str | None = field(default=None, repr=False) @dataclass(frozen=True, slots=True) @@ -1640,6 +1643,7 @@ class LLMProvider(ABC): context_window_tokens=( provider_context.context_window_tokens ), + session_id=provider_context.session_id, ) if stripped is not None or stripped_context is not None: logger.warning( diff --git a/nanobot/providers/conversation_state.py b/nanobot/providers/conversation_state.py index 462fc316f..6d891603a 100644 --- a/nanobot/providers/conversation_state.py +++ b/nanobot/providers/conversation_state.py @@ -42,9 +42,11 @@ class ProviderConversationStateController: model: str | None, messages: list[dict[str, Any]], state: ProviderConversationState | None = None, + session_id: str | None = None, ) -> None: self._provider = provider self._model = model + self._session_id = session_id self._state = ( state if state is not None @@ -60,9 +62,12 @@ class ProviderConversationStateController: context_window_tokens: int | None, ) -> ProviderCallContext | None: """Return typed provider context for a request that does not resume state.""" - if context_window_tokens is None: + if context_window_tokens is None and self._session_id is None: return None - return ProviderCallContext(context_window_tokens=context_window_tokens) + return ProviderCallContext( + context_window_tokens=context_window_tokens, + session_id=self._session_id, + ) def prepare_request( self, @@ -112,6 +117,7 @@ class ProviderConversationStateController: if independent_context is not None else None ), + session_id=self._session_id, ) def observe_response( diff --git a/nanobot/providers/fallback_provider.py b/nanobot/providers/fallback_provider.py index 24a55227a..12f22f2b1 100644 --- a/nanobot/providers/fallback_provider.py +++ b/nanobot/providers/fallback_provider.py @@ -186,6 +186,7 @@ class FallbackProvider(LLMProvider): return ProviderCallContext( conversation_state=provider_context.conversation_state, context_window_tokens=context_window_tokens, + session_id=provider_context.session_id, ) def _primary_available(self) -> bool: @@ -541,6 +542,7 @@ class FallbackProvider(LLMProvider): fallback_kwargs["provider_context"] = ProviderCallContext( conversation_state=state, context_window_tokens=context_window_tokens, + session_id=provider_context.session_id, ) if fallback.reasoning_effort is None: fallback_kwargs.pop("reasoning_effort", None) diff --git a/nanobot/providers/openai_codex_provider.py b/nanobot/providers/openai_codex_provider.py index d8cd4cd82..91d2dfffd 100644 --- a/nanobot/providers/openai_codex_provider.py +++ b/nanobot/providers/openai_codex_provider.py @@ -103,6 +103,7 @@ class OpenAICodexProvider(LLMProvider): provider=self._responses_state_provider(), model=_strip_model_prefix(model), ) + session_id = provider_context.session_id if provider_context is not None else None body: dict[str, Any] = { "model": _strip_model_prefix(model), @@ -111,10 +112,11 @@ class OpenAICodexProvider(LLMProvider): "instructions": system_prompt, "input": input_items, "text": {"verbosity": "medium"}, - "prompt_cache_key": _prompt_cache_key(messages[:2]), "tool_choice": tool_choice or "auto", "parallel_tool_calls": True, } + if session_id: + body["prompt_cache_key"] = _prompt_cache_key(session_id) body["include"] = ["reasoning.encrypted_content"] reasoning_options = _build_reasoning_options(reasoning_effort) if replayed and "gpt-5.6" in _strip_model_prefix(model).lower(): @@ -496,9 +498,8 @@ async def _request_codex( return result -def _prompt_cache_key(messages: list[dict[str, Any]]) -> str: - raw = json.dumps(messages, ensure_ascii=True, sort_keys=True) - return hashlib.sha256(raw.encode("utf-8")).hexdigest() +def _prompt_cache_key(session_id: str) -> str: + return hashlib.sha256(session_id.encode("utf-8")).hexdigest() def _friendly_error(status_code: int, raw: str) -> str: diff --git a/tests/agent/test_runner_fallback.py b/tests/agent/test_runner_fallback.py index fd4fcdaa3..ee2805bec 100644 --- a/tests/agent/test_runner_fallback.py +++ b/tests/agent/test_runner_fallback.py @@ -384,12 +384,16 @@ class TestFallbackOnPrimaryError: messages=[{"role": "user", "content": "hi"}], model="gpt-5.6", max_tokens=10_000, - provider_context=ProviderCallContext(context_window_tokens=50_000), + provider_context=ProviderCallContext( + context_window_tokens=50_000, + session_id="webui:cache-test", + ), ) primary_context = primary.context_calls[0] assert primary_context is not None assert primary_context.context_window_tokens == 200_000 + assert primary_context.session_id == "webui:cache-test" assert resolve_compact_threshold( primary_context.context_window_tokens, 10_000, diff --git a/tests/agent/test_runner_runtime_identity.py b/tests/agent/test_runner_runtime_identity.py index 8458125d7..a0c186a16 100644 --- a/tests/agent/test_runner_runtime_identity.py +++ b/tests/agent/test_runner_runtime_identity.py @@ -8,6 +8,7 @@ from nanobot.providers.base import ( GenerationSettings, LLMProvider, LLMResponse, + ProviderCallContext, ToolCallRequest, ) from nanobot.utils.llm_runtime import LLMRuntime @@ -22,6 +23,7 @@ async def test_active_run_keeps_provider_captured_at_admission() -> None: first_calls = 0 second_calls = 0 request_temperatures: list[float] = [] + request_session_ids: list[str | None] = [] selected_runtime = LLMRuntime.capture( first_provider, "captured-model", @@ -33,6 +35,9 @@ async def test_active_run_keeps_provider_captured_at_admission() -> None: nonlocal first_calls, selected_runtime first_calls += 1 request_temperatures.append(kwargs["temperature"]) + provider_context = kwargs["provider_context"] + assert isinstance(provider_context, ProviderCallContext) + request_session_ids.append(provider_context.session_id) selected_runtime = LLMRuntime.capture( second_provider, "future-model", @@ -63,9 +68,11 @@ async def test_active_run_keeps_provider_captured_at_admission() -> None: runtime=selected_runtime, max_iterations=2, max_tool_result_chars=AgentDefaults().max_tool_result_chars, + session_key="webui:cache-test", )) assert first_calls == 2 assert second_calls == 0 assert request_temperatures == [0.2, 0.2] + assert request_session_ids == ["webui:cache-test", "webui:cache-test"] assert selected_runtime.provider is second_provider diff --git a/tests/providers/test_conversation_state.py b/tests/providers/test_conversation_state.py index 157fc2c80..7c62159e8 100644 --- a/tests/providers/test_conversation_state.py +++ b/tests/providers/test_conversation_state.py @@ -9,6 +9,7 @@ import pytest from nanobot.providers.base import ( LLMProvider, LLMResponse, + ProviderCallContext, ProviderConversationState, ToolCallRequest, ) @@ -289,3 +290,23 @@ def test_independent_request_exposes_context_without_capability_check() -> None: assert provider_context.conversation_state is None assert provider_context.context_window_tokens == 200_000 provider.supports_native_compaction.assert_not_called() + + +def test_independent_request_exposes_session_id_without_token_budget() -> None: + provider = _provider(compact=False) + messages = [{"role": "user", "content": "hello"}] + controller = ProviderConversationStateController( + provider=provider, + model="gpt-5.6", + messages=messages, + session_id="webui:cache-test", + ) + + provider_context = controller.prepare_request( + messages, + context_window_tokens=None, + ) + + assert provider_context == ProviderCallContext( + session_id="webui:cache-test", + ) diff --git a/tests/providers/test_openai_codex_provider.py b/tests/providers/test_openai_codex_provider.py index 57f37a094..e7c2b05a9 100644 --- a/tests/providers/test_openai_codex_provider.py +++ b/tests/providers/test_openai_codex_provider.py @@ -260,8 +260,8 @@ async def test_codex_request_uses_configured_proxy(monkeypatch) -> None: @pytest.mark.asyncio -async def test_codex_prompt_cache_key_uses_stable_conversation_prefix(monkeypatch) -> None: - bodies: list[dict] = [] +async def test_codex_omits_prompt_cache_key_without_session_id(monkeypatch) -> None: + bodies: list[dict[str, Any]] = [] _mock_codex_token(monkeypatch) @@ -289,25 +289,40 @@ async def test_codex_prompt_cache_key_uses_stable_conversation_prefix(monkeypatc {"role": "assistant", "content": "first answer"}, ], ) - await provider.chat( - [ - {"role": "system", "content": "You are nanobot."}, - {"role": "user", "content": "first request"}, - {"role": "assistant", "content": "first answer"}, - {"role": "user", "content": "follow up"}, - ], - ) - await provider.chat( - [ - {"role": "system", "content": "You are nanobot."}, - {"role": "user", "content": "different request"}, - {"role": "assistant", "content": "first answer"}, - ], - ) + + assert "prompt_cache_key" not in bodies[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]] = [] + _mock_codex_token(monkeypatch) + + async def fake_request(_url, _headers, body, **_kwargs): + bodies.append(body) + return provider_base.LLMResponse(content="ok") + + monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request) + provider = OpenAICodexProvider() + + for session_id, first_request in ( + ("session-a", "first request"), + ("session-a", "different visible prefix"), + ("session-b", "first request"), + ): + await provider.chat( + [ + {"role": "system", "content": "You are nanobot."}, + {"role": "user", "content": first_request}, + ], + provider_context=provider_base.ProviderCallContext( + session_id=session_id, + ), + ) assert bodies[0]["prompt_cache_key"] == bodies[1]["prompt_cache_key"] assert bodies[0]["prompt_cache_key"] != bodies[2]["prompt_cache_key"] - assert all("service_tier" not in body for body in bodies) @pytest.mark.asyncio diff --git a/tests/providers/test_provider_retry.py b/tests/providers/test_provider_retry.py index 5f603de91..7c3bf7af8 100644 --- a/tests/providers/test_provider_retry.py +++ b/tests/providers/test_provider_retry.py @@ -431,13 +431,17 @@ async def test_image_retry_discards_provider_state_with_images( response = await provider.chat_with_retry( messages=messages, - provider_context=ProviderCallContext(conversation_state=state), + provider_context=ProviderCallContext( + conversation_state=state, + session_id="webui:cache-test", + ), ) assert response.content == "ok, no image" retry_context = provider.contexts[-1] assert isinstance(retry_context, ProviderCallContext) assert retry_context.conversation_state is None + assert retry_context.session_id == "webui:cache-test" public_content = messages[0]["content"] if isinstance(public_content, list): assert all(block.get("type") != "image_url" for block in public_content)