Prewarm more Guardian sampling WebSockets (#40403)

## What changed

- Increase the Guardian sampler's initial WebSocket pool from 2 connections to 8.
- Keep pool growth and recovery tests focused on the two-connection scenarios they exercise while making the remaining tests derive their expectations from the initial pool size.

GitOrigin-RevId: e2eabcdde5828ac7df333271f3253e74d696f396

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

@@ -65,6 +65,7 @@ use crate::async_scorer::config::CLASSIFICATION_OUTPUT_INSTRUCTIONS;
 use crate::async_scorer::config::DEFAULT_MODEL_CONTEXT_ITEM_TOKENS;
 use crate::async_scorer::config::DEFAULT_PARENT_COMPACTION_TOKENS;
 use crate::async_scorer::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC;
+use crate::async_scorer::sampler::INITIAL_WEBSOCKET_CONNECTIONS;
 use crate::async_scorer::sampler::MODEL;
 use crate::async_scorer::transcript::truncate_entry;
 
@@ -98,12 +99,10 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> {
     ];
     // Keep the sampled connection open for another request so only auth
     // invalidation, not a server close, forces the next handshake.
-    let server = responses::start_websocket_server(vec![
+    let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
-        Vec::new(),
+    connections.push(vec![events.clone(), events.clone()]);
-        vec![events.clone(), events.clone()],
+    connections.push(vec![events]);
-        vec![events],
+    let server = responses::start_websocket_server(connections).await;
-    ])
-    .await;
     let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("original"));
     auth_manager
         .set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new("original"))))
@@ -180,25 +179,27 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> {
         );
     }
 
+    let mut expected_authorizations =
+        vec![Some("Bearer original".to_owned()); INITIAL_WEBSOCKET_CONNECTIONS];
+    expected_authorizations.push(Some("Bearer refreshed".to_owned()));
     assert_eq!(
         server
             .handshakes()
             .iter()
             .map(|handshake| handshake.header("authorization"))
             .collect::<Vec<_>>(),
-        vec![
+        expected_authorizations
-            Some("Bearer original".to_owned()),
-            Some("Bearer original".to_owned()),
-            Some("Bearer refreshed".to_owned()),
-        ]
     );
+
+    let mut expected_requests = vec![0; INITIAL_WEBSOCKET_CONNECTIONS - 1];
+    expected_requests.extend([1, 1]);
     assert_eq!(
         server
             .connections()
             .iter()
             .map(Vec::len)
             .collect::<Vec<_>>(),
-        vec![0, 1, 1]
+        expected_requests
     );
     Ok(())
 }
@@ -587,7 +588,9 @@ async fn sample_configured_conversation_history(
         "total_tokens": 150,
     });
     let events = vec![ev_assistant_message("sample", "high"), completed];
-    let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
+    let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
+    connections.push(vec![events]);
+    let server = responses::start_websocket_server(connections).await;
     let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
         "http://{}/v1",
         server.uri().trim_start_matches("ws://")
@@ -661,7 +664,10 @@ async fn sample_configured_conversation_history(
 
     let request = tokio::time::timeout(
         Duration::from_secs(5),
-        server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0),
+        server.wait_for_request(
+            /*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
+            /*request_index*/ 0,
+        ),
     )
     .await?;
     Ok((request.body_json(), test, registry))
@@ -823,7 +829,9 @@ async fn contributor_fails_closed_when_luna_classification_fails() -> Result<()>
         ev_assistant_message("sample", "invalid"),
         ev_completed("response-invalid"),
     ];
-    let server = responses::start_websocket_server(vec![Vec::new(), vec![invalid_score]]).await;
+    let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
+    connections.push(vec![invalid_score]);
+    let server = responses::start_websocket_server(connections).await;
     let mut config = fixture.test.config.clone();
     config.model_provider = ModelProviderInfo::create_openai_provider(Some(format!(
         "http://{}/v1",
@@ -1647,7 +1655,8 @@ async fn contributor_skips_models_requiring_managed_guardian_review() -> Result<
         .build_with_auto_env(&thread_server)
         .await?;
 
-    let server = responses::start_websocket_server(vec![Vec::new(), Vec::new()]).await;
+    let server =
+        responses::start_websocket_server(vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS]).await;
     let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
         "http://{}/v1",
         server.uri().trim_start_matches("ws://")
@@ -2028,7 +2037,9 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
         ev_assistant_message("sample", "low"),
         ev_completed("response-1"),
     ];
-    let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
+    let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
+    connections.push(vec![events]);
+    let server = responses::start_websocket_server(connections).await;
     let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
         "http://{}/v1",
         server.uri().trim_start_matches("ws://")
@@ -2109,7 +2120,10 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
 
     let request = tokio::time::timeout(
         Duration::from_secs(5),
-        server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0),
+        server.wait_for_request(
+            /*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
+            /*request_index*/ 0,
+        ),
     )
     .await?
     .body_json();

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

@@ -46,7 +46,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;
-const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2;
+pub(super) const INITIAL_WEBSOCKET_CONNECTIONS: usize = 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

@@ -35,6 +35,7 @@ use tokio::net::TcpListener;
 use tokio::net::TcpStream;
 
 use super::CLASSIFICATION_TOKEN_USAGE_METRIC;
+use super::INITIAL_WEBSOCKET_CONNECTIONS;
 use super::LunaSampler;
 use super::LunaSamplerConfig;
 use super::LunaSamplingRequest;
@@ -88,7 +89,20 @@ fn assert_connection_metadata(server: &responses::WebSocketTestServer) -> Result
     Ok(thread_id)
 }
 
+#[derive(Clone, Copy)]
+enum ProxyPrewarmLimit {
+    AllConnections,
+    StopAfter { ready_connections: usize },
+}
+
 async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result<String> {
+    proxy_websocket_servers_with_prewarm_limit(servers, ProxyPrewarmLimit::AllConnections).await
+}
+
+async fn proxy_websocket_servers_with_prewarm_limit(
+    servers: &[&responses::WebSocketTestServer],
+    prewarm_limit: ProxyPrewarmLimit,
+) -> Result<String> {
     let listener = TcpListener::bind("127.0.0.1:0").await?;
     let address = listener.local_addr()?;
     let targets = servers
@@ -96,7 +110,15 @@ async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) ->
         .map(|server| server.uri().trim_start_matches("ws://").to_owned())
         .collect::<Vec<_>>();
     tokio::spawn(async move {
-        for target in targets {
+        for (index, target) in targets.into_iter().enumerate() {
+            if let ProxyPrewarmLimit::StopAfter { ready_connections } = prewarm_limit
+                && index == ready_connections
+            {
+                let Ok((connection, _)) = listener.accept().await else {
+                    return;
+                };
+                drop(connection);
+            }
             let Ok((mut incoming, _)) = listener.accept().await else {
                 return;
             };
@@ -171,7 +193,9 @@ async fn sampler_records_token_usage_after_returning_an_early_classification() -
         ev_output_text_delta("low"),
         ev_completed_with_tokens("response-1", /*total_tokens*/ 37),
     ];
-    let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
+    let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
+    connections.push(vec![events]);
+    let server = responses::start_websocket_server(connections).await;
     let metrics = Arc::new(RecordingMetrics::default());
     let mut config = sampler_config(format!(
         "http://{}/v1",
@@ -240,7 +264,13 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications
     let refreshed =
         responses::start_websocket_server(vec![vec![scripted_requests[1].clone()]]).await;
     let server = responses::start_websocket_server(vec![scripted_requests]).await;
-    let base_url = proxy_websocket_servers(&[&idle_server, &server, &refreshed]).await?;
+    let base_url = proxy_websocket_servers_with_prewarm_limit(
+        &[&idle_server, &server, &refreshed],
+        ProxyPrewarmLimit::StopAfter {
+            ready_connections: 2,
+        },
+    )
+    .await?;
     let manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-api-key"));
     manager
         .set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new(
@@ -391,8 +421,9 @@ async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Re
             ev_assistant_message("sample", "low"),
             ev_completed("response-1"),
         ];
-        let server =
+        let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
-            responses::start_websocket_server(vec![Vec::new(), vec![events.clone(), events]]).await;
+        connections.push(vec![events.clone(), events]);
+        let server = responses::start_websocket_server(connections).await;
         let mut config = sampler_config(format!(
             "http://{}/v1",
             server.uri().trim_start_matches("ws://")
@@ -412,7 +443,10 @@ async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Re
         assert_eq!(sampler.sample(request).await?, "low");
 
         let request = server
-            .wait_for_request(/*connection_index*/ 1, /*request_index*/ 0)
+            .wait_for_request(
+                /*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
+                /*request_index*/ 0,
+            )
             .await
             .body_json();
         let input = request["input"].as_array().expect("input items");
@@ -430,7 +464,10 @@ async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Re
             switched_request.parent_compaction_hash = Some("incompatible".to_owned());
             assert_eq!(sampler.sample(switched_request).await?, "low");
             let switched_request = server
-                .wait_for_request(/*connection_index*/ 1, /*request_index*/ 1)
+                .wait_for_request(
+                    /*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
+                    /*request_index*/ 1,
+                )
                 .await
                 .body_json();
             assert_eq!(switched_request["input"][2]["role"], "user");
@@ -509,7 +546,9 @@ async fn sampler_keeps_first_classification_token_when_later_output_disagrees()
         ev_assistant_message("sample", "lowhigh"),
         ev_completed("response-1"),
     ];
-    let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
+    let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
+    connections.push(vec![events]);
+    let server = responses::start_websocket_server(connections).await;
     let sampler = LunaSampler::connect(sampler_config(format!(
         "http://{}/v1",
         server.uri().trim_start_matches("ws://")
@@ -589,7 +628,13 @@ async fn sampler_grows_its_pool_for_overlapping_requests() -> Result<()> {
     let second = responses::start_websocket_server(vec![response("response-2")]).await;
     let third = responses::start_websocket_server(vec![response("response-3")]).await;
     let sampler = LunaSampler::connect(sampler_config(
-        proxy_websocket_servers(&[&first, &second, &third]).await?,
+        proxy_websocket_servers_with_prewarm_limit(
+            &[&first, &second, &third],
+            ProxyPrewarmLimit::StopAfter {
+                ready_connections: 2,
+            },
+        )
+        .await?,
     ))
     .await?;
 
@@ -661,7 +706,11 @@ async fn sampler_replaces_scored_drains_before_unfinished_classifications() -> R
         ]]])
         .await,
     );
-    let server_refs = servers.iter().collect::<Vec<_>>();
+    let server_refs = servers[2..INITIAL_WEBSOCKET_CONNECTIONS]
+        .iter()
+        .chain(servers[..2].iter())
+        .chain(servers[INITIAL_WEBSOCKET_CONNECTIONS..].iter())
+        .collect::<Vec<_>>();
     let sampler = Arc::new(
         LunaSampler::connect(sampler_config(proxy_websocket_servers(&server_refs).await?)).await?,
     );
@@ -785,7 +834,13 @@ async fn sampler_assigns_a_fresh_identity_when_replacing_aged_connections() -> R
     let second = responses::start_websocket_server(vec![response.clone()]).await;
     let replacement = responses::start_websocket_server(vec![response]).await;
     let sampler = LunaSampler::connect(sampler_config(
-        proxy_websocket_servers(&[&first, &second, &replacement]).await?,
+        proxy_websocket_servers_with_prewarm_limit(
+            &[&first, &second, &replacement],
+            ProxyPrewarmLimit::StopAfter {
+                ready_connections: 2,
+            },
+        )
+        .await?,
     ))
     .await?;
 
@@ -830,7 +885,13 @@ async fn sampler_reconnects_after_transient_service_failures() -> Result<()> {
     ]]])
     .await;
     let sampler = LunaSampler::connect(sampler_config(
-        proxy_websocket_servers(&[&first, &second, &recovered]).await?,
+        proxy_websocket_servers_with_prewarm_limit(
+            &[&first, &second, &recovered],
+            ProxyPrewarmLimit::StopAfter {
+                ready_connections: 2,
+            },
+        )
+        .await?,
     ))
     .await?;
 
@@ -865,7 +926,13 @@ async fn sampler_limits_transient_recovery_attempts() -> Result<()> {
     let third = responses::start_websocket_server(unavailable()).await;
     let unused = responses::start_websocket_server(unavailable()).await;
     let sampler = LunaSampler::connect(sampler_config(
-        proxy_websocket_servers(&[&first, &second, &third, &unused]).await?,
+        proxy_websocket_servers_with_prewarm_limit(
+            &[&first, &second, &third, &unused],
+            ProxyPrewarmLimit::StopAfter {
+                ready_connections: 2,
+            },
+        )
+        .await?,
     ))
     .await?;