From 76542bcea73785e17d918f9c27e2d9e7303969c5 Mon Sep 17 00:00:00 2001 From: kosta Date: Wed, 1 Jul 2026 21:10:23 -0400 Subject: [PATCH] Align Codex cache affinity header --- agent/transports/codex.py | 11 ++++--- .../agent/transports/test_codex_transport.py | 31 +++++++++++++++++-- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/agent/transports/codex.py b/agent/transports/codex.py index 56374b875335f..39876aa900397 100644 --- a/agent/transports/codex.py +++ b/agent/transports/codex.py @@ -253,9 +253,10 @@ def build_kwargs( # prompt_cache_key is content-addressed from the static prefix # (instructions + tools), NOT session_id — recurring cron jobs carry a # per-fire timestamp in session_id (cron__) that made every run - # cache-cold. session_id is left untouched for transcript isolation and - # the cache-scope routing headers below. Falls back to session_id when - # there is no static content to hash. + # cache-cold. session_id is left untouched for transcript isolation; + # x-client-request-id below follows this cache key so the Codex backend + # routes stable prefixes to the same cache shard. Falls back to + # session_id when there is no static content to hash. cache_key = _content_cache_key(instructions, response_tools) or session_id # xAI Responses takes prompt_cache_key in extra_body (set further # down); GitHub Models opts out of cache-key routing entirely. @@ -342,7 +343,9 @@ def build_kwargs( } ) merged_extra_headers["session_id"] = cache_scope_id - merged_extra_headers["x-client-request-id"] = cache_scope_id + merged_extra_headers["x-client-request-id"] = str( + kwargs.get("prompt_cache_key") or cache_key or cache_scope_id + ) kwargs["extra_headers"] = merged_extra_headers max_tokens = params.get("max_tokens") diff --git a/tests/agent/transports/test_codex_transport.py b/tests/agent/transports/test_codex_transport.py index 6389890519a06..058fb8a6279bc 100644 --- a/tests/agent/transports/test_codex_transport.py +++ b/tests/agent/transports/test_codex_transport.py @@ -195,8 +195,13 @@ def test_codex_backend_sets_cache_routing_headers(self, transport): ) headers = kw.get("extra_headers", {}) + # For ChatGPT/Codex, keep the raw session_id header for transcript + # identity, but route x-client-request-id to the same content-addressed + # value as prompt_cache_key so repeated stable prefixes get cache + # affinity across resumed turns and recurring jobs. assert headers.get("session_id") == "conv-codex-1" - assert headers.get("x-client-request-id") == "conv-codex-1" + assert headers.get("x-client-request-id") == kw["prompt_cache_key"] + assert headers.get("x-client-request-id") != "conv-codex-1" def test_codex_backend_no_headers_without_session_id(self, transport): messages = [{"role": "user", "content": "Hi"}] @@ -224,8 +229,30 @@ def test_codex_backend_preserves_caller_extra_headers(self, transport): headers = kw.get("extra_headers", {}) assert headers.get("x-test") == "1" + # For ChatGPT/Codex, keep the raw session_id header for transcript + # identity, but route x-client-request-id to the same content-addressed + # value as prompt_cache_key so repeated stable prefixes get cache + # affinity across resumed turns and recurring jobs. assert headers.get("session_id") == "conv-codex-1" - assert headers.get("x-client-request-id") == "conv-codex-1" + assert headers.get("x-client-request-id") == kw["prompt_cache_key"] + assert headers.get("x-client-request-id") != "conv-codex-1" + + def test_codex_backend_x_client_request_id_follows_prompt_cache_override(self, transport): + messages = [{"role": "user", "content": "Hi"}] + + kw = transport.build_kwargs( + model="gpt-5.4", + messages=messages, + tools=[], + session_id="conv-codex-1", + is_codex_backend=True, + request_overrides={"prompt_cache_key": "manual-cache-key"}, + ) + + headers = kw.get("extra_headers", {}) + assert kw["prompt_cache_key"] == "manual-cache-key" + assert headers.get("session_id") == "conv-codex-1" + assert headers.get("x-client-request-id") == "manual-cache-key" def test_non_codex_responses_preserves_caller_extra_headers(self, transport): messages = [{"role": "user", "content": "Hi"}]