feature: upstream cancel (#19524)

Co-authored-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

.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:

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

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<Result<Bytes, BenchErr>> {
+    let chunk = Bytes::from(vec![0u8; chunk_size]);
+    (0..n).map(|_| Ok::<_, BenchErr>(chunk.clone())).collect()
+}
+
+fn make_worker() -> Arc<dyn Worker> {
+    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::<Result<Bytes, BenchErr>>();
+            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);

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};

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<HeaderMap>,
 }
 
+/// 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::<BreakerOutcomesRecorded>()
+                            .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::<Value>(&error_body) {
+                Ok(error_body) => match serde_json::from_slice::<Value>(&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<Value>,
         return_logprob: bool,
-        decode_url: Option<String>,
         headers: Option<HeaderMap>,
         prefill: Arc<dyn Worker>,
         decode: Arc<dyn Worker>,
@@ -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(),
             );

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<dyn Worker>,
         is_stream: bool,
         load_guard: Option<WorkerLoadGuard>,
     ) -> 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;

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;

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<dyn crate::core::Worker>,
     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::<Result<Bytes, io::Error>>();
 
@@ -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<dyn crate::core::Worker>,
     headers: Option<&HeaderMap>,
     req: StreamingRequest,
     active_mcp: &Arc<smg_mcp::McpManager>,
@@ -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(&current_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!("<body read failed: {e}>"));
+                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,

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()

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<Item = Result<Bytes, E>>` 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<E = reqwest::Error> {
+    inner: Pin<Box<dyn Stream<Item = Result<Bytes, E>> + Send + 'static>>,
+    worker: Arc<dyn Worker>,
+    log_url: String,
+    terminal: Terminal,
+}
+
+impl<E> BreakerTrackedStream<E> {
+    pub fn new<S>(inner: S, worker: Arc<dyn Worker>, log_url: String) -> Self
+    where
+        S: Stream<Item = Result<Bytes, E>> + 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<E: Display> Stream for BreakerTrackedStream<E> {
+    type Item = Result<Bytes, E>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        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<E> Drop for BreakerTrackedStream<E> {
+    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<dyn Worker> {
+        Arc::new(BasicWorkerBuilder::new("http://test-worker").build())
+    }
+
+    fn breaker_counters(w: &Arc<dyn Worker>) -> (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::<Result<Bytes, TestErr>>();
+        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::<Bytes, _>(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::<Result<Bytes, TestErr>>();
+        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::<Bytes, _>(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));
+    }
+}

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::<f32>() < 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<Arc<RwLock<MockWorkerConfig>>>) -> 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::<Result<Event, std::io::Error>>(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::<Result<Event, std::io::Error>>(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::<Result<Event, std::io::Error>>(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::<String>();
+                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<Mutex<HashMap<u16, usize>>> = OnceLock::new();
+
+fn get_slow_stream_config() -> &'static Mutex<HashMap<u16, usize>> {
+    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<usize> {
+    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<Mutex<HashMap<u16, usize>>> = OnceLock::new();
+
+fn get_stream_error_after_config() -> &'static Mutex<HashMap<u16, usize>> {
+    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<usize> {
+    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<Mutex<HashMap<u16, u16>>> = OnceLock::new();
+
+fn get_fail_status_code_config() -> &'static Mutex<HashMap<u16, u16>> {
+    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<u16> {
+    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<Mutex<HashMap<u16, StreamTrackingState>>> = OnceLock::new();
+
+fn get_stream_tracker() -> &'static Mutex<HashMap<u16, StreamTrackingState>> {
+    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<Mutex<HashMap<u16, Arc<Notify>>>> = OnceLock::new();
+
+fn get_stream_finish_notifier_map() -> &'static Mutex<HashMap<u16, Arc<Notify>>> {
+    STREAM_FINISH_NOTIFIERS.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+fn get_stream_finish_notifier(port: u16) -> Arc<Notify> {
+    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<Notify>);
+
+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<StreamTrackingState> {
+    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<StreamTrackingState> {
+    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<Mutex<HashMap<u16, HashSet<String>>>> = OnceLock::new();
 

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;