diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index 99063164da..9101334258 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -8,9 +8,11 @@ use codex_models_manager::bundled_models_response; use serde_json::Value; use serde_json::json; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::sync::mpsc; use wiremock::Mock; use wiremock::MockServer; use wiremock::Request; @@ -54,6 +56,8 @@ const CALENDAR_LIST_EVENTS_RESOURCE_URI: &str = "connector://calendar/tools/cale pub const DOCUMENT_EXTRACT_TEXT_RESOURCE_URI: &str = "connector://calendar/tools/calendar_extract_text"; +type AppsStartupInitializeGate = Arc>>>; + #[derive(Clone)] pub struct AppsTestServer { pub chatgpt_base_url: String, @@ -63,6 +67,7 @@ pub struct AppsTestServer { pub struct AppsTestServerStartupControl { initialize_attempts: Arc, remaining_initialize_failures: Arc, + successful_initialize_gate: AppsStartupInitializeGate, } impl AppsTestServerStartupControl { @@ -74,6 +79,16 @@ impl AppsTestServerStartupControl { pub fn initialize_attempts(&self) -> usize { self.initialize_attempts.load(Ordering::SeqCst) } + + /// Holds one successful startup without turning recovery into a fresh runtime refresh. + pub fn hold_next_successful_initialize(&self) -> mpsc::Sender<()> { + let (release, gate) = mpsc::channel(); + *self + .successful_initialize_gate + .lock() + .expect("Apps initialization gate lock should not be poisoned") = Some(gate); + release + } } #[derive(Clone, Copy)] @@ -177,6 +192,7 @@ impl AppsTestServer { let control = AppsTestServerStartupControl { initialize_attempts: Arc::new(AtomicUsize::new(0)), remaining_initialize_failures: Arc::new(AtomicUsize::new(0)), + successful_initialize_gate: Arc::new(Mutex::new(None)), }; mount_streamable_http_json_rpc_with_startup_control( server, @@ -188,6 +204,7 @@ impl AppsTestServer { AppsTestToolsListBehavior::AlwaysAvailable, Some(Arc::clone(&control.initialize_attempts)), Some(Arc::clone(&control.remaining_initialize_failures)), + Some(Arc::clone(&control.successful_initialize_gate)), ) .await; Ok(( @@ -413,6 +430,7 @@ async fn mount_streamable_http_json_rpc_at_path( tools_list_behavior, /*initialize_attempts*/ None, /*remaining_initialize_failures*/ None, + /*successful_initialize_gate*/ None, ) .await; } @@ -428,6 +446,7 @@ async fn mount_streamable_http_json_rpc_with_startup_control( tools_list_behavior: AppsTestToolsListBehavior, initialize_attempts: Option>, remaining_initialize_failures: Option>, + successful_initialize_gate: Option, ) { Mock::given(method("POST")) .and(path_regex(mcp_path_regex)) @@ -439,6 +458,7 @@ async fn mount_streamable_http_json_rpc_with_startup_control( tools_list_behavior, initialize_attempts, remaining_initialize_failures, + successful_initialize_gate, }) .mount(server) .await; @@ -452,6 +472,7 @@ struct CodexAppsJsonRpcResponder { tools_list_behavior: AppsTestToolsListBehavior, initialize_attempts: Option>, remaining_initialize_failures: Option>, + successful_initialize_gate: Option, } impl Respond for CodexAppsJsonRpcResponder { @@ -491,6 +512,20 @@ impl Respond for CodexAppsJsonRpcResponder { "error": "simulated non-retryable Apps MCP startup failure", })); } + let gate = self.successful_initialize_gate.as_ref().and_then(|gate| { + gate.lock() + .expect("Apps initialization gate lock should not be poisoned") + .take() + }); + // This responder's dedicated mock server stays paused while the test controls + // when this existing client's background retry may finish. + if let Some(gate) = gate + && gate.recv().is_err() + { + return ResponseTemplate::new(500).set_body_json(json!({ + "error": "Apps initialization gate was dropped before release", + })); + } let id = body.get("id").cloned().unwrap_or(Value::Null); let protocol_version = body .pointer("/params/protocolVersion") diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 8a3c687ab4..0c33196ccc 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -62,6 +62,10 @@ struct CoalescingMcpContributor { observed_markers: Mutex>, } +struct AppsMcpServerContributor { + url: String, +} + impl CoalescingMcpContributor { fn new() -> Self { Self { @@ -110,6 +114,26 @@ impl McpServerContributor for CoalescingMcpContributor { } } +impl McpServerContributor for AppsMcpServerContributor { + fn id(&self) -> &'static str { + "deferred_apps_recovery_test" + } + + fn contribute<'a>( + &'a self, + _context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let config = serde_json::from_value(json!({ "url": self.url })) + .expect("test Apps MCP server config should be valid"); + vec![McpServerContribution::Set { + name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + config: Box::new(config), + }] + }) + } +} + fn format_labeled_requests_snapshot( scenario: &str, sections: &[(&str, &ResponsesRequest)], @@ -587,10 +611,17 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let (apps_server, startup_control) = - AppsTestServer::mount_with_startup_control(&server).await?; - // The exact step is captured once before the first request; its background retry can recover. + // Wiremock has one server thread, so isolate the held MCP response from model/app discovery. + let apps_server_mock = responses::start_mock_server().await; + let (gated_apps_server, startup_control) = + AppsTestServer::mount_with_startup_control(&apps_server_mock).await?; + let apps_server = AppsTestServer::mount_searchable(&server).await?; startup_control.fail_next_initialize_attempts(/*attempts*/ 1); + let release_apps_recovery = startup_control.hold_next_successful_initialize(); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(Arc::new(AppsMcpServerContributor { + url: format!("{}/api/codex/ps/mcp", gated_apps_server.chatgpt_base_url), + })); let call_id = "pause-for-apps"; let response = mount_sse_sequence( &server, @@ -627,6 +658,7 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur ) .await; let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone()) + .with_extensions(Arc::new(extensions.build())) .with_config(|config| { config .features @@ -651,33 +683,13 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur thread_settings: Default::default(), }) .await?; - let request = tokio::time::timeout(Duration::from_secs(3), async { - let mut request = None; - let mut apps_ready = false; - while request.is_none() || !apps_ready { - let event = test - .codex - .next_event() - .await - .expect("event stream should stay open"); - match event.msg { - EventMsg::RequestUserInput(next_request) => request = Some(next_request), - EventMsg::McpStartupUpdate(update) - if update.server == CODEX_APPS_MCP_SERVER_NAME - && matches!( - update.status, - codex_protocol::protocol::McpStartupStatus::Ready - ) => - { - apps_ready = true; - } - _ => {} - } - } - request.expect("request user input event") + let EventMsg::RequestUserInput(request) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::RequestUserInput(_)) }) .await - .expect("Apps should recover before the second sampling request"); + else { + unreachable!("wait_for_event should return the request-user-input event") + }; let initial_requests = response.requests(); assert_eq!(initial_requests.len(), 1); @@ -700,6 +712,19 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur "Calendar namespace should not be advertised before recovery: {initial_tools_state}" ); + release_apps_recovery + .send(()) + .expect("background Apps recovery should still be waiting"); + wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::McpStartupUpdate(update) + if update.server == CODEX_APPS_MCP_SERVER_NAME + && matches!(update.status, codex_protocol::protocol::McpStartupStatus::Ready) + ) + }) + .await; + test.codex .submit(Op::UserInputAnswer { id: request.turn_id, @@ -757,6 +782,7 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur }), "recovered request should advertise tool_search: {recovered_body}" ); + assert_eq!(startup_control.initialize_attempts(), 2); insta::assert_snapshot!( "deferred_tools_recover_during_sampling", format_labeled_requests_snapshot(