diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 42c9fecd1a..1abc89c1b7 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -409,12 +409,23 @@ pub fn create_text_param_for_request( }) } +/// Distinguishes draining a live response from abandoning obsolete work. +#[derive(Clone, Copy)] +pub enum ResponseInterrupt { + /// Drain normally, without validating or waiting for interruption acknowledgements. + Drain, + /// Skip undispatched work or retire its transport at the cleanup deadline. + /// Undispatched work ends without a completion event and leaves the socket reusable. + Cancel { deadline: tokio::time::Instant }, +} + pub struct ResponseStream { pub rx_event: mpsc::Receiver>, /// Server-assigned `x-request-id` response header, when present. pub upstream_request_id: Option, /// Requests a graceful interrupt. Keep consuming events through completion. - pub interrupt: Option>, + /// Only `Cancel` validates and waits for an interruption acknowledgement. + pub interrupt: Option>, } impl Stream for ResponseStream { diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 5418b1c235..77ed1ea60d 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -1,5 +1,6 @@ use crate::auth::SharedAuthProvider; use crate::common::ResponseEvent; +use crate::common::ResponseInterrupt; use crate::common::ResponseStream; use crate::common::ResponsesWsRequest; use crate::common::SafetyBufferingTreatment; @@ -25,6 +26,7 @@ use http::StatusCode; use serde::Deserialize; use serde_json::Value; use serde_json::map::Map as JsonMap; +use std::future::Future; use std::sync::Arc; use std::sync::OnceLock; use std::time::Duration; @@ -293,14 +295,38 @@ impl ResponsesWebsocketConnection { .send(Ok(ResponseEvent::ServerReasoningIncluded(true))) .await; } + // Abandoned work has a cleanup deadline; ordinary steering preserves the response. + let interrupt = async move { + match rx_interrupt.await { + Ok(interrupt) => interrupt, + Err(_) => std::future::pending().await, + } + } + .shared(); + let interrupt_timeout = async { + match interrupt.clone().await { + ResponseInterrupt::Drain => std::future::pending().await, + ResponseInterrupt::Cancel { deadline } => { + tokio::time::sleep_until(deadline).await; + } + } + }; + tokio::pin!(interrupt_timeout); let mut guard = tokio::select! { biased; _ = tx_event.closed() => return, + () = &mut interrupt_timeout => { + let _ = tx_event.send(Err(ApiError::Stream( + "timed out draining interrupted response".into(), + ))).await; + return; + } guard = stream.lock() => guard, }; if tx_event.is_closed() { return; } + let mut completion = None; let result = { let Some(ws_stream) = guard.as_mut() else { let _ = tx_event @@ -321,8 +347,12 @@ impl ResponsesWebsocketConnection { telemetry, turn_state.as_deref(), &timing_log_context, - rx_interrupt, + interrupt.clone(), + &mut completion, ) => result, + () = &mut interrupt_timeout => Err(ApiError::Stream( + "timed out draining interrupted response".into(), + )), _ = tx_event.closed() => Err(ApiError::Stream( "response event consumer dropped".to_string(), )), @@ -335,7 +365,15 @@ impl ResponsesWebsocketConnection { let failed_stream = guard.take(); drop(guard); drop(failed_stream); - let _ = tx_event.send(Err(err)).await; + // Failed cancellation retires the socket while preserving any received + // completion and its usage for the abandoned classifier's cleanup. + let event = completion.map_or_else(|| Err(err), Ok); + let _ = tx_event.send(event).await; + } else if let Some(completed) = completion { + // The handshake is finished. Consumer backpressure must not let the + // interruption deadline discard completion or its usage accounting. + drop(guard); + let _ = tx_event.send(Ok(completed)).await; } } .instrument(current_span), @@ -662,6 +700,11 @@ fn map_wrapped_websocket_error_event( })) } +#[derive(Deserialize)] +struct InterruptAcknowledgement { + response_id: String, +} + #[allow(clippy::too_many_arguments)] async fn run_websocket_response_stream( ws_stream: &mut WsStream, @@ -671,10 +714,18 @@ async fn run_websocket_response_stream( telemetry: Option>, turn_state: Option<&OnceLock>, timing_log_context: &ResponsesWebsocketTimingLogContext, - interrupt: oneshot::Receiver<()>, + interrupt: impl Future + Clone + Unpin, + completion: &mut Option, ) -> Result<(), ApiError> { let mut last_server_model: Option = None; let mut safety_buffering_treatment = SafetyBufferingTreatment::default(); + if matches!( + interrupt.clone().now_or_never(), + Some(ResponseInterrupt::Cancel { .. }) + ) { + // No request reached the backend, so this socket is still reusable. + return Ok(()); + } send_websocket_request( ws_stream, request_text, @@ -685,15 +736,18 @@ async fn run_websocket_response_stream( .await?; // A response owns its interrupt, and create must be sent before interrupt. - let mut interrupt = interrupt.fuse(); + let mut interrupt_request = interrupt.clone().fuse(); + let mut interruption = None; let mut response_id = None; + let mut awaiting_interrupt_ack = false; loop { let poll_start = Instant::now(); let response = tokio::select! { response = tokio::time::timeout(idle_timeout, ws_stream.next()) => { response.map_err(|_| ApiError::Stream("idle timeout waiting for websocket".into())) } - Ok(()) = &mut interrupt, if response_id.is_some() => { + request = &mut interrupt_request, if response_id.is_some() => { + interruption = Some(request); send_websocket_request( ws_stream, serde_json::json!({ @@ -706,6 +760,7 @@ async fn run_websocket_response_stream( timing_log_context.connection_reused, ) .await?; + awaiting_interrupt_ack = matches!(request, ResponseInterrupt::Cancel { .. }); continue; } }; @@ -743,6 +798,34 @@ async fn run_websocket_response_stream( continue; } }; + if matches!(interruption, Some(ResponseInterrupt::Cancel { .. })) + && matches!( + event.kind(), + "response.interrupt.accepted" | "response.interrupt.failed" + ) + { + let acknowledgement: InterruptAcknowledgement = serde_json::from_str(&text) + .map_err(|_| { + ApiError::Stream("invalid interruption acknowledgement".into()) + })?; + if !awaiting_interrupt_ack + || response_id.as_ref() != Some(&acknowledgement.response_id) + { + return Err(ApiError::Stream( + "unexpected interruption acknowledgement".into(), + )); + } + if event.kind() == "response.interrupt.failed" { + return Err(ApiError::Stream( + "response interruption was rejected".into(), + )); + } + awaiting_interrupt_ack = false; + if completion.is_some() { + break; + } + continue; + } emit_responses_websocket_timing_event( event.kind(), text.as_str(), @@ -778,7 +861,16 @@ async fn run_websocket_response_stream( } continue; } - if let Some(model) = event.response_model() + let model = event.response_model(); + // Preserve terminal usage before any accompanying notification can block. + let event = match process_responses_event(event) { + Ok(Some(event @ ResponseEvent::Completed { .. })) => { + *completion = Some(event); + Ok(None) + } + event => event, + }; + if let Some(model) = model && last_server_model.as_deref() != Some(model.as_str()) { let _ = tx_event @@ -816,22 +908,36 @@ async fn run_websocket_response_stream( "response event consumer dropped".to_string(), )); } - match process_responses_event(event) { + match event { Ok(Some(event)) => { if let ResponseEvent::Created { response_id: id } = &event { response_id.clone_from(id); } - let is_completed = matches!(event, ResponseEvent::Completed { .. }); let _ = tx_event.send(Ok(event)).await; - if is_completed { - break; - } } Ok(None) => {} Err(error) => { return Err(error.into_api_error()); } } + if completion.is_some() { + if awaiting_interrupt_ack { + continue; + } + if matches!( + interrupt.clone().now_or_never(), + Some(ResponseInterrupt::Cancel { .. }) + ) { + break; + } + // Ordinary responses deliver completion before releasing the socket. + // Preserve it if cancellation wins while delivery is backpressured. + let permit = tx_event.reserve().await; + if let (Ok(permit), Some(event)) = (permit, completion.take()) { + permit.send(Ok(event)); + } + break; + } } Message::Binary(_) => { return Err(ApiError::Stream("unexpected binary websocket event".into())); diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index b7511d9e50..3afbcc6ef1 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -40,6 +40,7 @@ pub use crate::common::ReasoningSummaryDelivery; pub use crate::common::ResponseCreateWsRequest; pub use crate::common::ResponseEvent; pub use crate::common::ResponseId; +pub use crate::common::ResponseInterrupt; pub use crate::common::ResponseStream; pub use crate::common::ResponsesApiRequest; pub use crate::common::ResponsesApiTools; diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 7a870f4151..a19431c5de 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -117,7 +117,7 @@ fn normalize_image_detail(detail: &mut Option, model_info: &ModelIn pub struct ResponseStream { pub(crate) rx_event: mpsc::Receiver>, - pub(crate) interrupt: Option>, + pub(crate) interrupt: Option>, /// Signals the mapper task that the consumer stopped polling before the /// provider stream reached its own terminal event. pub(crate) consumer_dropped: CancellationToken, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 70a6f0bcaf..64b42248bc 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2689,7 +2689,7 @@ async fn try_run_sampling_request( Ok(Err(_)) => { if let Some(interrupt) = stream.interrupt.take() { if step_context.settings.model_info.use_responses_lite { - let _ = interrupt.send(()); + let _ = interrupt.send(codex_api::ResponseInterrupt::Drain); } // Drain the response normally before reusing its connection and history. preempt = CancellationToken::new(); diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index e8ae9aa96c..806b1093b9 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -1239,16 +1239,23 @@ async fn injected_response_item_reopens_turn_after_final_answer() { server.shutdown().await; } -#[test_case(false; "completed_reasoning")] -#[test_case(true; "discarded_reasoning")] +#[test_case(false, Some("response.interrupt.accepted"); "completed_reasoning")] +#[test_case(true, Some("response.interrupt.accepted"); "discarded_reasoning")] +#[test_case(false, Some("response.interrupt.failed"); "rejected_interrupt")] +#[test_case(false, None; "missing_acknowledgement")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn steer_interrupts_and_drains_websocket(discard_partial: bool) -> anyhow::Result<()> { +async fn steer_interrupts_and_drains_websocket( + discard_partial: bool, + acknowledgement: Option<&str>, +) -> anyhow::Result<()> { core_test_support::skip_if_no_network!(Ok(())); let mut interrupted = ev_completed_with_tokens("resp-interrupted", /*total_tokens*/ 42); - interrupted["type"] = json!("response.incomplete"); - interrupted["response"]["status"] = json!("incomplete"); - interrupted["response"]["incomplete_details"] = json!({"reason": "interrupted"}); + if acknowledgement == Some("response.interrupt.accepted") { + interrupted["type"] = json!("response.incomplete"); + interrupted["response"]["status"] = json!("incomplete"); + interrupted["response"]["incomplete_details"] = json!({"reason": "interrupted"}); + } let server = responses::start_websocket_server(vec![vec![ vec![ev_response_created("warm-1"), ev_completed("warm-1")], vec![ @@ -1256,13 +1263,15 @@ async fn steer_interrupts_and_drains_websocket(discard_partial: bool) -> anyhow: ev_reasoning_item_added("reason-1", &["thinking"]), ], // Finish the active response only after receiving response.interrupt. - vec![ - json!({ - "type": "response.interrupt.accepted", - "response_id": "resp-interrupted", - "sequence_number": 3, + [ + acknowledgement.map(|kind| { + json!({ + "type": kind, + "response_id": "resp-interrupted", + "sequence_number": 3, + }) }), - if discard_partial { + Some(if discard_partial { json!({ "type": "response.output_item.interrupted", "response_id": "resp-interrupted", @@ -1272,9 +1281,12 @@ async fn steer_interrupts_and_drains_websocket(discard_partial: bool) -> anyhow: }) } else { ev_reasoning_item("reason-1", &["thinking"], &[]) - }, - interrupted, - ], + }), + Some(interrupted), + ] + .into_iter() + .flatten() + .collect(), vec![ ev_response_created("resp-follow-up"), ev_completed("resp-follow-up"), 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 9f97141cdf..a910a4c9f5 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/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; 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 index c561f94e95..e5197f7e42 100644 --- 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 @@ -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; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs new file mode 100644 index 0000000000..6c4930acbe --- /dev/null +++ b/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(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index 89021097e3..ce27e83c20 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/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; }; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/response.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/response.rs new file mode 100644 index 0000000000..065cea23ea --- /dev/null +++ b/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>, +} + +impl PendingResponse { + pub(super) fn new( + stream: ResponseStream, + lease: ConnectionLease, + metrics: Option>, + ) -> Self { + Self { + response: Some((stream, lease)), + metrics, + } + } + + pub(super) async fn next(&mut self) -> Option> { + 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; + }); + } +} 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 d1ceb537ba..cdf812bb32 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 @@ -142,7 +142,9 @@ pub(in crate::async_scorer) enum ProxyPrewarmLimit { StopAfter { ready_connections: usize }, } -async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result { +pub(super) async fn proxy_websocket_servers( + servers: &[&responses::WebSocketTestServer], +) -> Result { 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 { +pub(super) async fn connect_sampler(config: LunaSamplerConfig) -> Result { 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>); +pub(super) struct RecordingMetrics(pub(super) Mutex>); 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(