Prioritize new Guardian classifications under load (#38596)

## What changed

- Expand the Guardian sampling pool from 8 to 16 WebSocket connections.
- When the pool is full, supersede the oldest request that has already produced a score before replacing an unfinished classification. Treat superseded classifications as a no-op in the extension.
- Stop response WebSocket work when its event consumer is dropped, including while waiting for the connection lock or draining a completed sample.
- Allow retries across both initially warmed connections for retryable stream failures.

## Testing

- Add a concurrent sampler test that fills the pool and verifies scored drains are replaced before an unfinished classification.

GitOrigin-RevId: 9e4df7b01516094726ff3e11878f8ca742d21b03
This commit is contained in:
jif
2026-08-14 15:49:42 +00:00
committed by copyberry
parent 23094236ac
commit 742edd6c16
4 changed files with 214 additions and 39 deletions

View File

@@ -282,7 +282,14 @@ impl ResponsesWebsocketConnection {
.send(Ok(ResponseEvent::ServerReasoningIncluded(true)))
.await;
}
let mut guard = stream.lock().await;
let mut guard = tokio::select! {
biased;
_ = tx_event.closed() => return,
guard = stream.lock() => guard,
};
if tx_event.is_closed() {
return;
}
let result = {
let Some(ws_stream) = guard.as_mut() else {
let _ = tx_event
@@ -293,16 +300,21 @@ impl ResponsesWebsocketConnection {
return;
};
run_websocket_response_stream(
ws_stream,
tx_event.clone(),
request_text,
idle_timeout,
telemetry,
turn_state.as_deref(),
&timing_log_context,
)
.await
tokio::select! {
biased;
result = run_websocket_response_stream(
ws_stream,
tx_event.clone(),
request_text,
idle_timeout,
telemetry,
turn_state.as_deref(),
&timing_log_context,
) => result,
_ = tx_event.closed() => Err(ApiError::Stream(
"response event consumer dropped".to_string(),
)),
}
};
if let Err(err) = result {