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
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
+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 {