From e5589843a3c7a8360f53bb07fcecae6dab92403b Mon Sep 17 00:00:00 2001 From: Douglas Yang Date: Mon, 18 May 2026 01:20:09 -0700 Subject: [PATCH] feature: upstream cancel (#19524) Co-authored-by: Kangyan Zhou Co-authored-by: Claude Opus 4.7 (1M context) --- .pre-commit-config.yaml | 6 + sgl-model-gateway/Cargo.toml | 5 + .../benches/streaming_utils_bench.rs | 125 + sgl-model-gateway/src/core/mod.rs | 2 +- .../src/routers/http/pd_router.rs | 210 +- sgl-model-gateway/src/routers/http/router.rs | 87 +- sgl-model-gateway/src/routers/mod.rs | 1 + .../src/routers/openai/responses/streaming.rs | 383 ++- .../src/routers/openai/router.rs | 59 +- .../src/routers/streaming_utils.rs | 271 ++ sgl-model-gateway/tests/common/mock_worker.rs | 551 +++- sgl-model-gateway/tests/reliability/mod.rs | 1 + .../tests/reliability/upstream_cancel_test.rs | 2714 +++++++++++++++++ 13 files changed, 4206 insertions(+), 209 deletions(-) create mode 100644 sgl-model-gateway/benches/streaming_utils_bench.rs create mode 100644 sgl-model-gateway/src/routers/streaming_utils.rs create mode 100644 sgl-model-gateway/tests/reliability/upstream_cancel_test.rs diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7fea2c91d..988d30a06 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -101,6 +101,12 @@ repos: pass_filenames: false always_run: true stages: [pre-commit] + - id: rustfmt-sgl-model-gateway + name: rustfmt sgl-model-gateway (nightly) + entry: bash -c 'rustup component add --toolchain nightly rustfmt >/dev/null && cd sgl-model-gateway && cargo +nightly fmt -- --check' + language: system + files: ^sgl-model-gateway/.*\.rs$ + pass_filenames: false - repo: https://github.com/lycheeverse/lychee.git rev: lychee-v0.22.0 hooks: diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 23b4c7cae..648d647c1 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -164,6 +164,11 @@ name = "manual_policy_benchmark" harness = false path = "benches/manual_policy_benchmark.rs" +[[bench]] +name = "streaming_utils_bench" +harness = false +path = "benches/streaming_utils_bench.rs" + [profile.release] opt-level = "z" # Optimize for size lto = "fat" # Full LTO for smaller binaries diff --git a/sgl-model-gateway/benches/streaming_utils_bench.rs b/sgl-model-gateway/benches/streaming_utils_bench.rs new file mode 100644 index 000000000..9e4e12eca --- /dev/null +++ b/sgl-model-gateway/benches/streaming_utils_bench.rs @@ -0,0 +1,125 @@ +//! Microbench for `BreakerTrackedStream` per-poll overhead. +//! +//! The cancel feature wraps every upstream streaming response body in +//! `BreakerTrackedStream`, which sits in the per-chunk hot path. This bench +//! drains a synthetic in-memory chunk stream both bare and wrapped so the +//! delta isolates the wrapper's cost (state machine + Pin dispatch + +//! terminal-state bookkeeping) with no network noise. +//! +//! Run via `cargo bench --bench streaming_utils_bench`. For A/B against +//! main, use Criterion's baseline machinery: +//! `cargo bench --bench streaming_utils_bench -- --save-baseline main` +//! on main, then `--baseline main` on the feature branch. + +use std::{fmt, sync::Arc}; + +use bytes::Bytes; +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use futures_util::StreamExt; +use smg::{ + core::{BasicWorkerBuilder, Worker}, + routers::streaming_utils::BreakerTrackedStream, +}; +use tokio::runtime::Runtime; + +/// Minimal `Display`-able error type — keeps the wrapper generic so we +/// don't have to fabricate `reqwest::Error` instances. +#[derive(Debug)] +struct BenchErr; + +impl fmt::Display for BenchErr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("bench") + } +} + +/// Build `n` `Ok` chunks of `chunk_size` zero bytes. `Bytes::clone` is +/// cheap (refcount bump), so each call reuses the same allocation. +fn make_chunks(n: usize, chunk_size: usize) -> Vec> { + let chunk = Bytes::from(vec![0u8; chunk_size]); + (0..n).map(|_| Ok::<_, BenchErr>(chunk.clone())).collect() +} + +fn make_worker() -> Arc { + Arc::new(BasicWorkerBuilder::new("http://bench-worker").build()) +} + +/// Drain a bare `stream::iter` — the reference cost without the +/// `BreakerTrackedStream` wrapper. Whatever delta the `wrapped` bench +/// shows on top of this is the per-chunk overhead of the wrapper. +fn bench_iter_baseline(c: &mut Criterion) { + let rt = Runtime::new().expect("runtime"); + let mut group = c.benchmark_group("streaming_utils_baseline_iter"); + const CHUNK_SIZE: usize = 1024; + + for &n in &[16usize, 64, 256] { + group.throughput(Throughput::Elements(n as u64)); + group.bench_function(BenchmarkId::from_parameter(n), |b| { + b.iter(|| { + rt.block_on(async { + let chunks = make_chunks(n, CHUNK_SIZE); + let mut stream = futures_util::stream::iter(chunks); + while let Some(item) = stream.next().await { + black_box(item.expect("bench chunk")); + } + }); + }); + }); + } + group.finish(); +} + +/// Drain a `BreakerTrackedStream` over the same synthetic chunks. This is +/// the cancel-feature hot path: every upstream byte chunk on a streaming +/// request is polled through this wrapper. +fn bench_tracked_clean(c: &mut Criterion) { + let rt = Runtime::new().expect("runtime"); + let mut group = c.benchmark_group("streaming_utils_tracked_clean"); + const CHUNK_SIZE: usize = 1024; + + for &n in &[16usize, 64, 256] { + group.throughput(Throughput::Elements(n as u64)); + group.bench_function(BenchmarkId::from_parameter(n), |b| { + b.iter(|| { + rt.block_on(async { + let chunks = make_chunks(n, CHUNK_SIZE); + let worker = make_worker(); + let mut stream = BreakerTrackedStream::new( + futures_util::stream::iter(chunks), + worker, + "http://bench".to_string(), + ); + while let Some(item) = stream.next().await { + black_box(item.expect("bench chunk")); + } + }); + }); + }); + } + group.finish(); +} + +/// `mark_completed` is the PD-streaming `[DONE]`-sentinel fast-path: the +/// caller pre-marks the wrapper completed and stops polling. Times the +/// allocation + mark + drop sequence so changes to the `Terminal` +/// state-machine or `Drop` impl surface here too. +fn bench_tracked_mark_completed_drop(c: &mut Criterion) { + c.bench_function("streaming_utils_tracked_mark_completed_drop", |b| { + b.iter(|| { + let worker = make_worker(); + let inner = futures_util::stream::pending::>(); + let mut tracked = BreakerTrackedStream::new(inner, worker, "http://bench".to_string()); + tracked.mark_completed(); + black_box(&tracked); + drop(tracked); + }); + }); +} + +criterion_group!( + benches, + bench_iter_baseline, + bench_tracked_clean, + bench_tracked_mark_completed_drop, +); +criterion_main!(benches); diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs index 055812c80..405f53466 100644 --- a/sgl-model-gateway/src/core/mod.rs +++ b/sgl-model-gateway/src/core/mod.rs @@ -28,7 +28,7 @@ pub mod worker_registry; pub mod worker_service; // Re-export commonly used types for convenience -pub use circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; +pub use circuit_breaker::{CircuitBreaker, CircuitBreakerConfig, CircuitState}; pub use error::{WorkerError, WorkerResult}; pub use job_queue::{Job, JobQueue, JobQueueConfig}; pub use model_card::{ModelCard, ProviderType}; diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 45e801f36..4796c47ab 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -40,7 +40,9 @@ use crate::{ routers::{ error, grpc::utils::{error_type_from_status, route_to_endpoint}, - header_utils, RouterTrait, + header_utils, + streaming_utils::BreakerTrackedStream, + RouterTrait, }, }; @@ -65,6 +67,15 @@ struct PDRequestContext<'a> { headers: Option, } +/// Marker placed on a `Response` by paths inside +/// `execute_dual_dispatch_internal` that have already recorded prefill and +/// decode breaker outcomes against the workers' actual per-side results +/// (rather than the final response status). The outer dispatcher reads this +/// and skips its own status-based `record_outcome` calls so a decode-only +/// transport failure can't be misattributed to a healthy prefill. +#[derive(Clone, Copy)] +struct BreakerOutcomesRecorded; + impl PDRouter { async fn proxy_to_first_prefill_worker( &self, @@ -341,6 +352,7 @@ impl PDRouter { Err(e) => return Self::handle_serialization_error(e), }; + let ctx_is_stream = context.is_stream; let response = self .execute_dual_dispatch_internal( headers, @@ -353,9 +365,24 @@ impl PDRouter { .await; let status = response.status(); - let not_error = status.is_success() || status.is_client_error(); - prefill.record_outcome(not_error); - decode.record_outcome(not_error); + let outcomes_already_recorded = response + .extensions() + .get::() + .is_some(); + if !outcomes_already_recorded { + let not_error = status.is_success() || status.is_client_error(); + // Prefill is always non-streaming and fully read before + // we get here, so its outcome is final. + prefill.record_outcome(not_error); + // Decode for a streaming request is still mid-flight at + // this point; the `BreakerTrackedStream` wrapped around + // its byte stream records the outcome on drop. Skip the + // eager success record to avoid masking "200-then-broken" + // decode workers. + if !ctx_is_stream { + decode.record_outcome(not_error); + } + } // Record worker errors for server errors (5xx) if status.is_server_error() { @@ -428,13 +455,24 @@ impl PDRouter { // Handle streaming error response let response_headers = header_utils::preserve_response_headers(res.headers()); let error_payload = match res.bytes().await { - Ok(error_body) => { - if let Ok(error_json) = serde_json::from_slice::(&error_body) { + Ok(error_body) => match serde_json::from_slice::(&error_body) { + Ok(error_json) => { json!({ "message": error_json, "status": status.as_u16() }) - } else { - json!({ "message": String::from_utf8_lossy(&error_body).to_string(), "status": status.as_u16() }) } - } + Err(parse_err) => { + let body_text = String::from_utf8_lossy(&error_body).to_string(); + let preview: String = body_text.chars().take(256).collect(); + tracing::warn!( + "Failed to parse decode error body as JSON from {}: {} \ + (status={}, body preview: {:?})", + decode.url(), + parse_err, + status.as_u16(), + preview + ); + json!({ "message": body_text, "status": status.as_u16() }) + } + }, Err(e) => { json!({ "message": format!("Decode server error: {}", e), "status": status.as_u16() }) } @@ -446,13 +484,11 @@ impl PDRouter { ); let error_stream = tokio_stream::once(Ok(axum::body::Bytes::from(sse_data))); - let decode_url = decode.url().to_string(); self.create_streaming_response( error_stream, status, None, context.return_logprob, - Some(decode_url), Some(response_headers), prefill, decode, @@ -595,9 +631,39 @@ impl PDRouter { status ); - return self + // Per-worker breaker attribution before the synthetic 5xx + // response takes over. Prefill ran concurrently in the + // `tokio::join!`: tick it based on its actual response + // status, not on the decode-driven failure. For + // non-streaming the response carries no tracked stream + // so record decode's outcome here too — but treat 4xx + // as a client fault rather than a worker fault, matching + // the legacy outer-dispatcher rule and the streaming + // `BreakerTrackedStream` pre-mark in + // `create_streaming_response`. For streaming + // `handle_decode_error_response` wraps the synthetic + // error SSE in a `BreakerTrackedStream` that ticks + // decode on drop, so skip to avoid double-counting. + // Mark the response so the outer dispatcher skips its + // status-derived `record_outcome`. + let prefill_ok = match &prefill_result { + Ok(r) => { + let s = r.status(); + s.is_success() || s.is_client_error() + } + Err(_) => false, + }; + prefill.record_outcome(prefill_ok); + if !context.is_stream { + let decode_ok = status.is_success() || status.is_client_error(); + decode.record_outcome(decode_ok); + } + + let mut response = self .handle_decode_error_response(res, &context, prefill, decode) .await; + response.extensions_mut().insert(BreakerOutcomesRecorded); + return response; } // Process prefill response @@ -644,7 +710,6 @@ impl PDRouter { status, prefill_logprobs, context.return_logprob, - None, Some(response_headers), prefill, decode, @@ -688,7 +753,33 @@ impl PDRouter { error = %e, "Decode request failed" ); - error::bad_gateway("decode_server_error", format!("Decode server error: {}", e)) + // Decode failed at TCP/transport level. No tracked + // stream will ever wrap a response (streaming path) and + // we shortcut past the outer non-streaming + // `record_outcome` too — so record decode failure + // directly. Prefill ran concurrently in the + // `tokio::join!`: record its real per-worker outcome + // (success on a 2xx/4xx send, failure on transport + // error) so the decode-driven 502 doesn't penalise a + // healthy prefill. Mark the response so the outer + // dispatcher skips its status-derived `record_outcome` + // and we don't double-count. + decode.record_outcome(false); + let prefill_ok = match &prefill_result { + Ok(res) => { + let s = res.status(); + s.is_success() || s.is_client_error() + } + Err(_) => false, + }; + prefill.record_outcome(prefill_ok); + + let mut response = error::bad_gateway( + "decode_server_error", + format!("Decode server error: {}", e), + ); + response.extensions_mut().insert(BreakerOutcomesRecorded); + response } } } @@ -837,7 +928,6 @@ impl PDRouter { status: StatusCode, prefill_logprobs: Option, return_logprob: bool, - decode_url: Option, headers: Option, prefill: Arc, decode: Arc, @@ -846,33 +936,80 @@ impl PDRouter { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + // Uses select! to race stream.next() against tx.closed() so that + // when the client disconnects the upstream HTTP connection is dropped + // promptly, allowing the engine to abort the request. + // `biased;` drains a ready upstream chunk before observing client + // disconnect, so a chunk already produced by reqwest reaches the + // client (and the logprob merger) before we tear the loop down. + // + // The upstream stream is wrapped in `BreakerTrackedStream` so the + // decode worker's circuit breaker is updated once on drop: success + // on clean completion (`[DONE]` sentinel or `None`), failure on + // stream error, neither on client disconnect. PD's pre-PR semantics + // treated 4xx (client error) as not-a-worker-fault, so we only + // pre-mark the wrapper as Errored on 5xx — `handle_decode_error_response` + // synthesizes a single-chunk SSE error envelope that would otherwise + // stream cleanly to None and record a spurious success. + let mut tracked = + BreakerTrackedStream::new(stream, Arc::clone(&decode), decode.url().to_string()); + if !(status.is_success() || status.is_client_error()) { + tracked.mark_errored(); + } + let decode_for_log = decode.clone(); tokio::spawn(async move { - futures_util::pin_mut!(stream); - while let Some(chunk_result) = stream.next().await { - match chunk_result { - Ok(chunk) => { - let is_done = memmem::find(&chunk, b"data: [DONE]").is_some(); + loop { + tokio::select! { + biased; + chunk_result = tracked.next() => { + match chunk_result { + Some(Ok(chunk)) => { + let is_done = memmem::find(&chunk, b"data: [DONE]").is_some(); - let result = if return_logprob && prefill_logprobs.is_some() { - Self::merge_streaming_logprobs(prefill_logprobs.clone(), &chunk) - .unwrap_or(chunk) - } else { - chunk - }; + let result = if return_logprob && prefill_logprobs.is_some() { + Self::merge_streaming_logprobs(prefill_logprobs.clone(), &chunk) + .unwrap_or(chunk) + } else { + chunk + }; - if tx.send(Ok(result)).is_err() { - break; - } + // Mark the wrapper completed before the client + // send: upstream finished cleanly regardless of + // whether the client is still listening, and + // the worker deserves the success tick either + // way. `mark_completed` is a no-op once Errored + // is set, so the synthetic-error path is unaffected. + if is_done { + tracked.mark_completed(); + } - if is_done { - break; + if tx.send(Ok(result)).is_err() { + tracing::debug!( + "Receiver dropped (likely client disconnect), \ + cancelling upstream PD stream" + ); + break; + } + + if is_done { + break; + } + } + Some(Err(e)) => { + // BreakerTrackedStream already logged the error + // and marked the terminal state as Errored so + // the worker's circuit breaker will tick on drop. + let _ = tx.send(Err(format!("Stream error: {}", e))); + break; + } + None => break, } } - Err(e) => { - if let Some(ref url) = decode_url { - error!("Stream error from decode server {}: {}", url, e); - } - let _ = tx.send(Err(format!("Stream error: {}", e))); + _ = tx.closed() => { + tracing::info!( + "Client disconnected, cancelling upstream PD stream from {}", + decode_for_log.url() + ); break; } } @@ -1549,7 +1686,6 @@ mod tests { None, false, None, - None, prefill_ref.clone(), decode_ref.clone(), ); diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index ffb036d91..029c9d1a0 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -9,7 +9,6 @@ use axum::{ }; use futures_util::{stream, StreamExt}; use reqwest::Client; -use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::{debug, error}; use crate::{ @@ -38,7 +37,9 @@ use crate::{ routers::{ error::{self, extract_error_code_from_response}, grpc::utils::{error_type_from_status, route_to_endpoint}, - header_utils, RouterTrait, + header_utils, + streaming_utils::BreakerTrackedStream, + RouterTrait, }, }; @@ -307,20 +308,21 @@ impl Router { let headers = Some(&headers_with_trace); let response = self - .send_typed_request( - headers, - typed_req, - route, - worker.url(), - is_stream, - load_guard, - ) + .send_typed_request(headers, typed_req, route, &worker, is_stream, load_guard) .await; events::RequestReceivedEvent {}.emit(); let status = response.status(); - worker.record_outcome(status.is_success()); + // For streaming responses, the wrapped body (`BreakerTrackedStream`) + // records the circuit-breaker outcome once the stream actually + // terminates (success on clean end, failure on mid-stream error). + // Recording it eagerly here based on the initial status code would + // mask "200-then-broken" workers — every request would tick a + // success before the stream had a chance to error out. + if !is_stream { + worker.record_outcome(status.is_success()); + } // Record worker errors for server errors (5xx) if status.is_server_error() { @@ -487,13 +489,12 @@ impl Router { headers: Option<&HeaderMap>, typed_req: &T, route: &'static str, - worker_url: &str, + worker: &Arc, is_stream: bool, load_guard: Option, ) -> Response { - // Get the worker once and reuse for API key and load tracking - let worker = self.worker_registry.get_by_url(worker_url); - let api_key = worker.as_ref().and_then(|w| w.api_key().clone()); + let worker_url = worker.url(); + let api_key = worker.api_key().clone(); // Static key string to avoid per-request allocations const DP_RANK_KEY: &str = "data_parallel_rank"; @@ -570,6 +571,19 @@ impl Router { worker_url, route, e ); + // For streaming requests the caller skips the eager + // `record_outcome` on the assumption that a + // `BreakerTrackedStream` will tick the breaker on drop — + // but no tracked stream is installed when send() fails + // before any response stream exists. Record the failure + // here so a worker flapping at the TCP layer doesn't + // stay permanently selectable. Non-streaming requests + // are already covered by the caller's + // `worker.record_outcome(status.is_success())`, so + // gating on `is_stream` avoids double-counting. + if is_stream { + worker.record_outcome(false); + } return convert_reqwest_error(e); } }; @@ -602,29 +616,26 @@ impl Router { // Ensure we set the correct content-type for SSE response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); - let stream = res.bytes_stream(); - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - - // Spawn task to forward stream - tokio::spawn(async move { - let mut stream = stream; - while let Some(chunk) = stream.next().await { - match chunk { - Ok(bytes) => { - if tx.send(Ok(bytes)).is_err() { - break; - } - } - Err(e) => { - let _ = tx.send(Err(format!("Stream error: {}", e))); - break; - } - } - } - }); - - let stream = UnboundedReceiverStream::new(rx); - let body = Body::from_stream(stream); + // Pass the reqwest byte stream straight through as the response body. + // Dropping the response body drops this stream, which closes the + // upstream HTTP connection and lets the engine abort generation — + // no spawned task or channel needed. `BreakerTrackedStream` + // updates the worker's circuit breaker exactly once on drop: + // success on clean end, failure on stream error, neither on + // client disconnect. For non-2xx responses we pre-mark the + // wrapper as Errored — otherwise the small error body would + // stream cleanly to `None` and Drop would record a spurious + // success (and the streaming branch also skips the eager + // `record_outcome` above). + let mut tracked = BreakerTrackedStream::new( + res.bytes_stream(), + worker.clone(), + worker_url.to_string(), + ); + if !status.is_success() { + tracked.mark_errored(); + } + let body = Body::from_stream(tracked); let mut response = Response::new(body); *response.status_mut() = status; diff --git a/sgl-model-gateway/src/routers/mod.rs b/sgl-model-gateway/src/routers/mod.rs index 11d34a68b..11cb05bbb 100644 --- a/sgl-model-gateway/src/routers/mod.rs +++ b/sgl-model-gateway/src/routers/mod.rs @@ -32,6 +32,7 @@ pub mod openai; pub mod parse; pub mod persistence_utils; pub mod router_manager; +pub mod streaming_utils; pub mod tokenize; pub use factory::RouterFactory; diff --git a/sgl-model-gateway/src/routers/openai/responses/streaming.rs b/sgl-model-gateway/src/routers/openai/responses/streaming.rs index 2c7da80af..c84a7a608 100644 --- a/sgl-model-gateway/src/routers/openai/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/streaming.rs @@ -14,7 +14,7 @@ use axum::{ http::{header::CONTENT_TYPE, HeaderMap, HeaderValue, StatusCode}, response::{IntoResponse, Response}, }; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; use futures_util::StreamExt; use serde_json::{json, Value}; use tokio::sync::mpsc; @@ -44,6 +44,7 @@ use crate::{ mcp_utils::{ensure_request_mcp_client, McpLoopConfig}, openai::context::{RequestContext, StreamingEventContext, StreamingRequest}, persistence_utils::persist_conversation_items, + streaming_utils::BreakerTrackedStream, }, }; @@ -191,6 +192,18 @@ fn send_sse_event( tx.send(Ok(Bytes::from(block))).is_ok() } +/// Append a fully-formed SSE block to a `\n\n` terminator without going +/// through `format!()` — the `/responses` streaming path forwards a chunk +/// for every model token, so avoiding the intermediate `String` allocation +/// matters under load. +#[inline] +fn sse_block_to_bytes(block: &str) -> Bytes { + let mut buf = BytesMut::with_capacity(block.len() + 2); + buf.extend_from_slice(block.as_bytes()); + buf.extend_from_slice(b"\n\n"); + buf.freeze() +} + /// Transform fc_* item IDs to mcp_* format #[inline] fn transform_fc_to_mcp_id(item_id: &str) -> String { @@ -484,7 +497,7 @@ pub(super) fn send_final_response_event( /// Simple pass-through streaming without MCP interception pub(super) async fn handle_simple_streaming_passthrough( client: &reqwest::Client, - circuit_breaker: &crate::core::CircuitBreaker, + worker: Arc, headers: Option<&HeaderMap>, req: StreamingRequest, ) -> Response { @@ -499,7 +512,7 @@ pub(super) async fn handle_simple_streaming_passthrough( let response = match request_builder.send().await { Ok(resp) => resp, Err(err) => { - circuit_breaker.record_failure(); + worker.circuit_breaker().record_failure(); return ( StatusCode::BAD_GATEWAY, format!("Failed to forward request to OpenAI: {}", err), @@ -513,18 +526,28 @@ pub(super) async fn handle_simple_streaming_passthrough( StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); if !status.is_success() { - circuit_breaker.record_failure(); + worker.circuit_breaker().record_failure(); let error_body = response .text() .await .unwrap_or_else(|err| format!("Failed to read upstream error body: {}", err)); return (status_code, error_body).into_response(); } - - circuit_breaker.record_success(); + // Do NOT record_success here at status-OK time — the spawned forwarder + // below records the actual stream outcome on termination (success on + // clean end / `[DONE]`, failure on mid-stream error). Recording success + // eagerly here would mask 200-then-broken workers and double-count when + // the stream completes normally. let preserved_headers = preserve_response_headers(response.headers()); - let mut upstream_stream = response.bytes_stream(); + // Wrap upstream in `BreakerTrackedStream` so the breaker tick is decided + // once on drop: success on clean `None`, failure on `Some(Err)`, neither + // on early drop (client disconnected before upstream terminated). + let upstream_stream = BreakerTrackedStream::new( + response.bytes_stream(), + Arc::clone(&worker), + req.url.clone(), + ); let (tx, rx) = mpsc::unbounded_channel::>(); @@ -533,86 +556,168 @@ pub(super) async fn handle_simple_streaming_passthrough( let persist_needed = original_request.conversation.is_some(); let previous_response_id = req.previous_response_id; let storage = req.storage; + let upstream_url = req.url; + + // When persistence is needed (should_store || persist_needed), we must + // continue consuming upstream even after client disconnect to accumulate + // the full response — the upstream HTTP request is intentionally NOT + // cancelled in that branch. When neither is needed, we use select! to + // race stream.next() against tx.closed() so the upstream HTTP connection + // is dropped promptly when the client disconnects. + let need_persistence = should_store || persist_needed; tokio::spawn(async move { - let mut accumulator = StreamingResponseAccumulator::new(); - let mut upstream_failed = false; - let mut receiver_connected = true; let mut chunk_processor = ChunkProcessor::new(); - while let Some(chunk_result) = upstream_stream.next().await { - match chunk_result { - Ok(chunk) => { - chunk_processor.push_chunk(&chunk); + if need_persistence { + // Persistence path: keep consuming upstream even after client + // disconnect. The wrapper's terminal state at drop encodes the + // breaker outcome; `upstream_failed` is just a local flag used + // to skip persistence on stream error. + let mut upstream_stream = upstream_stream; + let mut accumulator = StreamingResponseAccumulator::new(); + let mut upstream_failed = false; + let mut receiver_connected = true; + while let Some(chunk_result) = upstream_stream.next().await { + match chunk_result { + Ok(chunk) => { + chunk_processor.push_chunk(&chunk); - while let Some(raw_block) = chunk_processor.next_block() { - let block_cow = match rewrite_streaming_block( - &raw_block, - &original_request, - previous_response_id.as_deref(), - ) { - Some(modified) => Cow::Owned(modified), - None => Cow::Borrowed(raw_block.as_str()), - }; + while let Some(raw_block) = chunk_processor.next_block() { + let block_cow = match rewrite_streaming_block( + &raw_block, + &original_request, + previous_response_id.as_deref(), + ) { + Some(modified) => Cow::Owned(modified), + None => Cow::Borrowed(raw_block.as_str()), + }; - if should_store || persist_needed { accumulator.ingest_block(&block_cow); - } - if receiver_connected { - let chunk_to_send = format!("{}\n\n", block_cow); - if tx.send(Ok(Bytes::from(chunk_to_send))).is_err() { + if receiver_connected + && tx.send(Ok(sse_block_to_bytes(&block_cow))).is_err() + { receiver_connected = false; + tracing::debug!( + "Client disconnected during /responses persistence \ + stream from {}; continuing to drain upstream for \ + storage", + upstream_url + ); } } - - if !receiver_connected && !should_store { - break; - } } - - if !receiver_connected && !should_store { + Err(err) => { + upstream_failed = true; + // BreakerTrackedStream already marked terminal=Errored + // and logged; just forward the error to the client. + let io_err = io::Error::other(err); + let _ = tx.send(Err(io_err)); break; } } - Err(err) => { - upstream_failed = true; - let io_err = io::Error::other(err); - let _ = tx.send(Err(io_err)); - break; - } } - } - if (should_store || persist_needed) && !upstream_failed { - if chunk_processor.has_remaining() { - accumulator.ingest_block(&chunk_processor.take_remaining()); - } - let encountered_error = accumulator.encountered_error().cloned(); - if let Some(mut response_json) = accumulator.into_final_response() { - patch_response_with_request_metadata( - &mut response_json, - &original_request, - previous_response_id.as_deref(), + if upstream_failed { + warn!( + "Skipping /responses persistence due to upstream stream error from {}: \ + store={} conversation={:?}", + upstream_url, should_store, original_request.conversation ); - - // Always persist conversation items and response (even without conversation) - if let Err(err) = persist_conversation_items( - storage.conversation.clone(), - storage.conversation_item.clone(), - storage.response.clone(), - &response_json, - &original_request, - ) - .await - { - warn!("Failed to persist conversation items (stream): {}", err); - } - } else if let Some(error_payload) = encountered_error { - warn!("Upstream streaming error payload: {}", error_payload); } else { - warn!("Streaming completed without a final response payload"); + if chunk_processor.has_remaining() { + accumulator.ingest_block(&chunk_processor.take_remaining()); + } + let encountered_error = accumulator.encountered_error().cloned(); + if let Some(mut response_json) = accumulator.into_final_response() { + patch_response_with_request_metadata( + &mut response_json, + &original_request, + previous_response_id.as_deref(), + ); + + // Always persist conversation items and response (even without conversation) + if let Err(err) = persist_conversation_items( + storage.conversation.clone(), + storage.conversation_item.clone(), + storage.response.clone(), + &response_json, + &original_request, + ) + .await + { + warn!("Failed to persist conversation items (stream): {}", err); + } + } else if let Some(error_payload) = encountered_error { + warn!("Upstream streaming error payload: {}", error_payload); + } else { + warn!("Streaming completed without a final response payload"); + } } + // upstream_stream dropped here → breaker tick fires (success or + // failure depending on terminal state set during the loop). + } else { + // No persistence: use select! to cancel upstream on client disconnect. + // `biased;` drains a ready upstream chunk before observing client + // disconnect, so a chunk already produced by reqwest reaches the + // client before we tear the loop down. The `BreakerTrackedStream` + // wrapper records the breaker outcome on drop based on the + // terminal it observed: clean `None` → success, `Some(Err)` → + // failure, dropped while still active (client disconnect or + // `tx.send` failure) → neither. + futures_util::pin_mut!(upstream_stream); + 'outer: loop { + tokio::select! { + biased; + chunk_result = upstream_stream.next() => { + match chunk_result { + Some(Ok(chunk)) => { + chunk_processor.push_chunk(&chunk); + + while let Some(raw_block) = chunk_processor.next_block() { + let block_cow = match rewrite_streaming_block( + &raw_block, + &original_request, + previous_response_id.as_deref(), + ) { + Some(modified) => Cow::Owned(modified), + None => Cow::Borrowed(raw_block.as_str()), + }; + + if tx.send(Ok(sse_block_to_bytes(&block_cow))).is_err() { + tracing::debug!( + "Receiver dropped (likely client disconnect), \ + cancelling upstream /responses stream from {}", + upstream_url + ); + break 'outer; + } + } + } + Some(Err(err)) => { + // BreakerTrackedStream already marked terminal=Errored. + let io_err = io::Error::other(err); + let _ = tx.send(Err(io_err)); + break 'outer; + } + None => { + // BreakerTrackedStream already marked terminal=Completed. + break 'outer; + } + } + } + _ = tx.closed() => { + tracing::info!( + "Client disconnected, cancelling upstream /responses stream from {}", + upstream_url + ); + break 'outer; + } + } + } + // upstream_stream dropped here → breaker tick fires based on + // terminal state. } }); @@ -632,9 +737,20 @@ pub(super) async fn handle_simple_streaming_passthrough( response } -/// Handle streaming WITH MCP tool call interception and execution +/// Handle streaming WITH MCP tool call interception and execution. +/// +/// Note on cancellation vs persistence: +/// Unlike `handle_simple_streaming_passthrough`, this path does *not* keep +/// consuming the upstream after a client disconnect even when `should_store` +/// or `persist_needed` is set. Tool interception is multi-iteration and +/// drives MCP tool execution between iterations, which is unbounded work. +/// We deliberately give up on persistence when the client is gone so we +/// don't keep workers and external MCP services busy on results no one +/// will read. Callers that need guaranteed persistence on disconnect +/// should not enable MCP tools for that request. pub(super) async fn handle_streaming_with_tool_interception( client: &reqwest::Client, + worker: Arc, headers: Option<&HeaderMap>, req: StreamingRequest, active_mcp: &Arc, @@ -658,6 +774,7 @@ pub(super) async fn handle_streaming_with_tool_interception( let payload_clone = payload.clone(); let active_mcp_clone = Arc::clone(active_mcp); let server_keys_clone = server_keys.clone(); + let worker_for_breaker = worker; // Spawn the streaming loop task tokio::spawn(async move { @@ -695,6 +812,18 @@ pub(super) async fn handle_streaming_with_tool_interception( }; loop { + // Check if the client has already disconnected before making a new + // upstream request (e.g. between tool-call iterations). This avoids + // issuing a fresh upstream HTTP request whose response we'd just + // discard. Mid-stream disconnects are caught by the inner select!. + if tx.is_closed() { + tracing::info!( + "Client disconnected before next tool-call iteration to {}, cancelling", + url_clone + ); + return; + } + // Make streaming request let mut request_builder = client_clone.post(&url_clone).json(¤t_payload); if let Some(ref h) = headers_opt { @@ -702,9 +831,30 @@ pub(super) async fn handle_streaming_with_tool_interception( } request_builder = request_builder.header("Accept", "text/event-stream"); - let response = match request_builder.send().await { + // Race send() against tx.closed() so that a client disconnect + // while we're still waiting for upstream response headers also + // cancels this in-flight HTTP request (not just mid-stream). + let response = tokio::select! { + biased; + res = request_builder.send() => res, + _ = tx.closed() => { + tracing::info!( + "Client disconnected before /responses tool-interception \ + upstream response from {}, cancelling", + url_clone + ); + return; + } + }; + let response = match response { Ok(r) => r, Err(e) => { + worker_for_breaker.circuit_breaker().record_failure(); + tracing::error!( + "Failed to send /responses tool-interception request to {}: {}", + url_clone, + e + ); let error_event = format!( "event: error\ndata: {{\"error\": {{\"message\": \"{}\"}}}}\n\n", e @@ -715,15 +865,33 @@ pub(super) async fn handle_streaming_with_tool_interception( }; if !response.status().is_success() { + worker_for_breaker.circuit_breaker().record_failure(); let status = response.status(); - let body = response.text().await.unwrap_or_default(); + let body = response + .text() + .await + .unwrap_or_else(|e| format!("")); + tracing::error!( + "Upstream /responses tool-interception non-success from {}: status={} \ + body={}", + url_clone, + status, + body + ); let error_event = format!("event: error\ndata: {{\"error\": {{\"message\": \"Upstream error {}: {}\"}}}}\n\n", status, body); let _ = tx.send(Ok(Bytes::from(error_event))); return; } - // Stream events and check for tool calls - let mut upstream_stream = response.bytes_stream(); + // Stream events and check for tool calls. + // Uses select! to race stream.next() against tx.closed() so that + // when the client disconnects the upstream HTTP connection is dropped + // promptly, allowing the engine to abort the request. + // `biased;` drains a ready upstream chunk before observing client + // disconnect, so a chunk already produced by reqwest reaches both + // the client and the chunk_processor before we tear the loop down. + let upstream_stream = response.bytes_stream(); + futures_util::pin_mut!(upstream_stream); let mut handler = StreamingToolHandler::with_starting_index(next_output_index); if let Some(ref id) = preserved_response_id { handler.original_response_id = Some(id.clone()); @@ -732,9 +900,23 @@ pub(super) async fn handle_streaming_with_tool_interception( let mut tool_calls_detected = false; let mut seen_in_progress = false; - while let Some(chunk_result) = upstream_stream.next().await { + loop { + let chunk_result = tokio::select! { + biased; + chunk = upstream_stream.next() => chunk, + _ = tx.closed() => { + // Client disconnected — drop the stream and exit + tracing::info!( + "Client disconnected, cancelling upstream /responses \ + tool-interception stream from {}", + url_clone + ); + return; + } + }; + match chunk_result { - Ok(chunk) => { + Some(Ok(chunk)) => { chunk_processor.push_chunk(&chunk); while let Some(raw_block) = chunk_processor.next_block() { @@ -838,11 +1020,25 @@ pub(super) async fn handle_streaming_with_tool_interception( break; } } - Err(e) => { + Some(Err(e)) => { + worker_for_breaker.circuit_breaker().record_failure(); + tracing::error!( + "Upstream /responses tool-interception stream error from {}: {}", + url_clone, + e + ); + if should_store || persist_needed { + warn!( + "Skipping /responses persistence (tool-interception) due to \ + upstream stream error from {}: store={} conversation={:?}", + url_clone, should_store, original_request.conversation + ); + } let error_event = format!("event: error\ndata: {{\"error\": {{\"message\": \"Stream error: {}\"}}}}\n\n", e); let _ = tx.send(Ok(Bytes::from(error_event))); return; } + None => break, } } @@ -906,8 +1102,15 @@ pub(super) async fn handle_streaming_with_tool_interception( err ); } + } else if should_store || persist_needed { + warn!( + "/responses tool-interception stream ended with no final \ + response payload; persistence skipped (store={} conversation={:?})", + should_store, original_request.conversation + ); } + worker_for_breaker.circuit_breaker().record_success(); let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); return; } @@ -925,10 +1128,24 @@ pub(super) async fn handle_streaming_with_tool_interception( }; if state.total_calls > effective_limit { + // Reaching the tool-call iteration limit is a request-level + // shape problem (the user asked for unbounded work or the + // model is in a tight tool-call loop); the upstream worker + // is not unhealthy. Record a success to reflect that the + // last upstream call returned cleanly. + worker_for_breaker.circuit_breaker().record_success(); warn!( "Reached tool call limit during streaming: {}", effective_limit ); + if should_store || persist_needed { + warn!( + "Skipping /responses persistence (tool-interception): \ + max_tool_calls limit reached, no final response built. \ + store={} conversation={:?}", + should_store, original_request.conversation + ); + } let error_event = "event: error\ndata: {\"error\": {\"message\": \"Exceeded max_tool_calls limit\"}}\n\n".to_string(); let _ = tx.send(Ok(Bytes::from(error_event))); let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); @@ -965,6 +1182,18 @@ pub(super) async fn handle_streaming_with_tool_interception( // Continue loop to make next streaming request } Err(e) => { + // Upstream stream finished cleanly (we got pending tool + // calls out of it); the gateway-side payload build failed. + // Credit the worker for the upstream call, mirroring the + // tool-iteration-limit path above. + worker_for_breaker.circuit_breaker().record_success(); + if should_store || persist_needed { + warn!( + "Skipping /responses persistence (tool-interception): \ + build_resume_payload failed ({}); store={} conversation={:?}", + e, should_store, original_request.conversation + ); + } let error_event = format!("event: error\ndata: {{\"error\": {{\"message\": \"Failed to build resume payload: {}\"}}}}\n\n", e); let _ = tx.send(Ok(Bytes::from(error_event))); let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); @@ -986,7 +1215,6 @@ pub(super) async fn handle_streaming_with_tool_interception( /// Main entry point for streaming responses pub async fn handle_streaming_response(ctx: RequestContext) -> Response { let worker = ctx.worker().expect("Worker not selected").clone(); - let circuit_breaker = worker.circuit_breaker(); let headers = ctx.headers().cloned(); let original_body = ctx.responses_request(); let mcp_manager = ctx.components.mcp_manager().expect("MCP manager required"); @@ -1011,7 +1239,7 @@ pub async fn handle_streaming_response(ctx: RequestContext) -> Response { if active_mcp.is_none() { return handle_simple_streaming_passthrough( &client, - circuit_breaker, + Arc::clone(&worker), headers.as_ref(), req, ) @@ -1023,6 +1251,7 @@ pub async fn handle_streaming_response(ctx: RequestContext) -> Response { // MCP is active - transform tools and set up interception handle_streaming_with_tool_interception( &client, + Arc::clone(&worker), headers.as_ref(), req, &active_mcp, diff --git a/sgl-model-gateway/src/routers/openai/router.rs b/sgl-model-gateway/src/routers/openai/router.rs index 6c920bc14..f41d679aa 100644 --- a/sgl-model-gateway/src/routers/openai/router.rs +++ b/sgl-model-gateway/src/routers/openai/router.rs @@ -13,10 +13,8 @@ use axum::{ Json, }; use data_connector::{ConversationId, ListParams, ResponseId, SortOrder}; -use futures_util::{future::join_all, StreamExt}; +use futures_util::future::join_all; use serde_json::{json, to_value, Value}; -use tokio::sync::mpsc; -use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::warn; use super::{ @@ -42,7 +40,10 @@ use crate::{ ResponsesGetParams, ResponsesRequest, }, }, - routers::header_utils::{apply_provider_headers, extract_auth_header}, + routers::{ + header_utils::{apply_provider_headers, extract_auth_header}, + streaming_utils::BreakerTrackedStream, + }, }; pub struct OpenAIRouter { @@ -596,16 +597,15 @@ impl crate::routers::RouterTrait for OpenAIRouter { let status = StatusCode::from_u16(resp.status().as_u16()) .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - // Record circuit breaker failure for error status codes - if !status.is_success() { - worker.circuit_breaker().record_failure(); - } - if !is_streaming { + // Non-streaming: record the breaker outcome inline + // once the body is fully read. + if !status.is_success() { + worker.circuit_breaker().record_failure(); + } let content_type = resp.headers().get(CONTENT_TYPE).cloned(); match resp.bytes().await { Ok(body) => { - // Only record success after body is fully read if status.is_success() { worker.circuit_breaker().record_success(); } @@ -626,30 +626,23 @@ impl crate::routers::RouterTrait for OpenAIRouter { } } } else { - // Streaming response - record success when stream starts - if status.is_success() { - worker.circuit_breaker().record_success(); + // Streaming response: pass the reqwest byte stream + // through `BreakerTrackedStream`, which records the + // circuit-breaker outcome exactly once on drop (success + // on clean end, failure on stream error, neither on + // client disconnect). For non-2xx responses we pre-mark + // the wrapper as Errored — otherwise the small error + // body would stream cleanly to `None` and Drop would + // record a spurious success. + let mut tracked = BreakerTrackedStream::new( + resp.bytes_stream(), + Arc::clone(&worker), + url.clone(), + ); + if !status.is_success() { + tracked.mark_errored(); } - let stream = resp.bytes_stream(); - let (tx, rx) = mpsc::unbounded_channel(); - tokio::spawn(async move { - let mut s = stream; - while let Some(chunk) = s.next().await { - match chunk { - Ok(bytes) => { - if tx.send(Ok(bytes)).is_err() { - break; - } - } - Err(e) => { - let _ = tx.send(Err(format!("Stream error: {}", e))); - break; - } - } - } - }); - let mut response = - Response::new(Body::from_stream(UnboundedReceiverStream::new(rx))); + let mut response = Response::new(Body::from_stream(tracked)); *response.status_mut() = status; response .headers_mut() diff --git a/sgl-model-gateway/src/routers/streaming_utils.rs b/sgl-model-gateway/src/routers/streaming_utils.rs new file mode 100644 index 000000000..94e3fc867 --- /dev/null +++ b/sgl-model-gateway/src/routers/streaming_utils.rs @@ -0,0 +1,271 @@ +//! Shared streaming helpers used by the HTTP / OpenAI / PD routers. +//! +//! The main type here is [`BreakerTrackedStream`], a `Stream` adapter that +//! records circuit-breaker outcomes based on how the upstream stream +//! terminates: +//! +//! - **Clean end** (`Poll::Ready(None)`) → `record_success`. +//! - **Upstream transport error** (`Poll::Ready(Some(Err(_)))`) → `record_failure`. +//! - **Caller drops the stream while still active** (client disconnect) → no +//! breaker call. The HTTP response already shipped a 200 status; whether +//! the worker is healthy is unknown from this signal alone. +//! +//! Callers can override the terminal state in two cases: +//! +//! - [`BreakerTrackedStream::mark_completed`] — for routers that detect +//! end-of-stream via an in-band sentinel (e.g. PD's `data: [DONE]`) +//! before the underlying byte stream returns `None`. Note that +//! `Completed` is *not* absorbing — a subsequent `poll_next` returning +//! `Err` will still escalate the terminal to `Errored`. Callers should +//! stop polling after `mark_completed`. +//! - [`BreakerTrackedStream::mark_errored`] — for routers wrapping content +//! that already represents a worker failure (e.g. a non-2xx response +//! body or a synthetic SSE error envelope built from a 5xx). `Errored` +//! is absorbing: once set, it persists regardless of later polls. + +use std::{ + fmt::Display, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +use bytes::Bytes; +use futures_util::Stream; +use tracing::error; + +use crate::core::Worker; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Terminal { + /// Stream is still in flight, or was dropped before terminating. + Active, + /// Stream returned `None` or the caller marked it complete. + Completed, + /// Stream yielded an `Err` item. + Errored, +} + +/// Wraps a `Stream>` so that the circuit breaker on +/// `worker` is updated exactly once on drop: +/// - completed → `record_success` +/// - errored → `record_failure` +/// - dropped while still active → neither (client disconnected; the worker +/// is innocent from our point of view). +/// +/// `E` defaults to `reqwest::Error` to match the common producer +/// (`response.bytes_stream()`), but can be any `Display + Send + 'static` +/// error type — useful for tests that construct streams with simpler +/// error types. +#[must_use = "BreakerTrackedStream must be polled to completion (or pre-marked) \ + and then dropped for the circuit breaker to record an outcome; \ + discarding it immediately records nothing"] +pub struct BreakerTrackedStream { + inner: Pin> + Send + 'static>>, + worker: Arc, + log_url: String, + terminal: Terminal, +} + +impl BreakerTrackedStream { + pub fn new(inner: S, worker: Arc, log_url: String) -> Self + where + S: Stream> + Send + 'static, + { + Self { + inner: Box::pin(inner), + worker, + log_url, + terminal: Terminal::Active, + } + } + + /// Mark the stream as cleanly completed. Use this from callers that + /// detect end-of-stream via an in-band sentinel (e.g. `data: [DONE]`) + /// before the underlying byte stream returns `None`. + /// + /// Has no effect once the wrapper is in any non-Active state. `Completed` + /// is *not* absorbing — a later `poll_next` returning `Err` will still + /// escalate the terminal to `Errored`. Callers should stop polling + /// after calling this. + pub fn mark_completed(&mut self) { + if self.terminal == Terminal::Active { + self.terminal = Terminal::Completed; + } + } + + /// Pre-tag the stream as terminally errored. Use this from callers + /// constructing a wrapper around content that already represents a + /// failed worker outcome (e.g. a non-2xx response body or a + /// synthetic error envelope) so Drop records `record_failure` even + /// if the underlying stream terminates cleanly. `Errored` is + /// absorbing — once set it stays set regardless of later events. + pub fn mark_errored(&mut self) { + self.terminal = Terminal::Errored; + } +} + +impl Stream for BreakerTrackedStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(b))) => Poll::Ready(Some(Ok(b))), + Poll::Ready(Some(Err(e))) => { + error!("Upstream stream error from worker {}: {}", self.log_url, e); + self.terminal = Terminal::Errored; + Poll::Ready(Some(Err(e))) + } + Poll::Ready(None) => { + if self.terminal == Terminal::Active { + self.terminal = Terminal::Completed; + } + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for BreakerTrackedStream { + fn drop(&mut self) { + match self.terminal { + Terminal::Completed => self.worker.circuit_breaker().record_success(), + Terminal::Errored => self.worker.circuit_breaker().record_failure(), + // Client disconnected before we knew the worker's verdict. + // Leaving the breaker untouched is the correct default — we + // got a 200 header and some bytes; nothing said the worker + // is unhealthy. + Terminal::Active => {} + } + } +} + +#[cfg(test)] +mod tests { + use std::{fmt, sync::Arc}; + + use bytes::Bytes; + use futures_util::StreamExt; + + use super::BreakerTrackedStream; + use crate::core::{BasicWorkerBuilder, Worker}; + + /// Lightweight error type for tests — keeps the wrapper generic so we + /// don't need to fabricate `reqwest::Error` instances. + #[derive(Debug)] + struct TestErr(&'static str); + + impl fmt::Display for TestErr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.0) + } + } + + fn worker() -> Arc { + Arc::new(BasicWorkerBuilder::new("http://test-worker").build()) + } + + fn breaker_counters(w: &Arc) -> (u64, u64) { + let cb = w.circuit_breaker(); + (cb.total_successes(), cb.total_failures()) + } + + #[tokio::test] + async fn drop_while_active_records_nothing() { + let w = worker(); + let inner = futures_util::stream::pending::>(); + let tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + drop(tracked); + assert_eq!(breaker_counters(&w), (0, 0)); + } + + #[tokio::test] + async fn clean_stream_records_one_success() { + let w = worker(); + let inner = futures_util::stream::iter(vec![ + Ok::<_, TestErr>(Bytes::from_static(b"a")), + Ok(Bytes::from_static(b"b")), + ]); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + while tracked.next().await.is_some() {} + drop(tracked); + assert_eq!(breaker_counters(&w), (1, 0)); + } + + #[tokio::test] + async fn stream_error_records_one_failure() { + let w = worker(); + let inner = + futures_util::stream::iter(vec![Ok(Bytes::from_static(b"a")), Err(TestErr("boom"))]); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + while tracked.next().await.is_some() {} + drop(tracked); + assert_eq!(breaker_counters(&w), (0, 1)); + } + + #[tokio::test] + async fn errored_is_absorbing_across_subsequent_polls() { + let w = worker(); + let inner = futures_util::stream::iter(vec![ + Err::(TestErr("boom")), + Ok(Bytes::from_static(b"after")), + ]); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + while tracked.next().await.is_some() {} + drop(tracked); + assert_eq!(breaker_counters(&w), (0, 1)); + } + + #[tokio::test] + async fn mark_completed_then_drop_records_success() { + let w = worker(); + let inner = futures_util::stream::pending::>(); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + tracked.mark_completed(); + drop(tracked); + assert_eq!(breaker_counters(&w), (1, 0)); + } + + #[tokio::test] + async fn mark_errored_then_clean_end_still_records_failure() { + let w = worker(); + let inner = futures_util::stream::iter(vec![Ok::<_, TestErr>(Bytes::from_static(b"a"))]); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + tracked.mark_errored(); + while tracked.next().await.is_some() {} + drop(tracked); + assert_eq!(breaker_counters(&w), (0, 1)); + } + + #[tokio::test] + async fn mark_completed_does_not_overwrite_errored() { + let w = worker(); + let inner = futures_util::stream::iter(vec![Err::(TestErr("boom"))]); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + while tracked.next().await.is_some() {} + tracked.mark_completed(); + drop(tracked); + assert_eq!(breaker_counters(&w), (0, 1)); + } + + // PD's [DONE] handler calls mark_completed before the underlying byte + // stream finishes; a trailing transport error must still flip the + // terminal to Errored so the breaker records failure. + #[tokio::test] + async fn mark_completed_then_later_err_escalates_to_failure() { + let w = worker(); + let inner = futures_util::stream::iter(vec![ + Ok::<_, TestErr>(Bytes::from_static(b"data: [DONE]\n\n")), + Err(TestErr("trailing transport error")), + ]); + let mut tracked = BreakerTrackedStream::new(inner, Arc::clone(&w), "u".into()); + // Caller observes the [DONE] chunk and pre-marks completed... + assert!(tracked.next().await.is_some()); + tracked.mark_completed(); + // ...but a trailing poll surfaces a transport error. + assert!(matches!(tracked.next().await, Some(Err(_)))); + drop(tracked); + assert_eq!(breaker_counters(&w), (0, 1)); + } +} diff --git a/sgl-model-gateway/tests/common/mock_worker.rs b/sgl-model-gateway/tests/common/mock_worker.rs index 19e863ea1..f44809a68 100755 --- a/sgl-model-gateway/tests/common/mock_worker.rs +++ b/sgl-model-gateway/tests/common/mock_worker.rs @@ -20,7 +20,7 @@ use axum::{ }; use futures_util::stream::{self, StreamExt}; use serde_json::json; -use tokio::sync::RwLock; +use tokio::sync::{Notify, RwLock}; use uuid::Uuid; /// Configuration for mock worker behavior @@ -158,6 +158,15 @@ async fn should_fail(config: &MockWorkerConfig) -> bool { rand::random::() < config.fail_rate } +/// Pick the HTTP status used when `should_fail` triggers. Defaults to 500 +/// for backwards compatibility; tests can override via +/// [`set_fail_status_code`] to exercise 4xx/non-5xx failure paths. +fn fail_status_code(port: u16) -> StatusCode { + get_fail_status_code_for_port(port) + .and_then(|s| StatusCode::from_u16(s).ok()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR) +} + async fn health_handler(State(config): State>>) -> Response { let config = config.read().await; @@ -303,7 +312,7 @@ async fn generate_handler( if should_fail(&config).await { return ( - StatusCode::INTERNAL_SERVER_ERROR, + fail_status_code(config.port), [("x-worker-id", worker_id)], Json(json!({ "error": "Random failure for testing" @@ -324,6 +333,71 @@ async fn generate_handler( if is_stream { let stream_delay = config.response_delay_ms; + if let Some(num_chunks) = get_slow_stream_chunks_for_port(config.port) { + let port = config.port; + let delay_ms = stream_delay; + let error_after = get_stream_error_after_for_port(port); + init_stream_tracking(port, num_chunks); + + let (tx, rx) = + tokio::sync::mpsc::channel::>(MOCK_STREAM_BUFFER); + tokio::spawn(async move { + let _exit_guard = install_stream_exit_notifier(port); + let timestamp_start = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + for i in 0..num_chunks { + if let Some(n) = error_after { + if i == n { + let _ = tx + .send(Err(std::io::Error::other( + "simulated upstream worker crash", + ))) + .await; + return; + } + } + if delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await; + } + let data = json!({ + "text": format!("chunk-{} ", i), + "meta_info": { + "prompt_tokens": 10, + "completion_tokens": (i + 1) as u64, + "completion_tokens_wo_jump_forward": (i + 1) as u64, + "input_token_logprobs": null, + "output_token_logprobs": null, + "first_token_latency": delay_ms as f64 / 1000.0, + "time_to_first_token": delay_ms as f64 / 1000.0, + "time_per_output_token": 0.01, + "start_time": timestamp_start, + "finish_reason": null + }, + "stage": "mid" + }); + if tx + .send(Ok(Event::default().data(data.to_string()))) + .await + .is_err() + { + return; + } + record_chunk_sent(port); + } + let _ = tx.send(Ok(Event::default().data("[DONE]"))).await; + mark_stream_completed(port); + }); + + let stream = tokio_stream::wrappers::ReceiverStream::new(rx); + return ( + [("x-worker-id", worker_id)], + Sse::new(stream).keep_alive(KeepAlive::default()), + ) + .into_response(); + } + // Check if it's a batch request let is_batch = payload.get("text").and_then(|t| t.as_array()).is_some(); @@ -442,28 +516,99 @@ async fn chat_completions_handler( if is_stream { let request_id = format!("chatcmpl-{}", Uuid::new_v4()); - let stream = stream::once(async move { - let chunk = json!({ - "id": request_id, - "object": "chat.completion.chunk", - "created": timestamp, - "model": "mock-model", - "choices": [{ - "index": 0, - "delta": { - "content": "This is a mock chat response." - }, - "finish_reason": null - }] + // Check for slow streaming mode (used by upstream cancel tests). + // Reads from the global SLOW_STREAM_CONFIG (set via set_slow_stream_chunks) + // rather than the payload, because the gateway deserializes/re-serializes + // the request body and drops unknown fields. + let slow_chunks = get_slow_stream_chunks_for_port(config.port); + + if let Some(num_chunks) = slow_chunks { + let port = config.port; + let delay_ms = config.response_delay_ms; + let error_after = get_stream_error_after_for_port(port); + + init_stream_tracking(port, num_chunks); + + // Small bounded capacity gives a bit of slack between the producer + // task and the SSE consumer; on receiver drop, send().await + // returns Err and the loop exits regardless of capacity. + let (tx, rx) = + tokio::sync::mpsc::channel::>(MOCK_STREAM_BUFFER); + + tokio::spawn(async move { + let _exit_guard = install_stream_exit_notifier(port); + for i in 0..num_chunks { + if let Some(n) = error_after { + if i == n { + // Inject a transport-level error to exercise the + // gateway's `Some(Err(_))` arm. + let _ = tx + .send(Err(std::io::Error::other( + "simulated upstream worker crash", + ))) + .await; + return; + } + } + if delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await; + } + let chunk = json!({ + "id": &request_id, + "object": "chat.completion.chunk", + "created": timestamp, + "model": "mock-model", + "choices": [{ + "index": 0, + "delta": { + "content": format!("chunk-{} ", i) + }, + "finish_reason": null + }] + }); + if tx + .send(Ok(Event::default().data(chunk.to_string()))) + .await + .is_err() + { + // Client disconnected, stream was cancelled + return; + } + record_chunk_sent(port); + } + // Send [DONE] + let _ = tx.send(Ok(Event::default().data("[DONE]"))).await; + mark_stream_completed(port); }); - Ok::<_, Infallible>(Event::default().data(chunk.to_string())) - }) - .chain(stream::once(async { Ok(Event::default().data("[DONE]")) })); + let stream = tokio_stream::wrappers::ReceiverStream::new(rx); + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() + } else { + let stream = stream::once(async move { + let chunk = json!({ + "id": request_id, + "object": "chat.completion.chunk", + "created": timestamp, + "model": "mock-model", + "choices": [{ + "index": 0, + "delta": { + "content": "This is a mock chat response." + }, + "finish_reason": null + }] + }); - Sse::new(stream) - .keep_alive(KeepAlive::default()) - .into_response() + Ok::<_, Infallible>(Event::default().data(chunk.to_string())) + }) + .chain(stream::once(async { Ok(Event::default().data("[DONE]")) })); + + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() + } } else { Json(json!({ "id": format!("chatcmpl-{}", Uuid::new_v4()), @@ -793,8 +938,13 @@ async fn responses_handler( Sse::new(stream) .keep_alive(KeepAlive::default()) .into_response() - } else if has_tools && has_function_output { - // Second turn: emit streaming text response + } else if has_tools + && has_function_output + && get_slow_stream_chunks_for_port(config.port).is_none() + { + // Second turn: emit streaming text response. + // If slow-stream is configured, fall through to the slow-stream + // branch below so cancel tests can disconnect mid second-turn. let rid = request_id.clone(); let msg_id = format!( "msg_{}", @@ -941,6 +1091,126 @@ async fn responses_handler( Sse::new(stream) .keep_alive(KeepAlive::default()) .into_response() + } else if let Some(num_chunks) = get_slow_stream_chunks_for_port(config.port) { + // Slow-stream mode for /responses cancel tests. Mirrors the + // chat-completions slow-stream path so the same set_slow_stream_chunks + // helper drives both endpoints. + let port = config.port; + let delay_ms = config.response_delay_ms; + let error_after = get_stream_error_after_for_port(port); + let rid = request_id.clone(); + let msg_id = format!( + "msg_{}", + Uuid::new_v4().to_string().split('-').next().unwrap() + ); + + init_stream_tracking(port, num_chunks); + + let (tx, rx) = + tokio::sync::mpsc::channel::>(MOCK_STREAM_BUFFER); + + tokio::spawn(async move { + let _exit_guard = install_stream_exit_notifier(port); + // Emit response.created and response.in_progress so the + // gateway's /responses persistence accumulator has the + // structural events it expects. + let created = Event::default().event("response.created").data( + json!({ + "type": "response.created", + "response": { + "id": rid.clone(), + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "status": "in_progress" + } + }) + .to_string(), + ); + if tx.send(Ok(created)).await.is_err() { + return; + } + let in_progress = Event::default().event("response.in_progress").data( + json!({ + "type": "response.in_progress", + "response": { + "id": rid.clone(), + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "status": "in_progress" + } + }) + .to_string(), + ); + if tx.send(Ok(in_progress)).await.is_err() { + return; + } + + for i in 0..num_chunks { + if let Some(n) = error_after { + if i == n { + let _ = tx + .send(Err(std::io::Error::other( + "simulated upstream worker crash", + ))) + .await; + return; + } + } + if delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await; + } + let delta = Event::default().event("response.output_text.delta").data( + json!({ + "type": "response.output_text.delta", + "output_index": 0, + "content_index": 0, + "item_id": msg_id.clone(), + "delta": format!("chunk-{} ", i) + }) + .to_string(), + ); + if tx.send(Ok(delta)).await.is_err() { + return; + } + record_chunk_sent(port); + } + + let aggregated_text = (0..num_chunks) + .map(|i| format!("chunk-{} ", i)) + .collect::(); + let completed = Event::default().event("response.completed").data( + json!({ + "type": "response.completed", + "response": { + "id": rid, + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "status": "completed", + "output": [{ + "id": msg_id, + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": aggregated_text + }] + }] + } + }) + .to_string(), + ); + let _ = tx.send(Ok(completed)).await; + let _ = tx.send(Ok(Event::default().data("[DONE]"))).await; + mark_stream_completed(port); + }); + + let stream = tokio_stream::wrappers::ReceiverStream::new(rx); + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() } else { // Default streaming response let stream = stream::once(async move { @@ -1195,6 +1465,241 @@ async fn responses_cancel_handler( } } +// --- Slow-stream configuration (for upstream cancel tests) --- +// Configured via a global map keyed by worker port so that tests +// can enable slow streaming WITHOUT relying on the request payload +// (the gateway deserializes/re-serializes the body, dropping unknown fields). + +static SLOW_STREAM_CONFIG: OnceLock>> = OnceLock::new(); + +fn get_slow_stream_config() -> &'static Mutex> { + SLOW_STREAM_CONFIG.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Configure a worker (by port) to send `num_chunks` chunks with +/// `response_delay_ms` between each when handling a streaming request. +/// Call this before making the request through the gateway. +pub fn set_slow_stream_chunks(port: u16, num_chunks: usize) { + let mut map = get_slow_stream_config().lock().unwrap(); + map.insert(port, num_chunks); +} + +/// Clear slow-stream configuration for a worker port. +pub fn clear_slow_stream_chunks(port: u16) { + let mut map = get_slow_stream_config().lock().unwrap(); + map.remove(&port); +} + +fn get_slow_stream_chunks_for_port(port: u16) -> Option { + let map = get_slow_stream_config().lock().unwrap(); + map.get(&port).copied() +} + +// --- Stream error injection (for upstream cancel + error tests) --- +// When set for `port`, the slow-stream producer emits an io::Error to the +// SSE stream after the configured number of successfully-sent chunks. +// reqwest will surface this as a transport error, which exercises the +// gateway's `Some(Err(_))` arm. + +static STREAM_ERROR_AFTER: OnceLock>> = OnceLock::new(); + +fn get_stream_error_after_config() -> &'static Mutex> { + STREAM_ERROR_AFTER.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Configure a worker (by port) to abort its SSE stream with an error +/// after sending `n` chunks. Must be combined with +/// [`set_slow_stream_chunks`] to take effect. +pub fn set_stream_error_after_chunks(port: u16, n: usize) { + let mut map = get_stream_error_after_config().lock().unwrap(); + map.insert(port, n); +} + +/// Clear error-injection configuration for a worker port. +pub fn clear_stream_error_after_chunks(port: u16) { + let mut map = get_stream_error_after_config().lock().unwrap(); + map.remove(&port); +} + +fn get_stream_error_after_for_port(port: u16) -> Option { + let map = get_stream_error_after_config().lock().unwrap(); + map.get(&port).copied() +} + +// --- Failure-status override (for breaker attribution tests) --- +// When set for `port`, `should_fail`-triggered failures return this HTTP +// status instead of the default 500. Lets a test pin breaker semantics for +// the 4xx-from-worker case (the gateway treats 4xx as "not a worker fault") +// without having to fabricate a separate mock worker. + +static FAIL_STATUS_CODE: OnceLock>> = OnceLock::new(); + +fn get_fail_status_code_config() -> &'static Mutex> { + FAIL_STATUS_CODE.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Configure a worker (by port) to return `status` when `fail_rate` +/// triggers a failure response, instead of the default 500. +pub fn set_fail_status_code(port: u16, status: u16) { + let mut map = get_fail_status_code_config().lock().unwrap(); + map.insert(port, status); +} + +/// Clear failure-status override for a worker port. +pub fn clear_fail_status_code(port: u16) { + let mut map = get_fail_status_code_config().lock().unwrap(); + map.remove(&port); +} + +fn get_fail_status_code_for_port(port: u16) -> Option { + let map = get_fail_status_code_config().lock().unwrap(); + map.get(&port).copied() +} + +// --- Stream cancellation tracking (for upstream cancel tests) --- + +/// Tracks the state of a streaming response for cancel verification. +#[derive(Clone, Debug)] +pub struct StreamTrackingState { + pub total_chunks: usize, + pub chunks_sent: usize, + pub completed: bool, +} + +static STREAM_CANCEL_TRACKER: OnceLock>> = OnceLock::new(); + +fn get_stream_tracker() -> &'static Mutex> { + STREAM_CANCEL_TRACKER.get_or_init(|| Mutex::new(HashMap::new())) +} + +// Per-port `Notify` fired when the worker's producer task exits (either +// because its outbound `send().await` failed — i.e. the gateway dropped +// the upstream connection — or because the stream completed naturally). +// Tests await this notification instead of polling counters, so cancel +// assertions don't depend on timing windows. +static STREAM_FINISH_NOTIFIERS: OnceLock>>> = OnceLock::new(); + +fn get_stream_finish_notifier_map() -> &'static Mutex>> { + STREAM_FINISH_NOTIFIERS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn get_stream_finish_notifier(port: u16) -> Arc { + let mut map = get_stream_finish_notifier_map().lock().unwrap(); + map.entry(port) + .or_insert_with(|| Arc::new(Notify::new())) + .clone() +} + +/// Bound on the per-stream mpsc buffer used by every slow-stream producer +/// task in this mock worker. Tests assert that `chunks_sent` after a cancel +/// grew by at most this many over the pre-drop snapshot, on the theory that +/// anything more means the gateway did not propagate the disconnect +/// upstream. +pub const MOCK_STREAM_BUFFER: usize = 4; + +/// RAII guard that fires the per-port finish notifier on drop, so the +/// notification fires whether the producer task exits normally or returns +/// early on `tx.send(...).await.is_err()`. +#[must_use = "StreamExitNotifier must be bound to a local (typically `_exit_guard`) \ + and held until the producer task ends — dropping it immediately fires \ + the notifier early, causing `wait_for_stream_finish` to return before \ + the producer has actually exited"] +pub struct StreamExitNotifier(Arc); + +impl Drop for StreamExitNotifier { + fn drop(&mut self) { + self.0.notify_one(); + } +} + +/// Install the exit notifier inside a producer task. Hold the returned +/// guard until the task ends (typically by binding it to `_exit_guard`). +#[must_use = "the returned guard fires the exit notifier on drop; bind it to a local \ + (e.g. `let _exit_guard = install_stream_exit_notifier(port);`) so it lives \ + for the producer task's lifetime"] +pub fn install_stream_exit_notifier(port: u16) -> StreamExitNotifier { + StreamExitNotifier(get_stream_finish_notifier(port)) +} + +/// Reset the stream tracker for a given port before starting a new test. +/// Also replaces the finish notifier so any unconsumed permit from a +/// previous test doesn't satisfy this test's wait immediately. +pub fn reset_stream_tracker(port: u16) { + let mut map = get_stream_tracker().lock().unwrap(); + map.remove(&port); + let mut nmap = get_stream_finish_notifier_map().lock().unwrap(); + nmap.insert(port, Arc::new(Notify::new())); +} + +/// Get the stream tracking state for a given port. +pub fn get_stream_tracking_state(port: u16) -> Option { + let map = get_stream_tracker().lock().unwrap(); + map.get(&port).cloned() +} + +/// Wait until the worker's producer task for `port` exits — either because +/// the gateway dropped the upstream connection (`send().await` failed) or +/// because the stream finished naturally. Returns the final tracking state. +/// The `timeout` is a safety net for hung tests; a healthy run returns the +/// instant the producer task drops its exit guard. +/// +/// **Precondition:** call [`reset_stream_tracker`] before issuing the +/// gateway request whose producer you intend to wait on. The reset +/// installs a fresh `Notify` so a stale permit left by a previous test +/// on the same port can't satisfy this wait immediately. +pub async fn wait_for_stream_finish( + port: u16, + timeout: tokio::time::Duration, +) -> Option { + let notifier = get_stream_finish_notifier(port); + if tokio::time::timeout(timeout, notifier.notified()) + .await + .is_err() + { + // A hung producer would silently look like a successful cancel + // (chunks_sent < total_chunks, completed=false) if we just + // returned what we have. Panic instead so the test fails loudly. + panic!( + "wait_for_stream_finish timed out after {:?} for port {} — \ + producer task never fired its exit notifier. Last tracker \ + state: {:?}", + timeout, + port, + get_stream_tracking_state(port) + ); + } + get_stream_tracking_state(port) +} + +// Initialize tracking for a new stream. `map.insert` overwrites any prior +// entry for this port, so callers don't need to reset first; we still expose +// `reset_stream_tracker` so tests can opt into removing the entry entirely. +fn init_stream_tracking(port: u16, total_chunks: usize) { + let mut map = get_stream_tracker().lock().unwrap(); + map.insert( + port, + StreamTrackingState { + total_chunks, + chunks_sent: 0, + completed: false, + }, + ); +} + +fn record_chunk_sent(port: u16) { + let mut map = get_stream_tracker().lock().unwrap(); + if let Some(state) = map.get_mut(&port) { + state.chunks_sent += 1; + } +} + +fn mark_stream_completed(port: u16) { + let mut map = get_stream_tracker().lock().unwrap(); + if let Some(state) = map.get_mut(&port) { + state.completed = true; + } +} + // --- Simple in-memory response store per worker port (for tests) --- static RESP_STORE: OnceLock>>> = OnceLock::new(); diff --git a/sgl-model-gateway/tests/reliability/mod.rs b/sgl-model-gateway/tests/reliability/mod.rs index 5ca376857..06d093bdc 100644 --- a/sgl-model-gateway/tests/reliability/mod.rs +++ b/sgl-model-gateway/tests/reliability/mod.rs @@ -4,3 +4,4 @@ pub mod circuit_breaker_test; pub mod fault_tolerance_test; pub mod rate_limiting_test; pub mod retries_test; +pub mod upstream_cancel_test; diff --git a/sgl-model-gateway/tests/reliability/upstream_cancel_test.rs b/sgl-model-gateway/tests/reliability/upstream_cancel_test.rs new file mode 100644 index 000000000..2b915516e --- /dev/null +++ b/sgl-model-gateway/tests/reliability/upstream_cancel_test.rs @@ -0,0 +1,2714 @@ +//! Upstream request cancellation tests +//! +//! Verifies that when a client disconnects mid-stream, the gateway +//! terminates the upstream request to the backend worker promptly +//! (via the `tokio::select!` / `tx.closed()` mechanism in the router). + +use std::{sync::Arc, time::Duration}; + +use axum::{ + body::Body, + extract::Request, + http::{header::CONTENT_TYPE, StatusCode}, +}; +use http_body_util::BodyExt; +use serde_json::json; +use smg::config::RouterConfig; +use tower::ServiceExt; + +use crate::common::{ + mock_worker::{ + clear_fail_status_code, clear_slow_stream_chunks, clear_stream_error_after_chunks, + get_stream_tracking_state, reset_stream_tracker, set_fail_status_code, + set_slow_stream_chunks, set_stream_error_after_chunks, wait_for_stream_finish, + StreamTrackingState, MOCK_STREAM_BUFFER, + }, + AppTestContext, TestRouterConfig, TestWorkerConfig, +}; + +/// Read up to `max_chunks` data frames from a streaming response body. +async fn read_n_chunks(body: &mut Body, max_chunks: usize) -> usize { + let mut chunks_read = 0; + while chunks_read < max_chunks { + match body.frame().await { + Some(Ok(frame)) if frame.is_data() => { + chunks_read += 1; + } + Some(Ok(_)) => continue, + _ => break, + } + } + chunks_read +} + +/// Read up to `max_chunks` data frames, returning the count and accumulated +/// bytes so callers can parse the SSE payload (e.g. to capture a `response.id` +/// before dropping the body). +async fn read_n_chunks_with_bytes(body: &mut Body, max_chunks: usize) -> (usize, Vec) { + let mut chunks_read = 0; + let mut buf: Vec = Vec::new(); + while chunks_read < max_chunks { + match body.frame().await { + Some(Ok(frame)) if frame.is_data() => { + if let Ok(data) = frame.into_data() { + buf.extend_from_slice(&data); + } + chunks_read += 1; + } + Some(Ok(_)) => continue, + _ => break, + } + } + (chunks_read, buf) +} + +/// Extract the `id` field from the first `response.created` SSE event in `buf`. +/// Returns `None` if the event hasn't arrived yet (caller should read more). +fn extract_response_id_from_sse(buf: &[u8]) -> Option { + let s = std::str::from_utf8(buf).ok()?; + let data_line = s + .lines() + .find(|l| l.starts_with("data:") && l.contains("\"response.created\""))?; + let json_str = data_line.trim_start_matches("data:").trim(); + let value: serde_json::Value = serde_json::from_str(json_str).ok()?; + value + .get("response") + .and_then(|r| r.get("id")) + .and_then(|id| id.as_str()) + .map(|s| s.to_string()) +} + +/// Safety timeout for `wait_for_stream_finish` — the worker notifies the +/// instant its producer task exits, so a healthy run returns well before +/// this. The 3s budget is just a guard against a hung test. +const STREAM_FINISH_TIMEOUT: Duration = Duration::from_secs(3); + +async fn assert_cancelled_before_completion(port: u16) -> StreamTrackingState { + let state = wait_for_stream_finish(port, STREAM_FINISH_TIMEOUT) + .await + .unwrap_or_else(|| { + panic!( + "Stream tracking state should exist for worker port {}", + port + ) + }); + + assert!( + !state.completed, + "Stream should NOT have completed - gateway should have cancelled it. \ + Chunks sent: {}, total: {}", + state.chunks_sent, state.total_chunks + ); + assert!( + state.chunks_sent < state.total_chunks, + "Worker should have sent fewer chunks than total ({} < {}). \ + Stream was not cancelled in time.", + state.chunks_sent, + state.total_chunks + ); + state +} + +#[cfg(test)] +mod upstream_cancel_tests { + use super::*; + + /// Test that the gateway cancels the upstream stream when the client + /// disconnects before consuming all chunks. + /// + /// Setup: + /// - Mock worker sends 20 chunks with 50ms delay between each (~1s total). + /// - Client reads a few chunks then drops the response body. + /// + /// Expectation: + /// - The mock worker stops producing once the gateway closes its + /// upstream connection. We assert that by waiting on the worker's + /// exit notifier (fired when its producer task drops, either via + /// send-failure or natural completion) and snapshotting + /// `chunks_sent` before/after the drop — proving the worker + /// actually halted instead of just being slower than our fixed sleep. + #[tokio::test] + async fn test_streaming_cancel_on_client_disconnect() { + let worker_port = 20250; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = TestRouterConfig::round_robin(4250); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 50)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Tell me a long story"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let chunks_read = read_n_chunks(&mut body, 3).await; + assert!( + chunks_read > 0, + "Should have read at least one chunk before disconnecting" + ); + + // Snapshot the worker counter the moment we drop, then wait for + // the producer task to fire its exit notifier. If cancel propagation + // is broken the producer keeps running until total_chunks and the + // counter ends up at `total_chunks`. + let snapshot = get_stream_tracking_state(worker_port) + .map(|s| s.chunks_sent) + .unwrap_or(0); + drop(body); + + let final_state = assert_cancelled_before_completion(worker_port).await; + assert!( + final_state.chunks_sent <= snapshot + MOCK_STREAM_BUFFER, + "Chunks_sent grew by more than the channel buffer ({} -> {}); \ + gateway likely did not propagate cancel.", + snapshot, + final_state.chunks_sent + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// Test that a fully consumed stream is NOT cancelled prematurely — + /// the worker sends all chunks and completes normally. + #[tokio::test] + async fn test_streaming_completes_when_client_consumes_all() { + let worker_port = 20251; + let total_chunks: usize = 5; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = TestRouterConfig::round_robin(4251); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 10)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Short response"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let _full_body = resp.into_body().collect().await.unwrap().to_bytes(); + + let state = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT) + .await + .expect("Stream tracking state should exist"); + assert!( + state.completed, + "Stream should have completed when client consumed all chunks. \ + Chunks sent: {}, total: {}", + state.chunks_sent, state.total_chunks + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// Test that a non-streaming request is not affected by cancel logic. + #[tokio::test] + async fn test_non_streaming_request_unaffected() { + let config = TestRouterConfig::round_robin(4252); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20252)]).await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": false + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body_bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let body: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap(); + assert_eq!( + body.get("object").and_then(|v| v.as_str()), + Some("chat.completion"), + "Non-streaming response should be a complete chat.completion object" + ); + + ctx.shutdown().await; + } + + /// Cancel before the worker emits any chunk. Catches a select! that + /// only wakes `tx.closed()` after the first `stream.next()` resolves. + #[tokio::test] + async fn test_streaming_cancel_before_first_chunk() { + let worker_port = 20253; + let total_chunks: usize = 10; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = TestRouterConfig::round_robin(4253); + let ctx = AppTestContext::new_with_config( + // 200ms per-chunk delay; the very first chunk takes the + // full 200ms because the worker sleeps before emitting. + config, + vec![TestWorkerConfig::slow(worker_port, 200)], + ) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Drop immediately, before pulling any frame. + drop(resp.into_body()); + + let state = assert_cancelled_before_completion(worker_port).await; + assert!( + state.chunks_sent <= 4, + "Worker should have sent very few chunks (≤ buffer capacity), \ + saw {} of {}", + state.chunks_sent, + state.total_chunks + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// Disconnecting *after* the stream completes naturally must be a + /// no-op — no panic, no spurious "cancel" log, completed=true stays. + #[tokio::test] + async fn test_streaming_cancel_after_done_is_noop() { + let worker_port = 20254; + let total_chunks: usize = 3; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = TestRouterConfig::round_robin(4254); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 5)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Drain entire body, then drop after a short pause. + let body = resp.into_body().collect().await.unwrap().to_bytes(); + assert!(!body.is_empty()); + tokio::time::sleep(Duration::from_millis(50)).await; + drop(body); + + let state = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT) + .await + .expect("tracking state"); + assert!(state.completed, "Stream should have completed cleanly"); + assert_eq!(state.chunks_sent, state.total_chunks); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// Cancelling one client request must not affect a concurrent request + /// that's hitting a *different* upstream replica. + #[tokio::test] + async fn test_cancel_one_request_does_not_affect_concurrent() { + let worker_a = 20255; + let worker_b = 20256; + let total_a: usize = 20; + let total_b: usize = 5; + + reset_stream_tracker(worker_a); + reset_stream_tracker(worker_b); + set_slow_stream_chunks(worker_a, total_a); + set_slow_stream_chunks(worker_b, total_b); + + let config = TestRouterConfig::round_robin(4255); + let ctx = AppTestContext::new_with_config( + config, + vec![ + TestWorkerConfig::slow(worker_a, 50), + TestWorkerConfig::slow(worker_b, 10), + ], + ) + .await; + + // Two parallel requests. Round-robin should send them to different + // workers. We don't strictly need to know which got which, but we + // assume the FIRST request lands on worker_a — that's the one we + // cancel — and we await the SECOND to completion. + let app = ctx.create_app().await; + let app2 = app.clone(); + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Long"}], + "stream": true + }); + let body_str = Arc::new(serde_json::to_string(&payload).unwrap()); + + let body_str_a = Arc::clone(&body_str); + let h_cancel = tokio::spawn(async move { + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from((*body_str_a).clone())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let mut body = resp.into_body(); + let _ = read_n_chunks(&mut body, 2).await; + drop(body); + }); + + // Small stagger so round-robin index advances. + tokio::time::sleep(Duration::from_millis(20)).await; + + let body_str_b = Arc::clone(&body_str); + let h_consume = tokio::spawn(async move { + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from((*body_str_b).clone())) + .unwrap(); + let resp = app2.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let _ = resp.into_body().collect().await.unwrap(); + }); + + h_cancel.await.unwrap(); + h_consume.await.unwrap(); + + // Worker that should have completed normally. + let state_b = wait_for_stream_finish(worker_b, STREAM_FINISH_TIMEOUT) + .await + .expect("worker_b tracking state"); + // Worker that should have been cancelled. + let state_a = wait_for_stream_finish(worker_a, STREAM_FINISH_TIMEOUT) + .await + .expect("worker_a tracking state"); + + // We don't know which worker got which request because round-robin + // is shared across the run, so accept either ordering. + let (cancelled, completed) = if state_a.completed { + (state_b, state_a) + } else { + (state_a, state_b) + }; + assert!( + !cancelled.completed, + "Cancelled stream should not have completed (sent {}/{})", + cancelled.chunks_sent, cancelled.total_chunks + ); + assert!( + completed.completed, + "Concurrent stream should have completed (sent {}/{})", + completed.chunks_sent, completed.total_chunks + ); + + clear_slow_stream_chunks(worker_a); + clear_slow_stream_chunks(worker_b); + ctx.shutdown().await; + } + + /// Mid-stream worker error must surface as `Stream error: …` to the + /// client, must NOT be silently swallowed as cancel, and must trigger + /// the gateway's error log path. We assert (a) the client sees the + /// error frame and (b) `chunks_sent` reflects the partial output the + /// worker sent before erroring. + #[tokio::test] + async fn test_streaming_worker_error_propagates_not_cancel() { + let worker_port = 20257; + let total_chunks: usize = 10; + let error_after: usize = 3; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + set_stream_error_after_chunks(worker_port, error_after); + + let config = TestRouterConfig::round_robin(4256); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 10)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Collecting the body should yield a transport-level error AFTER + // the first few chunks. axum surfaces the upstream error as a + // failed body.collect(), so we just stitch frames manually. + let mut body = resp.into_body(); + let mut combined = Vec::::new(); + let mut saw_transport_err = false; + loop { + match body.frame().await { + Some(Ok(frame)) => { + if let Ok(data) = frame.into_data() { + combined.extend_from_slice(&data); + } + } + Some(Err(_)) => { + saw_transport_err = true; + break; + } + None => break, + } + } + let combined_str = String::from_utf8_lossy(&combined); + assert!( + saw_transport_err + || combined_str.contains("Stream error") + || combined_str.contains("simulated upstream worker crash"), + "Client should observe a stream error event or transport error; got: {}", + combined_str + ); + + // Worker should have sent some chunks but not all, and the stream + // should NOT be marked completed (we crashed before [DONE]). + let state = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT) + .await + .expect("tracking state"); + assert!(!state.completed, "Errored stream should not be completed"); + assert!( + state.chunks_sent <= error_after, + "Worker reported {} chunks_sent, expected ≤ {}", + state.chunks_sent, + error_after + ); + + clear_stream_error_after_chunks(worker_port); + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// PD-disagg streaming cancel: client disconnects mid-decode-stream. + /// The decode worker's slow-stream tracker is what proves cancel + /// actually reached the upstream — prefill is fully drained + /// synchronously by `process_prefill_response`, so it's expected + /// to complete regardless. + #[tokio::test] + async fn test_pd_streaming_cancel_on_client_disconnect() { + let prefill_port = 20258; + let decode_port = 20259; + let total_chunks: usize = 20; + + reset_stream_tracker(decode_port); + set_slow_stream_chunks(decode_port, total_chunks); + + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4257) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::prefill(prefill_port), { + // Use a slow decode worker (50ms per chunk). + let mut w = TestWorkerConfig::decode(decode_port); + w.response_delay_ms = 50; + w + }], + ) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "text": "PD streaming test", + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let chunks_read = read_n_chunks(&mut body, 3).await; + assert!( + chunks_read > 0, + "Should have read at least one chunk before disconnecting" + ); + + let snapshot = get_stream_tracking_state(decode_port) + .map(|s| s.chunks_sent) + .unwrap_or(0); + drop(body); + + let final_state = assert_cancelled_before_completion(decode_port).await; + assert!( + final_state.chunks_sent <= snapshot + MOCK_STREAM_BUFFER, + "Decode chunks_sent grew by more than the buffer ({} -> {}); \ + gateway likely did not propagate cancel through PD path.", + snapshot, + final_state.chunks_sent + ); + + clear_slow_stream_chunks(decode_port); + ctx.shutdown().await; + } + + /// /v1/responses with no persistence (`store: false`, no conversation): + /// client disconnect must propagate to the upstream worker. + #[tokio::test] + async fn test_responses_streaming_cancel_no_persistence() { + let worker_port = 20260; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4258) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 50)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "input": "Tell me a story", + "stream": true, + "store": false + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + // First two events are response.created/response.in_progress, then + // chunks start. Read a few to ensure we're past the bootstrap. + let chunks_read = read_n_chunks(&mut body, 3).await; + assert!(chunks_read > 0); + + let snapshot = get_stream_tracking_state(worker_port) + .map(|s| s.chunks_sent) + .unwrap_or(0); + drop(body); + + let final_state = assert_cancelled_before_completion(worker_port).await; + assert!( + final_state.chunks_sent <= snapshot + MOCK_STREAM_BUFFER, + "/responses chunks_sent grew by more than the buffer ({} -> {}); \ + gateway likely did not propagate cancel through /responses path.", + snapshot, + final_state.chunks_sent + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// /v1/responses with persistence (`store: true`): the upstream is + /// intentionally NOT cancelled on client disconnect — the gateway + /// keeps consuming so the response can be persisted. We assert the + /// worker reaches `completed = true` after the client disconnects. + #[tokio::test] + async fn test_responses_streaming_persistence_drains_after_disconnect() { + let worker_port = 20261; + let total_chunks: usize = 6; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4259) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 20)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "input": "Tell me a story", + "stream": true, + "store": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + // Read enough frames to capture the response.created event so we + // can later look up the stored response by its id. + let (_, captured) = read_n_chunks_with_bytes(&mut body, 3).await; + let response_id = + extract_response_id_from_sse(&captured).expect("response.created event with id"); + drop(body); + + // With persistence the gateway keeps reading; worker should + // eventually mark the stream completed despite client gone. + let state = wait_for_stream_finish(worker_port, Duration::from_secs(5)) + .await + .expect("tracking state"); + assert!( + state.completed, + "Persistence path should drain upstream to completion despite \ + client disconnect (sent {}/{})", + state.chunks_sent, state.total_chunks + ); + assert_eq!(state.chunks_sent, state.total_chunks); + + // Draining is necessary but not sufficient — also verify the + // gateway actually called persist_conversation_items and the + // response landed in storage. Poll briefly because persistence + // happens after the upstream loop exits. + let storage = ctx.app_context.response_storage.clone(); + let stored = { + use data_connector::ResponseId; + let id = ResponseId::from(response_id.clone()); + let mut found = None; + for _ in 0..20 { + if let Ok(Some(r)) = storage.get_response(&id).await { + found = Some(r); + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + found.unwrap_or_else(|| { + panic!( + "Response {} should have been persisted after client \ + disconnect on store=true /responses stream", + response_id + ) + }) + }; + assert_eq!(stored.id.0, response_id); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// Dual of `_drains_after_disconnect`: with `store=true` and a mid-stream + /// upstream error, the gateway must NOT persist a torn response. The + /// commit `714b62f24` warn-log on the persistence-skipped path is the + /// observable signal; here we assert the stronger property that no row + /// lands in storage. + #[tokio::test] + async fn test_responses_simple_streaming_error_skips_persistence() { + let worker_port = 20266; + let total_chunks: usize = 10; + let error_after: usize = 3; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + set_stream_error_after_chunks(worker_port, error_after); + let _guard = StreamInjectionGuard(worker_port); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4266) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 20)]) + .await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + let (_, f_pre) = breaker_counts(&worker); + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "input": "Tell me a story", + "stream": true, + "store": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Drain the body so the response.created event lands in `captured` + // and the upstream error is observed (no client disconnect — the + // skip is driven by the error, not by a cancel). + let mut body = resp.into_body(); + let mut captured: Vec = Vec::new(); + while let Some(Ok(frame)) = body.frame().await { + if frame.is_data() { + if let Ok(data) = frame.into_data() { + captured.extend_from_slice(&data); + } + } + } + drop(body); + let response_id = + extract_response_id_from_sse(&captured).expect("response.created event with id"); + + // Wait for the gateway-side producer task to exit so persistence + // (or its skip) and breaker tick have run. + let _ = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT).await; + + // Co-assert that the breaker tick fired. Without this, a regression + // that silently swallowed the error (record nothing, persist nothing) + // would also produce an empty storage and pass the lookup below. + let (_, f_post) = breaker_counts(&worker); + assert!( + f_post > f_pre, + "/responses simple: mid-stream upstream error must record a \ + breaker failure. failures {}→{}", + f_pre, + f_post + ); + + let storage = ctx.app_context.response_storage.clone(); + use data_connector::ResponseId; + let id = ResponseId::from(response_id.clone()); + if let Ok(Some(_)) = storage.get_response(&id).await { + panic!( + "Response {} should NOT have been persisted after \ + mid-stream upstream error on store=true /responses", + response_id + ); + } + + ctx.shutdown().await; + } + + /// OpenAI-mode (non-/responses) chat-completions cancel: this exercises + /// the `OpenAIRouter` impl, which is a separate codepath from + /// `http::Router`, so the same cancel semantics need their own test. + #[tokio::test] + async fn test_openai_router_streaming_cancel() { + let worker_port = 20262; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4260) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 50)]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Tell me a story"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let chunks_read = read_n_chunks(&mut body, 3).await; + assert!(chunks_read > 0); + + let snapshot = get_stream_tracking_state(worker_port) + .map(|s| s.chunks_sent) + .unwrap_or(0); + drop(body); + + let final_state = assert_cancelled_before_completion(worker_port).await; + assert!( + final_state.chunks_sent <= snapshot + MOCK_STREAM_BUFFER, + "OpenAI-mode chunks_sent grew by more than the buffer ({} -> {})", + snapshot, + final_state.chunks_sent + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// Tool-interception streaming cancel: when the client disconnects mid + /// second-turn (after the gateway has executed an MCP tool call and + /// re-issued an upstream request), the gateway must drop that second + /// upstream connection promptly. This guards the explicit policy + /// documented at `streaming.rs:712-722` ("don't keep workers and + /// external MCP services busy on results no one will read") against + /// silent regression — the inner `select! { ... _ = tx.closed() }` + /// in `handle_streaming_with_tool_interception` is the load-bearing + /// piece. + #[tokio::test] + async fn test_tool_interception_streaming_cancel_on_client_disconnect() { + use smg::routers::{RouterFactory, RouterTrait}; + + use crate::common::{ + mock_mcp_server::MockMCPServer, + mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType}, + }; + + let worker_port = 20263; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let mut mcp = MockMCPServer::start().await.expect("start mcp"); + let mcp_yaml = format!( + "servers:\n - name: mock\n protocol: streamable\n url: {}\n", + mcp.url() + ); + let dir = tempfile::tempdir().expect("tmpdir"); + let cfg_path = dir.path().join("mcp.yaml"); + std::fs::write(&cfg_path, mcp_yaml).expect("write mcp cfg"); + + let mut worker = MockWorker::new(MockWorkerConfig { + port: worker_port, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 50, + fail_rate: 0.0, + }); + let worker_url = worker.start().await.expect("start worker"); + // Allow the mock worker's HTTP listener to bind before the router + // probes its health. + tokio::time::sleep(Duration::from_millis(200)).await; + + let router_cfg = RouterConfig::builder() + .openai_mode(vec![worker_url]) + .random_policy() + .host("127.0.0.1") + .port(4263) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = crate::common::create_test_context_with_mcp_config( + router_cfg, + cfg_path.to_str().unwrap(), + ) + .await; + let router: Arc = + Arc::from(RouterFactory::create_router(&ctx).await.expect("router")); + let app = crate::common::test_app::create_test_app_with_context(router, ctx); + + let payload = json!({ + "model": "mock-model", + "input": "search something", + "stream": true, + "store": false, + "tools": [{ + "type": "mcp", + "server_label": "mock", + "server_url": mcp.url() + }] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + // Pull a few frames so the body keeps draining while the gateway + // works through turn 1 (tool call) and starts turn 2 (slow text). + let _ = read_n_chunks(&mut body, 8).await; + + // Wait until the slow second-turn upstream request has actually + // started producing chunks — only the slow-stream branch in the + // mock initialises the tracker, so seeing chunks_sent>0 here + // means we're inside the second upstream request. + let mut waited_ms: u64 = 0; + loop { + if let Some(s) = get_stream_tracking_state(worker_port) { + if s.chunks_sent > 0 { + break; + } + } + if waited_ms >= 5000 { + panic!( + "Second-turn upstream stream never started producing chunks \ + within 5s — tool-interception path did not reach select!" + ); + } + tokio::time::sleep(Duration::from_millis(50)).await; + waited_ms += 50; + } + + let snapshot = get_stream_tracking_state(worker_port) + .map(|s| s.chunks_sent) + .unwrap_or(0); + drop(body); + + let final_state = assert_cancelled_before_completion(worker_port).await; + assert!( + final_state.chunks_sent <= snapshot + MOCK_STREAM_BUFFER, + "Tool-interception second-turn chunks_sent grew by more than \ + the buffer ({} -> {}) — cancel did not propagate to upstream", + snapshot, + final_state.chunks_sent + ); + + clear_slow_stream_chunks(worker_port); + worker.stop().await; + mcp.stop().await; + } + + /// Tool-interception path: client disconnects WHILE the gateway is + /// still waiting for the upstream's response headers (inside the + /// `request_builder.send().await` future, not yet streaming). + /// + /// The mock sleeps 1500ms before returning headers; the test drops + /// the response body ~50ms after the gateway has dispatched the + /// request. With the `tokio::select! { res = send() => …, _ = tx.closed() => return }` + /// guard in place, the gateway aborts the send before the mock ever + /// reaches its slow-stream init — so `get_stream_tracking_state` + /// stays at `None`. If the guard regresses to a plain + /// `request_builder.send().await`, the mock would complete its sleep, + /// initialise the tracker, and `get_stream_tracking_state` would + /// return `Some(...)`. + #[tokio::test] + async fn test_tool_interception_cancel_during_send() { + use smg::routers::{RouterFactory, RouterTrait}; + + use crate::common::{ + mock_mcp_server::MockMCPServer, + mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType}, + }; + + let worker_port = 20264; + // slow_stream is configured so that IF the mock ever gets past + // the pre-response delay, the tracker is populated and the test + // would observe the regression. + let total_chunks: usize = 5; + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let mut mcp = MockMCPServer::start().await.expect("start mcp"); + let mcp_yaml = format!( + "servers:\n - name: mock\n protocol: streamable\n url: {}\n", + mcp.url() + ); + let dir = tempfile::tempdir().expect("tmpdir"); + let cfg_path = dir.path().join("mcp.yaml"); + std::fs::write(&cfg_path, mcp_yaml).expect("write mcp cfg"); + + let mut worker = MockWorker::new(MockWorkerConfig { + port: worker_port, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + // 1500ms pre-response delay: long enough that the test's + // ~50ms drop reliably races *inside* the send().await window. + response_delay_ms: 1500, + fail_rate: 0.0, + }); + let worker_url = worker.start().await.expect("start worker"); + tokio::time::sleep(Duration::from_millis(200)).await; + + let router_cfg = RouterConfig::builder() + .openai_mode(vec![worker_url]) + .random_policy() + .host("127.0.0.1") + .port(4264) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = crate::common::create_test_context_with_mcp_config( + router_cfg, + cfg_path.to_str().unwrap(), + ) + .await; + let router: Arc = + Arc::from(RouterFactory::create_router(&ctx).await.expect("router")); + let app = crate::common::test_app::create_test_app_with_context(router, ctx); + + let payload = json!({ + "model": "mock-model", + "input": "search something", + "stream": true, + "store": false, + "tools": [{ + "type": "mcp", + "server_label": "mock", + "server_url": mcp.url() + }] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Drop immediately, while the mock is still inside its 1500ms + // pre-response sleep. The gateway's spawned task is parked in + // `select! { res = send() => ..., _ = tx.closed() => return }`. + let body = resp.into_body(); + tokio::time::sleep(Duration::from_millis(50)).await; + drop(body); + + // Give the gateway and mock plenty of time to process the cancel + // and finish their respective sleeps. 2500ms > 1500ms ensures + // that even if the select! guard regressed, the mock would have + // long since reached `init_stream_tracking` by the time we check. + tokio::time::sleep(Duration::from_millis(2500)).await; + + assert!( + get_stream_tracking_state(worker_port).is_none(), + "Mock worker initialised the slow-stream tracker, which means \ + its handler completed the pre-response sleep — i.e. the gateway \ + waited for upstream headers instead of cancelling send().await \ + on client disconnect. Tracker state: {:?}", + get_stream_tracking_state(worker_port) + ); + + clear_slow_stream_chunks(worker_port); + worker.stop().await; + mcp.stop().await; + } + + /// Tool-interception path with `store=true`: when the second-turn + /// upstream errors mid-stream, the gateway must NOT persist a torn + /// response. Mirrors `test_responses_simple_streaming_error_skips_persistence` + /// for the MCP-interception branch (streaming.rs:1023-1033). + #[tokio::test] + async fn test_tool_interception_streaming_error_skips_persistence() { + use data_connector::ResponseId; + use smg::routers::{RouterFactory, RouterTrait}; + + use crate::common::{ + mock_mcp_server::MockMCPServer, + mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType}, + }; + + let worker_port = 20267; + let total_chunks: usize = 10; + let error_after: usize = 2; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + set_stream_error_after_chunks(worker_port, error_after); + let _guard = StreamInjectionGuard(worker_port); + + let mut mcp = MockMCPServer::start().await.expect("start mcp"); + let mcp_yaml = format!( + "servers:\n - name: mock\n protocol: streamable\n url: {}\n", + mcp.url() + ); + let dir = tempfile::tempdir().expect("tmpdir"); + let cfg_path = dir.path().join("mcp.yaml"); + std::fs::write(&cfg_path, mcp_yaml).expect("write mcp cfg"); + + let mut worker = MockWorker::new(MockWorkerConfig { + port: worker_port, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 20, + fail_rate: 0.0, + }); + let worker_url = worker.start().await.expect("start worker"); + tokio::time::sleep(Duration::from_millis(200)).await; + + let router_cfg = RouterConfig::builder() + .openai_mode(vec![worker_url.clone()]) + .random_policy() + .host("127.0.0.1") + .port(4267) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = crate::common::create_test_context_with_mcp_config( + router_cfg, + cfg_path.to_str().unwrap(), + ) + .await; + let storage = ctx.response_storage.clone(); + let router: Arc = + Arc::from(RouterFactory::create_router(&ctx).await.expect("router")); + let pinned_worker = ctx + .worker_registry + .get_by_url(&worker_url) + .expect("worker should be registered after router create"); + let (_, f_pre) = breaker_counts(&pinned_worker); + let app = crate::common::test_app::create_test_app_with_context(router, ctx); + + let payload = json!({ + "model": "mock-model", + "input": "search something", + "stream": true, + "store": true, + "tools": [{ + "type": "mcp", + "server_label": "mock", + "server_url": mcp.url() + }] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Drain fully so the second-turn upstream error is observed (no + // client disconnect — the skip is driven by the error). + let mut body = resp.into_body(); + let mut captured: Vec = Vec::new(); + while let Some(Ok(frame)) = body.frame().await { + if frame.is_data() { + if let Ok(data) = frame.into_data() { + captured.extend_from_slice(&data); + } + } + } + drop(body); + let response_id = + extract_response_id_from_sse(&captured).expect("response.created event with id"); + + // Wait for the second-turn producer to exit. + let state = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT) + .await + .expect("second-turn producer to exit within timeout"); + + // Pin that the second-turn slow_stream branch actually fired with the + // error injection: if a future mock refactor sent turn 2 down the + // happy-path JSON branch, no chunks would have streamed and the + // persistence-skip assertion below would still pass for the wrong + // reason. + assert!( + state.chunks_sent >= error_after && !state.completed, + "Second-turn must have entered slow_stream and errored after \ + {} chunks (got chunks_sent={}, completed={})", + error_after, + state.chunks_sent, + state.completed, + ); + + // Co-assert the breaker tick fired so a "skip-persistence + record + // nothing" regression can't silently pass this test. + let (_, f_post) = breaker_counts(&pinned_worker); + assert!( + f_post > f_pre, + "tool-interception: second-turn upstream error must record a \ + breaker failure. failures {}→{}", + f_pre, + f_post + ); + + let id = ResponseId::from(response_id.clone()); + if let Ok(Some(_)) = storage.get_response(&id).await { + panic!( + "Response {} should NOT have been persisted after \ + second-turn upstream error on store=true tool-interception path", + response_id + ); + } + + worker.stop().await; + mcp.stop().await; + } + + /// After enough consecutive mid-stream upstream errors, the + /// `BreakerTrackedStream` drop should record failures often enough + /// that the worker's circuit breaker opens. This locks in the + /// contract that mid-stream errors are *not* silently swallowed — + /// regressing to "log only, no breaker tick" would leave a + /// 200-then-broken worker permanently selectable. + #[tokio::test] + async fn test_streaming_errors_trip_circuit_breaker() { + use smg::config::CircuitBreakerConfig; + + let worker_port = 20265; + let total_chunks: usize = 10; + let error_after: usize = 1; + let failure_threshold = 3; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + set_stream_error_after_chunks(worker_port, error_after); + + let config = TestRouterConfig::round_robin_with_circuit_breaker( + 4265, + CircuitBreakerConfig { + failure_threshold, + success_threshold: 2, + timeout_duration_secs: 30, + window_duration_secs: 60, + }, + ); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 10)]) + .await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let app = ctx.create_app().await; + + // Drain each request fully so the BreakerTrackedStream sees + // `Some(Err(...))` and tags the terminal state as Errored before + // Drop fires `record_failure`. + for _ in 0..failure_threshold { + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + // Don't assert on status — once the breaker trips, the gateway + // returns 503 instead of dispatching. Both outcomes count + // toward the test as long as the breaker opens by the end. + let mut body = resp.into_body(); + while body.frame().await.is_some() {} + } + + let worker = ctx + .app_context + .worker_registry + .get_by_url(&worker_url) + .expect("worker should be registered"); + let breaker = worker.circuit_breaker(); + assert!( + !matches!(breaker.state(), smg::core::CircuitState::Closed), + "Circuit breaker should NOT be Closed after {} streaming errors; \ + state = {:?}, consecutive_failures = {}, total_failures = {}, \ + total_successes = {}", + failure_threshold, + breaker.state(), + breaker.consecutive_failures(), + breaker.total_failures(), + breaker.total_successes(), + ); + + clear_slow_stream_chunks(worker_port); + clear_stream_error_after_chunks(worker_port); + ctx.shutdown().await; + } + + // -------- Breaker accounting tests -------- + // + // For single-upstream-call streaming paths (http chat, OpenAI chat, PD + // generate, /responses simple), the worker's circuit breaker is ticked + // exactly once per request based on the upstream's actual termination: + // success on clean end, failure on mid-stream error, neither on client + // disconnect. The tool-interception path (/responses with MCP) is the + // documented exception — it ticks once per upstream HTTP call inside + // the tool loop, so a 3-iteration loop can tick up to 3 times. + + /// Returns `(total_successes, total_failures)` for the given worker. + /// + /// Callers MUST capture the `Arc` once at test start (via + /// `worker_registry.get_by_url(...).unwrap()`) and reuse it for every + /// snapshot. Looking up by URL each time is unsafe: any path that + /// re-registers a worker (e.g. the admin `UpdateWorkerPropertiesStep` + /// workflow) replaces the registry's `Arc` with a freshly-built worker + /// that has a fresh `CircuitBreaker`. Two `get_by_url` calls bracketing + /// a request can therefore return handles to two different breakers, + /// making counter deltas vacuous. + fn breaker_counts(worker: &Arc) -> (u64, u64) { + let breaker = worker.circuit_breaker(); + (breaker.total_successes(), breaker.total_failures()) + } + + /// Capture the worker for a given URL once at test start. See + /// `breaker_counts` for why repeated `get_by_url` lookups are unsafe. + fn pin_worker(ctx: &AppTestContext, worker_url: &str) -> Arc { + ctx.app_context + .worker_registry + .get_by_url(worker_url) + .expect("worker should be registered") + } + + /// RAII cleanup for per-port stream injection state. Tests that + /// configure `set_slow_stream_chunks` / `set_stream_error_after_chunks` + /// must use this — without it, a panicking assertion would leave the + /// global injection map populated and poison any future test that + /// reuses the same port. + struct StreamInjectionGuard(u16); + impl Drop for StreamInjectionGuard { + fn drop(&mut self) { + clear_stream_error_after_chunks(self.0); + clear_slow_stream_chunks(self.0); + } + } + + /// http chat: client disconnect mid-stream must NOT tick the breaker. + /// Guards `BreakerTrackedStream`'s drop-while-Active path. + #[tokio::test] + async fn test_disconnect_does_not_move_breaker_http_chat() { + let worker_port = 20270; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = TestRouterConfig::round_robin(4270); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 50)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "long"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let _ = read_n_chunks(&mut body, 3).await; + drop(body); + + // Wait until the upstream producer exits so the body Drop has + // run and any breaker tick has landed. + let _ = assert_cancelled_before_completion(worker_port).await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + (s_post - s_pre, f_post - f_pre), + (0, 0), + "http chat: client disconnect must not move breaker. \ + successes {}→{}, failures {}→{}", + s_pre, + s_post, + f_pre, + f_post + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// OpenAIRouter chat: client disconnect mid-stream must NOT tick the + /// breaker. Same `BreakerTrackedStream` drop-while-Active story as the + /// http chat path. + #[tokio::test] + async fn test_disconnect_does_not_move_breaker_openai_chat() { + let worker_port = 20271; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4271) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 50)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "long"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let _ = read_n_chunks(&mut body, 3).await; + drop(body); + + let _ = assert_cancelled_before_completion(worker_port).await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + (s_post - s_pre, f_post - f_pre), + (0, 0), + "openai chat: client disconnect must not move breaker. \ + successes {}→{}, failures {}→{}", + s_pre, + s_post, + f_pre, + f_post + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// PD streaming client disconnect: + /// - decode breaker must show zero delta (`BreakerTrackedStream` drops + /// Active → no tick). + /// - prefill breaker must show exactly +1 success (prefill is fully + /// drained before decode streaming starts, so `record_outcome(true)` + /// fires for the 2xx prefill regardless of what the client does to + /// the decode stream). + #[tokio::test] + async fn test_disconnect_does_not_move_breaker_pd_decode() { + let prefill_port = 20272; + let decode_port = 20273; + let total_chunks: usize = 20; + + reset_stream_tracker(decode_port); + set_slow_stream_chunks(decode_port, total_chunks); + + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4272) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::prefill(prefill_port), { + let mut w = TestWorkerConfig::decode(decode_port); + w.response_delay_ms = 50; + w + }], + ) + .await; + let app = ctx.create_app().await; + let decode_url = format!("http://127.0.0.1:{}", decode_port); + let prefill_url = format!("http://127.0.0.1:{}", prefill_port); + let decode_worker = pin_worker(&ctx, &decode_url); + let prefill_worker = pin_worker(&ctx, &prefill_url); + + let (s_pre_decode, f_pre_decode) = breaker_counts(&decode_worker); + let (s_pre_prefill, f_pre_prefill) = breaker_counts(&prefill_worker); + + let payload = json!({ + "text": "PD streaming breaker test", + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let _ = read_n_chunks(&mut body, 3).await; + drop(body); + + let _ = assert_cancelled_before_completion(decode_port).await; + + let (s_post_decode, f_post_decode) = breaker_counts(&decode_worker); + assert_eq!( + (s_post_decode - s_pre_decode, f_post_decode - f_pre_decode), + (0, 0), + "PD decode: client disconnect must not move breaker. \ + successes {}→{}, failures {}→{}", + s_pre_decode, + s_post_decode, + f_pre_decode, + f_post_decode + ); + + let (s_post_prefill, f_post_prefill) = breaker_counts(&prefill_worker); + assert_eq!( + ( + s_post_prefill - s_pre_prefill, + f_post_prefill - f_pre_prefill + ), + (1, 0), + "PD prefill: 2xx must record exactly one success regardless of \ + client decode-stream disconnect. successes {}→{}, failures {}→{}", + s_pre_prefill, + s_post_prefill, + f_pre_prefill, + f_post_prefill + ); + + clear_slow_stream_chunks(decode_port); + ctx.shutdown().await; + } + + /// /v1/responses simple (no-persist): client disconnect must not + /// move the worker's circuit breaker — recording neither success + /// nor failure on a request the client abandoned mid-stream. + #[tokio::test] + async fn test_disconnect_does_not_move_breaker_responses_simple() { + let worker_port = 20274; + let total_chunks: usize = 20; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4274) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 50)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "input": "Tell me a story", + "stream": true, + "store": false + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let mut body = resp.into_body(); + let _ = read_n_chunks(&mut body, 3).await; + drop(body); + + let _ = assert_cancelled_before_completion(worker_port).await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + (s_post - s_pre, f_post - f_pre), + (0, 0), + "/responses simple: client disconnect must not move breaker. \ + successes {}→{}, failures {}→{}", + s_pre, + s_post, + f_pre, + f_post + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// /v1/responses simple, mid-stream error: the spawned forwarder in + /// `handle_simple_streaming_passthrough` must record a failure on + /// the worker's circuit breaker when an upstream stream errors + /// after headers — otherwise a "200 OK then broken pipe" worker + /// would never trip the breaker. + #[tokio::test] + async fn test_responses_simple_mid_stream_error_records_failure() { + let worker_port = 20275; + let total_chunks: usize = 10; + let error_after: usize = 2; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + set_stream_error_after_chunks(worker_port, error_after); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4275) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 10)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + + let (_, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "input": "Tell me a story", + "stream": true, + "store": false + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Drain the body so we observe the mid-stream error. + let mut body = resp.into_body(); + while body.frame().await.is_some() {} + drop(body); + + // Wait for the producer task to exit so any breaker tick is + // observable. + let _ = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT).await; + + let (_, f_post) = breaker_counts(&worker); + assert!( + f_post > f_pre, + "/responses simple: mid-stream upstream error must record \ + at least one failure on the breaker. failures {}→{}", + f_pre, + f_post + ); + + clear_stream_error_after_chunks(worker_port); + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// OpenAIRouter chat, mid-stream upstream error: breaker MUST record + /// at least one failure. `Some(Err(_))` → terminal = Errored → + /// Drop ticks `record_failure`. Single-request analogue of + /// `test_streaming_errors_trip_circuit_breaker`. + #[tokio::test] + async fn test_openai_chat_mid_stream_error_records_failure() { + let worker_port = 20276; + let total_chunks: usize = 10; + let error_after: usize = 2; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + set_stream_error_after_chunks(worker_port, error_after); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4276) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 10)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + + let (_, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }); + + reset_stream_tracker(worker_port); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + let _ = resp.status(); + let mut body = resp.into_body(); + while body.frame().await.is_some() {} + drop(body); + let _ = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT).await; + + let (_s_post, f_post) = breaker_counts(&worker); + assert!( + f_post > f_pre, + "openai chat: mid-stream upstream error must record at \ + least one failure on the breaker. failures {}→{}", + f_pre, + f_post + ); + + clear_stream_error_after_chunks(worker_port); + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + // -------- 5xx-streaming and happy-path success coverage -------- + + /// http::Router streaming 5xx must record `record_failure`, not success. + /// Guards the `mark_errored()` pre-tag on the streaming branch — without + /// it, the small error body streams cleanly to `None` and Drop would + /// record a spurious success. + #[tokio::test] + async fn test_http_chat_streaming_5xx_records_failure() { + let worker_port = 20290; + let config = TestRouterConfig::round_robin(4290); + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::flaky(worker_port, 1.0)], + ) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "x"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let _ = resp.into_body().collect().await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + s_post - s_pre, + 0, + "http chat streaming 5xx must not record success" + ); + assert!( + f_post > f_pre, + "http chat streaming 5xx must record at least one failure. \ + failures {}→{}", + f_pre, + f_post + ); + + ctx.shutdown().await; + } + + /// OpenAIRouter streaming 5xx must record `record_failure`, not success. + /// Guards the `mark_errored()` pre-tag on the streaming branch of + /// `OpenAIRouter::route_chat_completions`. + #[tokio::test] + async fn test_openai_chat_streaming_5xx_records_failure() { + let worker_port = 20291; + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4291) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::flaky(worker_port, 1.0)], + ) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "x"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let _ = resp.into_body().collect().await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + s_post - s_pre, + 0, + "openai chat streaming 5xx must not record success" + ); + assert!( + f_post > f_pre, + "openai chat streaming 5xx must record at least one failure. \ + failures {}→{}", + f_pre, + f_post + ); + + ctx.shutdown().await; + } + + /// PD decode 5xx on streaming request: decode breaker records failure, + /// not success. Guards the `mark_errored()` pre-tag in + /// `PDRouter::create_streaming_response` — the synthetic single-Ok + /// SSE envelope built by `handle_decode_error_response` would otherwise + /// terminate cleanly and record success. + #[tokio::test] + async fn test_pd_decode_streaming_5xx_records_failure() { + let prefill_port = 20292; + let decode_port = 20293; + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4292) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::prefill(prefill_port), { + let mut w = TestWorkerConfig::decode(decode_port); + w.fail_rate = 1.0; + w + }], + ) + .await; + let app = ctx.create_app().await; + let decode_url = format!("http://127.0.0.1:{}", decode_port); + let prefill_url = format!("http://127.0.0.1:{}", prefill_port); + let decode = pin_worker(&ctx, &decode_url); + let prefill = pin_worker(&ctx, &prefill_url); + let (s_pre, f_pre) = breaker_counts(&decode); + let (_s_pre_prefill, f_pre_prefill) = breaker_counts(&prefill); + + let payload = json!({ "text": "x", "stream": true }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let _ = resp.into_body().collect().await; + + let (s_post, f_post) = breaker_counts(&decode); + assert_eq!( + s_post - s_pre, + 0, + "PD decode streaming 5xx must not record success" + ); + assert!( + f_post > f_pre, + "PD decode streaming 5xx must record at least one failure. \ + failures {}→{}", + f_pre, + f_post + ); + + // Healthy prefill must not be penalised when only decode returns + // 5xx. The outer dispatcher used to derive prefill's outcome + // from the synthetic 5xx response status returned by + // `handle_decode_error_response`, falsely failing prefill. + let (_s_post_prefill, f_post_prefill) = breaker_counts(&prefill); + assert_eq!( + f_post_prefill - f_pre_prefill, + 0, + "PD streaming: healthy prefill must not be penalised when only \ + decode returns 5xx. prefill failures {}→{}", + f_pre_prefill, + f_post_prefill + ); + + ctx.shutdown().await; + } + + /// PD non-streaming, decode 4xx: the decode breaker MUST NOT record a + /// failure. 4xx is a client-fault (malformed input, auth, etc.), not a + /// worker fault — the old outer dispatcher used `not_error = + /// is_success() || is_client_error()` and the streaming path's + /// `BreakerTrackedStream` pre-mark in `create_streaming_response` + /// still preserves that distinction. The early-record path added + /// for prefill misattribution must keep the same semantics for + /// decode, otherwise a client sending malformed payloads can open + /// the breaker on a healthy worker. + #[tokio::test] + async fn test_pd_decode_non_streaming_4xx_does_not_penalise_breaker() { + let prefill_port = 20313; + let decode_port = 20314; + + // Force decode's failure response to 400 (client error) instead + // of the default 500. + set_fail_status_code(decode_port, 400); + + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4313) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::prefill(prefill_port), { + let mut w = TestWorkerConfig::decode(decode_port); + w.fail_rate = 1.0; + w + }], + ) + .await; + let app = ctx.create_app().await; + let decode_url = format!("http://127.0.0.1:{}", decode_port); + let prefill_url = format!("http://127.0.0.1:{}", prefill_port); + let decode = pin_worker(&ctx, &decode_url); + let prefill = pin_worker(&ctx, &prefill_url); + let (_s_pre_decode, f_pre_decode) = breaker_counts(&decode); + let (_s_pre_prefill, f_pre_prefill) = breaker_counts(&prefill); + + // Non-streaming /generate request. + let payload = json!({ "text": "x" }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let _ = resp.into_body().collect().await; + + // Legacy semantics (preserved by the streaming path's + // `BreakerTrackedStream` pre-mark and the old outer + // `not_error = is_success() || is_client_error()` rule): a 4xx + // response is recorded as a non-fault outcome — it must NOT + // increment `total_failures`, otherwise repeated client-caused + // 400s could open the breaker on a healthy worker. Whether it + // increments `total_successes` is incidental; we only pin the + // load-bearing invariant (no failure tick). + let (_s_post_decode, f_post_decode) = breaker_counts(&decode); + assert_eq!( + f_post_decode - f_pre_decode, + 0, + "PD decode 4xx is a client fault, not a worker fault — the \ + decode breaker must not record a failure. failures {}→{}", + f_pre_decode, + f_post_decode, + ); + + // Prefill stayed healthy and must also not be penalised by a + // client-caused decode 4xx. + let (_s_post_prefill, f_post_prefill) = breaker_counts(&prefill); + assert_eq!( + f_post_prefill - f_pre_prefill, + 0, + "PD prefill must not be penalised by a decode 4xx. \ + failures {}→{}", + f_pre_prefill, + f_post_prefill, + ); + + clear_fail_status_code(decode_port); + ctx.shutdown().await; + } + + /// /v1/responses simple streaming 5xx must record `record_failure`, + /// not success. Guards the `record_failure()` in the non-success status + /// arm of `handle_simple_streaming_passthrough` and confirms the eager + /// post-status `record_success()` is gone. + #[tokio::test] + async fn test_responses_simple_streaming_5xx_records_failure() { + let worker_port = 20294; + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4294) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::flaky(worker_port, 1.0)], + ) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "input": "x", + "stream": true, + "store": false + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let _ = resp.into_body().collect().await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + s_post - s_pre, + 0, + "/responses simple 5xx must not record success" + ); + assert!( + f_post > f_pre, + "/responses simple 5xx must record at least one failure. \ + failures {}→{}", + f_pre, + f_post + ); + + ctx.shutdown().await; + } + + /// /v1/responses simple, clean stream: must record exactly one + /// success and no failures. Pins the absence of the old eager + /// `record_success()` at status-OK time (which would have produced + /// 2 successes — one eager, one on stream-end). + #[tokio::test] + async fn test_responses_simple_clean_stream_records_one_success() { + let worker_port = 20295; + let total_chunks: usize = 4; + + reset_stream_tracker(worker_port); + set_slow_stream_chunks(worker_port, total_chunks); + + let config = RouterConfig::builder() + .openai_mode(vec![format!("http://127.0.0.1:{}", worker_port)]) + .round_robin_policy() + .host("127.0.0.1") + .port(4295) + .max_payload_size(8 * 1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(32) + .queue_timeout_secs(5) + .build_unchecked(); + let ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(worker_port, 5)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + let (s_pre, f_pre) = breaker_counts(&worker); + + let payload = json!({ + "model": "mock-model", + "input": "x", + "stream": true, + "store": false + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let _ = resp.into_body().collect().await; + let _ = wait_for_stream_finish(worker_port, STREAM_FINISH_TIMEOUT).await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + (s_post - s_pre, f_post - f_pre), + (1, 0), + "/responses simple clean stream must record exactly 1 success \ + and 0 failures. successes {}→{}, failures {}→{}", + s_pre, + s_post, + f_pre, + f_post + ); + + clear_slow_stream_chunks(worker_port); + ctx.shutdown().await; + } + + /// PD generate, clean stream: decode breaker must record exactly one + /// success, prefill exactly one success, neither failure. Specifically + /// guards the PD streaming loop's `[DONE]` detection — `mark_completed()` + /// must transition the wrapper from Active to Completed so Drop ticks + /// `record_success`, not "Active = no tick". + #[tokio::test] + async fn test_pd_clean_stream_records_one_success() { + let prefill_port = 20296; + let decode_port = 20297; + let total_chunks: usize = 4; + + reset_stream_tracker(decode_port); + set_slow_stream_chunks(decode_port, total_chunks); + + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4296) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + let ctx = AppTestContext::new_with_config( + config, + vec![TestWorkerConfig::prefill(prefill_port), { + let mut w = TestWorkerConfig::decode(decode_port); + w.response_delay_ms = 5; + w + }], + ) + .await; + let app = ctx.create_app().await; + let decode_url = format!("http://127.0.0.1:{}", decode_port); + let prefill_url = format!("http://127.0.0.1:{}", prefill_port); + let decode = pin_worker(&ctx, &decode_url); + let prefill = pin_worker(&ctx, &prefill_url); + let (s_pre_decode, f_pre_decode) = breaker_counts(&decode); + let (s_pre_prefill, f_pre_prefill) = breaker_counts(&prefill); + + let payload = json!({ "text": "x", "stream": true }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let _ = resp.into_body().collect().await; + let _ = wait_for_stream_finish(decode_port, STREAM_FINISH_TIMEOUT).await; + + let (s_post_decode, f_post_decode) = breaker_counts(&decode); + assert_eq!( + (s_post_decode - s_pre_decode, f_post_decode - f_pre_decode), + (1, 0), + "PD decode clean stream must record exactly 1 success and 0 failures. \ + successes {}→{}, failures {}→{}", + s_pre_decode, + s_post_decode, + f_pre_decode, + f_post_decode + ); + let (s_post_prefill, f_post_prefill) = breaker_counts(&prefill); + assert_eq!( + ( + s_post_prefill - s_pre_prefill, + f_post_prefill - f_pre_prefill + ), + (1, 0), + "PD prefill clean stream must record exactly 1 success and 0 failures. \ + successes {}→{}, failures {}→{}", + s_pre_prefill, + s_post_prefill, + f_pre_prefill, + f_post_prefill + ); + + clear_slow_stream_chunks(decode_port); + ctx.shutdown().await; + } + + /// PD generate, prefill 5xx (decode never reached): prefill breaker + /// records failure, decode breaker untouched. Guards the prefill-only + /// failure attribution in the PD retry/dispatch path. + #[tokio::test] + async fn test_pd_prefill_5xx_records_failure() { + let prefill_port = 20298; + let decode_port = 20299; + + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4298) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + let ctx = AppTestContext::new_with_config( + config, + vec![ + { + let mut p = TestWorkerConfig::prefill(prefill_port); + p.fail_rate = 1.0; + p + }, + TestWorkerConfig::decode(decode_port), + ], + ) + .await; + let app = ctx.create_app().await; + let prefill_url = format!("http://127.0.0.1:{}", prefill_port); + let decode_url = format!("http://127.0.0.1:{}", decode_port); + let prefill = pin_worker(&ctx, &prefill_url); + let decode = pin_worker(&ctx, &decode_url); + let (s_pre_p, f_pre_p) = breaker_counts(&prefill); + let (s_pre_d, f_pre_d) = breaker_counts(&decode); + + let payload = json!({ "text": "x", "stream": true }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let _ = resp.into_body().collect().await; + + let (s_post_p, f_post_p) = breaker_counts(&prefill); + let (s_post_d, f_post_d) = breaker_counts(&decode); + assert!( + f_post_p > f_pre_p, + "PD prefill 5xx must record at least one failure. failures {}→{}", + f_pre_p, + f_post_p + ); + assert_eq!( + s_post_p - s_pre_p, + 0, + "PD prefill 5xx must not record success" + ); + assert_eq!( + (s_post_d - s_pre_d, f_post_d - f_pre_d), + (0, 0), + "PD decode must be untouched when prefill fails. \ + successes {}→{}, failures {}→{}", + s_pre_d, + s_post_d, + f_pre_d, + f_post_d + ); + + ctx.shutdown().await; + } + + /// http chat streaming, upstream connect failure BEFORE the + /// `BreakerTrackedStream` is constructed: the breaker MUST still record + /// a failure. Guards the pre-stream error arm in + /// `send_typed_request` — returning `convert_reqwest_error(e)` without + /// ticking the worker breaker would let a worker that's flapping at + /// the TCP layer remain selectable indefinitely (the streaming branch + /// skips the eager `record_outcome` on the assumption that a tracked + /// stream will fire on drop, but no tracked stream was ever installed + /// on this path). + #[tokio::test] + async fn test_http_chat_pre_stream_failure_records_breaker_streaming() { + use smg::config::RetryConfig; + + let worker_port = 20310; + + // max_retries=1 keeps the assertion exact: one attempt → one + // failure tick. Any larger value just multiplies the count. + let config = TestRouterConfig::round_robin_with_retry( + 4310, + RetryConfig { + max_retries: 1, + initial_backoff_ms: 10, + max_backoff_ms: 50, + backoff_multiplier: 1.0, + jitter_factor: 0.0, + }, + ); + let mut ctx = + AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(worker_port)]) + .await; + let app = ctx.create_app().await; + let worker_url = format!("http://127.0.0.1:{}", worker_port); + let worker = pin_worker(&ctx, &worker_url); + let (s_pre, f_pre) = breaker_counts(&worker); + + // Stop the worker AFTER startup health check has marked it + // healthy. The periodic health checker isn't spawned in + // AppTestContext setups (it's started in `server.rs`), so + // `is_healthy()` stays true and the worker remains selectable. + // The next streaming request will fail at TCP connect → reqwest + // returns Err → `convert_reqwest_error` synthesises a 5xx + // Response without any `BreakerTrackedStream` ever wrapping the + // body. + ctx.workers[0].stop().await; + + let payload = json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "x"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert!( + resp.status().is_server_error(), + "Expected 5xx after upstream connect failure, got {}", + resp.status() + ); + let _ = resp.into_body().collect().await; + + let (s_post, f_post) = breaker_counts(&worker); + assert_eq!( + (s_post - s_pre, f_post - f_pre), + (0, 1), + "http chat streaming: pre-stream upstream failure must record \ + exactly one breaker failure (no tracked stream was installed, \ + so the deferred-record path doesn't fire). \ + successes {}→{}, failures {}→{}", + s_pre, + s_post, + f_pre, + f_post + ); + + ctx.shutdown().await; + } + + /// PD streaming, decode connect failure BEFORE the + /// `BreakerTrackedStream` is constructed: decode breaker MUST record + /// a failure. Guards `pd_router.rs`'s pre-stream error arm — returning + /// `error::bad_gateway` without ticking the decode breaker would let a + /// decode worker that's flapping at the TCP layer remain selectable + /// indefinitely (the streaming branch skips the eager `record_outcome` + /// on the assumption that a tracked stream will fire on drop, but no + /// tracked stream was ever installed on this path). + #[tokio::test] + async fn test_pd_decode_pre_stream_failure_records_breaker_streaming() { + use smg::config::RetryConfig; + + let prefill_port = 20311; + let decode_port = 20312; + + let config = RouterConfig::builder() + .prefill_decode_mode( + vec![(format!("http://127.0.0.1:{}", prefill_port), None)], + vec![format!("http://127.0.0.1:{}", decode_port)], + ) + .round_robin_policy() + .host("127.0.0.1") + .port(4311) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .retry_config(RetryConfig { + max_retries: 1, + initial_backoff_ms: 10, + max_backoff_ms: 50, + backoff_multiplier: 1.0, + jitter_factor: 0.0, + }) + .build_unchecked(); + let mut ctx = AppTestContext::new_with_config( + config, + vec![ + TestWorkerConfig::prefill(prefill_port), + TestWorkerConfig::decode(decode_port), + ], + ) + .await; + let app = ctx.create_app().await; + let decode_url = format!("http://127.0.0.1:{}", decode_port); + let prefill_url = format!("http://127.0.0.1:{}", prefill_port); + let decode_worker = pin_worker(&ctx, &decode_url); + let prefill_worker = pin_worker(&ctx, &prefill_url); + let (s_pre_decode, f_pre_decode) = breaker_counts(&decode_worker); + let (_s_pre_prefill, f_pre_prefill) = breaker_counts(&prefill_worker); + + // Stop ONLY the decode worker (index 1; prefill was registered + // first). Prefill stays up so its half of the tokio::join! send + // succeeds — the test specifically exercises the + // "decode_result is Err" arm in `execute_dual_dispatch_internal`. + ctx.workers[1].stop().await; + + let payload = json!({ "text": "x", "stream": true }); + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert!( + resp.status().is_server_error(), + "Expected 5xx after decode connect failure, got {}", + resp.status() + ); + let _ = resp.into_body().collect().await; + + let (s_post_decode, f_post_decode) = breaker_counts(&decode_worker); + assert!( + f_post_decode > f_pre_decode, + "PD streaming: pre-stream decode failure must record at least \ + one breaker failure on the decode worker (no tracked stream \ + was installed). failures {}→{}", + f_pre_decode, + f_post_decode + ); + assert_eq!( + s_post_decode - s_pre_decode, + 0, + "PD streaming pre-stream decode failure must not record a \ + success on the decode breaker. successes {}→{}", + s_pre_decode, + s_post_decode + ); + + // Prefill stayed up and its `send()` returned 2xx. The decode + // connect failure must NOT be misattributed to prefill — the + // outer dispatcher used to record `prefill.record_outcome(false)` + // based on the final 502 response status, penalising a healthy + // worker for its peer's failure. + let (_s_post_prefill, f_post_prefill) = breaker_counts(&prefill_worker); + assert_eq!( + f_post_prefill - f_pre_prefill, + 0, + "PD streaming: healthy prefill must not be penalised when only \ + decode fails. prefill failures {}→{}", + f_pre_prefill, + f_post_prefill + ); + + ctx.shutdown().await; + } +}