diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index f6c2b6f580..af109d3624 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -53,6 +53,7 @@ use rmcp::transport::streamable_http_server::session::local::LocalSessionManager use serde_json::json; use tempfile::TempDir; use tokio::net::TcpListener; +use tokio::sync::Notify; use tokio::task::JoinHandle; use tokio::time::timeout; @@ -887,7 +888,7 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result connectors, tools, Duration::ZERO, - Duration::from_millis(150), + Duration::ZERO, ) .await?; @@ -926,6 +927,8 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result assert_eq!(initial_data.len(), 1); assert!(initial_data.iter().all(|app| app.is_accessible)); + let directory_gate = server_control.pause_next_directory_list(); + let tool_list_gate = server_control.pause_next_tools_list(); write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("chatgpt-token-invalid") @@ -944,6 +947,9 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result force_refetch: true, }) .await?; + timeout(DEFAULT_TIMEOUT, tool_list_gate.wait_started()).await?; + directory_gate.release(); + timeout(DEFAULT_TIMEOUT, directory_gate.wait_finished()).await?; let refetch_error: JSONRPCError = timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_error_message(RequestId::Integer(refetch_request)), @@ -951,7 +957,8 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result .await??; assert!(refetch_error.error.message.contains("failed to")); - tokio::time::sleep(Duration::from_millis(250)).await; + tool_list_gate.release(); + timeout(DEFAULT_TIMEOUT, tool_list_gate.wait_finished()).await?; let cached_request = mcp .send_apps_list_request(AppsListParams { @@ -1340,18 +1347,28 @@ struct AppsServerState { expected_bearer: String, expected_account_id: String, response: Arc>, + directory_gate: Arc>>, directory_delay: Duration, } #[derive(Clone)] struct AppListMcpServer { tools: Arc>>, + tools_gate: Arc>>, tools_delay: Duration, } impl AppListMcpServer { - fn new(tools: Arc>>, tools_delay: Duration) -> Self { - Self { tools, tools_delay } + fn new( + tools: Arc>>, + tools_gate: Arc>>, + tools_delay: Duration, + ) -> Self { + Self { + tools, + tools_gate, + tools_delay, + } } } @@ -1359,6 +1376,8 @@ impl AppListMcpServer { struct AppsServerControl { response: Arc>, tools: Arc>>, + directory_gate: Arc>>, + tools_gate: Arc>>, } impl AppsServerControl { @@ -1377,6 +1396,82 @@ impl AppsServerControl { .unwrap_or_else(std::sync::PoisonError::into_inner); *tools_guard = tools; } + + fn pause_next_directory_list(&self) -> RequestGate { + pause_next_request(&self.directory_gate) + } + + fn pause_next_tools_list(&self) -> RequestGate { + pause_next_request(&self.tools_gate) + } +} + +#[derive(Clone)] +struct RequestGate { + started: Arc, + release: Arc, + finished: Arc, +} + +impl RequestGate { + fn new() -> Self { + Self { + started: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + finished: Arc::new(Notify::new()), + } + } + + async fn wait_started(&self) { + self.started.notified().await; + } + + fn release(&self) { + self.release.notify_one(); + } + + async fn wait_finished(&self) { + self.finished.notified().await; + } + + async fn wait_for_release(&self) { + self.started.notify_one(); + self.release.notified().await; + } +} + +fn pause_next_request(gate_slot: &Arc>>) -> RequestGate { + let gate = RequestGate::new(); + let mut gate_slot = gate_slot + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *gate_slot = Some(gate.clone()); + gate +} + +fn take_next_request_gate(gate_slot: &Arc>>) -> Option { + gate_slot + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() +} + +struct RequestGateFinish { + finished: Arc, +} + +impl RequestGateFinish { + fn new(gate: &RequestGate) -> Self { + Self { + finished: Arc::clone(&gate.finished), + } + } +} + +impl Drop for RequestGateFinish { + fn drop(&mut self) { + self.finished.notify_one(); + } } impl ServerHandler for AppListMcpServer { @@ -1394,8 +1489,14 @@ impl ServerHandler for AppListMcpServer { ) -> impl std::future::Future> + Send + '_ { let tools = self.tools.clone(); + let tools_gate = self.tools_gate.clone(); let tools_delay = self.tools_delay; async move { + let gate = take_next_request_gate(&tools_gate); + let _finish = gate.as_ref().map(RequestGateFinish::new); + if let Some(gate) = &gate { + gate.wait_for_release().await; + } if tools_delay > Duration::ZERO { tokio::time::sleep(tools_delay).await; } @@ -1434,16 +1535,21 @@ async fn start_apps_server_with_delays_and_control( json!({ "apps": connectors, "next_token": null }), )); let tools = Arc::new(StdMutex::new(tools)); + let directory_gate = Arc::new(StdMutex::new(None)); + let tools_gate = Arc::new(StdMutex::new(None)); let state = AppsServerState { expected_bearer: "Bearer chatgpt-token".to_string(), expected_account_id: "account-123".to_string(), response: response.clone(), + directory_gate: directory_gate.clone(), directory_delay, }; let state = Arc::new(state); let server_control = AppsServerControl { response, tools: tools.clone(), + directory_gate, + tools_gate: tools_gate.clone(), }; let listener = TcpListener::bind("127.0.0.1:0").await?; @@ -1452,7 +1558,14 @@ async fn start_apps_server_with_delays_and_control( let mcp_service = StreamableHttpService::new( { let tools = tools.clone(); - move || Ok(AppListMcpServer::new(tools.clone(), tools_delay)) + let tools_gate = tools_gate.clone(); + move || { + Ok(AppListMcpServer::new( + tools.clone(), + tools_gate.clone(), + tools_delay, + )) + } }, Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), @@ -1479,6 +1592,12 @@ async fn list_directory_connectors( headers: HeaderMap, uri: Uri, ) -> Result { + let gate = take_next_request_gate(&state.directory_gate); + let _finish = gate.as_ref().map(RequestGateFinish::new); + if let Some(gate) = &gate { + gate.wait_for_release().await; + } + if state.directory_delay > Duration::ZERO { tokio::time::sleep(state.directory_delay).await; }