diff --git a/experimental/sgl-router/monitoring/README.md b/experimental/sgl-router/monitoring/README.md index d792ff871..b001b6fe1 100644 --- a/experimental/sgl-router/monitoring/README.md +++ b/experimental/sgl-router/monitoring/README.md @@ -25,6 +25,7 @@ The dashboard graphs every family the router emits: | `sgl_router_worker_requests_total` | Counter | Per-worker **dispatches** by `worker_url`, `model_id`, `mode`, `outcome` (recorded after dispatch; blind to pre-dispatch drops) | | `sgl_router_request_duration_seconds` | Histogram | End-to-end request latency by `model_id` | | `sgl_router_ttft_seconds` | Histogram | Time to first token (streaming) by `model_id` | +| `sgl_router_stream_outcome_total` | Counter | Streaming outcomes by `worker_url`, `model_id`, and `outcome` (`ok`, `stream_error_event`, `upstream_error`, or `client_disconnect`). Counts committed 2xx streams only — non-2xx responses are counted by status in `responses_total` | | `sgl_router_active_load` | Gauge | Per-worker prefill-token / decode-block load | | `sgl_router_workers` | Gauge | Registered worker count by `mode` | | `sgl_router_worker_health` | Gauge | Per-worker health (1=breaker admits, 0=open) | diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index b01423e9d..7477f6f3d 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -164,7 +164,7 @@ impl Proxy { /// for the full streaming lifetime — without which a long-running SSE /// response would under-report load. // Each parameter is a distinct, required input to a single upstream - // forward (target, breaker, path, headers, body, plus the two + // forward (target, breaker, path, headers, body, plus the // streaming-lifetime callbacks). Bundling them into a struct purely to // satisfy the arg-count heuristic would add indirection without clarity. #[allow(clippy::too_many_arguments)] @@ -177,6 +177,7 @@ impl Proxy { body: Bytes, stream_guards: Option>, on_first_byte: Option>, + on_stream_end: Option>, ) -> Result, ApiError> { if !breaker.allow() { return Err(ApiError::BreakerOpen { @@ -217,23 +218,30 @@ impl Proxy { // is recorded as a failure. For 5xx headers we record_failure // up front and skip the pump hook (the body we surface is the // error response — its stream completing is not a worker win). - let on_complete: Option> = + let caller_end_hook = if status.is_success() { + on_stream_end + } else { + None + }; + let on_complete: Option> = if status.is_server_error() { breaker.record_failure(); None } else { let breaker_for_hook = Arc::clone(breaker); - Some(Box::new(move |ok| { - if ok { + Some(Box::new(move |end| { + if end.transport_ok { breaker_for_hook.record_success(); } else { breaker_for_hook.record_failure(); } + if let Some(hook) = caller_end_hook { + hook(end); + } })) }; - // Only record TTFT for successful streams — a 4xx/5xx error body - // streaming back is not a generated token, so drop the hook for - // non-2xx responses. + // Only record TTFT for successful streams; error-body chunks are not + // generated tokens. let first_byte_hook = if status.is_success() { on_first_byte } else { diff --git a/experimental/sgl-router/src/proxy/sse.rs b/experimental/sgl-router/src/proxy/sse.rs index 01a83f01f..700c5836a 100644 --- a/experimental/sgl-router/src/proxy/sse.rs +++ b/experimental/sgl-router/src/proxy/sse.rs @@ -11,6 +11,54 @@ use bytes::Bytes; use futures::{FutureExt, StreamExt}; use tokio_stream::wrappers::ReceiverStream; +/// How the SSE pump ended, reported to the `on_complete` hook. +#[derive(Debug, Clone, Copy)] +pub struct StreamEnd { + /// No upstream stream error and no pump panic. + pub transport_ok: bool, + /// An SSE error event (`data: {"error"...}`) rode the stream. + pub saw_error_event: bool, + /// The client dropped the response body before upstream finished. + pub client_disconnect: bool, +} + +/// A `data:` line whose payload's first JSON key is `error` — tolerant of +/// SSE-legal framing variants (no space after `data:`, whitespace after `{`), +/// so the match is anchored to the spec rather than one serializer's bytes. +fn is_error_event_line(line: &[u8]) -> bool { + line.strip_prefix(b"data:") + .map(|p| p.trim_ascii_start()) + .and_then(|p| p.strip_prefix(b"{")) + .map(|p| p.trim_ascii_start()) + .is_some_and(|p| p.starts_with(b"\"error\"")) +} + +/// Line-start bytes that suffice to decide `is_error_event_line`. +const LINE_PROBE: usize = 32; + +/// Finds error events emitted after an SSE response commits a 200. +/// Line-anchored, so lookalike text inside event payloads cannot match. +#[derive(Default)] +struct ErrorEventScanner { + line_start: Vec, +} + +impl ErrorEventScanner { + fn feed(&mut self, chunk: &[u8]) -> bool { + let mut hit = false; + for (i, segment) in chunk.split(|&b| b == b'\n').enumerate() { + if i > 0 { + hit |= is_error_event_line(&self.line_start); + self.line_start.clear(); + } + let room = LINE_PROBE - self.line_start.len(); + self.line_start + .extend_from_slice(&segment[..segment.len().min(room)]); + } + hit + } +} + /// Bridge a byte stream into an axum Body that streams chunks unchanged. /// /// Spawns one tokio task per stream so the handler can return immediately. @@ -47,14 +95,8 @@ use tokio_stream::wrappers::ReceiverStream; /// guard scope). /// /// # Completion hook -/// When `on_complete` is `Some`, the closure runs exactly once when the -/// pump task finishes. The bool argument is `true` on clean stream end -/// (including a clean client disconnect after at least the headers -/// landed cleanly), `false` on upstream stream error or pump panic. -/// `forward_streaming_to` passes a closure that records the worker's -/// circuit-breaker outcome — without this hook, a worker that returns -/// 2xx headers and then drops the stream mid-flight would stay credited -/// as healthy. +/// When `on_complete` is `Some`, it runs exactly once when the pump task +/// finishes with the transport, SSE error-event, and client-disconnect state. /// /// # First-byte hook /// When `on_first_byte` is `Some`, the closure runs exactly once, the moment @@ -65,7 +107,7 @@ use tokio_stream::wrappers::ReceiverStream; pub fn bytes_stream_to_body( stream: S, stream_guards: Option>, - on_complete: Option>, + on_complete: Option>, on_first_byte: Option>, ) -> Body where @@ -75,11 +117,14 @@ where let (tx, rx) = tokio::sync::mpsc::channel(64); tokio::spawn(async move { let tx_for_panic = tx.clone(); - // Capture the pump's outcome so we can report it through `on_complete` - // AFTER `pump.catch_unwind()` settles. The closure inside owns - // `outcome_setter`; the outer scope reads `outcome_holder` once. - let outcome_holder = Arc::new(parking_lot::Mutex::new(true)); - let outcome_setter = Arc::clone(&outcome_holder); + let outcome = Arc::new(parking_lot::Mutex::new(StreamEnd { + transport_ok: true, + saw_error_event: false, + client_disconnect: false, + })); + let outcome_setter = Arc::clone(&outcome); + // `None` once an error event is found — the scan is done for good. + let mut scanner = Some(ErrorEventScanner::default()); let pump = AssertUnwindSafe(async move { // Hold the guards for the task's lifetime — dropped when this // block exits (stream done or client disconnect). Leading @@ -95,17 +140,19 @@ where std::io::Error::other(msg) }); let is_err_chunk = item.is_err(); - // Fire the time-to-first-token hook on the first successful - // chunk from upstream. `take()` makes it fire at most once; - // an error-first stream never produced a token, so it's left - // unfired (and dropped on task end). - if !is_err_chunk { - if let Some(hook) = on_first_byte.take() { - hook(); + match &item { + Ok(bytes) => { + // TTFT hook: at most once (`take()`); an error-first + // stream never produced a token, so it stays unfired. + if let Some(hook) = on_first_byte.take() { + hook(); + } + if scanner.as_mut().is_some_and(|scanner| scanner.feed(bytes)) { + outcome_setter.lock().saw_error_event = true; + scanner = None; + } } - } - if is_err_chunk { - *outcome_setter.lock() = false; + Err(_) => outcome_setter.lock().transport_ok = false, } if tx.send(item).await.is_err() { // Receiver dropped. If we were about to ship an upstream @@ -114,6 +161,7 @@ where // not a router-side fault. if !is_err_chunk { tracing::debug!("SSE client disconnected mid-stream"); + outcome_setter.lock().client_disconnect = true; } break; } @@ -139,8 +187,9 @@ where .await; } if let Some(hook) = on_complete { - let ok = !panicked && *outcome_holder.lock(); - hook(ok); + let mut end = *outcome.lock(); + end.transport_ok &= !panicked; + hook(end); } }); Body::from_stream(ReceiverStream::new(rx)) @@ -394,4 +443,140 @@ mod tests { "pump drained the entire upstream after client disconnect ({final_polls} polls); the break-on-tx.send-err path is dead" ); } + + #[test] + fn error_event_scanner_detects_engine_error_event() { + let mut scanner = ErrorEventScanner::default(); + assert!(!scanner.feed(b"data: {\"choices\": [{\"delta\": {\"content\": \"hi\"}}]}\n\n")); + assert!( + scanner.feed(b"data: {\"error\": {\"message\": \"queue is full\", \"code\": 503}}\n\n") + ); + } + + #[test] + fn error_event_scanner_detects_error_split_across_chunks() { + let mut scanner = ErrorEventScanner::default(); + assert!(!scanner.feed(b"data: {\"err")); + assert!(scanner.feed(b"or\": {\"code\": 503}}\n\n")); + } + + #[test] + fn error_event_scanner_ignores_error_text_inside_content() { + let mut scanner = ErrorEventScanner::default(); + assert!(!scanner.feed( + b"data: {\"choices\": [{\"delta\": {\"content\": \"data: {\\\"error\\\" is how it looks\"}}]}\n\n", + )); + assert!(!scanner.feed(b"data: [DONE]\n\n")); + } + + #[test] + fn error_event_scanner_accepts_sse_framing_variants() { + for event in [ + &b"data:{\"error\": {\"code\": 503}}\n\n"[..], + b"data: { \"error\": {\"code\": 503}}\n\n", + b"data: {\"error\": \"queue full\"}\n\n", + ] { + assert!( + ErrorEventScanner::default().feed(event), + "missed variant: {}", + String::from_utf8_lossy(event) + ); + } + } + + #[test] + fn error_event_scanner_bounds_line_buffer() { + let mut scanner = ErrorEventScanner::default(); + let big = vec![b'x'; 1 << 20]; + assert!(!scanner.feed(&big)); + assert_eq!(scanner.line_start.len(), LINE_PROBE); + } + + fn body_with_completion( + chunks: Vec>, + ) -> (Body, tokio::sync::oneshot::Receiver) { + let (tx, rx) = tokio::sync::oneshot::channel(); + let body = bytes_stream_to_body( + stream::iter(chunks), + None, + Some(Box::new(move |end| { + let _ = tx.send(end); + })), + None, + ); + (body, rx) + } + + async fn stream_end(rx: tokio::sync::oneshot::Receiver) -> StreamEnd { + rx.await.expect("completion hook dropped") + } + + #[tokio::test] + async fn completion_reports_error_event() { + let chunks = vec![ + Ok::(Bytes::from_static(b"data: {\"err")), + Ok(Bytes::from_static(b"or\": {\"code\": 503}}\n\n")), + ]; + let (body, completion) = body_with_completion(chunks); + let _ = body.collect().await.unwrap(); + let end = stream_end(completion).await; + assert!(end.transport_ok); + assert!(end.saw_error_event); + assert!(!end.client_disconnect); + } + + #[tokio::test] + async fn completion_reports_error_event_then_transport_error() { + let chunks = vec![ + Ok(Bytes::from_static( + b"data: {\"error\": {\"code\": 503}}\n\n", + )), + Err(std::io::Error::other("connection reset")), + ]; + let (body, completion) = body_with_completion(chunks); + let _ = body.collect().await; + let end = stream_end(completion).await; + assert!(!end.transport_ok); + assert!(end.saw_error_event); + } + + #[tokio::test] + async fn completion_reports_upstream_error() { + let chunks = vec![Err(std::io::Error::other("upstream failed"))]; + let (body, completion) = body_with_completion(chunks); + let _ = body.collect().await; + let end = stream_end(completion).await; + assert!(!end.transport_ok); + assert!(!end.saw_error_event); + assert!(!end.client_disconnect); + } + + #[tokio::test] + async fn completion_reports_clean_end() { + let chunks = vec![Ok::(Bytes::from_static( + b"data: [DONE]\n\n", + ))]; + let (body, completion) = body_with_completion(chunks); + let _ = body.collect().await.unwrap(); + let end = stream_end(completion).await; + assert!(end.transport_ok); + assert!(!end.saw_error_event); + assert!(!end.client_disconnect); + } + + #[tokio::test] + async fn completion_reports_client_disconnect() { + let chunks = std::iter::repeat_with(|| { + Ok::(Bytes::from_static(b"data: x\n\n")) + }) + .take(1000) + .collect(); + let (body, completion) = body_with_completion(chunks); + let mut stream = body.into_data_stream(); + let _ = stream.next().await; + drop(stream); + let end = stream_end(completion).await; + assert!(end.transport_ok); + assert!(end.client_disconnect); + } } diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index ea0d536a1..0c2c86e08 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -23,6 +23,7 @@ //! | `sgl_router_worker_requests_total` | Counter | `worker_url`, `model_id`, `mode`, `outcome` | //! | `sgl_router_request_duration_seconds` | Histogram | `model_id` | //! | `sgl_router_ttft_seconds` | Histogram | `model_id` | +//! | `sgl_router_stream_outcome_total` | Counter | `worker_url`, `model_id`, `outcome` | //! | `sgl_router_active_load` | Gauge | `worker_url`, `kind` | //! | `sgl_router_workers` | Gauge | `mode` | //! | `sgl_router_worker_health` | Gauge | `worker_url` | @@ -49,6 +50,7 @@ //! The exposition is text/plain; version=0.0.4 per the Prometheus spec. use crate::config::PolicyKind; +use crate::proxy::sse::StreamEnd; use parking_lot::Mutex; use std::collections::HashMap; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; @@ -104,6 +106,39 @@ impl RequestOutcome { } } +/// Final outcome of a 2xx SSE stream. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamOutcome { + /// Stream ended without errors. + Ok, + /// The engine sent a `data: {"error"...}` SSE event. + StreamErrorEvent, + /// The upstream byte stream failed. + UpstreamError, + /// The client disconnected before the stream finished. + ClientDisconnect, +} + +pub(crate) fn classify_stream_end(end: StreamEnd) -> StreamOutcome { + match (end.transport_ok, end.saw_error_event, end.client_disconnect) { + (false, _, _) => StreamOutcome::UpstreamError, + (_, true, _) => StreamOutcome::StreamErrorEvent, + (_, _, true) => StreamOutcome::ClientDisconnect, + _ => StreamOutcome::Ok, + } +} + +impl StreamOutcome { + fn as_str(self) -> &'static str { + match self { + Self::Ok => "ok", + Self::StreamErrorEvent => "stream_error_event", + Self::UpstreamError => "upstream_error", + Self::ClientDisconnect => "client_disconnect", + } + } +} + /// Worker dispatch mode label — narrowed to the three modes the policy /// resolver distinguishes. The `Plain` variant covers the non-PD case. #[derive(Debug, Clone, Copy)] @@ -233,6 +268,7 @@ pub struct MetricsRegistry { // on `worker_requests_total` / the worker gauges instead. request_duration: Mutex>, ttft_seconds: Mutex>, + stream_outcome_total: Mutex>>, active_load: Mutex>>, stale_requests_total: Mutex>>, decode_affinity_total: Mutex>>, @@ -271,6 +307,14 @@ struct EdgeResponseKey { status_code: u16, } +/// Labels for `sgl_router_stream_outcome_total`. +#[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Clone)] +struct StreamOutcomeKey { + worker_url: String, + model_id: String, + outcome: &'static str, +} + /// Per-worker state sampled from the [`crate::workers::WorkerRegistry`] at /// scrape time and rendered as the `sgl_router_workers` / /// `sgl_router_worker_*` gauge families. Built by the `/metrics` route from @@ -429,6 +473,22 @@ impl MetricsRegistry { hist.observe(seconds); } + /// Record the final outcome of a 2xx stream. + pub fn record_stream_outcome(&self, worker_url: &str, model_id: &str, outcome: StreamOutcome) { + let key = StreamOutcomeKey { + worker_url: worker_url.to_owned(), + model_id: model_id.to_owned(), + outcome: outcome.as_str(), + }; + let counter = self + .stream_outcome_total + .lock() + .entry(key) + .or_default() + .clone(); + counter.fetch_add(1, Ordering::Relaxed); + } + /// Bump the edge counter `responses_total{route,method,status_code}`. Called /// at the middleware, so it captures every outcome — incl. early-exit /// 400/413/503 that the old per-handler site skipped. @@ -684,6 +744,26 @@ impl MetricsRegistry { } drop(guard); + // Final outcomes observed after a 2xx stream's headers are committed. + out.push_str("# HELP sgl_router_stream_outcome_total Final outcome of a 2xx stream.\n"); + out.push_str("# TYPE sgl_router_stream_outcome_total counter\n"); + let guard = self.stream_outcome_total.lock(); + let mut entries: Vec<(&StreamOutcomeKey, u64)> = guard + .iter() + .map(|(k, v)| (k, v.load(Ordering::Relaxed))) + .collect(); + entries.sort(); + for (key, value) in entries { + out.push_str(&format!( + "sgl_router_stream_outcome_total{{worker_url=\"{}\",model_id=\"{}\",outcome=\"{}\"}} {}\n", + escape_label(&key.worker_url), + escape_label(&key.model_id), + key.outcome, + value, + )); + } + drop(guard); + // responses_total — edge, by route/method/status (incl. early-exit 400/413/503) out.push_str( "# HELP sgl_router_responses_total Responses returned at the router HTTP edge, by route, method and HTTP status code.\n", @@ -1011,6 +1091,13 @@ fn escape_label(s: &str) -> String { mod tests { use super::*; + fn assert_metric_line(output: &str, expected: &str) { + assert!( + output.lines().any(|line| line == expected), + "missing metric line `{expected}`; got:\n{output}" + ); + } + #[test] fn empty_registry_renders_only_help_lines() { let reg = MetricsRegistry::new(); @@ -1151,6 +1238,59 @@ mod tests { } } + #[test] + fn stream_outcome_precedence() { + use StreamOutcome::*; + + for (transport_ok, saw_error_event, client_disconnect, expected) in [ + (false, false, false, UpstreamError), + (false, false, true, UpstreamError), + (false, true, false, UpstreamError), + (false, true, true, UpstreamError), + (true, false, false, Ok), + (true, false, true, ClientDisconnect), + (true, true, false, StreamErrorEvent), + (true, true, true, StreamErrorEvent), + ] { + let end = StreamEnd { + transport_ok, + saw_error_event, + client_disconnect, + }; + assert_eq!(classify_stream_end(end), expected, "{end:?}"); + } + } + + #[test] + fn record_stream_outcome_emits_labelled_counter_lines() { + let reg = MetricsRegistry::new(); + reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::Ok); + reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::Ok); + reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::StreamErrorEvent); + reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::UpstreamError); + reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::ClientDisconnect); + let out = reg.render(); + for expected in [ + r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="ok"} 2"#, + r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="stream_error_event"} 1"#, + r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="upstream_error"} 1"#, + r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="client_disconnect"} 1"#, + ] { + assert_metric_line(&out, expected); + } + } + + #[test] + fn stream_outcome_absent_until_recorded() { + let reg = MetricsRegistry::new(); + let out = reg.render(); + assert!(out.contains("# TYPE sgl_router_stream_outcome_total counter")); + assert!( + !out.contains("sgl_router_stream_outcome_total{"), + "no series until an outcome is recorded; got:\n{out}", + ); + } + #[test] fn record_response_counts_by_route_method_status_code() { let reg = MetricsRegistry::new(); diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index ec311d0ce..1f1dadcf6 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -17,11 +17,12 @@ use crate::policies::{ request_tokens_for, ExternalPrefixSignal, PrefillProposal, ProposalKind, RequestTokens, SelectionContext, }; +use crate::proxy::sse::StreamEnd; use crate::server::app_context::AppContext; use crate::server::error::ApiError; use crate::server::metrics::{ - MetricsRegistry, PolicySelectionFailureReason, RequestOutcome, StaleRequestOutcome, - WorkerModeLabel, + classify_stream_end, MetricsRegistry, PolicySelectionFailureReason, RequestOutcome, + StaleRequestOutcome, WorkerModeLabel, }; use crate::workers::{LoadGuard, Worker}; use axum::body::Body; @@ -764,6 +765,16 @@ pub async fn chat_completions( start, }; + // Classifies a 2xx stream after its headers are committed. Takes the + // streaming worker's URL (Final D in PD mode). + let make_stream_end_hook = |worker_url: String| -> Box { + let metrics = Arc::clone(&ctx.metrics); + let model = metrics_model.clone(); + Box::new(move |end| { + metrics.record_stream_outcome(&worker_url, &model, classify_stream_end(end)); + }) + }; + // Forward the router-computed tokens to the engine as `input_ids` so it // skips re-tokenizing the same prompt — but only when they are // engine-equivalent (chat-encoder path) AND the request contains nothing @@ -904,6 +915,7 @@ pub async fn chat_completions( outgoing_body, Some(stream_guards), Some(make_ttft_hook()), + Some(make_stream_end_hook(decode_worker.url.clone())), ); tokio::select! { biased; @@ -939,6 +951,7 @@ pub async fn chat_completions( outgoing_body, Some(stream_guards), Some(make_ttft_hook()), + Some(make_stream_end_hook(worker.url.clone())), ); // Bias `fetch` over the cancellation branch: a successful // response that completes in the same poll as the token firing diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index eec019aa0..185a6bc64 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -950,6 +950,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() { body, None, None, + None, ) .await; @@ -1319,35 +1320,20 @@ async fn streaming_active_load_drops_on_client_disconnect() { Duration::from_millis(100), ) .await; - let ctx = build_ctx_with_worker(&worker.url); + let (ctx, body) = stream_chat(&worker.url).await; let active_load = Arc::clone(&ctx.active_load); - let app = build_router(ctx); - - let req = Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header("content-type", "application/json") - .body(Body::from( - serde_json::to_vec(&serde_json::json!({ - "model": "tiny", - "messages": [{"role": "user", "content": "hi"}], - "stream": true, - })) - .unwrap(), - )) - .unwrap(); - let res = app.oneshot(req).await.unwrap(); // Read one chunk to confirm the stream is live, then drop the body. use futures::StreamExt; - let mut data_stream = res.into_body().into_data_stream(); + let mut data_stream = body.into_data_stream(); let _first = data_stream.next().await; drop(data_stream); - // Wait long enough for the SSE pump to notice the receiver-drop and - // exit (per `bytes_stream_to_body_breaks_on_client_disconnect` test - // in sse.rs, that takes well under 200 ms). - tokio::time::sleep(Duration::from_millis(300)).await; + let expected = format!( + r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="client_disconnect"}} 1"#, + worker.url, + ); + wait_for_metric(&ctx, &expected).await; assert_eq!( active_load.inflight_count(), @@ -1477,3 +1463,111 @@ async fn non_streaming_error_path_drops_active_load_guard() { "error path must drop the active-load guard", ); } + +fn has_metric_line(metrics: &str, expected: &str) -> bool { + metrics.lines().any(|line| line == expected) +} + +/// Send a streaming request and wait for the expected metric. +async fn stream_chat_and_render( + worker_url: &str, + expected_metric: &str, +) -> (Arc, String) { + let (ctx, body) = stream_chat(worker_url).await; + body.collect().await.unwrap(); + let metrics = wait_for_metric(&ctx, expected_metric).await; + (ctx, metrics) +} + +async fn stream_chat(worker_url: &str) -> (Arc, Body) { + let ctx = build_ctx_with_worker(worker_url); + let app = build_router(ctx.clone()); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": true + }) + .to_string(), + )) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + (ctx, res.into_body()) +} + +async fn wait_for_metric(ctx: &AppContext, expected_metric: &str) -> String { + let deadline = std::time::Instant::now() + Duration::from_secs(2); + loop { + let metrics = ctx.metrics.render(); + if has_metric_line(&metrics, expected_metric) { + return metrics; + } + assert!( + std::time::Instant::now() < deadline, + "timed out waiting for `{expected_metric}`; got:\n{metrics}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +/// A post-200 SSE error event is classified without affecting routing health. +#[tokio::test] +async fn streaming_error_event_records_outcome_without_tripping_breaker() { + let worker = crate::common::mock_worker::MockWorker::start(vec![ + "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n", + "data: {\"error\": {\"message\": \"The request queue is full.\", \"code\": 503}}\n\n", + "data: [DONE]\n\n", + ]) + .await; + let expected = format!( + r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="stream_error_event"}} 1"#, + worker.url, + ); + let (ctx, metrics) = stream_chat_and_render(&worker.url, &expected).await; + assert!(has_metric_line( + &metrics, + r#"sgl_router_responses_total{route="/v1/chat/completions",method="POST",status_code="200"} 1"# + )); + assert!( + ctx.registry + .all() + .iter() + .all(|worker| worker.breaker.would_allow()), + "SSE error event must not trip the circuit breaker", + ); +} + +#[tokio::test] +async fn streaming_clean_completion_records_ok() { + let worker = crate::common::mock_worker::MockWorker::start(vec![ + "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n", + "data: [DONE]\n\n", + ]) + .await; + let expected = format!( + r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="ok"}} 1"#, + worker.url, + ); + stream_chat_and_render(&worker.url, &expected).await; +} + +#[tokio::test] +async fn streaming_error_event_then_transport_failure_records_upstream_error() { + let worker = crate::common::mock_worker::MockWorker::start_returning_partial_body( + StatusCode::OK, + b"data: {\"error\": {\"code\": 503}}\n\n", + ) + .await; + let (ctx, body) = stream_chat(&worker.url).await; + assert!(body.collect().await.is_err()); + let expected = format!( + r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="upstream_error"}} 1"#, + worker.url, + ); + wait_for_metric(&ctx, &expected).await; +}