From 66919805ea080053d1933b6b43afeb0d8bf70c91 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 13 Aug 2026 16:55:34 +0000 Subject: [PATCH] Pool Guardian sampling WebSocket connections (#38406) ## What changed - Prewarm two Responses WebSocket connections and grow the pool on demand, up to eight connections, so overlapping samples can run concurrently. - Lease each connection to a single sample and return healthy connections to the idle pool after completion. - Replace closed or 55-minute-old connections and retry a sample once after a retryable or stream error. - Keep the sampler available when the second prewarm connection cannot be opened. ## Testing - Cover connection reuse, partial prewarm success, concurrent pool growth, and retrying an expired connection on another warm connection. GitOrigin-RevId: 3ee7b05adc00a2135ba7cb35f6e4455199bf3784 --- codex-rs/ext/guardian-v2/src/sampler.rs | 253 ++++++++++++------ codex-rs/ext/guardian-v2/src/sampler_tests.rs | 191 ++++++++++++- 2 files changed, 357 insertions(+), 87 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs index b9e55e97c1..324882c9b9 100644 --- a/codex-rs/ext/guardian-v2/src/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -1,4 +1,8 @@ use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; use codex_api::ApiError; use codex_api::Reasoning; @@ -25,9 +29,14 @@ use codex_protocol::protocol::SessionSource; use http::HeaderValue; use serde_json::Value; use thiserror::Error; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; const MODEL: &str = "gpt-5.6-luna"; const MAX_OUTPUT_BYTES: usize = 8 * 1024; +const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2; +const MAX_WEBSOCKET_CONNECTIONS: usize = 8; +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"; @@ -86,49 +95,92 @@ pub enum LunaSamplerError { OutputTooLarge, } -/// A persistent, authenticated Responses WebSocket dedicated to Luna sampling. -pub struct LunaSampler { +struct PooledConnection { connection: ResponsesWebsocketConnection, - session_id: String, - thread_id: String, - service_tier: Option, + connected_at: Instant, +} + +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); + } +} + +/// A bounded pool of authenticated Responses WebSockets dedicated to Luna sampling. +pub struct LunaSampler { + config: LunaSamplerConfig, + idle_connections: Arc>>, + capacity: Arc, } impl LunaSampler { - /// Opens the WebSocket before any sample is requested. + /// Opens the initial WebSockets before any sample is requested. pub async fn connect(config: LunaSamplerConfig) -> Result { - let provider = config + let sampler = Self { + config, + idle_connections: Arc::new(Mutex::new(Vec::with_capacity(MAX_WEBSOCKET_CONNECTIONS))), + capacity: Arc::new(Semaphore::new(MAX_WEBSOCKET_CONNECTIONS)), + }; + for index in 0..INITIAL_WEBSOCKET_CONNECTIONS { + let connection = match sampler.open_connection().await { + Ok(connection) => connection, + Err(error) if index == 0 => return Err(error), + Err(_) => break, + }; + sampler + .idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(connection); + } + Ok(sampler) + } + + async fn open_connection(&self) -> Result { + let provider = self + .config .provider .api_provider() .await .map_err(LunaSamplerError::Provider)?; - let auth = config + let auth = self + .config .provider .api_auth_for_scope(ProviderAuthScope { - agent_identity_policy: config.agent_identity_policy, - session_source: config.session_source, + 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 mut headers = build_session_headers( - Some(config.session_id.clone()), - Some(config.thread_id.clone()), + Some(self.config.session_id.clone()), + Some(self.config.thread_id.clone()), ); headers.insert( "openai-beta", HeaderValue::from_static(RESPONSES_WEBSOCKETS_BETA), ); - if let Some(originator) = config.originator.as_deref() { + if let Some(originator) = self.config.originator.as_deref() { add_originator_header(&mut headers, originator); } - if let Ok(request_id) = HeaderValue::from_str(&config.thread_id) { + if let Ok(request_id) = HeaderValue::from_str(&self.config.thread_id) { headers.insert("x-client-request-id", request_id); } - let provider_info = config.provider.info(); - if config + let provider_info = self.config.provider.info(); + if self + .config .provider .auth() .await @@ -141,7 +193,7 @@ impl LunaSampler { && provider_info.auth.is_none() && provider_info.aws.is_none() { - let routing_hint = match config.service_tier.as_deref() { + let routing_hint = match self.config.service_tier.as_deref() { Some(tier) => format!("model={MODEL};tier={tier}"), None => format!("model={MODEL}"), }; @@ -152,7 +204,7 @@ impl LunaSampler { let client = ResponsesWebsocketClient::new(provider, auth); let connect = client.connect( - &config.http_client_factory, + &self.config.http_client_factory, headers, default_headers(), /*turn_state*/ None, @@ -163,19 +215,46 @@ impl LunaSampler { .map_err(|_| LunaSamplerError::ConnectionTimeout)? .map_err(LunaSamplerError::Api)?; - Ok(Self { + Ok(PooledConnection { connection, - session_id: config.session_id, - thread_id: config.thread_id, - service_tier: config.service_tier, + connected_at: Instant::now(), }) } - /// Sends one structured, tool-less request on the existing WebSocket. + 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.connected_at.elapsed() < MAX_WEBSOCKET_AGE + && !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, + }) + } + + /// Sends one structured, tool-less request on an exclusively leased WebSocket. pub async fn sample(&self, request: LunaSamplingRequest) -> Result { let metadata = HashMap::from([ - ("session_id".to_owned(), self.session_id.clone()), - ("thread_id".to_owned(), self.thread_id.clone()), + ("session_id".to_owned(), self.config.session_id.clone()), + ("thread_id".to_owned(), self.config.thread_id.clone()), ("turn_id".to_owned(), request.turn_id), (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), ]); @@ -219,8 +298,8 @@ impl LunaSampler { stream: true, stream_options: None, include: Vec::new(), - service_tier: self.service_tier.clone(), - prompt_cache_key: Some(format!("guardian-v2:{}", self.thread_id)), + service_tier: self.config.service_tier.clone(), + prompt_cache_key: Some(format!("guardian-v2:{}", self.config.thread_id)), text: create_text_param_for_request( /*verbosity*/ None, &Some(request.output_schema), @@ -228,66 +307,88 @@ impl LunaSampler { ), client_metadata: Some(metadata), }; - let mut stream = self - .connection - .stream_request( - ResponsesWsRequest::ResponseCreate((&request).into()), - /*connection_reused*/ true, - /*turn_state*/ None, - ) - .await - .map_err(LunaSamplerError::Api)?; + let mut retried = false; + 'retry: loop { + let lease = self.lease_connection().await?; + let mut stream = lease + .connection + .connection + .stream_request( + ResponsesWsRequest::ResponseCreate((&request).into()), + /*connection_reused*/ true, + /*turn_state*/ None, + ) + .await + .map_err(LunaSamplerError::Api)?; - let mut output = String::new(); - let mut deltas = String::new(); - while let Some(event) = stream.rx_event.recv().await { - match event.map_err(LunaSamplerError::Api)? { - ResponseEvent::OutputTextDelta(delta) => { - deltas.push_str(&delta); - if deltas.len() > MAX_OUTPUT_BYTES { - return Err(LunaSamplerError::OutputTooLarge); + let mut output = String::new(); + let mut deltas = String::new(); + while let Some(event) = stream.rx_event.recv().await { + let event = match event { + Ok(event) => event, + Err(error) + if !retried + && matches!( + error, + ApiError::Retryable { .. } | ApiError::Stream(_) + ) => + { + retried = true; + continue 'retry; } + Err(error) => return Err(LunaSamplerError::Api(error)), + }; + match event { + ResponseEvent::OutputTextDelta(delta) => { + deltas.push_str(&delta); + if deltas.len() > MAX_OUTPUT_BYTES { + return Err(LunaSamplerError::OutputTooLarge); + } - if serde_json::from_str::>(&deltas).is_ok() { - let mut remaining_events = stream.rx_event; - tokio::spawn(async move { - // The transport needs a live consumer until completion to keep its - // authenticated WebSocket available for the next sample. - while let Some(event) = remaining_events.recv().await { - if matches!(event, Ok(ResponseEvent::Completed { .. }) | Err(_)) { - break; + if serde_json::from_str::>(&deltas).is_ok() { + let mut remaining_events = stream.rx_event; + tokio::spawn(async move { + while let Some(event) = remaining_events.recv().await { + match event { + Ok(ResponseEvent::Completed { .. }) => { + lease.reuse(); + break; + } + Err(_) => break, + _ => {} + } } - } - }); - return Ok(deltas); - } - } - ResponseEvent::OutputItemDone(ResponseItem::Message { role, content, .. }) - if role == "assistant" => - { - for item in content { - if let ContentItem::OutputText { text } = item { - output.push_str(&text); + }); + return Ok(deltas); } } - } - ResponseEvent::Completed { .. } => { - if !output.is_empty() { - return Ok(output); + ResponseEvent::OutputItemDone(ResponseItem::Message { + role, content, .. + }) if role == "assistant" => { + for item in content { + if let ContentItem::OutputText { text } = item { + output.push_str(&text); + } + } } - if !deltas.is_empty() { - return Ok(deltas); + ResponseEvent::Completed { .. } => { + lease.reuse(); + if !output.is_empty() { + return Ok(output); + } + if !deltas.is_empty() { + return Ok(deltas); + } + return Err(LunaSamplerError::MissingOutput); } - return Err(LunaSamplerError::MissingOutput); + _ => {} + } + if output.len() > MAX_OUTPUT_BYTES || deltas.len() > MAX_OUTPUT_BYTES { + return Err(LunaSamplerError::OutputTooLarge); } - _ => {} - } - if output.len() > MAX_OUTPUT_BYTES || deltas.len() > MAX_OUTPUT_BYTES { - return Err(LunaSamplerError::OutputTooLarge); } + return Err(LunaSamplerError::MissingOutput); } - - Err(LunaSamplerError::MissingOutput) } } diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs index b2ee19e039..9d5b401a4f 100644 --- a/codex-rs/ext/guardian-v2/src/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -18,17 +18,75 @@ use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use serde_json::json; use std::time::Duration; +use tokio::net::TcpListener; +use tokio::net::TcpStream; use super::LunaSampler; use super::LunaSamplerConfig; use super::LunaSamplingRequest; +async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let targets = servers + .iter() + .map(|server| server.uri().trim_start_matches("ws://").to_owned()) + .collect::>(); + tokio::spawn(async move { + for target in targets { + let Ok((mut incoming, _)) = listener.accept().await else { + return; + }; + tokio::spawn(async move { + let Ok(mut outgoing) = TcpStream::connect(target).await else { + return; + }; + let _ = tokio::io::copy_bidirectional(&mut incoming, &mut outgoing).await; + }); + } + }); + Ok(format!("http://{address}/v1")) +} + +fn sampler_config(base_url: String) -> LunaSamplerConfig { + LunaSamplerConfig { + provider: create_model_provider( + ModelProviderInfo::create_openai_provider(Some(base_url)), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "test-api-key", + ))), + ), + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly, + session_source: SessionSource::Exec, + session_id: "session-1".to_owned(), + thread_id: "thread-1".to_owned(), + originator: Some("guardian-v2-test".to_owned()), + service_tier: None, + } +} + +fn sample_request(turn_id: &str) -> LunaSamplingRequest { + LunaSamplingRequest { + instructions: "Return a risk score.".to_owned(), + input: "The user requested a README summary.".to_owned(), + output_schema: json!({ + "type": "object", + "properties": { "score": { "type": "number" } }, + "required": ["score"], + "additionalProperties": false + }), + reasoning_effort: ReasoningEffort::None, + turn_id: turn_id.to_owned(), + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requests() -> Result<()> { skip_if_no_network!(Ok(())); - let server = responses::start_websocket_server(vec![vec![ + let scripted_requests = vec![ vec![ ev_output_text_delta(r#"{"score":0"#), ev_output_text_delta(".25}"), @@ -40,11 +98,12 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ ev_assistant_message("sample-2", r#"{"score":0.75}"#), ev_completed("response-2"), ], - ]]) - .await; - let base_url = server.uri().replacen("ws://", "http://", 1); + ]; + let idle_server = responses::start_websocket_server(vec![scripted_requests.clone()]).await; + let server = responses::start_websocket_server(vec![scripted_requests]).await; + let base_url = proxy_websocket_servers(&[&idle_server, &server]).await?; let provider = create_model_provider( - ModelProviderInfo::create_openai_provider(Some(format!("{base_url}/v1"))), + ModelProviderInfo::create_openai_provider(Some(base_url)), Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( "test-api-key", ))), @@ -94,6 +153,18 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ turn_id: "turn-1".to_owned(), }) .await?; + tokio::time::timeout(Duration::from_secs(2), async { + while sampler + .idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len() + < 2 + { + tokio::task::yield_now().await; + } + }) + .await?; let second = sampler .sample(LunaSamplingRequest { instructions: "Return a risk score.".to_owned(), @@ -138,7 +209,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ async fn sampler_returns_complete_json_before_terminal_response_events() -> Result<()> { skip_if_no_network!(Ok(())); - let server = responses::start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + let config = WebSocketConnectionConfig { requests: vec![vec![ ev_output_text_delta(r#"{"score":0"#), ev_output_text_delta(".25}"), @@ -146,11 +217,12 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu response_headers: Vec::new(), accept_delay: None, close_after_requests: false, - }]) - .await; - let base_url = server.uri().replacen("ws://", "http://", 1); + }; + let idle_server = responses::start_websocket_server_with_headers(vec![config.clone()]).await; + let server = responses::start_websocket_server_with_headers(vec![config]).await; + let base_url = proxy_websocket_servers(&[&idle_server, &server]).await?; let provider = create_model_provider( - ModelProviderInfo::create_openai_provider(Some(format!("{base_url}/v1"))), + ModelProviderInfo::create_openai_provider(Some(base_url)), Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( "test-api-key", ))), @@ -185,6 +257,103 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu .await??; assert_eq!(output, r#"{"score":0.25}"#); - server.shutdown().await; + drop(sampler); + tokio::join!(idle_server.shutdown(), server.shutdown()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn sampler_remains_available_when_second_prewarm_fails() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_websocket_server(vec![vec![vec![ + ev_assistant_message("response-1", r#"{"score":0.25}"#), + ev_completed("response-1"), + ]]]) + .await; + let sampler = + LunaSampler::connect(sampler_config(proxy_websocket_servers(&[&server]).await?)).await?; + + assert_eq!( + sampler.sample(sample_request("turn-1")).await?, + r#"{"score":0.25}"# + ); + assert_eq!(server.handshakes().len(), 1); + 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, r#"{"score":0.25}"#), + 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 = LunaSampler::connect(sampler_config( + proxy_websocket_servers(&[&first, &second, &third]).await?, + )) + .await?; + + let outputs = tokio::try_join!( + sampler.sample(sample_request("turn-1")), + sampler.sample(sample_request("turn-2")), + sampler.sample(sample_request("turn-3")), + )?; + + assert_eq!( + outputs, + ( + r#"{"score":0.25}"#.to_owned(), + r#"{"score":0.25}"#.to_owned(), + r#"{"score":0.25}"#.to_owned(), + ) + ); + for server in [&first, &second, &third] { + assert_eq!(server.single_connection().len(), 1); + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_retries_expired_websockets_on_another_warm_connection() -> Result<()> { + skip_if_no_network!(Ok(())); + + let healthy = responses::start_websocket_server(vec![vec![vec![ + ev_assistant_message("response-1", r#"{"score":0.25}"#), + ev_completed("response-1"), + ]]]) + .await; + let expired = responses::start_websocket_server(vec![vec![vec![json!({ + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "code": "websocket_connection_limit_reached", + "message": "Responses websocket connection limit reached (60 minutes)." + } + })]]]) + .await; + let sampler = LunaSampler::connect(sampler_config( + proxy_websocket_servers(&[&healthy, &expired]).await?, + )) + .await?; + + let output = sampler.sample(sample_request("turn-1")).await?; + + assert_eq!(output, r#"{"score":0.25}"#); + let expired_requests = expired.single_connection(); + let healthy_requests = healthy.single_connection(); + assert_eq!(expired_requests.len(), 1); + assert_eq!(healthy_requests.len(), 1); + assert_eq!( + expired_requests[0].body_json(), + healthy_requests[0].body_json() + ); Ok(()) }