Make app list force-refetch cache tests deterministic

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<StdMutex<serde_json::Value>>,
+    directory_gate: Arc<StdMutex<Option<RequestGate>>>,
     directory_delay: Duration,
 }
 
 #[derive(Clone)]
 struct AppListMcpServer {
     tools: Arc<StdMutex<Vec<Tool>>>,
+    tools_gate: Arc<StdMutex<Option<RequestGate>>>,
     tools_delay: Duration,
 }
 
 impl AppListMcpServer {
-    fn new(tools: Arc<StdMutex<Vec<Tool>>>, tools_delay: Duration) -> Self {
-        Self { tools, tools_delay }
+    fn new(
+        tools: Arc<StdMutex<Vec<Tool>>>,
+        tools_gate: Arc<StdMutex<Option<RequestGate>>>,
+        tools_delay: Duration,
+    ) -> Self {
+        Self {
+            tools,
+            tools_gate,
+            tools_delay,
+        }
     }
 }
 
@@ -1359,6 +1376,8 @@ impl AppListMcpServer {
 struct AppsServerControl {
     response: Arc<StdMutex<serde_json::Value>>,
     tools: Arc<StdMutex<Vec<Tool>>>,
+    directory_gate: Arc<StdMutex<Option<RequestGate>>>,
+    tools_gate: Arc<StdMutex<Option<RequestGate>>>,
 }
 
 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<Notify>,
+    release: Arc<Notify>,
+    finished: Arc<Notify>,
+}
+
+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<StdMutex<Option<RequestGate>>>) -> 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<StdMutex<Option<RequestGate>>>) -> Option<RequestGate> {
+    gate_slot
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+        .take()
+}
+
+struct RequestGateFinish {
+    finished: Arc<Notify>,
+}
+
+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<Output = Result<ListToolsResult, rmcp::ErrorData>> + 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<impl axum::response::IntoResponse, StatusCode> {
+    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;
     }