Synchronize Git enrichment tests explicitly (#40006)

## What changed

- Track Git enrichment completion with a watch channel and unblock waiters when enrichment finishes or is canceled.
- Let the test sync tool wait for the current turn's enrichment with a bounded timeout.
- Replace polling in metadata tests and wait for enrichment before asserting workspace metadata in the guardian integration test.

GitOrigin-RevId: 8d35f75a0ebfe412674e1e797c6a13c03b1ea373
This commit is contained in:
felixxia-oai
2026-08-21 20:42:06 +00:00
committed by copyberry
parent dbe9dac1ae
commit 51d8d12236
6 changed files with 99 additions and 59 deletions

View File

@@ -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<tokio::sync::Mutex<HashMap<String, BarrierState>>> = OnceLock::new();
@@ -47,6 +48,8 @@ struct TestSyncArgs {
sleep_after_ms: Option<u64>,
#[serde(default)]
barrier: Option<BarrierArgs>,
#[serde(default)]
wait_for_git_enrichment: bool,
}
fn default_timeout_ms() -> u64 {
@@ -80,7 +83,7 @@ impl TestSyncHandler {
&self,
invocation: ToolInvocation,
) -> Result<Box<dyn crate::tools::context::ToolOutput>, 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
{

View File

@@ -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 {

View File

@@ -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,
})

View File

@@ -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<Option<JoinHandle<()>>>,
git_enrichment_complete: watch::Sender<bool>,
}
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<Self>) {
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);
}
}

View File

@@ -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());

View File

@@ -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)