diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index be11146c4..d77e36870 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -28,6 +28,7 @@ use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::OnceLock; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use codex_api::CompactClient as ApiCompactClient; @@ -119,6 +120,9 @@ use codex_response_debug_context::telemetry_transport_error_message; pub const OPENAI_BETA_HEADER: &str = "OpenAI-Beta"; pub const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; pub const X_CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; +pub const X_CODEX_PARENT_THREAD_ID_HEADER: &str = "x-codex-parent-thread-id"; +pub const X_CODEX_WINDOW_ID_HEADER: &str = "x-codex-window-id"; +pub const X_OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent"; pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str = "x-responsesapi-include-timing-metrics"; const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; @@ -137,6 +141,7 @@ pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = struct ModelClientState { auth_manager: Option>, conversation_id: ThreadId, + window_generation: AtomicU64, provider: ModelProviderInfo, auth_env_telemetry: AuthEnvTelemetry, session_source: SessionSource, @@ -274,6 +279,7 @@ impl ModelClient { state: Arc::new(ModelClientState { auth_manager, conversation_id, + window_generation: AtomicU64::new(0), provider, auth_env_telemetry, session_source, @@ -303,6 +309,24 @@ impl ModelClient { self.state.auth_manager.clone() } + pub(crate) fn set_window_generation(&self, window_generation: u64) { + self.state + .window_generation + .store(window_generation, Ordering::Relaxed); + self.store_cached_websocket_session(WebsocketSession::default()); + } + + pub(crate) fn advance_window_generation(&self) { + self.state.window_generation.fetch_add(1, Ordering::Relaxed); + self.store_cached_websocket_session(WebsocketSession::default()); + } + + fn current_window_id(&self) -> String { + let conversation_id = self.state.conversation_id; + let window_generation = self.state.window_generation.load(Ordering::Relaxed); + format!("{conversation_id}:{window_generation}") + } + fn take_cached_websocket_session(&self) -> WebsocketSession { let mut cached_websocket_session = self .state @@ -401,7 +425,7 @@ impl ModelClient { text, }; - let mut extra_headers = self.build_subagent_headers(); + let mut extra_headers = self.build_responses_identity_headers(); extra_headers.extend(build_conversation_headers(Some( self.state.conversation_id.to_string(), ))); @@ -461,21 +485,56 @@ impl ModelClient { fn build_subagent_headers(&self) -> ApiHeaderMap { let mut extra_headers = ApiHeaderMap::new(); - if let SessionSource::SubAgent(sub) = &self.state.session_source { - let subagent = match sub { - SubAgentSource::Review => "review".to_string(), - SubAgentSource::Compact => "compact".to_string(), - SubAgentSource::MemoryConsolidation => "memory_consolidation".to_string(), - SubAgentSource::ThreadSpawn { .. } => "collab_spawn".to_string(), - SubAgentSource::Other(label) => label.clone(), - }; - if let Ok(val) = HeaderValue::from_str(&subagent) { - extra_headers.insert("x-openai-subagent", val); - } + if let Some(subagent) = subagent_header_value(&self.state.session_source) + && let Ok(val) = HeaderValue::from_str(&subagent) + { + extra_headers.insert(X_OPENAI_SUBAGENT_HEADER, val); } extra_headers } + fn build_responses_identity_headers(&self) -> ApiHeaderMap { + let mut extra_headers = self.build_subagent_headers(); + if let Some(parent_thread_id) = parent_thread_id_header_value(&self.state.session_source) + && let Ok(val) = HeaderValue::from_str(&parent_thread_id) + { + extra_headers.insert(X_CODEX_PARENT_THREAD_ID_HEADER, val); + } + if let Ok(val) = HeaderValue::from_str(&self.current_window_id()) { + extra_headers.insert(X_CODEX_WINDOW_ID_HEADER, val); + } + extra_headers + } + + fn build_ws_client_metadata( + &self, + turn_metadata_header: Option<&str>, + ) -> HashMap { + let mut client_metadata = HashMap::new(); + client_metadata.insert( + X_CODEX_WINDOW_ID_HEADER.to_string(), + self.current_window_id(), + ); + if let Some(subagent) = subagent_header_value(&self.state.session_source) { + client_metadata.insert(X_OPENAI_SUBAGENT_HEADER.to_string(), subagent); + } + if let Some(parent_thread_id) = parent_thread_id_header_value(&self.state.session_source) { + client_metadata.insert( + X_CODEX_PARENT_THREAD_ID_HEADER.to_string(), + parent_thread_id, + ); + } + if let Some(turn_metadata_header) = parse_turn_metadata_header(turn_metadata_header) + && let Ok(turn_metadata) = turn_metadata_header.to_str() + { + client_metadata.insert( + X_CODEX_TURN_METADATA_HEADER.to_string(), + turn_metadata.to_string(), + ); + } + client_metadata + } + /// Builds request telemetry for unary API calls (e.g., Compact endpoint). fn build_request_telemetry( session_telemetry: &SessionTelemetry, @@ -655,6 +714,7 @@ impl ModelClient { headers.insert("x-client-request-id", header_value); } headers.extend(build_conversation_headers(Some(conversation_id))); + headers.extend(self.build_responses_identity_headers()); headers.insert( OPENAI_BETA_HEADER, HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), @@ -678,7 +738,7 @@ impl Drop for ModelClientSession { } impl ModelClientSession { - fn reset_websocket_session(&mut self) { + pub(crate) fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; self.websocket_session.last_response_rx = None; @@ -769,11 +829,15 @@ impl ModelClientSession { ApiResponsesOptions { conversation_id: Some(conversation_id), session_source: Some(self.client.state.session_source.clone()), - extra_headers: build_responses_headers( - self.client.state.beta_features_header.as_deref(), - Some(&self.turn_state), - turn_metadata_header.as_ref(), - ), + extra_headers: { + let mut headers = build_responses_headers( + self.client.state.beta_features_header.as_deref(), + Some(&self.turn_state), + turn_metadata_header.as_ref(), + ); + headers.extend(self.client.build_responses_identity_headers()); + headers + }, compression, turn_state: Some(Arc::clone(&self.turn_state)), } @@ -1137,7 +1201,7 @@ impl ModelClientSession { )?; let mut ws_payload = ResponseCreateWsRequest { client_metadata: response_create_client_metadata( - build_ws_client_metadata(turn_metadata_header), + Some(self.client.build_ws_client_metadata(turn_metadata_header)), request_trace.as_ref(), ), ..ResponseCreateWsRequest::from(&request) @@ -1370,14 +1434,6 @@ fn parse_turn_metadata_header(turn_metadata_header: Option<&str>) -> Option) -> Option> { - let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header)?; - let turn_metadata = turn_metadata_header.to_str().ok()?.to_string(); - let mut client_metadata = HashMap::new(); - client_metadata.insert(X_CODEX_TURN_METADATA_HEADER.to_string(), turn_metadata); - Some(client_metadata) -} - /// Builds the extra headers attached to Responses API requests. /// /// These headers implement Codex-specific conventions: @@ -1409,6 +1465,34 @@ fn build_responses_headers( headers } +fn subagent_header_value(session_source: &SessionSource) -> Option { + let SessionSource::SubAgent(subagent_source) = session_source else { + return None; + }; + match subagent_source { + SubAgentSource::Review => Some("review".to_string()), + SubAgentSource::Compact => Some("compact".to_string()), + SubAgentSource::MemoryConsolidation => Some("memory_consolidation".to_string()), + SubAgentSource::ThreadSpawn { .. } => Some("collab_spawn".to_string()), + SubAgentSource::Other(label) => Some(label.clone()), + } +} + +fn parent_thread_id_header_value(session_source: &SessionSource) -> Option { + match session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) => Some(parent_thread_id.to_string()), + SessionSource::Cli + | SessionSource::VSCode + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::SubAgent(_) + | SessionSource::Unknown => None, + } +} + fn map_response_stream( api_stream: S, session_telemetry: SessionTelemetry, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 0b37e5fea..a4fb68120 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -2,6 +2,10 @@ use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; use super::UnauthorizedRecoveryExecution; +use super::X_CODEX_PARENT_THREAD_ID_HEADER; +use super::X_CODEX_TURN_METADATA_HEADER; +use super::X_CODEX_WINDOW_ID_HEADER; +use super::X_OPENAI_SUBAGENT_HEADER; use codex_api::api_bridge::CoreAuthProvider; use codex_app_server_protocol::AuthMode; use codex_model_provider_info::WireApi; @@ -80,11 +84,49 @@ fn build_subagent_headers_sets_other_subagent_label() { ))); let headers = client.build_subagent_headers(); let value = headers - .get("x-openai-subagent") + .get(X_OPENAI_SUBAGENT_HEADER) .and_then(|value| value.to_str().ok()); assert_eq!(value, Some("memory_consolidation")); } +#[test] +fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() { + let parent_thread_id = ThreadId::new(); + let client = test_model_client(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 2, + agent_path: None, + agent_nickname: None, + agent_role: None, + })); + + client.advance_window_generation(); + + let client_metadata = client.build_ws_client_metadata(Some(r#"{"turn_id":"turn-123"}"#)); + let conversation_id = client.state.conversation_id; + assert_eq!( + client_metadata, + std::collections::HashMap::from([ + ( + X_CODEX_WINDOW_ID_HEADER.to_string(), + format!("{conversation_id}:1"), + ), + ( + X_OPENAI_SUBAGENT_HEADER.to_string(), + "collab_spawn".to_string(), + ), + ( + X_CODEX_PARENT_THREAD_ID_HEADER.to_string(), + parent_thread_id.to_string(), + ), + ( + X_CODEX_TURN_METADATA_HEADER.to_string(), + r#"{"turn_id":"turn-123"}"#.to_string(), + ), + ]) + ); +} + #[tokio::test] async fn summarize_memories_returns_empty_for_empty_input() { let client = test_model_client(SessionSource::Cli); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 3149e75bd..504420217 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1531,6 +1531,17 @@ impl Session { ), ), }; + let window_generation = match &initial_history { + InitialHistory::Resumed(resumed_history) => u64::try_from( + resumed_history + .history + .iter() + .filter(|item| matches!(item, RolloutItem::Compacted(_))) + .count(), + ) + .unwrap_or(u64::MAX), + InitialHistory::New | InitialHistory::Forked(_) => 0, + }; let state_builder = match &initial_history { InitialHistory::Resumed(resumed) => metadata::builder_from_items( resumed.history.as_slice(), @@ -1920,6 +1931,9 @@ impl Session { ), environment: environment_manager.current().await?, }; + services + .model_client + .set_window_generation(window_generation); let js_repl = Arc::new(JsReplHandle::with_node_path( config.js_repl_node_path.clone(), config.js_repl_node_module_dirs.clone(), @@ -3514,6 +3528,7 @@ impl Session { self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item)]) .await; } + self.services.model_client.advance_window_generation(); } async fn persist_rollout_response_items(&self, items: &[ResponseItem]) { @@ -5733,16 +5748,20 @@ pub(crate) async fn run_turn( let model_info = turn_context.model_info.clone(); let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX); + let mut prewarmed_client_session = prewarmed_client_session; // TODO(ccunningham): Pre-turn compaction runs before context updates and the // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively // when they would push the thread over the compaction threshold. - if run_pre_sampling_compact(&sess, &turn_context) - .await - .is_err() - { - error!("Failed to run pre-sampling compact"); - return None; + let pre_sampling_compacted = match run_pre_sampling_compact(&sess, &turn_context).await { + Ok(pre_sampling_compacted) => pre_sampling_compacted, + Err(_) => { + error!("Failed to run pre-sampling compact"); + return None; + } + }; + if pre_sampling_compacted && let Some(mut client_session) = prewarmed_client_session.take() { + client_session.reset_websocket_session(); } let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); @@ -6055,6 +6074,7 @@ pub(crate) async fn run_turn( { return None; } + client_session.reset_websocket_session(); continue; } @@ -6215,9 +6235,9 @@ pub(crate) async fn run_turn( async fn run_pre_sampling_compact( sess: &Arc, turn_context: &Arc, -) -> CodexResult<()> { +) -> CodexResult { let total_usage_tokens_before_compaction = sess.get_total_token_usage().await; - maybe_run_previous_model_inline_compact( + let mut pre_sampling_compacted = maybe_run_previous_model_inline_compact( sess, turn_context, total_usage_tokens_before_compaction, @@ -6231,8 +6251,9 @@ async fn run_pre_sampling_compact( // Compact if the total usage tokens are greater than the auto compact limit if total_usage_tokens >= auto_compact_limit { run_auto_compact(sess, turn_context, InitialContextInjection::DoNotInject).await?; + pre_sampling_compacted = true; } - Ok(()) + Ok(pre_sampling_compacted) } /// Runs pre-sampling compaction against the previous model when switching to a smaller diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 2f5b82ccc..27cf06926 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -220,6 +220,7 @@ async fn run_compact_task_inner( }; sess.replace_compacted_history(new_history, reference_context_item, compacted_item) .await; + client_session.reset_websocket_session(); sess.recompute_token_usage(&turn_context).await; sess.emit_turn_item_completed(&turn_context, compaction_item) diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 47c56357e..b4a200570 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -129,10 +129,16 @@ async fn responses_stream_includes_subagent_header_on_review() { } let request = request_recorder.single_request(); + let expected_window_id = format!("{conversation_id}:0"); assert_eq!( request.header("x-openai-subagent").as_deref(), Some("review") ); + assert_eq!( + request.header("x-codex-window-id").as_deref(), + Some(expected_window_id.as_str()) + ); + assert_eq!(request.header("x-codex-parent-thread-id"), None); assert_eq!(request.header("x-codex-sandbox"), None); } diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 0efeac47d..85dbe5e91 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -154,3 +154,4 @@ mod user_shell_cmd; mod view_image; mod web_search; mod websocket_fallback; +mod window_headers; diff --git a/codex-rs/core/tests/suite/window_headers.rs b/codex-rs/core/tests/suite/window_headers.rs new file mode 100644 index 000000000..76f135d20 --- /dev/null +++ b/codex-rs/core/tests/suite/window_headers.rs @@ -0,0 +1,146 @@ +#![allow(clippy::expect_used)] + +use super::compact::COMPACT_WARNING_MESSAGE; +use anyhow::Result; +use codex_core::CodexThread; +use codex_core::compact::SUMMARIZATION_PROMPT; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::WarningEvent; +use codex_protocol::user_input::UserInput; +use core_test_support::responses::ResponsesRequest; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use std::sync::Arc; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_fork() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("msg-1", "first reply"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-2", "summary"), + ev_completed("resp-2"), + ]), + sse(vec![ev_completed("resp-3")]), + sse(vec![ev_completed("resp-4")]), + sse(vec![ev_completed("resp-5")]), + ], + ) + .await; + + let mut builder = test_codex().with_config(|config| { + config.model_provider.name = "Non-OpenAI Model provider".to_string(); + config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string()); + }); + let initial = builder.build(&server).await?; + let initial_thread = Arc::clone(&initial.codex); + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + + submit_user_turn(&initial_thread, "before compact").await?; + submit_compact_turn(&initial_thread).await?; + submit_user_turn(&initial_thread, "after compact").await?; + shutdown_thread(&initial_thread).await?; + + let resumed = builder + .resume(&server, initial.home.clone(), rollout_path.clone()) + .await?; + submit_user_turn(&resumed.codex, "after resume").await?; + shutdown_thread(&resumed.codex).await?; + + let forked = resumed + .thread_manager + .fork_thread( + /*snapshot*/ 0usize, + resumed.config.clone(), + rollout_path, + /*persist_extended_history*/ false, + /*parent_trace*/ None, + ) + .await?; + submit_user_turn(&forked.thread, "after fork").await?; + shutdown_thread(&forked.thread).await?; + + let requests = request_log.requests(); + assert_eq!(requests.len(), 5, "expected five model requests"); + + let (initial_thread_id, first_generation) = window_id_parts(&requests[0]); + let (compact_thread_id, compact_generation) = window_id_parts(&requests[1]); + let (after_compact_thread_id, after_compact_generation) = window_id_parts(&requests[2]); + let (after_resume_thread_id, after_resume_generation) = window_id_parts(&requests[3]); + let (after_fork_thread_id, after_fork_generation) = window_id_parts(&requests[4]); + + assert_eq!(first_generation, 0); + assert_eq!(compact_thread_id, initial_thread_id); + assert_eq!(compact_generation, 0); + assert_eq!(after_compact_thread_id, initial_thread_id); + assert_eq!(after_compact_generation, 1); + assert_eq!(after_resume_thread_id, initial_thread_id); + assert_eq!(after_resume_generation, 1); + assert_ne!(after_fork_thread_id, initial_thread_id); + assert_eq!(after_fork_generation, 0); + + Ok(()) +} + +async fn submit_user_turn(codex: &Arc, text: &str) -> Result<()> { + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + }) + .await?; + wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + Ok(()) +} + +async fn submit_compact_turn(codex: &Arc) -> Result<()> { + codex.submit(Op::Compact).await?; + let warning_event = wait_for_event(codex, |event| matches!(event, EventMsg::Warning(_))).await; + let EventMsg::Warning(WarningEvent { message }) = warning_event else { + panic!("expected warning event after compact"); + }; + assert_eq!(message, COMPACT_WARNING_MESSAGE); + wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + Ok(()) +} + +async fn shutdown_thread(codex: &Arc) -> Result<()> { + codex.submit(Op::Shutdown).await?; + wait_for_event(codex, |event| matches!(event, EventMsg::ShutdownComplete)).await; + Ok(()) +} + +fn window_id_parts(request: &ResponsesRequest) -> (String, u64) { + let window_id = request + .header("x-codex-window-id") + .expect("missing x-codex-window-id header"); + let (thread_id, generation) = window_id + .rsplit_once(':') + .unwrap_or_else(|| panic!("invalid window id header: {window_id}")); + let generation = generation + .parse::() + .unwrap_or_else(|err| panic!("invalid window generation in {window_id}: {err}")); + (thread_id.to_string(), generation) +}