github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 16:17:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
felixxia-oai
2026-10-13 15:24:00 +00:00
committed by copyberry
parent 31434401fa
commit ffc1402633
12 changed files with 466 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/common.rs
+12 -1
View File
@@ -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<Result<ResponseEvent, ApiError>>,
/// Server-assigned `x-request-id` response header, when present.
pub upstream_request_id: Option<String>,
/// Requests a graceful interrupt. Keep consuming events through completion.
pub interrupt: Option<oneshot::Sender<()>>,
/// Only `Cancel` validates and waits for an interruption acknowledgement.
pub interrupt: Option<oneshot::Sender<ResponseInterrupt>>,
}
impl Stream for ResponseStream {
codex-rs/codex-api/src/endpoint/responses_websocket.rs
+117 -11
View File
@@ -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<Arc<dyn WebsocketTelemetry>>,
turn_state: Option<&OnceLock<String>>,
timing_log_context: &ResponsesWebsocketTimingLogContext,
interrupt: oneshot::Receiver<()>,
interrupt: impl Future<Output = ResponseInterrupt> + Clone + Unpin,
completion: &mut Option<ResponseEvent>,
) -> Result<(), ApiError> {
let mut last_server_model: Option<String> = 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()));
codex-rs/codex-api/src/lib.rs
+1
View File
@@ -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;
codex-rs/core/src/client_common.rs
+1 -1
View File
@@ -117,7 +117,7 @@ fn normalize_image_detail(detail: &mut Option<ImageDetail>, model_info: &ModelIn
pub struct ResponseStream {
pub(crate) rx_event: mpsc::Receiver<Result<ResponseEvent>>,
pub(crate) interrupt: Option<oneshot::Sender<()>>,
pub(crate) interrupt: Option<oneshot::Sender<codex_api::ResponseInterrupt>>,
/// Signals the mapper task that the consumer stopped polling before the
/// provider stream reached its own terminal event.
pub(crate) consumer_dropped: CancellationToken,
codex-rs/core/src/session/turn.rs
+1 -1
View File
@@ -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();
codex-rs/core/tests/suite/pending_input.rs
+27 -15
View File
@@ -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"),
codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs
+1
View File
@@ -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
+8 -2
View File
@@ -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
+199
View File
@@ -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
+12 -13
View File
@@ -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
+79
View File
@@ -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
+8 -5
View File
@@ -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(