From 27c05a52e0cdf93ef7db5a96ccaf88d6b9095472 Mon Sep 17 00:00:00 2001 From: Brandon Zhang Date: Fri, 21 Aug 2026 03:26:54 +0000 Subject: [PATCH] Include context window IDs in response metadata (#39847) ## What changed - Add `context_window_id` to turn metadata for sampling and compaction requests, using the UUID for the model-visible context window. - Keep the ID stable for requests in the same context window and advance it after compaction, independently of the existing thread-generation `window_id`. - Reserve `context_window_id` so client metadata cannot override it. ## Testing - Cover HTTP and WebSocket metadata, remote compaction, token-budget context alignment, resume behavior, and post-compaction window changes. GitOrigin-RevId: 55cceea377ad509c8dd96ad668f2b5703ea46a63 --- .../tests/suite/v2/client_metadata.rs | 14 ++++++++ codex-rs/core/src/compact.rs | 12 +++---- codex-rs/core/src/compact_remote_request.rs | 12 +++---- .../core/src/compact_remote_v2_attempt.rs | 12 +++---- codex-rs/core/src/responses_metadata.rs | 10 ++++++ codex-rs/core/src/session/mod.rs | 7 +++- codex-rs/core/src/session/session.rs | 18 ++++++++++ codex-rs/core/src/session/turn.rs | 8 ++--- codex-rs/core/src/turn_metadata_tests.rs | 20 +++++++---- codex-rs/core/tests/suite/compact_remote.rs | 15 +++++++++ ...pi_auth_prompt_cache_key_request_diff.snap | 2 +- ...ce_tier_prompt_cache_key_request_diff.snap | 2 +- codex-rs/core/tests/suite/token_budget.rs | 18 ++++++++++ codex-rs/core/tests/suite/window_headers.rs | 33 +++++++++++++++++++ 14 files changed, 151 insertions(+), 32 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/client_metadata.rs b/codex-rs/app-server/tests/suite/v2/client_metadata.rs index b6d7090e34..3904bee770 100644 --- a/codex-rs/app-server/tests/suite/v2/client_metadata.rs +++ b/codex-rs/app-server/tests/suite/v2/client_metadata.rs @@ -71,6 +71,10 @@ async fn turn_start_forwards_client_metadata_to_responses_request_v2() -> Result let client_metadata = HashMap::from([ ("fiber_run_id".to_string(), "fiber-start-123".to_string()), ("origin".to_string(), "gaas".to_string()), + ( + "context_window_id".to_string(), + "client-supplied".to_string(), + ), ]); let turn_req = mcp .send_turn_start_request(TurnStartParams { @@ -109,6 +113,11 @@ async fn turn_start_forwards_client_metadata_to_responses_request_v2() -> Result metadata["window_id"].as_str(), request.header("x-codex-window-id").as_deref() ); + assert!( + metadata["context_window_id"] + .as_str() + .is_some_and(|window_id| uuid::Uuid::parse_str(window_id).is_ok()) + ); Ok(()) } @@ -586,6 +595,11 @@ async fn turn_start_forwards_client_metadata_to_responses_websocket_request_body metadata["window_id"].as_str(), request["client_metadata"]["x-codex-window-id"].as_str() ); + assert!( + metadata["context_window_id"] + .as_str() + .is_some_and(|window_id| uuid::Uuid::parse_str(window_id).is_ok()) + ); websocket_server.shutdown().await; Ok(()) diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index b58e99d9d9..00a79c5442 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -261,12 +261,12 @@ async fn run_compact_task_inner_impl( // Reuse one client session so turn-scoped state (sticky routing, websocket incremental // request tracking) // survives retries within this compact turn. - let window_id = sess.current_window_id().await; - let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( - sess.installation_id.clone(), - window_id, - CodexResponsesRequestKind::Compaction(compaction_metadata), - ); + let responses_metadata = sess + .responses_metadata( + turn_context.as_ref(), + CodexResponsesRequestKind::Compaction(compaction_metadata), + ) + .await; loop { // Clone is required because of the loop diff --git a/codex-rs/core/src/compact_remote_request.rs b/codex-rs/core/src/compact_remote_request.rs index 25f8689fd8..8a46e6ce18 100644 --- a/codex-rs/core/src/compact_remote_request.rs +++ b/codex-rs/core/src/compact_remote_request.rs @@ -68,12 +68,12 @@ pub(super) async fn run_remote_compact_attempt( output_schema: None, output_schema_strict: true, }; - let window_id = sess.current_window_id().await; - let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( - sess.installation_id.clone(), - window_id, - CodexResponsesRequestKind::Compaction(compaction_metadata), - ); + let responses_metadata = sess + .responses_metadata( + turn_context.as_ref(), + CodexResponsesRequestKind::Compaction(compaction_metadata), + ) + .await; let new_history = sess .services .model_client diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index d66bf519be..6b567801b1 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -84,12 +84,12 @@ pub(super) async fn run_remote_compact_v2_attempt( output_schema_strict: true, }; - let window_id = sess.current_window_id().await; - let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( - sess.installation_id.clone(), - window_id, - CodexResponsesRequestKind::Compaction(compaction_metadata), - ); + let responses_metadata = sess + .responses_metadata( + turn_context.as_ref(), + CodexResponsesRequestKind::Compaction(compaction_metadata), + ) + .await; let trace_attempt = compaction_trace.start_attempt(&serde_json::json!({ "model": turn_context.model_info.slug.as_str(), "instructions": prompt.base_instructions.text.as_str(), diff --git a/codex-rs/core/src/responses_metadata.rs b/codex-rs/core/src/responses_metadata.rs index 64264c2599..1d12a7060b 100644 --- a/codex-rs/core/src/responses_metadata.rs +++ b/codex-rs/core/src/responses_metadata.rs @@ -16,6 +16,7 @@ use http::HeaderMap as ApiHeaderMap; use http::HeaderValue; use serde::Serialize; use serde_json::Value; +use uuid::Uuid; use crate::client::X_CODEX_INSTALLATION_ID_HEADER; use crate::client::X_CODEX_PARENT_THREAD_ID_HEADER; @@ -29,6 +30,7 @@ pub(crate) const THREAD_ID_KEY: &str = "thread_id"; pub(crate) const AGENT_NAME_KEY: &str = "agent_name"; pub(crate) const TURN_ID_KEY: &str = "turn_id"; pub(crate) const WINDOW_ID_KEY: &str = "window_id"; +pub(crate) const CONTEXT_WINDOW_ID_KEY: &str = "context_window_id"; pub(crate) const REQUEST_KIND_KEY: &str = "request_kind"; pub(crate) const COMPACTION_KEY: &str = "compaction"; // Keep the removed inventory reserved so callers cannot reintroduce oversized metadata. @@ -59,6 +61,7 @@ const RESERVED_METADATA_KEYS: &[&str] = &[ AGENT_NAME_KEY, TURN_ID_KEY, WINDOW_ID_KEY, + CONTEXT_WINDOW_ID_KEY, X_CODEX_WINDOW_ID_HEADER, X_CODEX_TURN_METADATA_HEADER, X_CODEX_PARENT_THREAD_ID_HEADER, @@ -209,6 +212,7 @@ pub struct CodexResponsesMetadata { pub(crate) turn_id: Option, pub(crate) routing_hint: Option, pub(crate) window_id: String, + pub(crate) context_window_id: Option, pub(crate) request_kind: Option, pub(crate) forked_from_thread_id: Option, pub(crate) parent_thread_id: Option, @@ -243,6 +247,7 @@ impl CodexResponsesMetadata { turn_id: None, routing_hint: None, window_id, + context_window_id: None, request_kind: None, forked_from_thread_id: None, parent_thread_id: None, @@ -365,6 +370,9 @@ impl CodexResponsesMetadata { .then_some(self.turn_id.as_deref()) .flatten(), window_id: has_request_identity.then_some(self.window_id.as_str()), + context_window_id: has_request_identity + .then_some(self.context_window_id) + .flatten(), request_kind: request_kind_value, forked_from_thread_id: self.forked_from_thread_id, parent_thread_id: self.parent_thread_id, @@ -487,6 +495,8 @@ struct CodexTurnMetadataPayload<'a> { #[serde(default, skip_serializing_if = "Option::is_none")] window_id: Option<&'a str>, #[serde(default, skip_serializing_if = "Option::is_none")] + context_window_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] request_kind: Option<&'static str>, #[serde(default, skip_serializing_if = "Option::is_none")] forked_from_thread_id: Option, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 14700373a4..d76155848a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3769,10 +3769,15 @@ impl Session { } pub(crate) async fn current_window_id(&self) -> String { + self.current_window().await.0 + } + + pub(crate) async fn current_window(&self) -> (String, Uuid) { let state = self.state.lock().await; let thread_id = self.thread_id; let window_number = state.auto_compact_window_number(); - format!("{thread_id}:{window_number}") + let context_window_id = state.auto_compact_window_ids().window_id; + (format!("{thread_id}:{window_number}"), context_window_id) } pub(crate) async fn advance_auto_compact_window(&self) -> (u64, AutoCompactWindowIds) { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8684fc4588..a492b0cc2c 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -6,6 +6,8 @@ use crate::config::ConstraintError; use crate::environment_selection::ThreadEnvironments; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::hook_mcp_executor::CoreHookMcpExecutor; +use crate::responses_metadata::CodexResponsesMetadata; +use crate::responses_metadata::CodexResponsesRequestKind; use crate::shell_snapshot::ShellSnapshot; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; @@ -638,6 +640,22 @@ impl Session { state.session_configuration.originator.clone() } + pub(crate) async fn responses_metadata( + &self, + turn_context: &TurnContext, + request_kind: CodexResponsesRequestKind, + ) -> CodexResponsesMetadata { + let (window_id, context_window_id) = self.current_window().await; + CodexResponsesMetadata { + context_window_id: Some(context_window_id), + ..turn_context.turn_metadata_state.to_responses_metadata( + self.installation_id.clone(), + window_id, + request_kind, + ) + } + } + #[instrument(name = "session_init", level = "info", skip_all)] #[allow(clippy::too_many_arguments)] pub(crate) async fn new( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index e3666553d6..9561099eac 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -375,11 +375,9 @@ pub(crate) async fn run_turn( .instrument(trace_span!("run_turn.prepare_sampling_request_input")) .await; - let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( - sess.installation_id.clone(), - window_id, - CodexResponsesRequestKind::Turn, - ); + let responses_metadata = sess + .responses_metadata(turn_context.as_ref(), CodexResponsesRequestKind::Turn) + .await; run_sampling_request( Arc::clone(&sess), Arc::clone(&step_context), diff --git a/codex-rs/core/src/turn_metadata_tests.rs b/codex-rs/core/src/turn_metadata_tests.rs index 8ede6e492d..5a4fe6cb5f 100644 --- a/codex-rs/core/src/turn_metadata_tests.rs +++ b/codex-rs/core/src/turn_metadata_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::responses_metadata::AUTO_REVIEW_ENABLED_KEY; +use crate::responses_metadata::CONTEXT_WINDOW_ID_KEY; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::CompactionTurnMetadata; use crate::responses_metadata::INSTALLATION_ID_KEY; @@ -738,6 +739,10 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields( ), ("turn_id".to_string(), "client-supplied".to_string()), (WINDOW_ID_KEY.to_string(), "client-supplied".to_string()), + ( + CONTEXT_WINDOW_ID_KEY.to_string(), + "client-supplied".to_string(), + ), ("thread_source".to_string(), "client-supplied".to_string()), ("request_kind".to_string(), "client-supplied".to_string()), ( @@ -819,6 +824,7 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields( assert_eq!(json["turn_id"].as_str(), Some("turn-a")); assert!(json.get("request_kind").is_none()); assert!(json.get(WINDOW_ID_KEY).is_none()); + assert!(json.get(CONTEXT_WINDOW_ID_KEY).is_none()); assert_eq!( json["turn_started_at_unix_ms"].as_i64(), Some(1_700_000_000_123) @@ -954,12 +960,14 @@ fn turn_metadata_state_overlays_compaction_only_on_compaction_requests() { #[test] fn responses_api_metadata_rejects_reserved_keys() { - assert_eq!( - validate_extra_metadata( - BTreeMap::from([("thread_source".to_string(), "sdk".to_string())]).iter() - ), - Err("responses_api_metadata contains a reserved key") - ); + for reserved_key in ["thread_source", WINDOW_ID_KEY, CONTEXT_WINDOW_ID_KEY] { + assert_eq!( + validate_extra_metadata( + BTreeMap::from([(reserved_key.to_string(), "sdk".to_string())]).iter() + ), + Err("responses_api_metadata contains a reserved key") + ); + } } #[tokio::test] diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index bef9b2babd..95a4db0a96 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -865,6 +865,7 @@ async fn remote_compact_replaces_history_for_followups() -> Result<()> { compact_metadata["window_id"].as_str(), compact_request.header("x-codex-window-id").as_deref() ); + assert!(compact_metadata["context_window_id"].as_str().is_some()); assert_eq!( compact_metadata["compaction"], json!({ @@ -892,6 +893,10 @@ async fn remote_compact_replaces_history_for_followups() -> Result<()> { first_response_metadata["turn_id"], compact_metadata["turn_id"], "manual compaction should use its own turn id" ); + assert_eq!( + first_response_metadata["context_window_id"], compact_metadata["context_window_id"], + "remote compaction should retain the active model-visible context window" + ); assert_eq!( compact_body["tools"], first_response_request.body_json()["tools"], @@ -951,6 +956,10 @@ async fn remote_compact_replaces_history_for_followups() -> Result<()> { follow_up_metadata["window_id"], compact_metadata["window_id"], "the following user turn should use the new compacted context window" ); + assert_ne!( + follow_up_metadata["context_window_id"], compact_metadata["context_window_id"], + "the following user turn should expose the new model-visible context window" + ); let follow_up_body = follow_up_request.body_json().to_string(); assert!( follow_up_body.contains("\"type\":\"compaction\""), @@ -1742,6 +1751,12 @@ async fn remote_compact_v2_retries_failures_with_stream_retry_budget() -> Result for compact_request in &response_requests[1..=3] { assert_eq!("/v1/responses", compact_request.path()); + let compact_metadata: Value = serde_json::from_str( + &compact_request + .header("x-codex-turn-metadata") + .expect("v2 compact request should include turn metadata"), + )?; + assert!(compact_metadata["context_window_id"].as_str().is_some()); assert!( compact_request .body_json() diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_api_auth_prompt_cache_key_request_diff.snap b/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_api_auth_prompt_cache_key_request_diff.snap index 055014de0b..2dad0171e8 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_api_auth_prompt_cache_key_request_diff.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_api_auth_prompt_cache_key_request_diff.snap @@ -12,7 +12,7 @@ Scenario: After five varied API-key-auth turns, remote manual compaction omits s - "thread_id": "", - "turn_id": "", - "x-codex-installation-id": "", -- "x-codex-turn-metadata": "{\"installation_id\":\"\",\"session_id\":\"\",\"thread_id\":\"\",\"agent_name\":\"/root\",\"turn_id\":\"\",\"window_id\":\":0\",\"request_kind\":\"turn\",\"root_turn_id\":\"\",\"sandbox\":\"\",\"sandbox_mode\":\"danger-full-access\",\"auto_review_enabled\":false,\"node_repl_auto_review_required\":false,\"node_repl_disabled\":false,\"turn_started_at_unix_ms\":}", +- "x-codex-turn-metadata": "{\"installation_id\":\"\",\"session_id\":\"\",\"thread_id\":\"\",\"agent_name\":\"/root\",\"turn_id\":\"\",\"window_id\":\":0\",\"context_window_id\":\"\",\"request_kind\":\"turn\",\"root_turn_id\":\"\",\"sandbox\":\"\",\"sandbox_mode\":\"danger-full-access\",\"auto_review_enabled\":false,\"node_repl_auto_review_required\":false,\"node_repl_disabled\":false,\"turn_started_at_unix_ms\":}", - "x-codex-window-id": ":0" - }, - "include": [ diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_chatgpt_auth_service_tier_prompt_cache_key_request_diff.snap b/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_chatgpt_auth_service_tier_prompt_cache_key_request_diff.snap index 1345e96aed..c4f8bc6830 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_chatgpt_auth_service_tier_prompt_cache_key_request_diff.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__compact_remote__remote_manual_compact_chatgpt_auth_service_tier_prompt_cache_key_request_diff.snap @@ -12,7 +12,7 @@ Scenario: After five varied ChatGPT-auth turns, remote manual compaction reuses - "thread_id": "", - "turn_id": "", - "x-codex-installation-id": "", -- "x-codex-turn-metadata": "{\"installation_id\":\"\",\"session_id\":\"\",\"thread_id\":\"\",\"agent_name\":\"/root\",\"turn_id\":\"\",\"window_id\":\":0\",\"request_kind\":\"turn\",\"root_turn_id\":\"\",\"sandbox\":\"\",\"sandbox_mode\":\"danger-full-access\",\"auto_review_enabled\":false,\"node_repl_auto_review_required\":false,\"node_repl_disabled\":false,\"turn_started_at_unix_ms\":}", +- "x-codex-turn-metadata": "{\"installation_id\":\"\",\"session_id\":\"\",\"thread_id\":\"\",\"agent_name\":\"/root\",\"turn_id\":\"\",\"window_id\":\":0\",\"context_window_id\":\"\",\"request_kind\":\"turn\",\"root_turn_id\":\"\",\"sandbox\":\"\",\"sandbox_mode\":\"danger-full-access\",\"auto_review_enabled\":false,\"node_repl_auto_review_required\":false,\"node_repl_disabled\":false,\"turn_started_at_unix_ms\":}", - "x-codex-window-id": ":0" - }, - "include": [ diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 55c06944a0..9e2981edc7 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -1022,6 +1022,15 @@ async fn token_budget_context_uses_new_window_after_compaction( assert_eq!(initial_token_budget.len(), 1); let (initial_first_window_id, initial_previous_window_id, initial_window_id) = token_budget_window_ids(&initial_token_budget[0], "/root"); + let initial_turn_metadata: Value = serde_json::from_str( + &requests[0] + .header("x-codex-turn-metadata") + .expect("initial context window metadata"), + )?; + assert_eq!( + initial_turn_metadata["context_window_id"].as_str(), + Some(initial_window_id.as_str()) + ); let post_compaction_token_budget = token_budget_contexts(&requests[1]); assert_eq!(post_compaction_token_budget.len(), 1); let ( @@ -1029,6 +1038,15 @@ async fn token_budget_context_uses_new_window_after_compaction( post_compaction_previous_window_id, post_compaction_window_id, ) = token_budget_window_ids(&post_compaction_token_budget[0], "/root"); + let post_compaction_turn_metadata: Value = serde_json::from_str( + &requests[1] + .header("x-codex-turn-metadata") + .expect("post-compaction context window metadata"), + )?; + assert_eq!( + post_compaction_turn_metadata["context_window_id"].as_str(), + Some(post_compaction_window_id.as_str()) + ); assert_eq!(initial_previous_window_id, None); assert_eq!(initial_first_window_id, initial_window_id); assert_eq!(post_compaction_first_window_id, initial_first_window_id); diff --git a/codex-rs/core/tests/suite/window_headers.rs b/codex-rs/core/tests/suite/window_headers.rs index 273b4eaa9a..a0a2312e93 100644 --- a/codex-rs/core/tests/suite/window_headers.rs +++ b/codex-rs/core/tests/suite/window_headers.rs @@ -97,6 +97,39 @@ async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_fork( assert_ne!(after_fork_thread_id, initial_thread_id); assert_eq!(after_fork_generation, 0); + let metadata = requests + .iter() + .map(|request| { + let metadata = request + .header("x-codex-turn-metadata") + .expect("turn metadata header"); + serde_json::from_str::(&metadata).expect("valid turn metadata") + }) + .collect::>(); + for (request, metadata) in requests.iter().zip(&metadata) { + assert_eq!( + metadata["window_id"].as_str(), + request.header("x-codex-window-id").as_deref() + ); + assert!( + metadata["context_window_id"] + .as_str() + .is_some_and(|window_id| uuid::Uuid::parse_str(window_id).is_ok()) + ); + } + assert_eq!( + metadata[0]["context_window_id"], + metadata[1]["context_window_id"] + ); + assert_ne!( + metadata[1]["context_window_id"], + metadata[2]["context_window_id"] + ); + assert_eq!( + metadata[2]["context_window_id"], + metadata[3]["context_window_id"] + ); + Ok(()) }