Stabilize Guardian WebSocket tests (#41191)

## Why

Guardian tests could continue after a WebSocket handshake but before the
connection entered the idle pool. The shared request waiter could also miss a
notification delivered between checking the request log and registering its
wait.

## What changed

- Wait for Guardian prewarming to populate the idle connection pool instead of
  treating completed handshakes as sufficient.
- Use two initial Guardian connections in tests while retaining eight in
  non-test builds.
- Register for request-log notifications before inspecting the log to avoid a
  lost wakeup.

GitOrigin-RevId: 88f6a22c5401e03c24e6e114f2563fc641bfef80

codex-rs/core/tests/common/responses.rs

@@ -604,6 +604,7 @@ impl WebSocketTestServer {
         request_index: usize,
     ) -> WebSocketRequest {
         loop {
+            let notified = self.request_log_updated.notified();
             if let Some(request) = self
                 .connections
                 .lock()
@@ -614,7 +615,7 @@ impl WebSocketTestServer {
             {
                 return request;
             }
-            self.request_log_updated.notified().await;
+            notified.await;
         }
     }
 

codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs

@@ -148,11 +148,11 @@ async fn installed_extension_warms_connections_without_blocking_thread_start() -
 
     assert!(server.handshakes().is_empty());
     assert!(thread_store.get::<LunaSampler>().is_some());
-    assert!(
-        server
-            .wait_for_handshakes(INITIAL_WEBSOCKET_CONNECTIONS, PREWARM_TIMEOUT)
-            .await
-    );
+    thread_store
+        .get::<LunaSampler>()
+        .expect("Guardian v2 should initialize")
+        .wait_for_prewarm(PREWARM_TIMEOUT)
+        .await?;
     Ok(())
 }
 
@@ -203,11 +203,11 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> {
             thread_store,
         })
         .await;
-    assert!(
-        server
-            .wait_for_handshakes(INITIAL_WEBSOCKET_CONNECTIONS, PREWARM_TIMEOUT)
-            .await
-    );
+    thread_store
+        .get::<LunaSampler>()
+        .expect("Guardian v2 should initialize")
+        .wait_for_prewarm(PREWARM_TIMEOUT)
+        .await?;
     let progress = thread_store
         .get::<GuardianV2ScoreProgress>()
         .expect("Guardian v2 should initialize");
@@ -927,11 +927,11 @@ async fn sample_configured_conversation_history_with_source(
             thread_store,
         })
         .await;
-    assert!(
-        server
-            .wait_for_handshakes(INITIAL_WEBSOCKET_CONNECTIONS, PREWARM_TIMEOUT)
-            .await
-    );
+    thread_store
+        .get::<LunaSampler>()
+        .expect("Guardian v2 should initialize")
+        .wait_for_prewarm(PREWARM_TIMEOUT)
+        .await?;
     let turn_store = ExtensionData::new("turn-1");
     let tool_name = ToolName::plain("read_file");
     let tool_payload = ToolPayload::Function {
@@ -1166,11 +1166,14 @@ async fn contributor_fails_closed_when_luna_classification_fails() -> Result<()>
             thread_store: fixture.test.codex.thread_extension_data(),
         })
         .await;
-    assert!(
-        server
-            .wait_for_handshakes(INITIAL_WEBSOCKET_CONNECTIONS, PREWARM_TIMEOUT)
-            .await
-    );
+    fixture
+        .test
+        .codex
+        .thread_extension_data()
+        .get::<LunaSampler>()
+        .expect("Guardian v2 should initialize")
+        .wait_for_prewarm(PREWARM_TIMEOUT)
+        .await?;
 
     fixture.score_tool(ToolName::plain("read_file")).await;
     fixture.assert_fails_closed("elevated_risk").await
@@ -2607,11 +2610,11 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
             thread_store,
         })
         .await;
-    assert!(
-        server
-            .wait_for_handshakes(INITIAL_WEBSOCKET_CONNECTIONS, PREWARM_TIMEOUT)
-            .await
-    );
+    thread_store
+        .get::<LunaSampler>()
+        .expect("Guardian v2 should initialize")
+        .wait_for_prewarm(PREWARM_TIMEOUT)
+        .await?;
     let turn_store = ExtensionData::new("turn-1");
     let tool_name = ToolName::plain("read_file");
     let tool_payload = ToolPayload::Function {

codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs

@@ -51,7 +51,7 @@ pub(crate) const MODEL: &str = "gpt-5.6-luna";
 pub(crate) const CLASSIFICATION_TOKEN_USAGE_METRIC: &str =
     "codex.guardian_v2.classification.token_usage";
 const MAX_OUTPUT_BYTES: usize = 8 * 1024;
-pub(super) const INITIAL_WEBSOCKET_CONNECTIONS: usize = 8;
+pub(super) const INITIAL_WEBSOCKET_CONNECTIONS: usize = if cfg!(test) { 2 } else { 8 };
 const MAX_WEBSOCKET_CONNECTIONS: usize = 16;
 const MAX_SAMPLING_RETRIES: usize = 2;
 const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60);

codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs

@@ -42,6 +42,19 @@ use super::LunaSamplingRequest;
 use super::MAX_SAMPLING_RETRIES;
 use super::MAX_WEBSOCKET_CONNECTIONS;
 
+impl LunaSampler {
+    /// Waits for warm sockets to enter the client pool, beyond the server handshake.
+    pub(in crate::async_scorer) async fn wait_for_prewarm(&self, timeout: Duration) -> Result<()> {
+        tokio::time::timeout(timeout, async {
+            while self.idle_connections.lock().unwrap().len() < INITIAL_WEBSOCKET_CONNECTIONS {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await?;
+        Ok(())
+    }
+}
+
 fn assert_connection_metadata(server: &responses::WebSocketTestServer) -> Result<String> {
     let handshake = server.single_handshake();
     let thread_id = handshake.header("thread-id").expect("classifier thread ID");
@@ -112,6 +125,7 @@ async fn proxy_websocket_servers_with_prewarm_limit(
     tokio::spawn(async move {
         for (index, target) in targets.into_iter().enumerate() {
             if let ProxyPrewarmLimit::StopAfter { ready_connections } = prewarm_limit
+                && ready_connections < INITIAL_WEBSOCKET_CONNECTIONS
                 && index == ready_connections
             {
                 let Ok((connection, _)) = listener.accept().await else {