From 6c0a92420376514aa6f4fc2a47e853023aad99f7 Mon Sep 17 00:00:00 2001 From: pash-openai Date: Fri, 13 Feb 2026 12:48:29 -0800 Subject: [PATCH] turn metadata: per-turn non-blocking (#11677) --- codex-rs/core/src/client.rs | 3 +- codex-rs/core/src/codex.rs | 84 +++-- codex-rs/core/src/compact.rs | 2 +- codex-rs/core/src/git_info.rs | 46 ++- codex-rs/core/src/memories/phase1.rs | 2 +- codex-rs/core/src/tasks/mod.rs | 3 + codex-rs/core/src/tasks/regular.rs | 5 +- codex-rs/core/src/turn_metadata.rs | 302 +++++++++++++++--- codex-rs/core/tests/responses_headers.rs | 142 +++++--- .../core/tests/suite/client_websockets.rs | 6 +- codex-rs/core/tests/suite/turn_state.rs | 20 ++ 11 files changed, 457 insertions(+), 158 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ec0d1e54f..99f8464ac 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -633,7 +633,6 @@ impl ModelClientSession { &mut self, otel_manager: &OtelManager, model_info: &ModelInfo, - turn_metadata_header: Option<&str>, ) -> std::result::Result<(), ApiError> { if !self.client.responses_websocket_enabled(model_info) || self.client.websockets_disabled() { @@ -656,7 +655,7 @@ impl ModelClientSession { client_setup.api_provider, client_setup.api_auth, Some(Arc::clone(&self.turn_state)), - turn_metadata_header, + None, ) .await?; self.connection = Some(connection); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4fbfe7b50..59cbf57dd 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -38,8 +38,7 @@ use crate::stream_events_utils::handle_output_item_done; use crate::stream_events_utils::last_assistant_message_from_item; use crate::terminal; use crate::truncate::TruncationPolicy; -use crate::turn_metadata::build_turn_metadata_header; -use crate::turn_metadata::resolve_turn_metadata_header_with_timeout; +use crate::turn_metadata::TurnMetadataState; use crate::util::error_or_panic; use async_channel::Receiver; use async_channel::Sender; @@ -92,7 +91,6 @@ use rmcp::model::RequestId; use serde_json; use serde_json::Value; use tokio::sync::Mutex; -use tokio::sync::OnceCell; use tokio::sync::RwLock; use tokio::sync::oneshot; use tokio::sync::watch; @@ -558,7 +556,7 @@ pub(crate) struct TurnContext { pub(crate) truncation_policy: TruncationPolicy, pub(crate) js_repl: Arc, pub(crate) dynamic_tools: Vec, - turn_metadata_header: OnceCell>, + turn_metadata_state: Arc, } impl TurnContext { pub(crate) fn model_context_window(&self) -> Option { @@ -638,7 +636,7 @@ impl TurnContext { truncation_policy, js_repl: Arc::clone(&self.js_repl), dynamic_tools: self.dynamic_tools.clone(), - turn_metadata_header: self.turn_metadata_header.clone(), + turn_metadata_state: self.turn_metadata_state.clone(), } } @@ -689,39 +687,16 @@ impl TurnContext { }) } - async fn build_turn_metadata_header(&self) -> Option { - let sandbox = sandbox_tag(&self.sandbox_policy, self.windows_sandbox_level); - self.turn_metadata_header - .get_or_init(|| async { - build_turn_metadata_header(self.cwd.as_path(), Some(sandbox)).await - }) - .await - .clone() + pub fn current_turn_metadata_header(&self) -> Option { + self.turn_metadata_state.current_header_value() } - /// Resolves the per-turn metadata header under a shared timeout policy. - /// - /// This uses the same timeout helper as websocket startup prewarm so both turn execution and - /// background prewarm observe identical "timeout means best-effort fallback" behavior. - pub async fn resolve_turn_metadata_header(&self) -> Option { - resolve_turn_metadata_header_with_timeout( - self.build_turn_metadata_header(), - self.turn_metadata_header.get().cloned().flatten(), - ) - .await + pub fn spawn_turn_metadata_enrichment_task(self: &Arc) { + self.turn_metadata_state.spawn_git_enrichment_task(); } - /// Starts best-effort background computation of turn metadata. - /// - /// This warms the cached value used by [`TurnContext::resolve_turn_metadata_header`] so turns - /// and websocket prewarm are less likely to pay metadata construction latency on demand. - pub fn spawn_turn_metadata_header_task(self: &Arc) { - let context = Arc::clone(self); - tokio::spawn(async move { - trace!("Spawning turn metadata calculation task"); - context.build_turn_metadata_header().await; - trace!("Turn metadata calculation task completed"); - }); + pub fn cancel_turn_metadata_enrichment_task(&self) { + self.turn_metadata_state.cancel_git_enrichment_task(); } } @@ -948,6 +923,17 @@ impl Session { }); let cwd = session_configuration.cwd.clone(); + let turn_metadata_state = Arc::new(TurnMetadataState::new( + sub_id.clone(), + cwd.clone(), + Some( + sandbox_tag( + session_configuration.sandbox_policy.get(), + session_configuration.windows_sandbox_level, + ) + .to_string(), + ), + )); TurnContext { sub_id, config: per_turn_config.clone(), @@ -978,7 +964,7 @@ impl Session { truncation_policy: model_info.truncation_policy.into(), js_repl, dynamic_tools: session_configuration.dynamic_tools.clone(), - turn_metadata_header: OnceCell::new(), + turn_metadata_state, } } @@ -1273,17 +1259,10 @@ impl Session { let prewarm_model_info = models_manager .get_model_info(session_configuration.collaboration_mode.model(), &config) .await; - let prewarm_cwd = session_configuration.cwd.clone(); - let turn_metadata_header = resolve_turn_metadata_header_with_timeout( - async move { build_turn_metadata_header(prewarm_cwd.as_path(), None).await }, - None, - ) - .boxed(); let startup_regular_task = RegularTask::with_startup_prewarm( services.model_client.clone(), services.otel_manager.clone(), prewarm_model_info, - turn_metadata_header, ); state.set_startup_regular_task(startup_regular_task); @@ -1779,7 +1758,7 @@ impl Session { turn_context.final_output_json_schema = final_schema; } let turn_context = Arc::new(turn_context); - turn_context.spawn_turn_metadata_header_task(); + turn_context.spawn_turn_metadata_enrichment_task(); turn_context } @@ -4026,9 +4005,21 @@ async fn spawn_review_thread( let session_source = parent_turn_context.session_source.clone(); let per_turn_config = Arc::new(per_turn_config); + let review_turn_id = sub_id.to_string(); + let turn_metadata_state = Arc::new(TurnMetadataState::new( + review_turn_id.clone(), + parent_turn_context.cwd.clone(), + Some( + sandbox_tag( + &parent_turn_context.sandbox_policy, + parent_turn_context.windows_sandbox_level, + ) + .to_string(), + ), + )); let review_turn_context = TurnContext { - sub_id: sub_id.to_string(), + sub_id: review_turn_id, config: per_turn_config, auth_manager: auth_manager_for_context, model_info: model_info.clone(), @@ -4057,7 +4048,7 @@ async fn spawn_review_thread( js_repl: Arc::clone(&sess.js_repl), dynamic_tools: parent_turn_context.dynamic_tools.clone(), truncation_policy: model_info.truncation_policy.into(), - turn_metadata_header: parent_turn_context.turn_metadata_header.clone(), + turn_metadata_state, }; // Seed the child task with the review prompt as the initial user message. @@ -4067,6 +4058,7 @@ async fn spawn_review_thread( text_elements: Vec::new(), }]; let tc = Arc::new(review_turn_context); + tc.spawn_turn_metadata_enrichment_task(); sess.spawn_task(tc.clone(), input, ReviewTask::new()).await; // Announce entering review mode so UIs can switch modes. @@ -4298,7 +4290,6 @@ pub(crate) async fn run_turn( // many turns, from the perspective of the user, it is a single turn. let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); - let turn_metadata_header = turn_context.resolve_turn_metadata_header().await; // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse // one instance across retries within this turn. let mut client_session = @@ -4350,6 +4341,7 @@ pub(crate) async fn run_turn( }) .map(|user_message| user_message.message()) .collect::>(); + let turn_metadata_header = turn_context.current_turn_metadata_header(); match run_sampling_request( Arc::clone(&sess), Arc::clone(&turn_context), diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index ad33144bd..08932bf7a 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -84,7 +84,6 @@ async fn run_compact_task_inner( let max_retries = turn_context.provider.stream_max_retries(); let mut retries = 0; - let turn_metadata_header = turn_context.resolve_turn_metadata_header().await; let mut client_session = sess.services.model_client.new_session(); // Reuse one client session so turn-scoped state (sticky routing, websocket append tracking) // survives retries within this compact turn. @@ -110,6 +109,7 @@ async fn run_compact_task_inner( personality: turn_context.personality, ..Default::default() }; + let turn_metadata_header = turn_context.current_turn_metadata_header(); let attempt_result = drain_to_completed( &sess, turn_context.as_ref(), diff --git a/codex-rs/core/src/git_info.rs b/codex-rs/core/src/git_info.rs index dcabc3a34..d735d4dbd 100644 --- a/codex-rs/core/src/git_info.rs +++ b/codex-rs/core/src/git_info.rs @@ -150,6 +150,15 @@ pub async fn get_head_commit_hash(cwd: &Path) -> Option { } } +pub async fn get_has_changes(cwd: &Path) -> Option { + let output = run_git_command_with_timeout(&["status", "--porcelain"], cwd).await?; + if !output.status.success() { + return None; + } + + Some(!output.stdout.is_empty()) +} + fn parse_git_remote_urls(stdout: &str) -> Option> { let mut remotes = BTreeMap::new(); for line in stdout.lines() { @@ -254,7 +263,11 @@ pub async fn git_diff_to_remote(cwd: &Path) -> Option { /// Run a git command with a timeout to prevent blocking on large repositories async fn run_git_command_with_timeout(args: &[&str], cwd: &Path) -> Option { let mut command = Command::new("git"); - command.args(args).current_dir(cwd).kill_on_drop(true); + command + .env("GIT_OPTIONAL_LOCKS", "0") + .args(args) + .current_dir(cwd) + .kill_on_drop(true); let result = timeout(GIT_COMMAND_TIMEOUT, command.output()).await; match result { @@ -962,6 +975,37 @@ mod tests { assert_eq!(git_info.branch, Some("feature-branch".to_string())); } + #[tokio::test] + async fn test_get_has_changes_non_git_directory_returns_none() { + let temp_dir = TempDir::new().expect("Failed to create temp dir"); + assert_eq!(get_has_changes(temp_dir.path()).await, None); + } + + #[tokio::test] + async fn test_get_has_changes_clean_repo_returns_false() { + let temp_dir = TempDir::new().expect("Failed to create temp dir"); + let repo_path = create_test_git_repo(&temp_dir).await; + assert_eq!(get_has_changes(&repo_path).await, Some(false)); + } + + #[tokio::test] + async fn test_get_has_changes_with_tracked_change_returns_true() { + let temp_dir = TempDir::new().expect("Failed to create temp dir"); + let repo_path = create_test_git_repo(&temp_dir).await; + + fs::write(repo_path.join("test.txt"), "updated tracked file").expect("write tracked file"); + assert_eq!(get_has_changes(&repo_path).await, Some(true)); + } + + #[tokio::test] + async fn test_get_has_changes_with_untracked_change_returns_true() { + let temp_dir = TempDir::new().expect("Failed to create temp dir"); + let repo_path = create_test_git_repo(&temp_dir).await; + + fs::write(repo_path.join("new_file.txt"), "untracked").expect("write untracked file"); + assert_eq!(get_has_changes(&repo_path).await, Some(true)); + } + #[tokio::test] async fn test_get_git_working_tree_state_clean_repo() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); diff --git a/codex-rs/core/src/memories/phase1.rs b/codex-rs/core/src/memories/phase1.rs index 654e062e3..17e0359c9 100644 --- a/codex-rs/core/src/memories/phase1.rs +++ b/codex-rs/core/src/memories/phase1.rs @@ -198,7 +198,7 @@ async fn build_request_context(session: &Arc, config: &Config) -> Reque let turn_context = session.new_default_turn().await; RequestContext::from_turn_context( turn_context.as_ref(), - turn_context.resolve_turn_metadata_header().await, + turn_context.current_turn_metadata_header(), model, ) } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index e05912f5f..17961d38b 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -197,6 +197,8 @@ impl Session { turn_context: Arc, last_agent_message: Option, ) { + turn_context.cancel_turn_metadata_enrichment_task(); + let mut active = self.active_turn.lock().await; let mut pending_input = Vec::::new(); let mut should_clear_active_turn = false; @@ -260,6 +262,7 @@ impl Session { trace!(task_kind = ?task.kind, sub_id, "aborting running task"); task.cancellation_token.cancel(); + task.turn_context.cancel_turn_metadata_enrichment_task(); let session_task = task.task; select! { diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 8782d7ced..1428f7775 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -10,7 +10,6 @@ use async_trait::async_trait; use codex_otel::OtelManager; use codex_protocol::openai_models::ModelInfo; use codex_protocol::user_input::UserInput; -use futures::future::BoxFuture; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use tracing::Instrument; @@ -39,13 +38,11 @@ impl RegularTask { model_client: ModelClient, otel_manager: OtelManager, model_info: ModelInfo, - turn_metadata_header: BoxFuture<'static, Option>, ) -> Self { let prewarmed_session_task = tokio::spawn(async move { let mut client_session = model_client.new_session(); - let turn_metadata_header = turn_metadata_header.await; match client_session - .prewarm_websocket(&otel_manager, &model_info, turn_metadata_header.as_deref()) + .prewarm_websocket(&otel_manager, &model_info) .await { Ok(()) => Some(client_session), diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index e17a78804..dd4d71744 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -1,88 +1,284 @@ -//! Helpers for computing and resolving optional per-turn metadata headers. -//! -//! This module owns both metadata construction and the shared timeout policy used by -//! turn execution and startup websocket prewarm. Keeping timeout behavior centralized -//! ensures both call sites treat timeout as the same best-effort fallback condition. - use std::collections::BTreeMap; -use std::future::Future; use std::path::Path; -use std::time::Duration; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::RwLock; use serde::Serialize; -use tracing::warn; +use tokio::task::JoinHandle; use crate::git_info::get_git_remote_urls_assume_git_repo; use crate::git_info::get_git_repo_root; +use crate::git_info::get_has_changes; use crate::git_info::get_head_commit_hash; -pub(crate) const TURN_METADATA_HEADER_TIMEOUT: Duration = Duration::from_millis(250); +#[derive(Clone, Debug, Default)] +struct WorkspaceGitMetadata { + associated_remote_urls: Option>, + latest_git_commit_hash: Option, + has_changes: Option, +} -/// Resolves turn metadata with a shared timeout policy. -/// -/// On timeout, this logs a warning and returns the provided fallback header. -/// -/// Keeping this helper centralized avoids drift between turn-time metadata resolution and startup -/// websocket prewarm, both of which need identical timeout semantics. -pub(crate) async fn resolve_turn_metadata_header_with_timeout( - build_header: F, - fallback_on_timeout: Option, -) -> Option -where - F: Future>, -{ - match tokio::time::timeout(TURN_METADATA_HEADER_TIMEOUT, build_header).await { - Ok(header) => header, - Err(_) => { - warn!( - "timed out after {}ms while building turn metadata header", - TURN_METADATA_HEADER_TIMEOUT.as_millis() - ); - fallback_on_timeout - } +impl WorkspaceGitMetadata { + fn is_empty(&self) -> bool { + self.associated_remote_urls.is_none() + && self.latest_git_commit_hash.is_none() + && self.has_changes.is_none() } } -#[derive(Serialize)] +#[derive(Clone, Debug, Serialize, Default)] struct TurnMetadataWorkspace { #[serde(default, skip_serializing_if = "Option::is_none")] associated_remote_urls: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] latest_git_commit_hash: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + has_changes: Option, } -#[derive(Serialize)] -struct TurnMetadata { +impl From for TurnMetadataWorkspace { + fn from(value: WorkspaceGitMetadata) -> Self { + Self { + associated_remote_urls: value.associated_remote_urls, + latest_git_commit_hash: value.latest_git_commit_hash, + has_changes: value.has_changes, + } + } +} + +#[derive(Clone, Debug, Serialize, Default)] +pub(crate) struct TurnMetadataBag { + #[serde(default, skip_serializing_if = "Option::is_none")] + turn_id: Option, #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] workspaces: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] sandbox: Option, } -pub async fn build_turn_metadata_header(cwd: &Path, sandbox: Option<&str>) -> Option { - let repo_root = get_git_repo_root(cwd); +impl TurnMetadataBag { + fn to_header_value(&self) -> Option { + serde_json::to_string(self).ok() + } +} - let (latest_git_commit_hash, associated_remote_urls) = tokio::join!( +fn build_turn_metadata_bag( + turn_id: Option, + sandbox: Option, + repo_root: Option, + workspace_git_metadata: Option, +) -> TurnMetadataBag { + let mut workspaces = BTreeMap::new(); + if let (Some(repo_root), Some(workspace_git_metadata)) = (repo_root, workspace_git_metadata) + && !workspace_git_metadata.is_empty() + { + workspaces.insert(repo_root, workspace_git_metadata.into()); + } + + TurnMetadataBag { + turn_id, + workspaces, + sandbox, + } +} + +pub async fn build_turn_metadata_header(cwd: &Path, sandbox: Option<&str>) -> Option { + let repo_root = get_git_repo_root(cwd).map(|root| root.to_string_lossy().into_owned()); + + let (latest_git_commit_hash, associated_remote_urls, has_changes) = tokio::join!( get_head_commit_hash(cwd), - get_git_remote_urls_assume_git_repo(cwd) + get_git_remote_urls_assume_git_repo(cwd), + get_has_changes(cwd), ); - if latest_git_commit_hash.is_none() && associated_remote_urls.is_none() && sandbox.is_none() { + if latest_git_commit_hash.is_none() + && associated_remote_urls.is_none() + && has_changes.is_none() + && sandbox.is_none() + { return None; } - let mut workspaces = BTreeMap::new(); - if let Some(repo_root) = repo_root { - workspaces.insert( - repo_root.to_string_lossy().into_owned(), - TurnMetadataWorkspace { - associated_remote_urls, - latest_git_commit_hash, - }, + build_turn_metadata_bag( + None, + sandbox.map(ToString::to_string), + repo_root, + Some(WorkspaceGitMetadata { + associated_remote_urls, + latest_git_commit_hash, + has_changes, + }), + ) + .to_header_value() +} + +#[derive(Clone, Debug)] +pub(crate) struct TurnMetadataState { + cwd: PathBuf, + repo_root: Option, + base_metadata: TurnMetadataBag, + base_header: String, + enriched_header: Arc>>, + enrichment_task: Arc>>>, +} + +impl TurnMetadataState { + pub(crate) fn new(turn_id: String, cwd: PathBuf, sandbox: Option) -> Self { + let repo_root = get_git_repo_root(&cwd).map(|root| root.to_string_lossy().into_owned()); + let base_metadata = build_turn_metadata_bag(Some(turn_id), sandbox, None, None); + let base_header = base_metadata + .to_header_value() + .unwrap_or_else(|| "{}".to_string()); + + Self { + cwd, + repo_root, + base_metadata, + base_header, + enriched_header: Arc::new(RwLock::new(None)), + enrichment_task: Arc::new(Mutex::new(None)), + } + } + + pub(crate) fn current_header_value(&self) -> Option { + if let Some(header) = self + .enriched_header + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .cloned() + { + return Some(header); + } + Some(self.base_header.clone()) + } + + pub(crate) fn spawn_git_enrichment_task(&self) { + if self.repo_root.is_none() { + return; + } + + let mut task_guard = self + .enrichment_task + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if task_guard.is_some() { + return; + } + + let state = self.clone(); + *task_guard = Some(tokio::spawn(async move { + let workspace_git_metadata = state.fetch_workspace_git_metadata().await; + let Some(repo_root) = state.repo_root.clone() else { + return; + }; + + let enriched_metadata = build_turn_metadata_bag( + state.base_metadata.turn_id.clone(), + state.base_metadata.sandbox.clone(), + Some(repo_root), + Some(workspace_git_metadata), + ); + if enriched_metadata.workspaces.is_empty() { + return; + } + + if let Some(header_value) = enriched_metadata.to_header_value() { + *state + .enriched_header + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(header_value); + } + })); + } + + pub(crate) fn cancel_git_enrichment_task(&self) { + let mut task_guard = self + .enrichment_task + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(task) = task_guard.take() { + task.abort(); + } + } + + async fn fetch_workspace_git_metadata(&self) -> WorkspaceGitMetadata { + let (latest_git_commit_hash, associated_remote_urls, has_changes) = tokio::join!( + get_head_commit_hash(&self.cwd), + get_git_remote_urls_assume_git_repo(&self.cwd), + get_has_changes(&self.cwd), + ); + + WorkspaceGitMetadata { + associated_remote_urls, + latest_git_commit_hash, + has_changes, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use serde_json::Value; + use tempfile::TempDir; + use tokio::process::Command; + + #[tokio::test] + async fn build_turn_metadata_header_includes_has_changes_for_clean_repo() { + let temp_dir = TempDir::new().expect("temp dir"); + let repo_path = temp_dir.path().join("repo"); + std::fs::create_dir_all(&repo_path).expect("create repo"); + + Command::new("git") + .args(["init"]) + .current_dir(&repo_path) + .output() + .await + .expect("git init"); + Command::new("git") + .args(["config", "user.name", "Test User"]) + .current_dir(&repo_path) + .output() + .await + .expect("git config user.name"); + Command::new("git") + .args(["config", "user.email", "test@example.com"]) + .current_dir(&repo_path) + .output() + .await + .expect("git config user.email"); + + std::fs::write(repo_path.join("README.md"), "hello").expect("write file"); + Command::new("git") + .args(["add", "."]) + .current_dir(&repo_path) + .output() + .await + .expect("git add"); + Command::new("git") + .args(["commit", "-m", "initial"]) + .current_dir(&repo_path) + .output() + .await + .expect("git commit"); + + let header = build_turn_metadata_header(&repo_path, Some("none")) + .await + .expect("header"); + let parsed: Value = serde_json::from_str(&header).expect("valid json"); + let workspace = parsed + .get("workspaces") + .and_then(Value::as_object) + .and_then(|workspaces| workspaces.values().next()) + .cloned() + .expect("workspace"); + + assert_eq!( + workspace.get("has_changes").and_then(Value::as_bool), + Some(false) ); } - serde_json::to_string(&TurnMetadata { - workspaces, - sandbox: sandbox.map(ToString::to_string), - }) - .ok() } diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index d22e5c66e..71c8b4717 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -375,6 +375,15 @@ async fn responses_stream_includes_turn_metadata_header_for_git_workspace_e2e() .expect("x-codex-turn-metadata header should be present"); let initial_parsed: serde_json::Value = serde_json::from_str(&initial_header).expect("x-codex-turn-metadata should be valid JSON"); + let initial_turn_id = initial_parsed + .get("turn_id") + .and_then(serde_json::Value::as_str) + .expect("turn_id should be present") + .to_string(); + assert!( + !initial_turn_id.is_empty(), + "turn_id should not be empty in x-codex-turn-metadata" + ); assert_eq!( initial_parsed .get("sandbox") @@ -424,56 +433,95 @@ async fn responses_stream_includes_turn_metadata_header_for_git_workspace_e2e() .trim() .to_string(); - let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); - loop { - let request_recorder = responses::mount_sse_once(&server, response_body.clone()).await; - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - test.submit_turn("hello") - .await - .expect("submit post-git turn prompt"); + let first_response = responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_reasoning_item("rsn-1", &["thinking"], &[]), + responses::ev_shell_command_call("call-1", "echo turn-metadata"), + responses::ev_completed("resp-2"), + ]); + let follow_up_response = responses::sse(vec![ + responses::ev_response_created("resp-3"), + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-3"), + ]); + let request_log = responses::mount_response_sequence( + &server, + vec![ + responses::sse_response(first_response), + responses::sse_response(follow_up_response), + ], + ) + .await; - let maybe_metadata = request_recorder - .single_request() + test.submit_turn("hello") + .await + .expect("submit post-git turn prompt"); + + let requests = request_log.requests(); + assert_eq!(requests.len(), 2, "expected two requests in one turn"); + + let first_parsed: serde_json::Value = serde_json::from_str( + &requests[0] .header("x-codex-turn-metadata") - .and_then(|header_value| { - let parsed: serde_json::Value = serde_json::from_str(&header_value).ok()?; - let workspace = parsed - .get("workspaces") - .and_then(serde_json::Value::as_object) - .and_then(|workspaces| workspaces.values().next()) - .cloned()?; - Some((parsed, workspace)) - }); - let Some((parsed, workspace)) = maybe_metadata else { - if tokio::time::Instant::now() >= deadline { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - continue; - }; + .expect("first request should include turn metadata"), + ) + .expect("first metadata should be valid json"); + let second_parsed: serde_json::Value = serde_json::from_str( + &requests[1] + .header("x-codex-turn-metadata") + .expect("second request should include turn metadata"), + ) + .expect("second metadata should be valid json"); - assert_eq!( - parsed.get("sandbox").and_then(serde_json::Value::as_str), - Some("none") - ); - assert_eq!( - workspace - .get("latest_git_commit_hash") - .and_then(serde_json::Value::as_str), - Some(expected_head.as_str()) - ); - assert_eq!( - workspace - .get("associated_remote_urls") - .and_then(serde_json::Value::as_object) - .and_then(|remotes| remotes.get("origin")) - .and_then(serde_json::Value::as_str), - Some(expected_origin.as_str()) - ); - return; - } + let first_turn_id = first_parsed + .get("turn_id") + .and_then(serde_json::Value::as_str) + .expect("first turn_id should be present"); + let second_turn_id = second_parsed + .get("turn_id") + .and_then(serde_json::Value::as_str) + .expect("second turn_id should be present"); + assert_eq!( + first_turn_id, second_turn_id, + "requests should share turn_id" + ); + assert_ne!( + second_turn_id, + initial_turn_id.as_str(), + "post-git turn should have a new turn_id" + ); - panic!( - "x-codex-turn-metadata with git workspace info was never observed within 5s after git setup" + assert_eq!( + second_parsed + .get("sandbox") + .and_then(serde_json::Value::as_str), + Some("none") + ); + + let workspace = second_parsed + .get("workspaces") + .and_then(serde_json::Value::as_object) + .and_then(|workspaces| workspaces.values().next()) + .cloned() + .expect("second request should include git workspace metadata"); + assert_eq!( + workspace + .get("latest_git_commit_hash") + .and_then(serde_json::Value::as_str), + Some(expected_head.as_str()) + ); + assert_eq!( + workspace + .get("associated_remote_urls") + .and_then(serde_json::Value::as_object) + .and_then(|remotes| remotes.get("origin")) + .and_then(serde_json::Value::as_str), + Some(expected_origin.as_str()) + ); + assert_eq!( + workspace + .get("has_changes") + .and_then(serde_json::Value::as_bool), + Some(false) ); } diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index d436cac59..1db81ecf8 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -106,7 +106,7 @@ async fn responses_websocket_preconnect_reuses_connection() { let harness = websocket_harness(&server).await; let mut client_session = harness.client.new_session(); client_session - .prewarm_websocket(&harness.otel_manager, &harness.model_info, None) + .prewarm_websocket(&harness.otel_manager, &harness.model_info) .await .expect("websocket prewarm failed"); let prompt = prompt_with_input(vec![message_item("hello")]); @@ -131,7 +131,7 @@ async fn responses_websocket_preconnect_is_reused_even_with_header_changes() { let harness = websocket_harness(&server).await; let mut client_session = harness.client.new_session(); client_session - .prewarm_websocket(&harness.otel_manager, &harness.model_info, None) + .prewarm_websocket(&harness.otel_manager, &harness.model_info) .await .expect("websocket prewarm failed"); let prompt = prompt_with_input(vec![message_item("hello")]); @@ -172,7 +172,7 @@ async fn responses_websocket_prewarm_uses_model_preference_when_feature_disabled let harness = websocket_harness_with_options(&server, false, false, false, true).await; let mut client_session = harness.client.new_session(); client_session - .prewarm_websocket(&harness.otel_manager, &harness.model_info, None) + .prewarm_websocket(&harness.otel_manager, &harness.model_info) .await .expect("websocket prewarm failed"); diff --git a/codex-rs/core/tests/suite/turn_state.rs b/codex-rs/core/tests/suite/turn_state.rs index 833c2aad9..62bd97439 100644 --- a/codex-rs/core/tests/suite/turn_state.rs +++ b/codex-rs/core/tests/suite/turn_state.rs @@ -16,6 +16,7 @@ use core_test_support::responses::start_websocket_server_with_headers; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; +use serde_json::Value; const TURN_STATE_HEADER: &str = "x-codex-turn-state"; @@ -65,6 +66,25 @@ async fn responses_turn_state_persists_within_turn_and_resets_after() -> Result< ); assert_eq!(requests[2].header(TURN_STATE_HEADER), None); + let parse_turn_id = |header: Option| { + let value = header?; + let parsed: Value = serde_json::from_str(&value).ok()?; + parsed + .get("turn_id") + .and_then(Value::as_str) + .map(str::to_string) + }; + + let first_turn_id = parse_turn_id(requests[0].header("x-codex-turn-metadata")) + .expect("first request should include turn metadata turn_id"); + let second_turn_id = parse_turn_id(requests[1].header("x-codex-turn-metadata")) + .expect("follow-up request should include turn metadata turn_id"); + let third_turn_id = parse_turn_id(requests[2].header("x-codex-turn-metadata")) + .expect("new turn request should include turn metadata turn_id"); + + assert_eq!(first_turn_id, second_turn_id); + assert_ne!(second_turn_id, third_turn_id); + Ok(()) }