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
2026-10-18 17:18:52 +00:00 · 2026-08-13 17:00:59 +00:00
parent a70211249a
commit 66919805ea
2 changed files with 357 additions and 87 deletions
@@ -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<String>,
+    connected_at: Instant,
+}
+
+struct ConnectionLease {
+    connection: PooledConnection,
+    idle_connections: Arc<Mutex<Vec<PooledConnection>>>,
+    _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<Mutex<Vec<PooledConnection>>>,
+    capacity: Arc<Semaphore>,
 }

 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<Self, LunaSamplerError> {
-        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<PooledConnection, LunaSamplerError> {
+        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<ConnectionLease, LunaSamplerError> {
+        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<String, LunaSamplerError> {
        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::<serde_json::Map<String, Value>>(&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::<serde_json::Map<String, Value>>(&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)
    }
 }

@@ -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<String> {
+    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::<Vec<_>>();
+    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(())
 }