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(()) }