diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 171ec36756..dc06fbb0db 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -40,12 +40,14 @@ use crate::test_codex::ApplyPatchModelOutput; #[derive(Debug, Clone)] pub struct ResponseMock { requests: Arc>>, + requests_updated: Arc, } impl ResponseMock { fn new() -> Self { Self { requests: Arc::new(Mutex::new(Vec::new())), + requests_updated: Arc::new(Notify::new()), } } @@ -80,6 +82,20 @@ impl ResponseMock { .iter() .find_map(|req| req.function_call_output_text(call_id)) } + + pub async fn wait_for_function_call_output_text(&self, call_id: &str) -> String { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let requests_updated = self.requests_updated.notified(); + if let Some(output) = self.function_call_output_text(call_id) { + return output; + } + requests_updated.await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {call_id} output")) + } } #[derive(Debug, Clone)] @@ -579,14 +595,17 @@ impl Match for ModelsMock { impl Match for ResponseMock { fn matches(&self, request: &wiremock::Request) -> bool { - self.requests - .lock() - .unwrap() - .push(ResponsesRequest(request.clone())); + { + self.requests + .lock() + .unwrap() + .push(ResponsesRequest(request.clone())); + } // Enforce invariant checks on every request body captured by the mock. // Panic on orphan tool outputs or calls to catch regressions early. validate_request_body_invariants(request); + self.requests_updated.notify_waiters(); true } } diff --git a/codex-rs/core/tests/suite/request_permissions_tool.rs b/codex-rs/core/tests/suite/request_permissions_tool.rs index 4df6602cdc..ed5bd714f1 100644 --- a/codex-rs/core/tests/suite/request_permissions_tool.rs +++ b/codex-rs/core/tests/suite/request_permissions_tool.rs @@ -300,9 +300,9 @@ async fn approved_folder_write_request_permissions_unblocks_later_exec_without_s } let exec_output = responses - .function_call_output_text("exec-call") - .map(|output| json!({ "output": output })) - .unwrap_or_else(|| panic!("expected exec-call output")); + .wait_for_function_call_output_text("exec-call") + .await; + let exec_output = json!({ "output": exec_output }); let (exit_code, stdout) = parse_result(&exec_output); assert!(exit_code.is_none() || exit_code == Some(0)); assert!(stdout.contains("folder-grant-ok")); @@ -463,9 +463,9 @@ async fn apply_patch_after_request_permissions(strict_auto_review: bool) -> Resu } let patch_output = responses - .function_call_output_text("apply-patch-call") - .map(|output| json!({ "output": output })) - .unwrap_or_else(|| panic!("expected apply-patch-call output")); + .wait_for_function_call_output_text("apply-patch-call") + .await; + let patch_output = json!({ "output": patch_output }); let (exit_code, stdout) = parse_result(&patch_output); assert!(exit_code.is_none() || exit_code == Some(0)); assert!(