Interrupt and drain canceled classifier responses before socket reuse (#53698)

## Why

Canceled classifier requests need bounded cleanup that preserves received token usage and returns only safe WebSocket connections to the pool.

## What changed

- Own classifier response streams and connection leases together, interrupting dropped responses with a one-second cleanup deadline.
- Require a matching interruption acknowledgement and response completion before reusing a canceled response's socket. Retire sockets on rejection, invalid acknowledgements, or timeout while preserving any received completion and token usage.
- Skip requests canceled before dispatch and return their unused sockets to the pool.
- Distinguish cancellation from ordinary steering, which continues draining without requiring an interruption acknowledgement.

## Testing

Add regression coverage for cancellation, acknowledgement ordering and failures, cleanup timeouts, socket reuse, and token accounting under backpressure. Extend steering tests to cover rejected and missing acknowledgements.

GitOrigin-RevId: c4e7aac49d36cea2c2d0e9f2ba5f6147f67471dd

codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs

@@ -3,6 +3,7 @@
 
 mod connection_pool;
 mod execution;
+mod response;
 
 use super::request::PreparedRequest;
 use connection_pool::ConnectionPool;

codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs

@@ -576,8 +576,10 @@ impl ConnectionLease {
         }
     }
 
-    pub(super) fn reuse(self) {
-        if let Connection::Websocket(connection) = self.connection {
+    pub(super) async fn reuse(self) {
+        if let Connection::Websocket(connection) = self.connection
+            && !connection.connection.is_closed().await
+        {
             self.pool
                 .idle_connections
                 .lock()
@@ -590,3 +592,7 @@ impl ConnectionLease {
 #[cfg(test)]
 #[path = "connection_pool_tests.rs"]
 mod tests;
+
+#[cfg(test)]
+#[path = "connection_pool_cancellation_tests.rs"]
+mod cancellation_tests;

codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs

@@ -0,0 +1,199 @@
+//! Classifier cancellation preserves usage and reuses only unused or fully drained sockets.
+
+use super::*;
+use crate::async_scorer::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC;
+use crate::async_scorer::sampler::response::PendingResponse;
+use crate::async_scorer::sampler::tests::RecordingMetrics;
+use crate::async_scorer::sampler::tests::connect_sampler;
+use crate::async_scorer::sampler::tests::proxy_websocket_servers;
+use crate::async_scorer::sampler::tests::sample_request;
+use crate::async_scorer::sampler::tests::sampler_config;
+use anyhow::Result;
+use codex_api::ResponseEvent;
+use codex_api::ResponseInterrupt;
+use core_test_support::responses;
+use core_test_support::responses::ev_assistant_message;
+use core_test_support::responses::ev_completed;
+use core_test_support::responses::ev_completed_with_tokens;
+use core_test_support::skip_if_no_network;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn cancelled_classifiers_interrupt_and_reuse_only_completed_streams() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let accepted = json!({"type": "response.interrupt.accepted", "response_id": "stale"});
+    let rejected = json!({
+        "type": "response.interrupt.failed",
+        "response_id": "stale",
+        "error": {"code": "interrupt_not_supported", "message": "Interruption is disabled"},
+    });
+    let error = json!({"type": "error", "status": 400});
+    let completed = ev_completed_with_tokens("stale", /*total_tokens*/ 42);
+    // None means reuse; explicit failures must close before the one-second drain deadline.
+    let close_now = Some(Duration::from_millis(500));
+    let close_after_timeout = Some(Duration::from_secs(2));
+    for (interrupted, close_timeout, expected_tokens) in [
+        (vec![accepted.clone(), completed.clone()], None, Some(42)),
+        (vec![completed.clone(), accepted], None, Some(42)),
+        (vec![rejected.clone()], close_now, None),
+        (vec![completed.clone(), rejected], close_now, Some(42)),
+        (vec![completed.clone(), error.clone()], close_now, Some(42)),
+        (
+            vec![
+                json!({"type": "response.interrupt.accepted", "response_id": "another-response"}),
+                completed.clone(),
+            ],
+            close_now,
+            None,
+        ),
+        (vec![completed], close_after_timeout, Some(42)),
+        (vec![error], close_after_timeout, None),
+        (Vec::new(), close_after_timeout, None),
+    ] {
+        let server = responses::start_websocket_server(vec![vec![
+            vec![responses::ev_response_created("stale")],
+            interrupted,
+            vec![ev_assistant_message("fresh", "low"), ev_completed("fresh")],
+        ]])
+        .await;
+        // Offer one socket: another classification can use it only if cancellation returned it.
+        let metrics = Arc::new(RecordingMetrics::default());
+        let mut config = sampler_config(proxy_websocket_servers(&[&server]).await?);
+        config.metrics = Some(metrics.clone());
+        let sampler = Arc::new(connect_sampler(config).await?);
+        let pending_sampler = Arc::clone(&sampler);
+        let pending =
+            tokio::spawn(async move { pending_sampler.sample(sample_request("stale")).await });
+        tokio::time::timeout(
+            Duration::from_secs(2),
+            server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0),
+        )
+        .await?;
+        pending.abort();
+        assert!(pending.await.unwrap_err().is_cancelled());
+        let interrupt = tokio::time::timeout(
+            Duration::from_secs(2),
+            server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 1),
+        )
+        .await?;
+        assert_eq!(
+            interrupt.body_json(),
+            json!({
+                "type": "response.interrupt", "response_id": "stale", "mode": "discard_partial_items"
+            })
+        );
+        if let Some(close_timeout) = close_timeout {
+            assert!(
+                server
+                    .wait_for_closed_connections(/*expected*/ 1, close_timeout)
+                    .await
+            );
+        }
+        // Reacquire all slots to await cleanup and accounting, even after socket closure.
+        let permits = tokio::time::timeout(
+            Duration::from_secs(2),
+            Arc::clone(&sampler.connections.classifications)
+                .acquire_many_owned(u32::try_from(MAX_CONCURRENT_REQUESTS)?),
+        )
+        .await??;
+        assert_eq!(
+            sampler.connections.idle_connections.lock().unwrap().len(),
+            usize::from(close_timeout.is_none())
+        );
+        drop(permits);
+        if close_timeout.is_none() {
+            assert_eq!(sampler.sample(sample_request("fresh")).await?, "low");
+            let requests = server.single_connection();
+            assert_eq!(requests.len(), 3);
+            assert_eq!(
+                requests[2].body_json()["client_metadata"]["parent_turn_id"],
+                json!("fresh")
+            );
+        }
+        // A missing acknowledgement must not discard usage already received by the transport.
+        if let Some(total) = expected_tokens {
+            let expected = (
+                CLASSIFICATION_TOKEN_USAGE_METRIC.to_owned(),
+                total,
+                vec![("token_type".to_owned(), "total".to_owned())],
+            );
+            // Both branches above synchronize with lease release, after accounting.
+            assert!(metrics.0.lock().unwrap().contains(&expected));
+        }
+        drop(sampler);
+        server.shutdown().await;
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn cancellation_preserves_unused_sockets_and_backpressured_completion() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let mut completed = ev_completed_with_tokens("stale", /*total_tokens*/ 42);
+    completed["safety_buffering"] = json!({"use_cases": [], "reasons": []});
+    let output_events = 2_000;
+    let events = std::iter::once(responses::ev_response_created("stale"))
+        .chain(std::iter::repeat_n(
+            responses::ev_output_text_delta("x"),
+            output_events,
+        ))
+        .chain([completed])
+        .collect();
+    // Keep the server open after an interrupt so only the cleanup deadline closes it.
+    let server =
+        responses::start_websocket_server(vec![vec![events, Vec::new(), Vec::new()]]).await;
+    let sampler =
+        connect_sampler(sampler_config(proxy_websocket_servers(&[&server]).await?)).await?;
+    let pool = &sampler.connections;
+    let lease = pool.lease().await?;
+    let thread_id = lease.thread_id.clone();
+    let request = sampler
+        .execution(
+            &sample_request("stale"),
+            Vec::new(),
+            /*auth_owner_generation*/ None,
+        )
+        .request;
+    // Cancel before this current-thread runtime can dispatch, then reuse its only socket.
+    let stream = lease.stream_request(&request).await?;
+    drop(PendingResponse::new(stream, lease, /*metrics*/ None));
+    let cleanup = pool
+        .classifications
+        .acquire_many(u32::try_from(MAX_CONCURRENT_REQUESTS)?);
+    drop(tokio::time::timeout(Duration::from_secs(/*secs*/ 2), cleanup).await??);
+    let lease = pool.lease().await?;
+    assert_eq!(lease.thread_id, thread_id);
+    let mut stream = lease.stream_request(&request).await?;
+    // Consume Created and leave exactly a full channel of output before completion.
+    let consumed = output_events
+        .checked_sub(stream.rx_event.max_capacity())
+        .expect("fixture must fill the response channel")
+        + 1;
+    for _ in 0..consumed {
+        stream.rx_event.recv().await.expect("output")?;
+    }
+    let interrupt = stream.interrupt.take().expect("websocket interruption");
+    assert!(
+        interrupt
+            .send(ResponseInterrupt::Cancel {
+                deadline: Instant::now() + Duration::from_secs(/*secs*/ 1),
+            })
+            .is_ok()
+    );
+    assert!(
+        server
+            .wait_for_closed_connections(/*expected*/ 1, Duration::from_secs(/*secs*/ 2))
+            .await
+    );
+    let mut total_tokens = None;
+    while let Some(event) = stream.rx_event.recv().await {
+        if let ResponseEvent::Completed { token_usage, .. } = event? {
+            total_tokens = token_usage.map(|usage| usage.total_tokens);
+        }
+    }
+    assert_eq!(total_tokens, Some(42));
+    drop(lease);
+    server.shutdown().await;
+    Ok(())
+}

codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs

@@ -2,6 +2,7 @@
 //! Snapshot output returns immediately with a detached drain. Retained requests publish
 //! the same early result but keep ownership until completion to capture reusable output.
 //! Requests and retries stop when the account owner that started the classification changes.
+//! Dropped responses interrupt and drain before their transport can return to the pool.
 
 use super::CLASSIFICATION_TOKEN_USAGE_METRIC;
 use super::ConnectionPool;
@@ -9,6 +10,7 @@ use super::LunaSamplerConfig;
 use super::LunaSamplerError;
 use super::MAX_OUTPUT_BYTES;
 use super::connection_pool::RequestMode;
+use super::response::PendingResponse;
 use codex_api::ApiError;
 use codex_api::ResponseEvent;
 use codex_api::ResponsesApiRequest;
@@ -316,7 +318,7 @@ impl SamplingExecution {
                 client_metadata.insert("parent_response_id".to_owned(), parent_response_id.clone());
             }
             self.request.client_metadata = Some(client_metadata);
-            let mut stream = match tokio::select! {
+            let stream = match tokio::select! {
                 biased;
                 _ = &mut owner_changed => return Err(account_changed_error()),
                 _ = mode.superseded() => return Err(LunaSamplerError::Superseded),
@@ -335,6 +337,7 @@ impl SamplingExecution {
                 }
             };
 
+            let mut response = PendingResponse::new(stream, lease, self.config.metrics.clone());
             let mut output = String::new();
             while let Some(event) = tokio::select! {
                 biased;
@@ -347,7 +350,7 @@ impl SamplingExecution {
                         Err(LunaSamplerError::Superseded)
                     };
                 }
-                event = stream.rx_event.recv() => event,
+                event = response.next() => event,
             } {
                 ensure_account_owner()?;
                 let event = match event {
@@ -391,21 +394,15 @@ impl SamplingExecution {
                             }
                         };
                         scored.store(true, Ordering::Relaxed);
-                        let mut remaining_events = stream.rx_event;
-                        let metrics = self.config.metrics.clone();
                         tokio::spawn(async move {
                             while let Some(event) = tokio::select! {
                                 biased;
                                 _ = &mut superseded => None,
-                                event = remaining_events.recv() => event,
+                                event = response.next() => event,
                             } {
                                 match event {
                                     Ok(ResponseEvent::Completed { token_usage, .. }) => {
-                                        record_token_usage(
-                                            metrics.as_deref(),
-                                            token_usage.as_ref(),
-                                        );
-                                        lease.reuse();
+                                        response.reuse(token_usage.as_ref()).await;
                                         break;
                                     }
                                     Err(_) => break,
@@ -430,8 +427,7 @@ impl SamplingExecution {
                         }
                     }
                     ResponseEvent::Completed { token_usage, .. } => {
-                        record_token_usage(self.config.metrics.as_deref(), token_usage.as_ref());
-                        lease.reuse();
+                        response.reuse(token_usage.as_ref()).await;
                         if !output.is_empty() {
                             return Ok(output);
                         }
@@ -453,7 +449,10 @@ impl SamplingExecution {
     }
 }
 
-fn record_token_usage(metrics: Option<&dyn ExtensionMetrics>, token_usage: Option<&TokenUsage>) {
+pub(super) fn record_token_usage(
+    metrics: Option<&dyn ExtensionMetrics>,
+    token_usage: Option<&TokenUsage>,
+) {
     let (Some(metrics), Some(token_usage)) = (metrics, token_usage) else {
         return;
     };

codex-rs/ext/guardian-v2/src/async_scorer/sampler/response.rs

@@ -0,0 +1,79 @@
+//! Owns an in-flight classifier stream and its exclusive transport lease.
+//! Normal completion reuses the lease. Dropped work delegates interruption and
+//! bounded draining to the transport, which confirms whether the socket is reusable.
+//! Cancellation before dispatch also returns the unused socket to the pool.
+//! The lease bounds cleanup tasks with the same permits as active classifications.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_api::ApiError;
+use codex_api::ResponseEvent;
+use codex_api::ResponseInterrupt;
+use codex_api::ResponseStream;
+use codex_extension_api::ExtensionMetrics;
+use codex_protocol::protocol::TokenUsage;
+
+use super::connection_pool::ConnectionLease;
+use super::execution::record_token_usage;
+
+pub(super) struct PendingResponse {
+    response: Option<(ResponseStream, ConnectionLease)>,
+    metrics: Option<Arc<dyn ExtensionMetrics>>,
+}
+
+impl PendingResponse {
+    pub(super) fn new(
+        stream: ResponseStream,
+        lease: ConnectionLease,
+        metrics: Option<Arc<dyn ExtensionMetrics>>,
+    ) -> Self {
+        Self {
+            response: Some((stream, lease)),
+            metrics,
+        }
+    }
+
+    pub(super) async fn next(&mut self) -> Option<Result<ResponseEvent, ApiError>> {
+        self.response.as_mut()?.0.rx_event.recv().await
+    }
+
+    pub(super) async fn reuse(mut self, token_usage: Option<&TokenUsage>) {
+        if let Some((_, lease)) = self.response.take() {
+            record_token_usage(self.metrics.as_deref(), token_usage);
+            lease.reuse().await;
+        }
+    }
+}
+
+impl Drop for PendingResponse {
+    fn drop(&mut self) {
+        let Some((mut stream, lease)) = self.response.take() else {
+            return;
+        };
+        let Some(interrupt) = stream.interrupt.take() else {
+            // HTTP cancellation has no response-interrupt handshake.
+            return;
+        };
+        let Ok(runtime) = tokio::runtime::Handle::try_current() else {
+            return;
+        };
+        // Signal before cleanup is scheduled so the transport can skip undispatched work.
+        // Completion may already be queued even if the interrupt receiver has closed.
+        let _ = interrupt.send(ResponseInterrupt::Cancel {
+            deadline: tokio::time::Instant::now() + Duration::from_secs(/*secs*/ 1),
+        });
+        let metrics = self.metrics.clone();
+        runtime.spawn(async move {
+            while let Some(Ok(event)) = stream.rx_event.recv().await {
+                if let ResponseEvent::Completed { token_usage, .. } = event {
+                    record_token_usage(metrics.as_deref(), token_usage.as_ref());
+                    break;
+                }
+            }
+            // The transport retires failed streams before reporting an error. Clean EOF
+            // also leaves an unused socket reusable when cancellation preceded dispatch.
+            lease.reuse().await;
+        });
+    }
+}

codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs

@@ -142,7 +142,9 @@ pub(in crate::async_scorer) enum ProxyPrewarmLimit {
     StopAfter { ready_connections: usize },
 }
 
-async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result<String> {
+pub(super) async fn proxy_websocket_servers(
+    servers: &[&responses::WebSocketTestServer],
+) -> Result<String> {
     proxy_websocket_servers_with_prewarm_limit(servers, ProxyPrewarmLimit::AllConnections).await
 }
 
@@ -230,7 +232,7 @@ pub(in crate::async_scorer) fn sampler_config(base_url: String) -> LunaSamplerCo
     }
 }
 
-async fn connect_sampler(config: LunaSamplerConfig) -> Result<LunaSampler> {
+pub(super) async fn connect_sampler(config: LunaSamplerConfig) -> Result<LunaSampler> {
     let sampler = LunaSampler::new(config);
     sampler.prewarm().await;
     Ok(sampler)
@@ -339,7 +341,7 @@ async fn classifier_preserves_api_key_cyber_program_policy() -> Result<()> {
 type RecordedMetric = (String, i64, Vec<(String, String)>);
 
 #[derive(Default)]
-struct RecordingMetrics(Mutex<Vec<RecordedMetric>>);
+pub(super) struct RecordingMetrics(pub(super) Mutex<Vec<RecordedMetric>>);
 
 impl ExtensionMetrics for RecordingMetrics {
     fn histogram_with_boundaries(
@@ -1154,8 +1156,9 @@ async fn sampler_uses_http_with_a_fresh_identity_when_warm_connections_expire()
         ev_assistant_message("response-1", "low"),
         ev_completed("response-1"),
     ];
-    let first = responses::start_websocket_server(vec![vec![response.clone()]]).await;
-    let second = responses::start_websocket_server(vec![vec![response.clone()]]).await;
+    // Keep the sockets open until the test expires their pooled connections.
+    let first = responses::start_websocket_server(vec![vec![response.clone(), Vec::new()]]).await;
+    let second = responses::start_websocket_server(vec![vec![response.clone(), Vec::new()]]).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(