diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 56084f797b..bf9e2e59d3 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/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![ - Vec::new(), - vec![events.clone(), events.clone()], - vec![events], - ]) - .await; + let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1]; + connections.push(vec![events.clone(), events.clone()]); + connections.push(vec![events]); + let server = responses::start_websocket_server(connections).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![ - Some("Bearer original".to_owned()), - Some("Bearer original".to_owned()), - Some("Bearer refreshed".to_owned()), - ] + expected_authorizations ); + + 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![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(); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index b85234b981..cd13f43183 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/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); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index 5f86b907ec..d5a99ed84c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/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 { + 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 { 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::>(); 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 = - responses::start_websocket_server(vec![Vec::new(), vec![events.clone(), events]]).await; + let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1]; + 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::>(); + let server_refs = servers[2..INITIAL_WEBSOCKET_CONNECTIONS] + .iter() + .chain(servers[..2].iter()) + .chain(servers[INITIAL_WEBSOCKET_CONNECTIONS..].iter()) + .collect::>(); 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?;