diff --git a/codex-rs/core/src/tools/handlers/test_sync.rs b/codex-rs/core/src/tools/handlers/test_sync.rs index 08d983b616..e4b4a9e1ac 100644 --- a/codex-rs/core/src/tools/handlers/test_sync.rs +++ b/codex-rs/core/src/tools/handlers/test_sync.rs @@ -23,6 +23,7 @@ use codex_tools::ToolSpec; pub struct TestSyncHandler; const DEFAULT_TIMEOUT_MS: u64 = 1_000; +const GIT_ENRICHMENT_TIMEOUT: Duration = Duration::from_secs(10); static BARRIERS: OnceLock>> = OnceLock::new(); @@ -47,6 +48,8 @@ struct TestSyncArgs { sleep_after_ms: Option, #[serde(default)] barrier: Option, + #[serde(default)] + wait_for_git_enrichment: bool, } fn default_timeout_ms() -> u64 { @@ -80,7 +83,7 @@ impl TestSyncHandler { &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { - let ToolInvocation { payload, .. } = invocation; + let ToolInvocation { payload, turn, .. } = invocation; let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -103,6 +106,20 @@ impl TestSyncHandler { wait_on_barrier(barrier).await?; } + if args.wait_for_git_enrichment { + tokio::time::timeout( + GIT_ENRICHMENT_TIMEOUT, + turn.turn_metadata_state.wait_for_git_enrichment(), + ) + .await + .map_err(|_| { + FunctionCallError::RespondToModel(format!( + "test_sync_tool git enrichment wait timed out after {} seconds", + GIT_ENRICHMENT_TIMEOUT.as_secs() + )) + })?; + } + if let Some(delay) = args.sleep_after_ms && delay > 0 { diff --git a/codex-rs/core/src/tools/handlers/test_sync_spec.rs b/codex-rs/core/src/tools/handlers/test_sync_spec.rs index 4de8104034..33bd720c94 100644 --- a/codex-rs/core/src/tools/handlers/test_sync_spec.rs +++ b/codex-rs/core/src/tools/handlers/test_sync_spec.rs @@ -46,6 +46,13 @@ pub fn create_test_sync_tool() -> ToolSpec { Some(false.into()), ), ), + ( + "wait_for_git_enrichment".to_string(), + JsonSchema::boolean(Some( + "Wait for Git enrichment for the current turn to finish, subject to a timeout." + .to_string(), + )), + ), ]); ToolSpec::Function(ResponsesApiTool { diff --git a/codex-rs/core/src/tools/handlers/test_sync_spec_tests.rs b/codex-rs/core/src/tools/handlers/test_sync_spec_tests.rs index 55f4a1a0a6..adcb2a2eeb 100644 --- a/codex-rs/core/src/tools/handlers/test_sync_spec_tests.rs +++ b/codex-rs/core/src/tools/handlers/test_sync_spec_tests.rs @@ -57,6 +57,13 @@ fn test_sync_tool_matches_expected_spec() { "Delay before any other action. Defaults to no delay.".to_string(), )), ), + ( + "wait_for_git_enrichment".to_string(), + JsonSchema::boolean(Some( + "Wait for Git enrichment for the current turn to finish, subject to a timeout." + .to_string(), + )), + ), ]), /*required*/ None, Some(false.into())), output_schema: None, }) diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index 81d4912327..fe85efe685 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -10,6 +10,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use serde_json::Value; +use tokio::sync::watch; use tokio::task::JoinHandle; use crate::responses_metadata::AGENT_NAME_KEY; @@ -124,6 +125,7 @@ pub(crate) struct TurnMetadataState { root_turn_ambiguous: AtomicBool, user_input_requested_during_turn: AtomicBool, enrichment_task: Mutex>>, + git_enrichment_complete: watch::Sender, } impl TurnMetadataState { @@ -185,6 +187,7 @@ impl TurnMetadataState { root_turn_ambiguous: AtomicBool::new(false), user_input_requested_during_turn: AtomicBool::new(false), enrichment_task: Mutex::new(None), + git_enrichment_complete: watch::channel(/*init*/ true).0, } } @@ -405,9 +408,9 @@ impl TurnMetadataState { } pub(crate) fn spawn_git_enrichment_task(self: &Arc) { - if self.repo_root.is_none() { + let Some(repo_root) = self.repo_root.clone() else { return; - } + }; let mut task_guard = self .enrichment_task @@ -417,29 +420,32 @@ impl TurnMetadataState { return; } + self.git_enrichment_complete.send_replace(/*value*/ false); let state = Arc::clone(self); *task_guard = Some(tokio::spawn(async move { - let Some(repo_root) = state.repo_root.clone() else { - return; - }; let workspace_git_metadata = state.fetch_workspace_git_metadata(&repo_root).await; - if workspace_git_metadata.is_empty() { - return; + if !workspace_git_metadata.is_empty() { + let mut workspaces = BTreeMap::new(); + workspaces.insert( + repo_root.to_string_lossy().into_owned(), + workspace_git_metadata.into(), + ); + *state + .enriched_workspaces + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(workspaces); } - let mut workspaces = BTreeMap::new(); - workspaces.insert( - repo_root.to_string_lossy().into_owned(), - workspace_git_metadata.into(), - ); - *state - .enriched_workspaces - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(workspaces); + state.git_enrichment_complete.send_replace(/*value*/ true); })); } + pub(crate) async fn wait_for_git_enrichment(&self) { + let mut completion = self.git_enrichment_complete.subscribe(); + let _ = completion.wait_for(|complete| *complete).await; + } + pub(crate) fn cancel_git_enrichment_task(&self) { let mut task_guard = self .enrichment_task @@ -447,6 +453,7 @@ impl TurnMetadataState { .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(task) = task_guard.take() { task.abort(); + self.git_enrichment_complete.send_replace(/*value*/ true); } } diff --git a/codex-rs/core/src/turn_metadata_tests.rs b/codex-rs/core/src/turn_metadata_tests.rs index 5a4fe6cb5f..bb19d1adb6 100644 --- a/codex-rs/core/src/turn_metadata_tests.rs +++ b/codex-rs/core/src/turn_metadata_tests.rs @@ -126,22 +126,10 @@ async fn create_clean_git_repo(repo_name: &str) -> (TempDir, AbsolutePathBuf) { } async fn wait_for_git_enrichment(state: &TurnMetadataState) -> Value { - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let header = test_turn_metadata_header(state); - let json: Value = serde_json::from_str(&header).expect("json"); - if json - .get("workspaces") - .and_then(Value::as_object) - .is_some_and(|workspaces| !workspaces.is_empty()) - { - return json; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("git enrichment should complete") + tokio::time::timeout(Duration::from_secs(2), state.wait_for_git_enrichment()) + .await + .expect("git enrichment should complete"); + serde_json::from_str(&test_turn_metadata_header(state)).expect("json") } #[tokio::test] @@ -1106,6 +1094,9 @@ async fn turn_metadata_state_git_enrichment_cancellation_is_retryable_and_errors .expect("enrichment task lock") .is_none() ); + tokio::time::timeout(Duration::from_secs(2), state.wait_for_git_enrichment()) + .await + .expect("cancelled git enrichment should unblock waiters"); assert!(state.current_workspaces().is_empty()); state.spawn_git_enrichment_task(); @@ -1138,20 +1129,10 @@ async fn turn_metadata_state_git_enrichment_cancellation_is_retryable_and_errors &model_info_from_slug("gpt-5.4"), )); invalid_state.spawn_git_enrichment_task(); - tokio::time::timeout(Duration::from_secs(2), async { - loop { - if invalid_state - .enrichment_task - .lock() - .expect("enrichment task lock") - .as_ref() - .is_some_and(tokio::task::JoinHandle::is_finished) - { - return; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) + tokio::time::timeout( + Duration::from_secs(2), + invalid_state.wait_for_git_enrichment(), + ) .await .expect("failed git enrichment should complete"); assert!(invalid_state.current_workspaces().is_empty()); diff --git a/codex-rs/core/tests/suite/git_enrichment.rs b/codex-rs/core/tests/suite/git_enrichment.rs index 4a2e8e584a..53c6a67869 100644 --- a/codex-rs/core/tests/suite/git_enrichment.rs +++ b/codex-rs/core/tests/suite/git_enrichment.rs @@ -163,6 +163,15 @@ async fn guardian_prewarm_and_review_skip_redundant_git_enrichment() -> Result<( let server = start_websocket_server(vec![ vec![vec![ev_response_created("warm-1"), ev_completed("warm-1")]], vec![vec![ev_response_created("warm-2"), ev_completed("warm-2")]], + vec![vec![ + ev_response_created("wait-for-parent-git"), + ev_function_call( + "wait-for-parent-git", + "test_sync_tool", + r#"{"wait_for_git_enrichment":true}"#, + ), + ev_completed("wait-for-parent-git"), + ]], vec![vec![ ev_response_created("approval-request"), ev_function_call("approval-call", "exec_command", &tool_args), @@ -175,11 +184,13 @@ async fn guardian_prewarm_and_review_skip_redundant_git_enrichment() -> Result<( ]) .await; let cwd = repo.path().to_path_buf(); - let mut builder = test_codex().with_config(move |config| { - config.cwd = cwd.abs(); - config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); - config.approvals_reviewer = ApprovalsReviewer::AutoReview; - }); + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(move |config| { + config.cwd = cwd.abs(); + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + }); let test = builder.build_with_websocket_server(&server).await?; let (first, second) = tokio::time::timeout(Duration::from_secs(5), async { @@ -201,19 +212,29 @@ async fn guardian_prewarm_and_review_skip_redundant_git_enrichment() -> Result<( text_elements: Vec::new(), }])) .await?; - let (user_turn, guardian_turn) = tokio::time::timeout(Duration::from_secs(5), async { - tokio::join!( - server.wait_for_request(/*connection_index*/ 2, /*request_index*/ 0), - server.wait_for_request(/*connection_index*/ 3, /*request_index*/ 0) - ) - }) - .await?; + let (initial_user_turn, user_turn, guardian_turn) = + tokio::time::timeout(Duration::from_secs(15), async { + tokio::join!( + server.wait_for_request(/*connection_index*/ 2, /*request_index*/ 0), + server.wait_for_request(/*connection_index*/ 3, /*request_index*/ 0), + server.wait_for_request(/*connection_index*/ 4, /*request_index*/ 0) + ) + }) + .await?; + let initial_user_turn = initial_user_turn.body_json(); let user_turn = user_turn.body_json(); let guardian_turn = guardian_turn.body_json(); + let initial_user_turn_metadata = turn_metadata(&initial_user_turn)?; assert_eq!( turn_metadata(&user_turn)?["auto_review_enabled"].as_bool(), Some(true) ); + if let Some(workspaces) = initial_user_turn_metadata.get("workspaces") { + assert_eq!( + workspaces, + &expected_workspace(repo.path(), &head, /*has_changes*/ true) + ); + } assert_eq!( turn_metadata(&user_turn)?["workspaces"], expected_workspace(repo.path(), &head, /*has_changes*/ true)