From f4fb23f3d001d164a93b7835caaa5b4c3a1002fd Mon Sep 17 00:00:00 2001 From: Georgio Constantinou <210088133+rungmc357@users.noreply.github.com> Date: Thu, 2 Jul 2026 00:33:56 -0400 Subject: [PATCH] perf(transport): gate prompt cache keys by provider capability --- agent/transports/chat_completions.py | 67 +++++++ providers/base.py | 6 + .../agent/transports/test_chat_completions.py | 165 +++++++++++++++++- 3 files changed, 237 insertions(+), 1 deletion(-) diff --git a/agent/transports/chat_completions.py b/agent/transports/chat_completions.py index 07f76736730a1..645d868d92897 100644 --- a/agent/transports/chat_completions.py +++ b/agent/transports/chat_completions.py @@ -9,6 +9,7 @@ which has provider-specific conditionals for max_tokens defaults, reasoning configuration, temperature handling, and extra_body assembly. """ +import json from typing import Any, Dict from agent.lmstudio_reasoning import resolve_lmstudio_effort @@ -18,6 +19,56 @@ from agent.transports.base import ProviderTransport from agent.transports.types import NormalizedResponse, ToolCall, Usage +def _static_prompt_instructions(messages: list[dict[str, Any]]) -> str: + """Return the stable system/developer prefix used for cache routing. + + Chat Completions carries instructions in its message list rather than a + separate ``instructions`` field. Only a leading system/developer message + is static by contract; later messages are conversation state and must not + split a warm prefix bucket on every turn. + """ + if not messages or not isinstance(messages[0], dict): + return "" + first = messages[0] + if first.get("role") not in {"system", "developer"}: + return "" + content = first.get("content") + if isinstance(content, str): + return content + try: + return json.dumps(content, sort_keys=True, ensure_ascii=False, separators=(",", ":")) + except (TypeError, ValueError): + return str(content or "") + + +def _add_prompt_cache_key( + api_kwargs: dict[str, Any], + *, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + supports_prompt_cache_key: bool, +) -> None: + """Add a content-addressed key only for an explicitly capable endpoint.""" + if not supports_prompt_cache_key: + return + + # An explicit caller body field is authoritative too. Do not add a + # duplicate top-level field whose SDK merge precedence could overwrite it. + extra_body = api_kwargs.get("extra_body") + if "prompt_cache_key" in api_kwargs or ( + isinstance(extra_body, dict) and "prompt_cache_key" in extra_body + ): + return + + # Reuse the Responses transport's single authoritative hash algorithm so + # equivalent static prefixes route to the same cache bucket across modes. + from agent.transports.codex import _content_cache_key + + cache_key = _content_cache_key(_static_prompt_instructions(messages), tools) + if cache_key: + api_kwargs["prompt_cache_key"] = cache_key + + def _reasoning_config_for_model(model: str, reasoning_config: dict | None) -> dict | None: """Return the model's wire-compatible reasoning config.""" if not isinstance(reasoning_config, dict): @@ -327,6 +378,8 @@ class ChatCompletionsTransport(ProviderTransport): # Claude on OpenRouter/Nous max output anthropic_max_output: int | None extra_body_additions: dict | None + supports_prompt_cache_key: bool — explicit endpoint capability for + the top-level Chat Completions request field; defaults off. """ # Codex sanitization: drop reasoning_items / call_id / response_item_id. # Pass model so the Gemini thought_signature (extra_content) is kept for @@ -507,6 +560,13 @@ class ChatCompletionsTransport(ProviderTransport): if overrides: api_kwargs.update(overrides) + _add_prompt_cache_key( + api_kwargs, + messages=sanitized, + tools=api_kwargs.get("tools"), + supports_prompt_cache_key=bool(params.get("supports_prompt_cache_key")), + ) + return api_kwargs def _build_kwargs_from_profile(self, profile, model, sanitized, tools, params): @@ -649,6 +709,13 @@ class ChatCompletionsTransport(ProviderTransport): if extra_body: api_kwargs["extra_body"] = extra_body + _add_prompt_cache_key( + api_kwargs, + messages=sanitized, + tools=api_kwargs.get("tools"), + supports_prompt_cache_key=bool(profile.supports_prompt_cache_key), + ) + return api_kwargs def normalize_response(self, response: Any, **kwargs) -> NormalizedResponse: diff --git a/providers/base.py b/providers/base.py index 554e01e4f7c77..1349d579bb3bb 100644 --- a/providers/base.py +++ b/providers/base.py @@ -72,6 +72,12 @@ class ProviderProfile: # (e.g. Xiaomi MiMo, which returns 400 "text is not set"). supports_vision_tool_messages: bool = True + # True only when this provider's Chat Completions endpoint explicitly + # documents ``prompt_cache_key`` as an accepted request body field. This + # is deliberately opt-in: many OpenAI-compatible endpoints reject unknown + # top-level fields rather than ignoring them. + supports_prompt_cache_key: bool = False + # ── Model catalog ───────────────────────────────────────── # fallback_models: curated list shown in /model picker when live fetch fails. # Only agentic models that support tool calling should appear here. diff --git a/tests/agent/transports/test_chat_completions.py b/tests/agent/transports/test_chat_completions.py index 92e2aaf163ca5..c3f4c3ab0b12c 100644 --- a/tests/agent/transports/test_chat_completions.py +++ b/tests/agent/transports/test_chat_completions.py @@ -1,8 +1,12 @@ """Tests for the ChatCompletionsTransport.""" -import pytest +import json from types import SimpleNamespace +import httpx +import pytest +from openai import OpenAI + from agent.transports import get_transport from agent.transports.types import NormalizedResponse @@ -537,3 +541,162 @@ class TestChatCompletionsGeminiNativeExtraBodyStrip: eb = kw.get("extra_body") assert eb and "tags" in eb + def test_tags_pass_through_on_gemini_openai_compat(self, transport): + # /openai compat endpoint is not "native" — unchanged behavior. + kw = transport.build_kwargs( + "anthropic/claude-sonnet-4.6", + [{"role": "user", "content": "hi"}], + None, + provider_profile=self._nous_profile(), + base_url="https://generativelanguage.googleapis.com/v1beta/openai", + session_id="s1", + max_tokens=None, + ) + eb = kw.get("extra_body") + assert eb and "tags" in eb + + +class TestPromptCacheKeyCapability: + """Chat Completions cache routing is opt-in and body-safe.""" + + @staticmethod + def _messages(instructions="You are stable."): + return [ + {"role": "system", "content": instructions}, + {"role": "user", "content": "hello"}, + ] + + @staticmethod + def _tools(name="lookup"): + return [{ + "type": "function", + "function": { + "name": name, + "description": "Look something up.", + "parameters": {"type": "object", "properties": {}}, + }, + }] + + def _request_body(self, kwargs, *, stream=False): + captured = {} + + def handler(request): + captured.update(json.loads(request.content)) + if stream: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + 'data: {"id":"chatcmpl_1","object":"chat.completion.chunk",' + '"choices":[{"index":0,"delta":{"content":"ok"},' + '"finish_reason":null}]}\n\n' + "data: [DONE]\n\n" + ), + ) + return httpx.Response(200, json={ + "id": "chatcmpl_1", + "object": "chat.completion", + "created": 0, + "model": kwargs["model"], + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }], + }) + + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + client = OpenAI( + api_key="test-key", + base_url="https://cache-capable.test/v1", + http_client=http_client, + ) + result = client.chat.completions.create(**kwargs, stream=stream) + if stream: + list(result) + return captured + + def test_profile_capability_emits_content_key_in_nonstream_request_body(self, transport): + from providers.base import ProviderProfile + + kwargs = transport.build_kwargs( + model="cache-model", + messages=self._messages(), + tools=self._tools(), + session_id="cron_job_2026-07-15T10:00:00Z", + provider_profile=ProviderProfile( + name="cache-capable", supports_prompt_cache_key=True, + ), + ) + + body = self._request_body(kwargs) + + assert body["prompt_cache_key"].startswith("pck_") + assert body["prompt_cache_key"] == kwargs["prompt_cache_key"] + + def test_legacy_capability_emits_same_key_in_streaming_request_body(self, transport): + kwargs = transport.build_kwargs( + model="cache-model", + messages=self._messages(), + tools=self._tools(), + session_id="cron_job_2026-07-15T10:05:00Z", + supports_prompt_cache_key=True, + ) + + body = self._request_body(kwargs, stream=True) + + assert body["prompt_cache_key"] == kwargs["prompt_cache_key"] + + @pytest.mark.parametrize("provider", [None, "anthropic", "custom"]) + def test_default_off_never_leaks_unknown_body_field(self, transport, provider): + from providers import get_provider_profile + + kwargs = transport.build_kwargs( + model="strict-model", + messages=self._messages(), + tools=self._tools(), + session_id="cron_job_2026-07-15T10:00:00Z", + provider_profile=(get_provider_profile(provider) if provider else None), + ) + + body = self._request_body(kwargs) + + assert "prompt_cache_key" not in kwargs + assert "prompt_cache_key" not in body + + def test_explicit_top_level_and_extra_body_overrides_are_preserved(self, transport): + from providers.base import ProviderProfile + + profile = ProviderProfile(name="cache-capable", supports_prompt_cache_key=True) + top_level = transport.build_kwargs( + model="cache-model", messages=self._messages(), tools=self._tools(), + provider_profile=profile, + request_overrides={"prompt_cache_key": "caller-top-level"}, + ) + in_extra_body = transport.build_kwargs( + model="cache-model", messages=self._messages(), tools=self._tools(), + provider_profile=profile, + request_overrides={"extra_body": {"prompt_cache_key": "caller-extra-body"}}, + ) + + assert top_level["prompt_cache_key"] == "caller-top-level" + assert "prompt_cache_key" not in top_level.get("extra_body", {}) + assert "prompt_cache_key" not in in_extra_body + assert in_extra_body["extra_body"]["prompt_cache_key"] == "caller-extra-body" + + def test_cron_ids_share_static_prefix_key_and_content_changes_invalidate(self, transport): + def key(session_id, *, instructions="You are stable.", tool_name="lookup"): + return transport.build_kwargs( + model="cache-model", + messages=self._messages(instructions), + tools=self._tools(tool_name), + session_id=session_id, + supports_prompt_cache_key=True, + )["prompt_cache_key"] + + first = key("cron_job_2026-07-15T10:00:00Z") + second = key("cron_job_2026-07-15T10:05:00Z") + + assert first == second + assert first != key("cron_job_2026-07-15T10:05:00Z", instructions="You are different.") + assert first != key("cron_job_2026-07-15T10:05:00Z", tool_name="search")