From 7769bccbb2b4e9469a36b12510e73594fa03c5d5 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 7 Sep 2026 09:14:06 +0000 Subject: [PATCH] Avoid WebSocket connection waits in Guardian v2 classification (#43408) ## Why When no healthy pooled WebSocket is available, classification waits for a new connection to open and can stall on its handshake. ## What changed - Use HTTP streaming when no healthy idle WebSocket is ready, and replenish the pool in the background with a five-minute cooldown after connection timeouts. - Share concurrency limits and the sampling retry budget across both transports, and bound the wait for HTTP response headers. - Cancel superseded requests while awaiting response headers and stop SSE polling when the consumer closes. ## Testing Add coverage for HTTP fallback and recovery after cooldown, shared transport capacity, stalled response headers, and cancellation before headers or during body draining. Update sampler and extension tests to cover fallback after socket expiry, authentication changes, and transient failures. GitOrigin-RevId: 0b23ccc1bfb2dfdbc8d2791d097706cd4dd44f39 --- .../app-server/tests/suite/v2/guardian_v2.rs | 68 +-- codex-rs/codex-api/src/sse/responses.rs | 6 +- codex-rs/core/tests/common/responses.rs | 2 +- .../src/async_scorer/extension_tests.rs | 43 +- .../guardian-v2/src/async_scorer/sampler.rs | 282 ++---------- .../async_scorer/sampler/connection_pool.rs | 427 ++++++++++++++++++ .../sampler/connection_pool_tests.rs | 347 ++++++++++++++ .../src/async_scorer/sampler_tests.rs | 300 ++++++------ 8 files changed, 1002 insertions(+), 473 deletions(-) create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index bbc85476d6..507bc18110 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -371,7 +371,9 @@ async fn parent_response( State(state): State>, Json(request): Json, ) -> impl IntoResponse { - let events = if request + let events = if request["model"] == "gpt-5.6-luna" { + luna_response(&state, request).await + } else if request .pointer("/client_metadata/x-openai-subagent") .and_then(Value::as_str) == Some("guardian") @@ -550,6 +552,39 @@ async fn parent_response( ) } +async fn luna_response(state: &MockResponsesState, request: Value) -> Vec { + let is_root_sample = state.root_worker + && state + .root_thread_id + .lock() + .expect("root thread lock should not be poisoned") + .as_ref() + .is_some_and(|thread_id| { + request["prompt_cache_key"] == format!("guardian-v2:{thread_id}") + }); + if !is_root_sample { + state + .luna_requests + .lock() + .expect("Luna request lock should not be poisoned") + .push(request); + state.allow_luna.notified().await; + } + let classification = if state.invalid_classification { + "invalid" + } else if state.luna_score < 0.5 { + "low" + } else { + "high" + }; + vec![ + responses::ev_response_created("luna-score"), + responses::ev_output_text_delta(classification), + responses::ev_assistant_message("luna-score-message", classification), + responses::ev_completed("luna-score"), + ] +} + async fn luna_websocket( State(state): State>, websocket: WebSocketUpgrade, @@ -561,36 +596,7 @@ async fn luna_websocket( continue; }; let request: Value = serde_json::from_str(&text).expect("valid Luna request"); - let is_root_sample = state.root_worker - && state - .root_thread_id - .lock() - .expect("root thread lock should not be poisoned") - .as_ref() - .is_some_and(|thread_id| { - request["prompt_cache_key"] == format!("guardian-v2:{thread_id}") - }); - if !is_root_sample { - state - .luna_requests - .lock() - .expect("Luna request lock should not be poisoned") - .push(request); - state.allow_luna.notified().await; - } - let classification = if state.invalid_classification { - "invalid" - } else if state.luna_score < 0.5 { - "low" - } else { - "high" - }; - for event in [ - responses::ev_response_created("luna-score"), - responses::ev_output_text_delta(classification), - responses::ev_assistant_message("luna-score-message", classification), - responses::ev_completed("luna-score"), - ] { + for event in luna_response(&state, request).await { if socket .send(Message::Text(event.to_string().into())) .await diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 68b6d663b8..e4f5f04d62 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -578,7 +578,11 @@ async fn process_sse_with_treatment( loop { let start = Instant::now(); - let response = timeout(idle_timeout, stream.next()).await; + let response = tokio::select! { + biased; + _ = tx_event.closed() => return, + response = timeout(idle_timeout, stream.next()) => response, + }; if let Some(t) = telemetry.as_ref() { t.on_sse_poll(&response, start.elapsed()); } diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 5aeb094c33..9c4fdb07f5 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -1079,7 +1079,7 @@ where fn base_mock() -> (MockBuilder, ResponseMock) { let response_mock = ResponseMock::new(); let mock = Mock::given(method("POST")) - .and(path_regex(".*/(responses|guardian)$")) + .and(path_regex(".*/(responses|guardian|guardian-classifier)$")) .and(response_mock.clone()); (mock, response_mock) } 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 85daa63ea7..acbbc6161a 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 @@ -88,6 +88,8 @@ use crate::async_scorer::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC; use crate::async_scorer::sampler::INITIAL_WEBSOCKET_CONNECTIONS; use crate::async_scorer::sampler::LunaSampler; use crate::async_scorer::sampler::MODEL; +use crate::async_scorer::sampler::tests::ProxyPrewarmLimit; +use crate::async_scorer::sampler::tests::proxy_websocket_servers_with_http; use crate::async_scorer::transcript::MAX_MESSAGE_ENTRY_TOKENS; use crate::async_scorer::transcript::MAX_TOOL_ENTRY_TOKENS; use crate::async_scorer::transcript::truncate_entry; @@ -184,7 +186,7 @@ async fn installed_extension_warms_connections_without_blocking_thread_start() - } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> { +async fn installed_extension_uses_http_after_warm_socket_auth_expires() -> Result<()> { skip_if_no_network!(Ok(())); let thread_server = responses::start_mock_server().await; @@ -197,20 +199,25 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> { ev_completed("response-1"), ]; // Keep the sampled connection open for another request so only auth - // invalidation, not a server close, forces the next handshake. + // invalidation forces the next classification to use HTTP. 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 http = responses::start_mock_server().await; + let http_mock = responses::mount_sse_once(&http, responses::sse(events)).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")))) .await?; let mut config = test.config.clone(); - config.model_provider = ModelProviderInfo::create_openai_provider(Some(format!( - "http://{}/v1", - server.uri().trim_start_matches("ws://") - ))); + config.model_provider = ModelProviderInfo::create_openai_provider(Some( + proxy_websocket_servers_with_http( + &[&server; INITIAL_WEBSOCKET_CONNECTIONS], + ProxyPrewarmLimit::AllConnections, + Some(&http.uri()), + ) + .await?, + )); config.features.enable(Feature::GuardianV2)?; let mut builder = ExtensionRegistryBuilder::new(); super::install( @@ -292,34 +299,28 @@ 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::>(), - expected_authorizations + vec![Some("Bearer original".to_owned()); INITIAL_WEBSOCKET_CONNECTIONS] ); - - let mut expected_requests = vec![0; INITIAL_WEBSOCKET_CONNECTIONS - 1]; - expected_requests.extend([1, 1]); + assert_eq!(progress.latest_failed_tool_call.load(Ordering::Acquire), 0); + let http_request = http_mock.single_request(); assert_eq!( - server - .connections() - .iter() - .map(Vec::len) - .collect::>(), - expected_requests + http_request.header("authorization"), + Some("Bearer refreshed".to_owned()) ); - let requests = server + let mut requests = server .connections() .into_iter() .flatten() .map(|request| request.body_json()) .collect::>(); + assert_eq!(requests.len(), 1); + requests.push(http_request.body_json()); for request in &requests { responses::assert_parent_turn(request, Some("turn-1"))?; responses::assert_root_turn(request, /*expected*/ None)?; 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 a720d18fb1..993e2d48d7 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -1,11 +1,15 @@ +//! Builds tool-less risk requests and publishes the first classifier output. +//! Both transports share request identity, retry, cancellation, and output handling. + +mod connection_pool; + +use connection_pool::ConnectionPool; use std::collections::HashMap; use std::collections::VecDeque; use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; -use std::time::Duration; -use std::time::Instant; use codex_api::ApiError; use codex_api::Reasoning; @@ -13,40 +17,26 @@ use codex_api::ReasoningContext; use codex_api::ResponseEvent; use codex_api::ResponsesApiRequest; use codex_api::ResponsesEndpoint; -use codex_api::ResponsesWebsocketClient; -use codex_api::ResponsesWebsocketConnection; -use codex_api::ResponsesWsRequest; use codex_api::TransportError; -use codex_api::build_session_headers; use codex_extension_api::ContextualUserFragment; use codex_extension_api::ExtensionMetrics; use codex_http_client::HttpClientFactory; use codex_login::AgentIdentityAuthPolicy; -use codex_login::CodexAuth; use codex_login::UnauthorizedRecovery; -use codex_login::default_client::add_originator_header; -use codex_login::default_client::default_headers; -use codex_model_provider::AgentIdentitySessionFallback; -use codex_model_provider::ProviderAuthScope; use codex_model_provider::SharedModelProvider; use codex_protocol::ResponseItemId; -use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TokenUsage; -use http::HeaderValue; use http::StatusCode; use serde_json::json; use thiserror::Error; -use tokio::sync::OwnedSemaphorePermit; -use tokio::sync::Semaphore; use tokio::sync::oneshot; use uuid::Uuid; -use super::metrics::sampler_failure_reason; use super::trusted_skills::GuardianTrustedSkillsFragment; use super::trusted_tools::GuardianTrustedToolFragment; @@ -55,10 +45,8 @@ 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 = if cfg!(test) { 2 } else { 8 }; -const MAX_WEBSOCKET_CONNECTIONS: usize = 16; +const MAX_CONCURRENT_REQUESTS: usize = 16; const MAX_SAMPLING_RETRIES: usize = 2; -const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60); -const RESPONSES_WEBSOCKETS_BETA: &str = "responses_websockets=2026-02-06"; const RESPONSES_LITE_METADATA_KEY: &str = "ws_request_header_x_openai_internal_codex_responses_lite"; const TURN_METADATA_KEY: &str = "x-codex-turn-metadata"; @@ -89,7 +77,7 @@ pub struct LunaSamplerConfig { pub metrics: Option>, } -/// One tool-less Luna classification request over an already-open connection. +/// One tool-less Luna classification request. pub struct LunaSamplingRequest { /// ID of the response handling the classified tool. pub parent_response_id: Option, @@ -123,8 +111,8 @@ pub enum LunaSamplerError { /// The thread's provider or scoped credentials could not be resolved. #[error("could not resolve the Luna model provider: {0}")] Provider(#[source] CodexErr), - /// The Responses WebSocket could not be opened or streamed. - #[error("Luna Responses WebSocket failed: {0}")] + /// The Responses request could not be opened or streamed. + #[error("Luna Responses request failed: {0}")] Api(#[source] ApiError), /// The provider's WebSocket connect deadline elapsed. #[error("Luna Responses WebSocket connection timed out")] @@ -143,30 +131,6 @@ pub enum LunaSamplerError { IncompatibleCompaction, } -struct PooledConnection { - connection: ResponsesWebsocketConnection, - endpoint: ResponsesEndpoint, - // The bridge routes by thread ID, so each socket needs its own identity. - thread_id: String, - expires_at: Instant, - auth_changes: Option>, -} - -struct ConnectionLease { - connection: PooledConnection, - idle_connections: Arc>>, - _permit: OwnedSemaphorePermit, -} - -impl ConnectionLease { - fn reuse(self) { - self.idle_connections - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(self.connection); - } -} - struct ActiveRequest { supersede: oneshot::Sender<()>, scored: Arc, @@ -200,11 +164,10 @@ fn record_token_usage(metrics: Option<&dyn ExtensionMetrics>, token_usage: Optio } } -/// A bounded pool of authenticated Responses WebSockets dedicated to Luna sampling. +/// Runs bounded Luna classifications over pooled WebSockets or HTTP. pub struct LunaSampler { - config: LunaSamplerConfig, - idle_connections: Arc>>, - capacity: Arc, + config: Arc, + connections: Arc, active_requests: Mutex>, } @@ -219,186 +182,20 @@ impl LunaSampler { } pub(super) fn new(config: LunaSamplerConfig) -> Self { + let config = Arc::new(config); Self { + connections: ConnectionPool::new(Arc::clone(&config)), config, - idle_connections: Arc::new(Mutex::new(Vec::with_capacity(MAX_WEBSOCKET_CONNECTIONS))), - capacity: Arc::new(Semaphore::new(MAX_WEBSOCKET_CONNECTIONS)), - active_requests: Mutex::new(VecDeque::with_capacity(MAX_WEBSOCKET_CONNECTIONS)), + active_requests: Mutex::new(VecDeque::with_capacity(MAX_CONCURRENT_REQUESTS)), } } pub(super) async fn prewarm(&self) { - for _ in 0..INITIAL_WEBSOCKET_CONNECTIONS { - let idle_connections = self - .idle_connections - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .len(); - if idle_connections >= self.capacity.available_permits() { - break; - } - let connection = match self.open_connection().await { - Ok(connection) => connection, - Err(_) => break, - }; - self.idle_connections - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(connection); + if let Some(refill) = self.connections.replenish() { + let _ = refill.await; } } - async fn responses_endpoint(&self) -> ResponsesEndpoint { - let provider = self.config.provider.info(); - if self.config.free_guardian - && self - .config - .provider - .auth() - .await - .as_ref() - .is_some_and(CodexAuth::uses_codex_backend) - && provider.supports_codex_backend_routes() - && provider.requires_openai_auth - && provider.env_key.is_none() - && provider.experimental_bearer_token.is_none() - && provider.auth.is_none() - && provider.aws.is_none() - { - ResponsesEndpoint::GuardianClassifier - } else { - ResponsesEndpoint::Responses - } - } - - async fn open_connection(&self) -> Result { - let provider = self - .config - .provider - .api_provider() - .await - .map_err(LunaSamplerError::Provider)?; - let auth_manager = self.config.provider.auth_manager(); - let auth_changes = auth_manager.map(|manager| manager.auth_change_receiver()); - let auth = self - .config - .provider - .api_auth_for_scope(ProviderAuthScope { - agent_identity_policy: self.config.agent_identity_policy, - session_source: self.config.session_source.clone(), - agent_identity_session_fallback: AgentIdentitySessionFallback::default(), - }) - .await - .map_err(LunaSamplerError::Provider)? - .auth; - let thread_id = ThreadId::new().to_string(); - let mut headers = build_session_headers( - Some(self.config.session_id.clone()), - Some(thread_id.clone()), - ); - headers.insert("x-openai-subagent", HeaderValue::from_static("guardian")); - headers.insert( - "x-codex-window-id", - HeaderValue::from_str(&format!("{thread_id}:0")).map_err(|error| { - LunaSamplerError::Api(ApiError::Stream(format!( - "invalid classifier window ID: {error}" - ))) - })?, - ); - headers.insert( - "openai-beta", - HeaderValue::from_static(RESPONSES_WEBSOCKETS_BETA), - ); - headers.insert( - "x-openai-internal-codex-responses-lite", - HeaderValue::from_static("true"), - ); - if let Some(originator) = self.config.originator.as_deref() { - add_originator_header(&mut headers, originator); - } - if let Ok(request_id) = HeaderValue::from_str(&thread_id) { - headers.insert("x-client-request-id", request_id); - } - - let provider_info = self.config.provider.info(); - let endpoint = self.responses_endpoint().await; - let client = ResponsesWebsocketClient::new(provider, auth).with_endpoint(endpoint); - let connect = client.connect( - &self.config.http_client_factory, - headers, - default_headers(), - /*turn_state*/ None, - /*telemetry*/ None, - ); - let started_at = Instant::now(); - let result = tokio::time::timeout(provider_info.websocket_connect_timeout(), connect) - .await - .map_err(|_| LunaSamplerError::ConnectionTimeout) - .and_then(|result| result.map_err(LunaSamplerError::Api)); - if let Some(metrics) = self.config.metrics.as_deref() { - let outcome = if result.is_ok() { "success" } else { "failure" }; - let mut tags = vec![("endpoint", endpoint.path()), ("outcome", outcome)]; - if let Err(error) = &result { - tags.push(("failure_reason", sampler_failure_reason(error))); - } - metrics.histogram( - "codex.guardian_v2.connection.duration_ms", - i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX), - &tags, - ); - } - let connection = result?; - if auth_changes - .as_ref() - .is_some_and(|auth| auth.has_changed().unwrap_or(true)) - { - return Err(LunaSamplerError::Api(ApiError::Stream( - "authentication changed while connecting".into(), - ))); - } - - Ok(PooledConnection { - connection, - endpoint, - thread_id, - expires_at: Instant::now() + MAX_WEBSOCKET_AGE, - auth_changes, - }) - } - - async fn lease_connection(&self) -> Result { - let permit = Arc::clone(&self.capacity) - .acquire_owned() - .await - .map_err(|_| LunaSamplerError::ConnectionTimeout)?; - let connection = loop { - let idle = self - .idle_connections - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .pop(); - match idle { - Some(connection) - if connection - .auth_changes - .as_ref() - .is_none_or(|auth| !auth.has_changed().unwrap_or(true)) - && Instant::now() < connection.expires_at - && !connection.connection.is_closed().await => - { - break connection; - } - Some(_) => {} - None => break self.open_connection().await?, - } - }; - Ok(ConnectionLease { - connection, - idle_connections: Arc::clone(&self.idle_connections), - _permit: permit, - }) - } - async fn retry_after_failure( &self, error: &LunaSamplerError, @@ -428,10 +225,7 @@ impl LunaSampler { if !recovery.has_next() || recovery.next().await.is_err() { return false; } - self.idle_connections - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clear(); + self.connections.clear(); return true; } else { status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS @@ -460,7 +254,7 @@ impl LunaSampler { false } - /// Sends one tool-less classification request on an exclusively leased WebSocket. + /// Sends one tool-less classification request using an available transport. pub async fn sample(&self, request: LunaSamplingRequest) -> Result { if request.parent_compaction.is_some() && !self.supports_parent_compaction(request.parent_compaction_hash.as_deref()) @@ -577,7 +371,7 @@ impl LunaSampler { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); active_requests.retain(|request| !request.supersede.is_closed()); - if active_requests.len() == MAX_WEBSOCKET_CONNECTIONS { + if active_requests.len() == MAX_CONCURRENT_REQUESTS { let oldest_scored = active_requests .iter() .position(|request| request.scored.load(Ordering::Relaxed)) @@ -601,7 +395,7 @@ impl LunaSampler { let lease = match tokio::select! { biased; _ = &mut superseded => return Err(LunaSamplerError::Superseded), - lease = self.lease_connection() => lease, + lease = self.connections.lease() => lease, } { Ok(lease) => lease, Err(error) => { @@ -614,13 +408,12 @@ impl LunaSampler { return Err(error); } }; - request.service_tier = - if lease.connection.endpoint == ResponsesEndpoint::GuardianClassifier { - None - } else { - self.config.service_tier.clone() - }; - let thread_id = &lease.connection.thread_id; + request.service_tier = if lease.endpoint == ResponsesEndpoint::GuardianClassifier { + None + } else { + self.config.service_tier.clone() + }; + let thread_id = &lease.thread_id; let mut turn_metadata = json!({ "session_id": self.config.session_id, "thread_id": thread_id, @@ -644,22 +437,17 @@ impl LunaSampler { turn_metadata["root_turn_id"] = json!(root_turn_id); } client_metadata.insert(TURN_METADATA_KEY.to_owned(), turn_metadata.to_string()); - if lease.connection.endpoint == ResponsesEndpoint::GuardianClassifier + if lease.endpoint == ResponsesEndpoint::GuardianClassifier && let Some(parent_response_id) = &parent_response_id { client_metadata.insert("parent_response_id".to_owned(), parent_response_id.clone()); } request.client_metadata = Some(client_metadata); - let mut stream = match lease - .connection - .connection - .stream_request( - ResponsesWsRequest::ResponseCreate((&request).into()), - /*connection_reused*/ true, - /*turn_state*/ None, - ) - .await - { + let mut stream = match tokio::select! { + biased; + _ = &mut superseded => return Err(LunaSamplerError::Superseded), + stream = lease.stream_request(&request) => stream, + } { Ok(stream) => stream, Err(error) => { let error = LunaSamplerError::Api(error); @@ -767,4 +555,4 @@ impl LunaSampler { #[cfg(test)] #[path = "sampler_tests.rs"] -mod tests; +pub(super) mod tests; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs new file mode 100644 index 0000000000..13adc5f1bc --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs @@ -0,0 +1,427 @@ +//! Keeps WebSocket establishment off the classification path. +//! One background worker fills the pool. Opening timeouts pause replenishment; +//! requests with no healthy idle socket use HTTP under the same concurrency limit. + +use super::super::metrics::sampler_failure_reason; +use super::INITIAL_WEBSOCKET_CONNECTIONS; +use super::LunaSamplerConfig; +use super::LunaSamplerError; +use super::MAX_CONCURRENT_REQUESTS; +use codex_api::ApiError; +use codex_api::Provider; +use codex_api::ReqwestTransport; +use codex_api::ResponseStream; +use codex_api::ResponsesApiRequest; +use codex_api::ResponsesClient; +use codex_api::ResponsesEndpoint; +use codex_api::ResponsesOptions; +use codex_api::ResponsesWebsocketClient; +use codex_api::ResponsesWebsocketConnection; +use codex_api::ResponsesWsRequest; +use codex_api::SharedAuthProvider; +use codex_api::TransportError; +use codex_api::build_session_headers; +use codex_http_client::ClientRouteClass; +use codex_login::CodexAuth; +use codex_login::default_client::add_originator_header; +use codex_login::default_client::create_client_for_route_async; +use codex_login::default_client::default_headers; +use codex_model_provider::AgentIdentitySessionFallback; +use codex_model_provider::ProviderAuthScope; +use codex_protocol::ThreadId; +use http::HeaderMap; +use http::HeaderValue; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use tokio::sync::OnceCell; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio::task::JoinHandle; +use tokio::time::Instant; + +const CONNECT_COOLDOWN: Duration = Duration::from_secs(5 * 60); +const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60); +const RESPONSES_WEBSOCKETS_BETA: &str = "responses_websockets=2026-02-06"; + +pub(super) struct ConnectionPool { + config: Arc, + pub(super) idle_connections: Mutex>, + classifications: Arc, + // A socket holds its permit while opening, idle, leased, and draining. + sockets: Arc, + replenishing: Arc>, + retry_after: Mutex>, + http_transport: Mutex>)>>, +} + +pub(super) struct PooledConnection { + connection: ResponsesWebsocketConnection, + endpoint: ResponsesEndpoint, + // The bridge routes by thread ID, so each socket needs its own identity. + thread_id: String, + pub(super) expires_at: Instant, + auth_changes: Option>, + _permit: OwnedSemaphorePermit, +} + +enum Connection { + Websocket(PooledConnection), + Http(ResponsesClient), +} + +pub(super) struct ConnectionLease { + pub(super) thread_id: String, + pub(super) endpoint: ResponsesEndpoint, + connection: Connection, + pool: Arc, + _permit: OwnedSemaphorePermit, +} + +impl ConnectionPool { + pub(super) fn new(config: Arc) -> Arc { + Arc::new(Self { + config, + idle_connections: Mutex::new(Vec::new()), + classifications: Arc::new(Semaphore::new(MAX_CONCURRENT_REQUESTS)), + sockets: Arc::new(Semaphore::new(MAX_CONCURRENT_REQUESTS)), + replenishing: Arc::new(tokio::sync::Mutex::new(())), + retry_after: Mutex::new(None), + http_transport: Mutex::new(None), + }) + } + + pub(super) fn clear(&self) { + self.idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clear(); + } + + /// Never waits for another opener or sleeps through the cooldown. + pub(super) fn replenish(self: &Arc) -> Option> { + let guard = Arc::clone(&self.replenishing).try_lock_owned().ok()?; + if self + .retry_after + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_some_and(|deadline| Instant::now() < deadline) + { + return None; + } + let pool = Arc::clone(self); + Some(tokio::spawn(async move { + let _guard = guard; + loop { + if pool + .idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len() + >= INITIAL_WEBSOCKET_CONNECTIONS + { + break; + } + let Ok(permit) = Arc::clone(&pool.sockets).try_acquire_owned() else { + break; + }; + match pool.open_connection(permit).await { + Ok(connection) => { + *pool + .retry_after + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = None; + pool.idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(connection); + } + Err(error) => { + if matches!(error, LunaSamplerError::ConnectionTimeout) { + *pool + .retry_after + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some(Instant::now() + CONNECT_COOLDOWN); + } + break; + } + } + } + })) + } + + pub(super) async fn lease(self: &Arc) -> Result { + let permit = Arc::clone(&self.classifications) + .acquire_owned() + .await + .map_err(|error| LunaSamplerError::Api(ApiError::Stream(error.to_string())))?; + let connection = loop { + let idle = self + .idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .pop(); + match idle { + Some(connection) + if connection + .auth_changes + .as_ref() + .is_none_or(|auth| !auth.has_changed().unwrap_or(true)) + && Instant::now() < connection.expires_at + && !connection.connection.is_closed().await => + { + break Some(connection); + } + Some(_) => {} + None => break None, + } + }; + let (connection, thread_id, endpoint) = match connection { + Some(connection) => { + let thread_id = connection.thread_id.clone(); + let endpoint = connection.endpoint; + (Connection::Websocket(connection), thread_id, endpoint) + } + None => { + self.replenish(); + let (mut provider, auth) = self.client_setup().await?; + // Sampling owns the retry budget across both transports. + provider.retry.max_attempts = 0; + let endpoint = self.responses_endpoint().await; + let url = provider.url_for_path(endpoint.path()); + let transport = { + let mut cached = self + .http_transport + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some((cached_url, transport)) = cached.as_ref() + && cached_url == &url + { + Arc::clone(transport) + } else { + let transport = Arc::new(OnceCell::new()); + *cached = Some((url.clone(), Arc::clone(&transport))); + transport + } + }; + let transport = transport + .get_or_try_init(|| async { + let client = create_client_for_route_async( + self.config.http_client_factory.clone(), + url, + ClientRouteClass::Api, + ) + .await + .map_err(|error| { + LunaSamplerError::Api(ApiError::Transport(TransportError::Build( + error.to_string(), + ))) + })?; + Ok::<_, LunaSamplerError>(ReqwestTransport::from_http_client(client)) + }) + .await? + .clone(); + let client = + ResponsesClient::new(transport, provider, auth).with_endpoint(endpoint); + ( + Connection::Http(client), + ThreadId::new().to_string(), + endpoint, + ) + } + }; + Ok(ConnectionLease { + thread_id, + endpoint, + connection, + pool: Arc::clone(self), + _permit: permit, + }) + } + + async fn client_setup(&self) -> Result<(Provider, SharedAuthProvider), LunaSamplerError> { + let provider = self + .config + .provider + .api_provider() + .await + .map_err(LunaSamplerError::Provider)?; + let auth = self + .config + .provider + .api_auth_for_scope(ProviderAuthScope { + agent_identity_policy: self.config.agent_identity_policy, + session_source: self.config.session_source.clone(), + agent_identity_session_fallback: AgentIdentitySessionFallback::default(), + }) + .await + .map_err(LunaSamplerError::Provider)? + .auth; + Ok((provider, auth)) + } + + fn headers(&self, thread_id: &str) -> Result { + let mut headers = build_session_headers( + Some(self.config.session_id.clone()), + Some(thread_id.to_owned()), + ); + headers.insert("x-openai-subagent", HeaderValue::from_static("guardian")); + headers.insert( + "x-codex-window-id", + HeaderValue::from_str(&format!("{thread_id}:0")).map_err(|error| { + LunaSamplerError::Api(ApiError::Stream(format!( + "invalid classifier window ID: {error}" + ))) + })?, + ); + headers.insert( + "x-openai-internal-codex-responses-lite", + HeaderValue::from_static("true"), + ); + if let Some(originator) = self.config.originator.as_deref() { + add_originator_header(&mut headers, originator); + } + if let Ok(request_id) = HeaderValue::from_str(thread_id) { + headers.insert("x-client-request-id", request_id); + } + Ok(headers) + } + async fn responses_endpoint(&self) -> ResponsesEndpoint { + let provider = self.config.provider.info(); + if self.config.free_guardian + && self + .config + .provider + .auth() + .await + .as_ref() + .is_some_and(CodexAuth::uses_codex_backend) + && provider.supports_codex_backend_routes() + && provider.requires_openai_auth + && provider.env_key.is_none() + && provider.experimental_bearer_token.is_none() + && provider.auth.is_none() + && provider.aws.is_none() + { + ResponsesEndpoint::GuardianClassifier + } else { + ResponsesEndpoint::Responses + } + } + + async fn open_connection( + &self, + permit: OwnedSemaphorePermit, + ) -> Result { + let auth_manager = self.config.provider.auth_manager(); + let auth_changes = auth_manager.map(|manager| manager.auth_change_receiver()); + let (provider, auth) = self.client_setup().await?; + let thread_id = ThreadId::new().to_string(); + let mut headers = self.headers(&thread_id)?; + headers.insert( + "openai-beta", + HeaderValue::from_static(RESPONSES_WEBSOCKETS_BETA), + ); + + let provider_info = self.config.provider.info(); + let endpoint = self.responses_endpoint().await; + let client = ResponsesWebsocketClient::new(provider, auth).with_endpoint(endpoint); + let connect = client.connect( + &self.config.http_client_factory, + headers, + default_headers(), + /*turn_state*/ None, + /*telemetry*/ None, + ); + let started_at = Instant::now(); + let result = tokio::time::timeout(provider_info.websocket_connect_timeout(), connect) + .await + .map_err(|_| LunaSamplerError::ConnectionTimeout) + .and_then(|result| result.map_err(LunaSamplerError::Api)); + if let Some(metrics) = self.config.metrics.as_deref() { + let outcome = if result.is_ok() { "success" } else { "failure" }; + let mut tags = vec![("endpoint", endpoint.path()), ("outcome", outcome)]; + if let Err(error) = &result { + tags.push(("failure_reason", sampler_failure_reason(error))); + } + metrics.histogram( + "codex.guardian_v2.connection.duration_ms", + i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX), + &tags, + ); + } + let connection = result?; + if auth_changes + .as_ref() + .is_some_and(|auth| auth.has_changed().unwrap_or(true)) + { + return Err(LunaSamplerError::Api(ApiError::Stream( + "authentication changed while connecting".into(), + ))); + } + + Ok(PooledConnection { + connection, + endpoint, + thread_id, + expires_at: Instant::now() + MAX_WEBSOCKET_AGE, + auth_changes, + _permit: permit, + }) + } +} + +impl ConnectionLease { + pub(super) async fn stream_request( + &self, + request: &ResponsesApiRequest, + ) -> Result { + match &self.connection { + Connection::Websocket(connection) => { + connection + .connection + .stream_request( + ResponsesWsRequest::ResponseCreate(request.into()), + /*connection_reused*/ true, + /*turn_state*/ None, + ) + .await + } + Connection::Http(client) => { + // The SSE idle timeout starts after headers arrive. Bound that wait too. + tokio::time::timeout( + self.pool.config.provider.info().stream_idle_timeout(), + client.stream_request( + request.clone(), + ResponsesOptions { + session_id: Some(self.pool.config.session_id.clone()), + thread_id: Some(self.thread_id.clone()), + extra_headers: self + .pool + .headers(&self.thread_id) + .map_err(|error| ApiError::Stream(error.to_string()))?, + ..Default::default() + }, + ), + ) + .await + .map_err(|_| ApiError::Transport(TransportError::Timeout))? + } + } + } + + pub(super) fn reuse(self) { + if let Connection::Websocket(connection) = self.connection { + self.pool + .idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(connection); + } + } +} + +#[cfg(test)] +#[path = "connection_pool_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs new file mode 100644 index 0000000000..42ed27c7ee --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs @@ -0,0 +1,347 @@ +//! Exercises classifier fallback, cooldown, and cancellation against local servers. + +use super::super::LunaSampler; +use super::super::tests::sample_request; +use super::super::tests::sampler_config; +use super::*; +use anyhow::Result; +use codex_login::AuthManager; +use codex_model_provider::create_model_provider; +use codex_model_provider_info::ModelProviderInfo; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::net::TcpStream; + +// Routes HTTP normally while selectively stalling WebSocket handshakes. +struct Gateway { + url: String, + opens: Arc, + allowed_opens: Arc, + task: JoinHandle<()>, +} + +impl Gateway { + async fn new(http_url: &str, websocket_url: &str) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let url = format!("http://{}", listener.local_addr()?); + let http_target = http_url.trim_start_matches("http://").to_owned(); + let websocket_target = websocket_url.trim_start_matches("ws://").to_owned(); + let opens = Arc::new(AtomicUsize::new(0)); + let allowed_opens = Arc::new(AtomicUsize::new(0)); + let count = Arc::clone(&opens); + let allowed = Arc::clone(&allowed_opens); + let task = tokio::spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + while let Ok((mut incoming, _)) = listener.accept().await { + let http_target = http_target.clone(); + let websocket_target = websocket_target.clone(); + let count = Arc::clone(&count); + let allowed = Arc::clone(&allowed); + connections.spawn(async move { + let mut method = [0; 4]; + incoming.read_exact(&mut method).await?; + let target = if &method == b"GET " { + let attempt = count.fetch_add(/*val*/ 1, Ordering::SeqCst) + 1; + if attempt > allowed.load(Ordering::SeqCst) { + tokio::io::copy(&mut incoming, &mut tokio::io::sink()).await?; + return Ok::<_, std::io::Error>(()); + } + websocket_target + } else { + http_target + }; + let mut outgoing = TcpStream::connect(target).await?; + outgoing.write_all(&method).await?; + tokio::io::copy_bidirectional(&mut incoming, &mut outgoing).await?; + Ok(()) + }); + } + }); + Ok(Self { + url, + opens, + allowed_opens, + task, + }) + } +} + +impl Drop for Gateway { + fn drop(&mut self) { + self.task.abort(); + } +} + +#[tokio::test] +async fn cold_pool_uses_http_during_open_timeout_then_recovers_after_cooldown() -> Result<()> { + skip_if_no_network!(Ok(())); + for free_guardian in [false, true] { + let http = responses::start_mock_server().await; + let events = vec![ + responses::ev_output_text_delta("low"), + responses::ev_completed("score"), + ]; + let http_mock = + responses::mount_sse_sequence(&http, vec![responses::sse(events.clone()); 3]).await; + let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1]; + connections.push(vec![events.clone()]); + let ws = responses::start_websocket_server(connections).await; + let gateway = Gateway::new(&http.uri(), ws.uri()).await?; + let base_path = if free_guardian { + "/backend-api/codex" + } else { + "/v1" + }; + let base_url = format!("{}{base_path}", gateway.url); + let mut config = sampler_config(base_url.clone()); + config.provider = create_model_provider( + ModelProviderInfo::create_openai_provider(Some(base_url)), + Some(AuthManager::from_auth_for_testing(if free_guardian { + CodexAuth::create_dummy_chatgpt_auth_for_testing() + } else { + CodexAuth::from_api_key("test-api-key") + })), + ); + config.free_guardian = free_guardian; + config.service_tier = Some("priority".to_owned()); + let sampler = LunaSampler::new(config); + let opener = sampler.connections.replenish().unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while gateway.opens.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await?; + let mut request = sample_request("parent-turn"); + request.parent_response_id = Some("resp-parent".to_owned()); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), sampler.sample(request)).await??, + "low" + ); + assert_eq!( + sampler.sample(sample_request("parent-turn-2")).await?, + "low" + ); + assert_eq!(gateway.opens.load(Ordering::SeqCst), 1); + assert!(!opener.is_finished()); + assert!(sampler.connections.replenish().is_none()); + + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(15)).await; + opener.await?; + tokio::time::resume(); + assert!(sampler.connections.replenish().is_none()); + assert_eq!( + sampler.sample(sample_request("during-cooldown")).await?, + "low" + ); + assert_eq!(gateway.opens.load(Ordering::SeqCst), 1); + + let requests = http_mock.requests(); + let first = &requests[0]; + let expected_path = if free_guardian { + "/backend-api/codex/guardian-classifier" + } else { + "/v1/responses" + }; + assert_eq!(first.path(), expected_path); + assert!(first.header("authorization").is_some()); + let body = first.body_json(); + assert_eq!( + body["service_tier"].as_str(), + if free_guardian { + None + } else { + Some("priority") + } + ); + assert_eq!( + body["client_metadata"]["parent_response_id"].as_str(), + free_guardian.then_some("resp-parent") + ); + + gateway.allowed_opens.store(usize::MAX, Ordering::SeqCst); + tokio::time::pause(); + tokio::time::advance(CONNECT_COOLDOWN).await; + tokio::time::resume(); + sampler.connections.replenish().unwrap().await?; + assert_eq!( + sampler.sample(sample_request("after-cooldown")).await?, + "low" + ); + assert_eq!(http_mock.requests().len(), 3); + assert_eq!( + gateway.opens.load(Ordering::SeqCst), + 1 + INITIAL_WEBSOCKET_CONNECTIONS + ); + } + Ok(()) +} + +#[tokio::test] +async fn cooldown_preserves_healthy_sockets_and_both_transports_share_capacity() -> Result<()> { + skip_if_no_network!(Ok(())); + let http = responses::start_mock_server().await; + let ws = responses::start_websocket_server_with_headers(vec![ + responses::WebSocketConnectionConfig { + requests: Vec::new(), + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: false, + }, + ]) + .await; + let gateway = Gateway::new(&http.uri(), ws.uri()).await?; + gateway.allowed_opens.store(/*val*/ 1, Ordering::SeqCst); + let sampler = LunaSampler::new(sampler_config(format!("{}/v1", gateway.url))); + let opener = sampler.connections.replenish().unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while gateway.opens.load(Ordering::SeqCst) < 2 { + tokio::task::yield_now().await; + } + }) + .await?; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(15)).await; + opener.await?; + tokio::time::resume(); + let mut leases = Vec::new(); + leases.push(sampler.connections.lease().await?); + assert!(matches!(leases[0].connection, Connection::Websocket(_))); + for _ in 1..MAX_CONCURRENT_REQUESTS { + let lease = sampler.connections.lease().await?; + assert!(matches!(lease.connection, Connection::Http(_))); + leases.push(lease); + } + assert!( + tokio::time::timeout(Duration::from_millis(20), sampler.connections.lease()) + .await + .is_err() + ); + leases.pop(); + assert!(matches!( + sampler.connections.lease().await?.connection, + Connection::Http(_) + )); + assert_eq!(gateway.opens.load(Ordering::SeqCst), 2); + Ok(()) +} + +#[tokio::test] +async fn stalled_http_headers_exhaust_the_sampling_retry_budget() -> Result<()> { + skip_if_no_network!(Ok(())); + let http = responses::start_mock_server().await; + let mock = responses::mount_response_sequence( + &http, + vec![responses::sse_response(String::new()).set_delay(Duration::from_secs(60)); 3], + ) + .await; + let mut config = sampler_config(format!("{}/v1", http.uri())); + let mut provider = config.provider.info().clone(); + provider.stream_idle_timeout_ms = Some(1_000); + config.provider = create_model_provider(provider, config.provider.auth_manager()); + let sampler = LunaSampler::new(config); + + let result = tokio::time::timeout( + Duration::from_secs(10), + sampler.sample(sample_request("stalled-headers")), + ) + .await?; + assert!(matches!( + result, + Err(LunaSamplerError::Api(ApiError::Transport( + TransportError::Timeout + ))) + )); + assert_eq!(mock.requests().len(), 3); + assert_eq!( + sampler.connections.classifications.available_permits(), + MAX_CONCURRENT_REQUESTS + ); + Ok(()) +} + +#[tokio::test] +async fn supersession_closes_http_before_headers_and_while_draining_the_body() -> Result<()> { + skip_if_no_network!(Ok(())); + // The second case publishes a score, then leaves the response body open forever. + for response in [ + String::new(), + format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n{}", + responses::sse(vec![responses::ev_output_text_delta("low")]), + ), + ] { + let has_score = !response.is_empty(); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let sampler = Arc::new(LunaSampler::new(sampler_config(format!( + "http://{}/v1", + listener.local_addr()?, + )))); + let (received, mut requests) = tokio::sync::mpsc::unbounded_channel(); + let (closed, mut disconnects) = tokio::sync::mpsc::unbounded_channel(); + let mut tasks = tokio::task::JoinSet::new(); + tasks.spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + let mut request_id = 0; + while let Ok((mut socket, _)) = listener.accept().await { + let mut method = [0; 4]; + socket.read_exact(&mut method).await?; + if &method != b"POST" { + continue; + } + request_id += 1; + let received = received.clone(); + let closed = closed.clone(); + let response = response.clone(); + connections.spawn(async move { + socket.write_all(response.as_bytes()).await?; + let _ = received.send(request_id); + let result = tokio::io::copy(&mut socket, &mut tokio::io::sink()).await; + let _ = closed.send(request_id); + result + }); + } + Ok::<_, std::io::Error>(()) + }); + let mut samples = tokio::task::JoinSet::new(); + let first = Arc::clone(&sampler); + samples.spawn(async move { first.sample(sample_request("first")).await }); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), requests.recv()).await?, + Some(1) + ); + if has_score { + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), samples.join_next()) + .await? + .unwrap()??, + "low", + ); + assert!(disconnects.try_recv().is_err()); + } + for _ in 0..MAX_CONCURRENT_REQUESTS { + let sampler = Arc::clone(&sampler); + samples.spawn(async move { sampler.sample(sample_request("newer")).await }); + } + if !has_score { + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), samples.join_next()) + .await? + .unwrap()?, + Err(LunaSamplerError::Superseded) + )); + } + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), disconnects.recv()).await?, + Some(1) + ); + } + Ok(()) +} 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 4735acb17b..de0f44fb87 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 @@ -28,9 +28,9 @@ use serde_json::json; use std::collections::HashSet; use std::sync::Arc; use std::sync::Mutex; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tokio::net::TcpStream; use uuid::Uuid; @@ -41,14 +41,15 @@ use super::LunaSampler; use super::LunaSamplerConfig; use super::LunaSamplerError; use super::LunaSamplingRequest; -use super::MAX_SAMPLING_RETRIES; -use super::MAX_WEBSOCKET_CONNECTIONS; +use super::MAX_CONCURRENT_REQUESTS; 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 { + while self.connections.idle_connections.lock().unwrap().len() + < INITIAL_WEBSOCKET_CONNECTIONS + { tokio::task::yield_now().await; } }) @@ -119,7 +120,7 @@ fn assert_connection_metadata( } #[derive(Clone, Copy)] -enum ProxyPrewarmLimit { +pub(in crate::async_scorer) enum ProxyPrewarmLimit { AllConnections, StopAfter { ready_connections: usize }, } @@ -131,6 +132,14 @@ async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> async fn proxy_websocket_servers_with_prewarm_limit( servers: &[&responses::WebSocketTestServer], prewarm_limit: ProxyPrewarmLimit, +) -> Result { + proxy_websocket_servers_with_http(servers, prewarm_limit, /*http_url*/ None).await +} + +pub(in crate::async_scorer) async fn proxy_websocket_servers_with_http( + servers: &[&responses::WebSocketTestServer], + prewarm_limit: ProxyPrewarmLimit, + http_url: Option<&str>, ) -> Result { let listener = TcpListener::bind("127.0.0.1:0").await?; let address = listener.local_addr()?; @@ -138,24 +147,40 @@ async fn proxy_websocket_servers_with_prewarm_limit( .iter() .map(|server| server.uri().trim_start_matches("ws://").to_owned()) .collect::>(); + let http_target = http_url.map(|url| url.trim_start_matches("http://").to_owned()); 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 { - return; - }; - drop(connection); + let mut index = 0; + let mut failed_prewarm = false; + while let Ok((mut incoming, _)) = listener.accept().await { + let mut method = [0; 4]; + if incoming.read_exact(&mut method).await.is_err() { + continue; } - let Ok((mut incoming, _)) = listener.accept().await else { - return; + let target = if &method == b"GET " { + if let ProxyPrewarmLimit::StopAfter { ready_connections } = prewarm_limit + && ready_connections < INITIAL_WEBSOCKET_CONNECTIONS + && index == ready_connections + && !failed_prewarm + { + failed_prewarm = true; + continue; + } + let target = targets.get(index).cloned(); + index += 1; + target + } else { + http_target.clone() + }; + let Some(target) = target else { + continue; }; tokio::spawn(async move { let Ok(mut outgoing) = TcpStream::connect(target).await else { return; }; + if outgoing.write_all(&method).await.is_err() { + return; + } let _ = tokio::io::copy_bidirectional(&mut incoming, &mut outgoing).await; }); } @@ -163,7 +188,7 @@ async fn proxy_websocket_servers_with_prewarm_limit( Ok(format!("http://{address}/v1")) } -fn sampler_config(base_url: String) -> LunaSamplerConfig { +pub(super) fn sampler_config(base_url: String) -> LunaSamplerConfig { LunaSamplerConfig { provider: create_model_provider( ModelProviderInfo::create_openai_provider(Some(base_url)), @@ -190,7 +215,7 @@ async fn connect_sampler(config: LunaSamplerConfig) -> Result { Ok(sampler) } -fn sample_request(parent_turn_id: &str) -> LunaSamplingRequest { +pub(super) fn sample_request(parent_turn_id: &str) -> LunaSamplingRequest { LunaSamplingRequest { parent_response_id: None, instructions: "Return high for high risk or low for low risk.".to_owned(), @@ -469,6 +494,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications .await?; tokio::time::timeout(Duration::from_secs(2), async { while sampler + .connections .idle_connections .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -480,6 +506,8 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications }) .await?; manager.refresh_token_from_authority().await?; + sampler.connections.clear(); + sampler.prewarm().await; let second = sampler .sample(LunaSamplingRequest { parent_response_id: None, @@ -704,48 +732,6 @@ async fn sampler_keeps_first_classification_token_when_later_output_disagrees() Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn sampler_recovers_after_initial_prewarm_failures() -> Result<()> { - skip_if_no_network!(Ok(())); - - let server = responses::start_websocket_server(vec![vec![vec![ - ev_assistant_message("recovered", "low"), - ev_completed("recovered"), - ]]]) - .await; - let listener = TcpListener::bind("127.0.0.1:0").await?; - let address = listener.local_addr()?; - let target = server.uri().trim_start_matches("ws://").to_owned(); - let failed_connections = Arc::new(AtomicUsize::new(0)); - let observed_failures = Arc::clone(&failed_connections); - tokio::spawn(async move { - for _ in 0..=MAX_SAMPLING_RETRIES { - let Ok((connection, _)) = listener.accept().await else { - return; - }; - observed_failures.fetch_add(/*val*/ 1, Ordering::Relaxed); - drop(connection); - } - while let Ok((mut incoming, _)) = listener.accept().await { - let target = target.clone(); - tokio::spawn(async move { - let Ok(mut outgoing) = TcpStream::connect(target).await else { - return; - }; - let _ = tokio::io::copy_bidirectional(&mut incoming, &mut outgoing).await; - }); - } - }); - - let sampler = connect_sampler(sampler_config(format!("http://{address}/v1"))).await?; - assert_eq!(failed_connections.load(Ordering::Relaxed), 1); - - assert_eq!(sampler.sample(sample_request("turn-1")).await?, "low"); - assert_eq!(server.handshakes().len(), 1); - - Ok(()) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn sampler_remains_available_when_second_prewarm_fails() -> Result<()> { skip_if_no_network!(Ok(())); @@ -763,70 +749,6 @@ async fn sampler_remains_available_when_second_prewarm_fails() -> Result<()> { Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn sampler_grows_its_pool_for_overlapping_requests() -> Result<()> { - skip_if_no_network!(Ok(())); - - let response = |id: &str| vec![vec![ev_assistant_message(id, "low"), ev_completed(id)]]; - let first = responses::start_websocket_server(vec![response("response-1")]).await; - let second = responses::start_websocket_server(vec![response("response-2")]).await; - let third = responses::start_websocket_server(vec![response("response-3")]).await; - let sampler = connect_sampler(sampler_config( - proxy_websocket_servers_with_prewarm_limit( - &[&first, &second, &third], - ProxyPrewarmLimit::StopAfter { - ready_connections: 2, - }, - ) - .await?, - )) - .await?; - - let mut first_request = sample_request("turn-1"); - first_request.root_turn_id = Some("root-1".to_owned()); - let mut second_request = sample_request("turn-2"); - second_request.root_turn_id = Some("root-2".to_owned()); - let outputs = tokio::try_join!( - sampler.sample(first_request), - sampler.sample(second_request), - sampler.sample(sample_request("turn-3")), - )?; - - assert_eq!( - outputs, - ("low".to_owned(), "low".to_owned(), "low".to_owned(),) - ); - let mut thread_ids = HashSet::new(); - let mut turn_ids = HashSet::new(); - let mut parent_turn_ids = HashSet::new(); - for server in [&first, &second, &third] { - assert_eq!(server.single_connection().len(), 1); - let metadata = server.single_connection()[0].body_json()["client_metadata"].clone(); - let parent_turn_id = metadata["parent_turn_id"].as_str().expect("owning turn ID"); - let expected_root = match parent_turn_id { - "turn-1" => Some("root-1"), - "turn-2" => Some("root-2"), - "turn-3" => None, - other => panic!("unexpected owning turn: {other}"), - }; - assert!(thread_ids.insert(assert_connection_metadata( - server, - &[(parent_turn_id, expected_root)] - )?)); - assert!(turn_ids.insert(metadata["turn_id"].as_str().expect("turn ID").to_owned())); - parent_turn_ids.insert(parent_turn_id.to_owned()); - } - assert_eq!( - parent_turn_ids, - HashSet::from_iter([ - "turn-1".to_owned(), - "turn-2".to_owned(), - "turn-3".to_owned() - ]) - ); - Ok(()) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn sampler_replaces_scored_drains_before_unfinished_classifications() -> Result<()> { skip_if_no_network!(Ok(())); @@ -847,28 +769,42 @@ async fn sampler_replaces_scored_drains_before_unfinished_classifications() -> R accept_delay: None, close_after_requests: false, }; - let mut servers = Vec::with_capacity(MAX_WEBSOCKET_CONNECTIONS + 2); + let mut servers = Vec::with_capacity(MAX_CONCURRENT_REQUESTS + 2); servers.push(responses::start_websocket_server_with_headers(vec![scored_response]).await); servers.push(responses::start_websocket_server_with_headers(vec![stalled_response]).await); - for _ in 2..=MAX_WEBSOCKET_CONNECTIONS { + for _ in 2..MAX_CONCURRENT_REQUESTS { servers.push( responses::start_websocket_server_with_headers(vec![incomplete_response.clone()]).await, ); } - servers.push( - responses::start_websocket_server(vec![vec![vec![ - ev_assistant_message("newest", "high"), - ev_completed("newest"), - ]]]) - .await, - ); + let http = responses::start_mock_server().await; + let _http_mock = responses::mount_sse_sequence( + &http, + ["low", "high"] + .map(|score| { + responses::sse(vec![ + ev_assistant_message("overflow", score), + ev_completed("overflow"), + ]) + }) + .to_vec(), + ) + .await; let server_refs = servers[2..INITIAL_WEBSOCKET_CONNECTIONS] .iter() .chain(servers[..2].iter()) .chain(servers[INITIAL_WEBSOCKET_CONNECTIONS..].iter()) .collect::>(); let sampler = Arc::new( - connect_sampler(sampler_config(proxy_websocket_servers(&server_refs).await?)).await?, + connect_sampler(sampler_config( + proxy_websocket_servers_with_http( + &server_refs, + ProxyPrewarmLimit::AllConnections, + Some(&http.uri()), + ) + .await?, + )) + .await?, ); let oldest_sampler = Arc::clone(&sampler); @@ -888,7 +824,8 @@ async fn sampler_replaces_scored_drains_before_unfinished_classifications() -> R ) .await?; - for index in 0..MAX_WEBSOCKET_CONNECTIONS - 2 { + for index in 0..MAX_CONCURRENT_REQUESTS - 2 { + sampler.prewarm().await; assert_eq!( sampler .sample(sample_request(&format!("turn-{index}"))) @@ -993,44 +930,50 @@ async fn sampler_retries_expired_websockets_on_another_warm_connection() -> Resu } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sampler_assigns_a_fresh_identity_when_replacing_aged_connections() -> Result<()> { +async fn sampler_uses_http_with_a_fresh_identity_when_warm_connections_expire() -> Result<()> { skip_if_no_network!(Ok(())); - let response = vec![vec![ + let response = vec![ ev_assistant_message("response-1", "low"), ev_completed("response-1"), - ]]; - let first = responses::start_websocket_server(vec![response.clone()]).await; - let second = responses::start_websocket_server(vec![response.clone()]).await; - let replacement = responses::start_websocket_server(vec![response]).await; + ]; + let first = responses::start_websocket_server(vec![vec![response.clone()]]).await; + let second = responses::start_websocket_server(vec![vec![response.clone()]]).await; + let http = responses::start_mock_server().await; + let http_mock = responses::mount_sse_once(&http, responses::sse(response)).await; let sampler = connect_sampler(sampler_config( - proxy_websocket_servers_with_prewarm_limit( - &[&first, &second, &replacement], - ProxyPrewarmLimit::StopAfter { - ready_connections: 2, - }, + proxy_websocket_servers_with_http( + &[&first, &second], + ProxyPrewarmLimit::AllConnections, + Some(&http.uri()), ) .await?, )) .await?; assert_eq!(sampler.sample(sample_request("turn-1")).await?, "low"); + sampler.wait_for_prewarm(Duration::from_secs(2)).await?; { - let mut connections = sampler.idle_connections.lock().unwrap(); + let mut connections = sampler.connections.idle_connections.lock().unwrap(); for connection in connections.iter_mut() { - connection.expires_at = std::time::Instant::now(); + connection.expires_at = tokio::time::Instant::now(); } } assert_eq!(sampler.sample(sample_request("turn-2")).await?, "low"); + let request = http_mock.single_request(); + let http_thread_id = request + .header("thread-id") + .expect("HTTP classifier thread ID"); + ThreadId::from_string(&http_thread_id)?; let thread_ids = HashSet::from([ assert_connection_metadata(&first, &[])?, assert_connection_metadata(&second, &[("turn-1", None)])?, - assert_connection_metadata(&replacement, &[("turn-2", None)])?, + http_thread_id, ]); assert_eq!(thread_ids.len(), 3); + responses::assert_parent_turn(&request.body_json(), Some("turn-2"))?; assert_eq!(second.single_connection().len(), 1); - assert_eq!(replacement.single_connection().len(), 1); Ok(()) } @@ -1050,17 +993,20 @@ async fn sampler_reconnects_after_transient_service_failures() -> Result<()> { }; let first = responses::start_websocket_server(unavailable()).await; let second = responses::start_websocket_server(unavailable()).await; - let recovered = responses::start_websocket_server(vec![vec![vec![ - ev_assistant_message("recovered", "low"), - ev_completed("recovered"), - ]]]) + let http = responses::start_mock_server().await; + let recovered = responses::mount_sse_once( + &http, + responses::sse(vec![ + ev_assistant_message("recovered", "low"), + ev_completed("recovered"), + ]), + ) .await; let sampler = connect_sampler(sampler_config( - proxy_websocket_servers_with_prewarm_limit( - &[&first, &second, &recovered], - ProxyPrewarmLimit::StopAfter { - ready_connections: 2, - }, + proxy_websocket_servers_with_http( + &[&first, &second], + ProxyPrewarmLimit::AllConnections, + Some(&http.uri()), ) .await?, )) @@ -1069,9 +1015,13 @@ async fn sampler_reconnects_after_transient_service_failures() -> Result<()> { assert_eq!(sampler.sample(sample_request("turn-1")).await?, "low"); assert_eq!(first.single_connection().len(), 1); assert_eq!(second.single_connection().len(), 1); - assert_eq!(recovered.single_connection().len(), 1); + let request = recovered.single_request(); assert_eq!( - recovered.single_handshake().header("authorization"), + request.body_json()["client_metadata"]["turn_id"], + first.single_connection()[0].body_json()["client_metadata"]["turn_id"], + ); + assert_eq!( + request.header("authorization"), Some("Bearer test-api-key".to_owned()) ); @@ -1094,14 +1044,21 @@ async fn sampler_limits_transient_recovery_attempts() -> Result<()> { }; let first = responses::start_websocket_server(unavailable()).await; let second = responses::start_websocket_server(unavailable()).await; - let third = responses::start_websocket_server(unavailable()).await; - let unused = responses::start_websocket_server(unavailable()).await; + let http = responses::start_mock_server().await; + let third = responses::mount_sse_once( + &http, + responses::sse(vec![json!({ + "type": "response.failed", "response": { + "error": {"code": "internal_server_error", "message": "HTTP sampling failed"} + } + })]), + ) + .await; let sampler = connect_sampler(sampler_config( - proxy_websocket_servers_with_prewarm_limit( - &[&first, &second, &third, &unused], - ProxyPrewarmLimit::StopAfter { - ready_connections: 2, - }, + proxy_websocket_servers_with_http( + &[&first, &second], + ProxyPrewarmLimit::AllConnections, + Some(&http.uri()), ) .await?, )) @@ -1112,11 +1069,10 @@ async fn sampler_limits_transient_recovery_attempts() -> Result<()> { .await .expect_err("sampling should stop after the bounded retries"); - assert!(error.to_string().contains("503")); + assert!(error.to_string().contains("HTTP sampling failed")); assert_eq!(first.single_connection().len(), 1); assert_eq!(second.single_connection().len(), 1); - assert_eq!(third.single_connection().len(), 1); - assert!(unused.handshakes().is_empty()); + assert_eq!(third.requests().len(), 1); Ok(()) }