Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Use session IDs for prompt cache keys (#33035)
## What changed

- Default `prompt_cache_key` to the session ID from response metadata instead of the thread ID.
- Preserve explicit prompt cache key overrides.
- Verify that API-key-authenticated root and subagent requests use the same session-based cache key even though they have different thread IDs.

GitOrigin-RevId: cb1bf21a3cb8fa2c5083c5c6afd4829ff2adcf1f
  • Loading branch information
jif-oai authored and copyberry committed Jul 14, 2026
commit 4aa950d456c6c90174d3269d7eaab4a2823e5889
6 changes: 3 additions & 3 deletions codex-rs/core/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -466,10 +466,10 @@ impl ModelClient {
self
}

fn prompt_cache_key(&self) -> String {
fn prompt_cache_key(&self, responses_metadata: &CodexResponsesMetadata) -> String {
self.prompt_cache_key_override
.clone()
.unwrap_or_else(|| self.state.thread_id.to_string())
.unwrap_or_else(|| responses_metadata.session_id.clone())
}

/// Creates a fresh turn-scoped streaming session.
Expand Down Expand Up @@ -885,7 +885,7 @@ impl ModelClient {
&prompt.output_schema,
prompt.output_schema_strict,
);
let prompt_cache_key = Some(self.prompt_cache_key());
let prompt_cache_key = Some(self.prompt_cache_key(responses_metadata));
let service_tier = model_info.service_tier_for_request(service_tier);
let request = ResponsesApiRequest {
model: model_info.slug.clone(),
Expand Down
2 changes: 1 addition & 1 deletion codex-rs/core/tests/suite/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1224,7 +1224,7 @@ async fn includes_session_id_thread_id_and_model_headers_in_request() {
assert_eq!(request_authorization, "Bearer Test API Key");
assert_eq!(
request_body["prompt_cache_key"].as_str(),
Some(thread_id_string.as_str())
Some(session_id_string.as_str())
);
assert_codex_client_metadata(
&request_body,
Expand Down
1 change: 1 addition & 0 deletions codex-rs/core/tests/suite/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ mod pending_input;
mod permissions_messages;
mod personality;
mod plugins;
mod prompt_cache_key;
mod prompt_caching;
mod prompt_debug_tests;
mod quota_exceeded;
Expand Down
160 changes: 160 additions & 0 deletions codex-rs/core/tests/suite/prompt_cache_key.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
use std::time::Duration;

use anyhow::Result;
use anyhow::anyhow;
use codex_features::Feature;
use codex_login::CodexAuth;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_function_call_with_namespace;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::mount_sse_once_match;
use core_test_support::responses::sse;
use core_test_support::responses::start_mock_server;
use core_test_support::test_codex::test_codex;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;

const ROOT_PROMPT: &str = "delegate the cache audit";
const CHILD_TASK: &str = "inspect the repository";
const SPAWN_CALL_ID: &str = "spawn-worker";
const COLLABORATION_NAMESPACE: &str = "collaboration";

fn body_contains(request: &wiremock::Request, text: &str) -> bool {
serde_json::from_slice::<Value>(&request.body).is_ok_and(|body| body.to_string().contains(text))
}

fn request_has_input_type(request: &wiremock::Request, input_type: &str) -> bool {
serde_json::from_slice::<Value>(&request.body)
.ok()
.and_then(|body| body.get("input").and_then(Value::as_array).cloned())
.is_some_and(|items| {
items
.iter()
.any(|item| item.get("type").and_then(Value::as_str) == Some(input_type))
})
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn api_key_subagent_uses_session_id_as_prompt_cache_key() -> Result<()> {
let server = start_mock_server().await;
let spawn_args = serde_json::to_string(&json!({
"message": CHILD_TASK,
"task_name": "worker",
}))?;
let root_request = mount_sse_once_match(
&server,
|request: &wiremock::Request| {
body_contains(request, ROOT_PROMPT)
&& !request_has_input_type(request, "agent_message")
&& !body_contains(request, SPAWN_CALL_ID)
},
sse(vec![
ev_response_created("root-response-1"),
ev_function_call_with_namespace(
SPAWN_CALL_ID,
COLLABORATION_NAMESPACE,
"spawn_agent",
&spawn_args,
),
ev_completed("root-response-1"),
]),
)
.await;
let child_request = mount_sse_once_match(
&server,
|request: &wiremock::Request| {
body_contains(request, CHILD_TASK) && !body_contains(request, SPAWN_CALL_ID)
},
sse(vec![
ev_response_created("child-response"),
ev_assistant_message("child-message", "inspection complete"),
ev_completed("child-response"),
]),
)
.await;
mount_sse_once_match(
&server,
|request: &wiremock::Request| body_contains(request, SPAWN_CALL_ID),
sse(vec![
ev_response_created("root-response-2"),
ev_assistant_message("root-message", "worker finished"),
ev_completed("root-response-2"),
]),
)
.await;

let mut builder = test_codex()
.with_auth(CodexAuth::from_api_key("dummy"))
.with_config(|config| {
config
.features
.enable(Feature::Collab)
.expect("test config should allow feature update");
config
.features
.enable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
let expected_session_id = test.session_configured.session_id.to_string();
test.submit_turn(ROOT_PROMPT).await?;

let root_request = root_request
.requests()
.into_iter()
.next()
.expect("root request");
let root_thread_id = root_request.header("thread-id").expect("root thread ID");
let child_request = tokio::time::timeout(Duration::from_secs(10), async {
loop {
if let Some(request) = child_request.requests().into_iter().find(|request| {
request
.header("thread-id")
.is_some_and(|thread_id| thread_id != root_thread_id)
}) {
break request;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.map_err(|_| anyhow!("timed out waiting for the child request"))?;
let child_thread_id = child_request.header("thread-id").expect("child thread ID");

assert_eq!(
json!({
"differentThreadIds": root_thread_id != child_thread_id,
"root": {
"sessionId": root_request.header("session-id"),
"threadId": &root_thread_id,
"clientRequestId": root_request.header("x-client-request-id"),
"promptCacheKey": root_request.body_json()["prompt_cache_key"].clone(),
},
"child": {
"sessionId": child_request.header("session-id"),
"threadId": &child_thread_id,
"clientRequestId": child_request.header("x-client-request-id"),
"promptCacheKey": child_request.body_json()["prompt_cache_key"].clone(),
},
}),
json!({
"differentThreadIds": true,
"root": {
"sessionId": &expected_session_id,
"threadId": &root_thread_id,
"clientRequestId": &root_thread_id,
"promptCacheKey": &expected_session_id,
},
"child": {
"sessionId": &expected_session_id,
"threadId": &child_thread_id,
"clientRequestId": &child_thread_id,
"promptCacheKey": &expected_session_id,
},
})
);

Ok(())
}
Loading