From 3cef6960764ebe9563a24b6654c2f64eba832828 Mon Sep 17 00:00:00 2001 From: Nick DiZazzo Date: Tue, 18 Aug 2026 02:45:48 -0400 Subject: [PATCH 1/5] fix(skippy): stop orphaned generation on a stalled SSE channel mpsc::Sender::blocking_send waits indefinitely for buffer space. A consumer that stops draining without the channel ever being dropped -- a client gone before the server's own disconnect detection notices, e.g. behind a proxy that doesn't propagate the close -- can pin a generation worker, and the execution lane it holds, forever, even after the request is cancelled. This matches the report's "orphaned generation after client disconnect" defect: GPU stayed busy for 3+ minutes and other requests failed with a lane timeout. Replace the six ad hoc blocking_send call sites in run_generation_stream with one send_generation_event helper that polls tx.try_send instead: it keeps retrying while the request is live and the buffer is merely full, but gives up -- cancelling the request if it hasn't been already -- once the caller is cancelled, the receiver is gone, or the buffer has stayed full past the same window other requests already wait for an execution lane (GENERATION_ADMISSION_TIMEOUT, 10s). That last case covers the report's open question directly: it no longer matters whether the SSE body's Drop fires promptly on disconnect, since a stalled channel self-cancels on its own. Added two regression tests: one proves a stalled-but-cancelled send no longer hangs, the other proves a stalled send with no external cancel at all still self-cancels once past the stall timeout (the stall timeout is injectable for the test so it doesn't wait out the real 10s). --- crates/skippy-server/src/frontend/backend.rs | 92 ++++++++++++++++--- .../src/frontend/backend/tests.rs | 76 +++++++++++++++ 2 files changed, 154 insertions(+), 14 deletions(-) diff --git a/crates/skippy-server/src/frontend/backend.rs b/crates/skippy-server/src/frontend/backend.rs index b8b0e901ab..202281a707 100644 --- a/crates/skippy-server/src/frontend/backend.rs +++ b/crates/skippy-server/src/frontend/backend.rs @@ -55,6 +55,7 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::thread; use std::time::Duration; use std::time::Instant; use tokio::sync::OwnedSemaphorePermit; @@ -67,6 +68,67 @@ fn request_cancelled_error() -> OpenAiError { OpenAiError::cancelled("request cancelled") } +/// How long a full SSE event channel is treated as merely backed up before +/// the generation worker gives up on it. Reuses the window other requests +/// already wait for an execution lane, so a stalled consumer can never pin a +/// lane longer than everyone else already times out for. +const STREAM_SEND_STALL_TIMEOUT: Duration = GENERATION_ADMISSION_TIMEOUT; + +/// Poll interval while `send_generation_event` waits for buffer space. +const STREAM_SEND_RETRY_INTERVAL: Duration = Duration::from_millis(20); + +/// Forward one generation event to the SSE channel. +/// +/// `mpsc::Sender::blocking_send` waits indefinitely for buffer space. A +/// consumer that stops draining without the channel ever being dropped -- +/// a client gone before the server's own disconnect detection notices, e.g. +/// behind a proxy that doesn't propagate the close -- can then pin the +/// generation worker, and the execution lane it holds, forever even after +/// the request is cancelled. This polls instead: it keeps retrying while the +/// request is live and the buffer is merely full, but gives up -- cancelling +/// the request if it hasn't been already -- once the caller is cancelled, +/// the receiver is gone, or the buffer has stayed full past +/// `STREAM_SEND_STALL_TIMEOUT`. +fn send_generation_event( + tx: &mpsc::Sender>, + event: OpenAiResult, + context: &OpenAiRequestContext, +) -> Result<(), OpenAiError> { + send_generation_event_with_stall_timeout(tx, event, context, STREAM_SEND_STALL_TIMEOUT) +} + +/// `send_generation_event` with the stall timeout as a parameter, so tests +/// can exercise the self-cancelling path without waiting out the real +/// `STREAM_SEND_STALL_TIMEOUT`. +fn send_generation_event_with_stall_timeout( + tx: &mpsc::Sender>, + mut event: OpenAiResult, + context: &OpenAiRequestContext, + stall_timeout: Duration, +) -> Result<(), OpenAiError> { + let stalled_since = Instant::now(); + loop { + if context.is_cancelled() { + return Err(request_cancelled_error()); + } + event = match tx.try_send(event) { + Ok(()) => return Ok(()), + Err(mpsc::error::TrySendError::Closed(_)) => { + context.cancel(); + return Err(OpenAiError::backend("stream receiver dropped")); + } + Err(mpsc::error::TrySendError::Full(rejected)) => rejected, + }; + if stalled_since.elapsed() >= stall_timeout { + context.cancel(); + return Err(OpenAiError::backend( + "stream receiver stalled past the admission timeout", + )); + } + thread::sleep(STREAM_SEND_RETRY_INTERVAL); + } +} + fn should_emit_stream_usage(request_include_usage: bool, context: &OpenAiRequestContext) -> bool { request_include_usage || context.observes_stream_usage() } @@ -849,16 +911,13 @@ impl StageOpenAiBackend { vec![GenerationStreamEvent::Delta(chunk.to_string())] }; for event in events { - tx.blocking_send(Ok(event)).map_err(|_| { - context.cancel(); - OpenAiError::backend("stream receiver dropped") - })?; + send_generation_event(&tx, Ok(event), &context)?; } Ok(()) }, ); if context.is_cancelled() { - let _ = tx.blocking_send(Err(request_cancelled_error())); + let _ = send_generation_event(&tx, Err(request_cancelled_error()), &context); return; } match result { @@ -867,15 +926,14 @@ impl StageOpenAiBackend { match parser.finish(&output.text) { Ok(events) => { for event in events { - if tx.blocking_send(Ok(event)).is_err() { - context.cancel(); + if send_generation_event(&tx, Ok(event), &context).is_err() { return; } } parser.finish_reason(output.finish_reason) } Err(error) => { - let _ = tx.blocking_send(Err(error)); + let _ = send_generation_event(&tx, Err(error), &context); return; } } @@ -883,17 +941,23 @@ impl StageOpenAiBackend { output.finish_reason }; if should_emit_stream_usage(include_usage, &context) - && tx - .blocking_send(Ok(GenerationStreamEvent::Usage(output.usage()))) - .is_err() + && send_generation_event( + &tx, + Ok(GenerationStreamEvent::Usage(output.usage())), + &context, + ) + .is_err() { - context.cancel(); return; } - let _ = tx.blocking_send(Ok(GenerationStreamEvent::Done(finish_reason))); + let _ = send_generation_event( + &tx, + Ok(GenerationStreamEvent::Done(finish_reason)), + &context, + ); } Err(error) => { - let _ = tx.blocking_send(Err(error)); + let _ = send_generation_event(&tx, Err(error), &context); } } }); diff --git a/crates/skippy-server/src/frontend/backend/tests.rs b/crates/skippy-server/src/frontend/backend/tests.rs index 926e94662d..ec675ba6c2 100644 --- a/crates/skippy-server/src/frontend/backend/tests.rs +++ b/crates/skippy-server/src/frontend/backend/tests.rs @@ -521,3 +521,79 @@ fn internal_stream_usage_observation_preserves_client_wire_preference() { let observed = OpenAiRequestContext::new().with_stream_usage_observation(); assert!(should_emit_stream_usage(false, &observed)); } + +/// Reproduces the orphaned-generation report: a client can vanish (dropped +/// connection, or one that hasn't been noticed yet -- e.g. behind a proxy +/// that doesn't propagate the close) leaving the SSE receiver alive but +/// permanently undrained. `send_generation_event` must not let that pin the +/// generation worker, and the execution lane it holds, forever: once the +/// request is cancelled it must give up promptly even though the channel +/// stays full and the receiver is never dropped. +/// +/// This runs the send on its own thread and waits for a result over a +/// bounded `recv_timeout` rather than joining directly, so a regression back +/// to an unconditional blocking send fails this test instead of hanging the +/// suite. +#[test] +fn stalled_receiver_does_not_pin_the_generation_worker_forever() { + let (tx, rx) = mpsc::channel(1); + tx.try_send(Ok(GenerationStreamEvent::Delta("first".to_owned()))) + .expect("channel has room for the first event"); + let context = OpenAiRequestContext::new(); + + let sender_context = context.clone(); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let result = send_generation_event( + &tx, + Ok(GenerationStreamEvent::Delta("second".to_owned())), + &sender_context, + ); + // Keep `rx` alive without draining it until after the send settles, + // so a fix that works only because the channel closed doesn't pass. + drop(rx); + let _ = done_tx.send(result.is_err()); + }); + + // Give the sender thread a chance to observe the full channel before + // cancelling -- simulating cancellation arriving (e.g. from a + // connection-drop observer) after the worker is already stuck sending. + std::thread::sleep(Duration::from_millis(50)); + context.cancel(); + + let cancelled = done_rx + .recv_timeout(Duration::from_secs(5)) + .expect("a stalled send must be interrupted by cancellation, not block forever"); + assert!(cancelled, "cancelled send must return an error"); +} + +/// Covers the case the report actually flagged as unproven: nothing ever +/// calls `cancel()` -- the connection-drop observer (`CancelOnDropSseStream`) +/// simply never fires, e.g. because the client vanished behind a proxy that +/// kept the socket to mesh-llm open. A stalled, never-dropped, never-drained +/// receiver must still cause the send to give up and self-cancel once it has +/// been full for the stall timeout, so the lane isn't held indefinitely. +#[test] +fn stalled_receiver_self_cancels_after_the_stall_timeout_with_no_external_cancel() { + let (tx, rx) = mpsc::channel(1); + tx.try_send(Ok(GenerationStreamEvent::Delta("first".to_owned()))) + .expect("channel has room for the first event"); + let context = OpenAiRequestContext::new(); + + let result = send_generation_event_with_stall_timeout( + &tx, + Ok(GenerationStreamEvent::Delta("second".to_owned())), + &context, + Duration::from_millis(50), + ); + + assert!( + result.is_err(), + "a send stalled past the timeout must fail rather than hang" + ); + assert!( + context.is_cancelled(), + "a self-detected stall must cancel the request so the lane is freed" + ); + drop(rx); +} From 085d399a8c0601f1533f9af58531b35259318a94 Mon Sep 17 00:00:00 2001 From: Nick DiZazzo Date: Tue, 18 Aug 2026 04:53:14 -0400 Subject: [PATCH 2/5] fix(skippy): deliver terminal SSE frames and drop the send poll loop send_generation_event's leading `if context.is_cancelled() { return Err(...) }` made the cancellation-branch call site in run_generation_stream a guaranteed no-op: it built the cancellation error frame and then called a helper that, by construction, always returned instantly without enqueuing it. The parser.finish error frame and the outer generation-error frame were dropped the same way. On main, blocking_send delivered these unconditionally whenever the 16-slot buffer had room. Losing the Err frame changes stream_lifecycle classification: it drives lifecycle.failed(error), which marks backend_error and yields StreamDropOutcome::BackendError/ StreamTerminal; without it, drop_outcome() falls through to StreamDropOutcome::Cancelled instead, corrupting both client output and telemetry for a case the whole PR exists to handle correctly. Replace the two free functions with a StreamEventSender carrying the channel, a runtime handle, and a stall timeout. It exposes two methods instead of one: - send(), used only by the in-flight on_text_chunk callback, checks cancellation first (via `biased` select) and aborts immediately when the request is already cancelled -- preserving the old early-return semantics deterministically. - send_terminal(), used for every frame emitted after generation finishes, deliberately does not consult cancellation, so a cancelled-but-still-draining receiver still gets its terminal frame -- matching main's unconditional blocking_send and fixing the swallow above. Both race the send against a stall timeout via tokio::select! under `tokio::runtime::Handle::block_on`, the same block_on-from-inside- spawn_blocking pattern already used by prompting.rs for hook calls in this exact blocking generation worker. This replaces the previous try_send + thread::sleep(20ms) poll: a healthy consumer only slightly slower than decode hits Full routinely as ordinary backpressure, and polling could waste up to 20ms per event, capping throughput near 50 events/s under sustained backpressure. select! wakes the moment either the channel has room or cancellation fires. The stall timeout gets its own constant, STREAM_SEND_STALL_TIMEOUT, rather than aliasing GENERATION_ADMISSION_TIMEOUT: it bounds a single send, not a whole generation, and admission queueing and stream-stall tolerance are unrelated policies that must be retunable independently. Finally, StreamEventSender tracks whether the receiver has been proven unreachable (closed, or stalled past the timeout). Once set, both send() and send_terminal() fail fast without waiting again. This matters because a genuinely stalled consumer would otherwise be waited on twice: once by the in-flight send that discovers the stall, and again by the terminal frame that follows it, doubling the execution lane's hold to 2x the stall timeout and undermining the point of freeing it promptly. A request that was merely cancelled externally, with a receiver that is still alive and draining, is not affected by this short-circuit -- only proven-unreachable receivers are. Adapted the existing stall/cancellation tests to the new API (each test builds its own tokio::runtime::Runtime and passes rt.handle().clone() to the sender, rather than #[tokio::test], since Handle::block_on panics from inside a runtime worker thread) and added two new tests: one proving terminal frames now reach a cancelled but live receiver, and one proving a proven-unreachable receiver's terminal send returns well under the stall timeout instead of waiting it out a second time. --- crates/skippy-server/src/frontend/backend.rs | 201 ++++++++++++------ .../src/frontend/backend/tests.rs | 103 ++++++++- 2 files changed, 233 insertions(+), 71 deletions(-) diff --git a/crates/skippy-server/src/frontend/backend.rs b/crates/skippy-server/src/frontend/backend.rs index 202281a707..720fbe585b 100644 --- a/crates/skippy-server/src/frontend/backend.rs +++ b/crates/skippy-server/src/frontend/backend.rs @@ -53,9 +53,9 @@ use skippy_runtime::SamplingConfig; use std::collections::BTreeMap; use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use std::thread; use std::time::Duration; use std::time::Instant; use tokio::sync::OwnedSemaphorePermit; @@ -69,63 +69,142 @@ fn request_cancelled_error() -> OpenAiError { } /// How long a full SSE event channel is treated as merely backed up before -/// the generation worker gives up on it. Reuses the window other requests -/// already wait for an execution lane, so a stalled consumer can never pin a -/// lane longer than everyone else already times out for. -const STREAM_SEND_STALL_TIMEOUT: Duration = GENERATION_ADMISSION_TIMEOUT; - -/// Poll interval while `send_generation_event` waits for buffer space. -const STREAM_SEND_RETRY_INTERVAL: Duration = Duration::from_millis(20); +/// the generation worker gives up on it and frees its execution lane. +/// +/// Deliberately its own value rather than an alias of +/// `GENERATION_ADMISSION_TIMEOUT`: admission queueing and stream-stall +/// tolerance are unrelated policies, and retuning one must not silently +/// retune the other. It bounds a single send, not a whole generation. +const STREAM_SEND_STALL_TIMEOUT: Duration = Duration::from_secs(10); -/// Forward one generation event to the SSE channel. +/// Forwards generation events to the SSE channel without letting a consumer +/// that has stopped draining pin the generation worker, and the execution +/// lane it holds, forever. /// /// `mpsc::Sender::blocking_send` waits indefinitely for buffer space. A /// consumer that stops draining without the channel ever being dropped -- /// a client gone before the server's own disconnect detection notices, e.g. /// behind a proxy that doesn't propagate the close -- can then pin the /// generation worker, and the execution lane it holds, forever even after -/// the request is cancelled. This polls instead: it keeps retrying while the -/// request is live and the buffer is merely full, but gives up -- cancelling -/// the request if it hasn't been already -- once the caller is cancelled, -/// the receiver is gone, or the buffer has stayed full past -/// `STREAM_SEND_STALL_TIMEOUT`. -fn send_generation_event( - tx: &mpsc::Sender>, - event: OpenAiResult, - context: &OpenAiRequestContext, -) -> Result<(), OpenAiError> { - send_generation_event_with_stall_timeout(tx, event, context, STREAM_SEND_STALL_TIMEOUT) +/// the request is cancelled. This races each send against cancellation and +/// against `stall_timeout` via `tokio::select!` instead of polling, so a +/// consumer draining normally is woken the instant space appears rather than +/// after up to a 20 ms poll tick. +struct StreamEventSender { + tx: mpsc::Sender>, + runtime: tokio::runtime::Handle, + stall_timeout: Duration, + /// Set once the receiver is gone or has stayed full past the stall + /// timeout. Nothing further can reach the client, so later frames are + /// dropped rather than waited on again. + receiver_unreachable: AtomicBool, } -/// `send_generation_event` with the stall timeout as a parameter, so tests -/// can exercise the self-cancelling path without waiting out the real -/// `STREAM_SEND_STALL_TIMEOUT`. -fn send_generation_event_with_stall_timeout( - tx: &mpsc::Sender>, - mut event: OpenAiResult, - context: &OpenAiRequestContext, - stall_timeout: Duration, -) -> Result<(), OpenAiError> { - let stalled_since = Instant::now(); - loop { - if context.is_cancelled() { - return Err(request_cancelled_error()); +impl StreamEventSender { + fn new( + tx: mpsc::Sender>, + runtime: tokio::runtime::Handle, + stall_timeout: Duration, + ) -> Self { + Self { + tx, + runtime, + stall_timeout, + receiver_unreachable: AtomicBool::new(false), + } + } + + /// Mark the receiver unreachable and free the request's execution lane. + /// Called once nothing further sent on this channel could possibly reach + /// the client: the receiver dropped, or it stayed full past + /// `stall_timeout`. + fn mark_receiver_unreachable(&self, context: &OpenAiRequestContext) { + self.receiver_unreachable.store(true, Ordering::Release); + context.cancel(); + } + + /// Send one in-flight event (from the `on_text_chunk` callback). Checks + /// cancellation first so an already-cancelled request returns + /// immediately and deterministically, matching the pre-existing + /// early-return semantics. + fn send( + &self, + event: OpenAiResult, + context: &OpenAiRequestContext, + ) -> Result<(), OpenAiError> { + if self.receiver_unreachable.load(Ordering::Acquire) { + return Err(OpenAiError::backend("stream receiver unreachable")); } - event = match tx.try_send(event) { - Ok(()) => return Ok(()), - Err(mpsc::error::TrySendError::Closed(_)) => { - context.cancel(); - return Err(OpenAiError::backend("stream receiver dropped")); + let cancellation = context.cancellation_token(); + // `tokio::time::sleep` needs an entered runtime the instant it is + // called, not just when polled, so it must be constructed inside the + // `block_on`-driven future rather than before it. + self.runtime.block_on(async { + let send = self.tx.send(event); + let sleep = tokio::time::sleep(self.stall_timeout); + tokio::select! { + biased; + () = cancellation.cancelled() => Err(request_cancelled_error()), + result = send => match result { + Ok(()) => Ok(()), + Err(_) => { + self.mark_receiver_unreachable(context); + Err(OpenAiError::backend("stream receiver dropped")) + } + }, + () = sleep => { + self.mark_receiver_unreachable(context); + Err(OpenAiError::backend( + "stream receiver stalled without draining", + )) + } } - Err(mpsc::error::TrySendError::Full(rejected)) => rejected, - }; - if stalled_since.elapsed() >= stall_timeout { - context.cancel(); - return Err(OpenAiError::backend( - "stream receiver stalled past the admission timeout", - )); + }) + } + + /// Send one terminal frame -- everything emitted after generation + /// finishes: the cancellation error, a parser-finish error, the backend + /// error, usage, and `Done`. + /// + /// Deliberately does **not** consult cancellation. These are exactly the + /// frames the frontend lifecycle needs when the request *was* cancelled: + /// on `main`, `blocking_send` delivered them unconditionally. Skipping + /// them for a cancelled-but-still-live receiver would flip + /// `stream_lifecycle`'s terminal classification -- an `Err` frame drives + /// `lifecycle.failed(error)`, which marks `backend_error` and yields + /// `StreamDropOutcome::BackendError`/`StreamTerminal`; without it, + /// `drop_outcome()` falls through to `StreamDropOutcome::Cancelled` + /// instead. + /// + /// It does still refuse to wait on a receiver already proven + /// unreachable: the in-flight send that got us here may already have + /// waited out `stall_timeout` once, and waiting again would double the + /// execution lane's hold to `2 * stall_timeout`. + fn send_terminal(&self, event: OpenAiResult) -> Result<(), OpenAiError> { + if self.receiver_unreachable.load(Ordering::Acquire) { + return Err(OpenAiError::backend("stream receiver unreachable")); } - thread::sleep(STREAM_SEND_RETRY_INTERVAL); + // See the matching comment in `send`: the sleep future must be + // constructed inside the entered runtime `block_on` provides. + self.runtime.block_on(async { + let send = self.tx.send(event); + let sleep = tokio::time::sleep(self.stall_timeout); + tokio::select! { + result = send => match result { + Ok(()) => Ok(()), + Err(_) => { + self.receiver_unreachable.store(true, Ordering::Release); + Err(OpenAiError::backend("stream receiver dropped")) + } + }, + () = sleep => { + self.receiver_unreachable.store(true, Ordering::Release); + Err(OpenAiError::backend( + "stream receiver stalled without draining", + )) + } + } + }) } } @@ -876,6 +955,11 @@ impl StageOpenAiBackend { let chat_parse_metadata = prompt.chat_parse_metadata.clone(); let (tx, rx) = mpsc::channel(16); let hook_runtime = Some(tokio::runtime::Handle::current()); + let sender = StreamEventSender::new( + tx, + tokio::runtime::Handle::current(), + STREAM_SEND_STALL_TIMEOUT, + ); let mut chat_stream_parser = if let (true, Some(request), Some(metadata)) = (parse_chat_output, hook_request.clone(), chat_parse_metadata) { @@ -911,13 +995,13 @@ impl StageOpenAiBackend { vec![GenerationStreamEvent::Delta(chunk.to_string())] }; for event in events { - send_generation_event(&tx, Ok(event), &context)?; + sender.send(Ok(event), &context)?; } Ok(()) }, ); if context.is_cancelled() { - let _ = send_generation_event(&tx, Err(request_cancelled_error()), &context); + let _ = sender.send_terminal(Err(request_cancelled_error())); return; } match result { @@ -926,14 +1010,14 @@ impl StageOpenAiBackend { match parser.finish(&output.text) { Ok(events) => { for event in events { - if send_generation_event(&tx, Ok(event), &context).is_err() { + if sender.send_terminal(Ok(event)).is_err() { return; } } parser.finish_reason(output.finish_reason) } Err(error) => { - let _ = send_generation_event(&tx, Err(error), &context); + let _ = sender.send_terminal(Err(error)); return; } } @@ -941,23 +1025,16 @@ impl StageOpenAiBackend { output.finish_reason }; if should_emit_stream_usage(include_usage, &context) - && send_generation_event( - &tx, - Ok(GenerationStreamEvent::Usage(output.usage())), - &context, - ) - .is_err() + && sender + .send_terminal(Ok(GenerationStreamEvent::Usage(output.usage()))) + .is_err() { return; } - let _ = send_generation_event( - &tx, - Ok(GenerationStreamEvent::Done(finish_reason)), - &context, - ); + let _ = sender.send_terminal(Ok(GenerationStreamEvent::Done(finish_reason))); } Err(error) => { - let _ = send_generation_event(&tx, Err(error), &context); + let _ = sender.send_terminal(Err(error)); } } }); diff --git a/crates/skippy-server/src/frontend/backend/tests.rs b/crates/skippy-server/src/frontend/backend/tests.rs index ec675ba6c2..3001fac8c6 100644 --- a/crates/skippy-server/src/frontend/backend/tests.rs +++ b/crates/skippy-server/src/frontend/backend/tests.rs @@ -1,6 +1,8 @@ use super::*; use openai_frontend::ChatCompletionRequest; +use openai_frontend::FinishReason; use serde_json::json; +use tokio::runtime::Runtime; fn trusted_ids(session_id: &str) -> OpenAiGenerationIds { OpenAiGenerationIds::new_with_trust(OpenAiCacheHints::default(), Some(session_id), true) @@ -525,27 +527,29 @@ fn internal_stream_usage_observation_preserves_client_wire_preference() { /// Reproduces the orphaned-generation report: a client can vanish (dropped /// connection, or one that hasn't been noticed yet -- e.g. behind a proxy /// that doesn't propagate the close) leaving the SSE receiver alive but -/// permanently undrained. `send_generation_event` must not let that pin the -/// generation worker, and the execution lane it holds, forever: once the +/// permanently undrained. `StreamEventSender::send` must not let that pin +/// the generation worker, and the execution lane it holds, forever: once the /// request is cancelled it must give up promptly even though the channel /// stays full and the receiver is never dropped. /// /// This runs the send on its own thread and waits for a result over a /// bounded `recv_timeout` rather than joining directly, so a regression back /// to an unconditional blocking send fails this test instead of hanging the -/// suite. +/// suite. It uses the real `STREAM_SEND_STALL_TIMEOUT`, so cancellation -- +/// not the stall timeout -- must be what ends the wait. #[test] fn stalled_receiver_does_not_pin_the_generation_worker_forever() { let (tx, rx) = mpsc::channel(1); tx.try_send(Ok(GenerationStreamEvent::Delta("first".to_owned()))) .expect("channel has room for the first event"); let context = OpenAiRequestContext::new(); + let rt = Runtime::new().expect("tokio runtime for stall test"); + let sender = StreamEventSender::new(tx, rt.handle().clone(), STREAM_SEND_STALL_TIMEOUT); let sender_context = context.clone(); let (done_tx, done_rx) = std::sync::mpsc::channel(); std::thread::spawn(move || { - let result = send_generation_event( - &tx, + let result = sender.send( Ok(GenerationStreamEvent::Delta("second".to_owned())), &sender_context, ); @@ -572,19 +576,20 @@ fn stalled_receiver_does_not_pin_the_generation_worker_forever() { /// simply never fires, e.g. because the client vanished behind a proxy that /// kept the socket to mesh-llm open. A stalled, never-dropped, never-drained /// receiver must still cause the send to give up and self-cancel once it has -/// been full for the stall timeout, so the lane isn't held indefinitely. +/// been full for the (here, injected and short) stall timeout, so the lane +/// isn't held indefinitely. #[test] fn stalled_receiver_self_cancels_after_the_stall_timeout_with_no_external_cancel() { let (tx, rx) = mpsc::channel(1); tx.try_send(Ok(GenerationStreamEvent::Delta("first".to_owned()))) .expect("channel has room for the first event"); let context = OpenAiRequestContext::new(); + let rt = Runtime::new().expect("tokio runtime for stall test"); + let sender = StreamEventSender::new(tx, rt.handle().clone(), Duration::from_millis(50)); - let result = send_generation_event_with_stall_timeout( - &tx, + let result = sender.send( Ok(GenerationStreamEvent::Delta("second".to_owned())), &context, - Duration::from_millis(50), ); assert!( @@ -597,3 +602,83 @@ fn stalled_receiver_self_cancels_after_the_stall_timeout_with_no_external_cancel ); drop(rx); } + +/// Red->green for the swallowed-terminal-frame defect: on the pre-fix code, +/// the `run_generation_stream` cancellation branch checked +/// `context.is_cancelled()` before sending, so an already-cancelled request +/// caused the cancellation error frame -- and, by the same shape, the +/// `parser.finish` error frame and the outer generation error frame -- to be +/// silently dropped instead of enqueued. That flips +/// `stream_lifecycle`'s terminal classification: without the `Err` frame, +/// `drop_outcome()` falls through to `StreamDropOutcome::Cancelled` instead +/// of the `BackendError`/`StreamTerminal` path `lifecycle.failed(error)` +/// drives. `send_terminal` must deliver the frame to a receiver that is +/// merely cancelled but still alive and draining, while `send` (used only +/// for in-flight events) must still refuse to send once cancelled. +#[test] +fn terminal_frames_are_delivered_after_the_request_is_cancelled() { + let (tx, mut rx) = mpsc::channel(4); + let context = OpenAiRequestContext::new(); + context.cancel(); + let rt = Runtime::new().expect("tokio runtime for terminal-delivery test"); + let sender = StreamEventSender::new(tx, rt.handle().clone(), STREAM_SEND_STALL_TIMEOUT); + + sender + .send_terminal(Ok(GenerationStreamEvent::Done(FinishReason::Stop))) + .expect("terminal frames must still reach a live, cancelled-but-draining receiver"); + + let received = rx + .try_recv() + .expect("the terminal frame must be enqueued, not silently swallowed"); + assert!(matches!( + received, + Ok(GenerationStreamEvent::Done(FinishReason::Stop)) + )); + + let send_result = sender.send( + Ok(GenerationStreamEvent::Delta("late".to_owned())), + &context, + ); + assert!( + send_result.is_err(), + "the cancellation check is bypassed only for terminal frames, not in-flight ones" + ); +} + +/// Bounds the double-wait hazard: once an in-flight send has already proven +/// the receiver unreachable (stalled past the timeout, here injected short), +/// a subsequent terminal send must not wait out the same stall timeout a +/// second time -- that would double the execution lane's hold to +/// `2 * stall_timeout` and defeat the point of freeing it promptly. +#[test] +fn terminal_frames_are_dropped_once_the_receiver_is_proven_unreachable() { + let (tx, rx) = mpsc::channel(1); + tx.try_send(Ok(GenerationStreamEvent::Delta("first".to_owned()))) + .expect("channel has room for the first event"); + let context = OpenAiRequestContext::new(); + let rt = Runtime::new().expect("tokio runtime for double-wait test"); + let sender = StreamEventSender::new(tx, rt.handle().clone(), Duration::from_millis(50)); + + let stalled = sender.send( + Ok(GenerationStreamEvent::Delta("second".to_owned())), + &context, + ); + assert!( + stalled.is_err(), + "the in-flight send must self-cancel once the receiver proves unreachable" + ); + + let started = Instant::now(); + let terminal = sender.send_terminal(Ok(GenerationStreamEvent::Done(FinishReason::Stop))); + let elapsed = started.elapsed(); + + assert!( + terminal.is_err(), + "a proven-unreachable receiver must not be handed a terminal frame either" + ); + assert!( + elapsed < Duration::from_millis(25), + "terminal send must short-circuit instead of waiting out the stall timeout again, took {elapsed:?}" + ); + drop(rx); +} From 3ea188eb4b9f1b8b01a81862b5869f942edb1b26 Mon Sep 17 00:00:00 2001 From: Nick DiZazzo Date: Tue, 18 Aug 2026 04:53:34 -0400 Subject: [PATCH 3/5] chore(openai-frontend): delete the uncompiled router/stream_lifecycle.rs No `mod` declaration in the crate reaches src/router/stream_lifecycle.rs: lib.rs's `pub mod router;` resolves to src/router.rs (which exists on its own, 34k), and the crate's only `mod stream_lifecycle;` (lib.rs:15) resolves to src/stream_lifecycle.rs. src/router/ contained only this one orphaned file, so it never compiled and has silently drifted from the live stream_lifecycle.rs -- diffing the two shows real divergence in imports and terminal-result handling. It sits in this PR's blast radius and is a trap for readers who assume it's live code; delete it. --- .../src/router/stream_lifecycle.rs | 160 ------------------ 1 file changed, 160 deletions(-) delete mode 100644 crates/openai-frontend/src/router/stream_lifecycle.rs diff --git a/crates/openai-frontend/src/router/stream_lifecycle.rs b/crates/openai-frontend/src/router/stream_lifecycle.rs deleted file mode 100644 index cc86dd5ad4..0000000000 --- a/crates/openai-frontend/src/router/stream_lifecycle.rs +++ /dev/null @@ -1,160 +0,0 @@ -use std::{ - convert::Infallible, - pin::Pin, - sync::{ - Arc, Mutex, - atomic::{AtomicBool, Ordering}, - }, - task::{Context, Poll}, -}; - -use axum::response::{ - IntoResponse, Response, - sse::{Event, KeepAlive, Sse}, -}; -use futures_util::Stream; -use mesh_llm_events::logging::events::TokenUsage; - -use crate::{ - backend::CancellationToken, - common::Usage, - errors::OpenAiError, - lifecycle::{ - OpenAiLifecycleContext, OpenAiLifecycleEvent, OpenAiLifecycleObserver, OpenAiTerminalResult, - }, -}; - -use super::{authoritative_usage, failure_for_status}; - -#[derive(Clone, Copy)] -pub(super) struct StreamingResponse; - -pub(super) fn sse_response( - events: S, - cancellation: CancellationToken, - lifecycle: StreamLifecycle, -) -> Response -where - S: Stream> + Send + 'static, -{ - let mut response = Sse::new(CancelOnDropSseStream::new(events, cancellation, lifecycle)) - .keep_alive(KeepAlive::default()) - .into_response(); - response.extensions_mut().insert(StreamingResponse); - response -} - -#[derive(Clone)] -pub(super) struct StreamLifecycle { - observer: Option>, - context: OpenAiLifecycleContext, - terminal: Arc, - usage: Arc>>, -} - -impl StreamLifecycle { - pub(super) fn new( - observer: Option>, - context: OpenAiLifecycleContext, - ) -> Self { - Self { - observer, - context, - terminal: Arc::new(AtomicBool::new(false)), - usage: Arc::new(Mutex::new(None)), - } - } - - pub(super) fn observe_usage(&self, usage: &Usage) { - if let Some(usage) = authoritative_usage(usage) { - *self.usage.lock().expect("stream usage lock poisoned") = Some(usage); - } - } - - fn completed(&self) { - self.observe_terminal(OpenAiLifecycleEvent::StreamTerminal { - context: self.context.clone(), - result: self - .usage - .lock() - .expect("stream usage lock poisoned") - .map_or( - OpenAiTerminalResult::Completed { status_code: 200 }, - |usage| OpenAiTerminalResult::CompletedWithUsage { - status_code: 200, - usage, - }, - ), - }); - } - - pub(super) fn failed(&self, error: &OpenAiError) { - self.observe_terminal(OpenAiLifecycleEvent::StreamTerminal { - context: self.context.clone(), - result: OpenAiTerminalResult::Failed { - status_code: error.status().as_u16(), - failure: failure_for_status(error.status()), - }, - }); - } - - fn dropped(&self, cancelled: bool) { - let event = if cancelled { - OpenAiLifecycleEvent::StreamCancelled { - context: self.context.clone(), - } - } else { - OpenAiLifecycleEvent::StreamDropped { - context: self.context.clone(), - } - }; - self.observe_terminal(event); - } - - fn observe_terminal(&self, event: OpenAiLifecycleEvent) { - if self.terminal.swap(true, Ordering::AcqRel) { - return; - } - if let Some(observer) = &self.observer { - observer.observe(&event); - } - } -} - -struct CancelOnDropSseStream { - inner: Pin> + Send + 'static>>, - cancellation: CancellationToken, - lifecycle: StreamLifecycle, -} - -impl CancelOnDropSseStream { - fn new(inner: S, cancellation: CancellationToken, lifecycle: StreamLifecycle) -> Self - where - S: Stream> + Send + 'static, - { - Self { - inner: Box::pin(inner), - cancellation, - lifecycle, - } - } -} - -impl Stream for CancelOnDropSseStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let poll = self.inner.as_mut().poll_next(cx); - if matches!(poll, Poll::Ready(None)) { - self.lifecycle.completed(); - } - poll - } -} - -impl Drop for CancelOnDropSseStream { - fn drop(&mut self) { - self.lifecycle.dropped(self.cancellation.is_cancelled()); - self.cancellation.cancel(); - } -} From e4d5e5f4d684cf785d1eadf8111779ca46b8e5d8 Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Wed, 19 Aug 2026 09:51:27 +1000 Subject: [PATCH 4/5] fix(skippy): log stalled and dropped SSE consumers with the request id When a generation worker frees its execution lane because the SSE consumer stalled past the timeout or dropped the receiver, emit a diagnostic naming the request and which failure occurred, so an operator can tell a client-initiated cancellation apart from a stalled consumer that pinned a lane. The request id is captured on the sender so terminal-frame delivery (which has no request context) can attribute its own stall/drop. Also derive the terminal short-circuit test's timing bound from the injected stall timeout instead of a fixed 25ms wall-clock number, giving it a wide margin on a loaded CI runner while keeping the assertion coupled to the timeout it is guarding against. Co-authored-by: Michael Neale Signed-off-by: Michael Neale --- crates/skippy-server/src/frontend/backend.rs | 24 +++++++++++ .../src/frontend/backend/tests.rs | 40 ++++++++++++++++--- 2 files changed, 58 insertions(+), 6 deletions(-) diff --git a/crates/skippy-server/src/frontend/backend.rs b/crates/skippy-server/src/frontend/backend.rs index 720fbe585b..4d0e770014 100644 --- a/crates/skippy-server/src/frontend/backend.rs +++ b/crates/skippy-server/src/frontend/backend.rs @@ -94,6 +94,11 @@ struct StreamEventSender { tx: mpsc::Sender>, runtime: tokio::runtime::Handle, stall_timeout: Duration, + /// Request identifier carried for diagnostics so a stalled or dropped + /// consumer can be attributed to the exact request that held (and then + /// freed) an execution lane. `send_terminal` has no request context of + /// its own, so the id is captured once at construction. + request_id: String, /// Set once the receiver is gone or has stayed full past the stall /// timeout. Nothing further can reach the client, so later frames are /// dropped rather than waited on again. @@ -105,11 +110,13 @@ impl StreamEventSender { tx: mpsc::Sender>, runtime: tokio::runtime::Handle, stall_timeout: Duration, + request_id: String, ) -> Self { Self { tx, runtime, stall_timeout, + request_id, receiver_unreachable: AtomicBool::new(false), } } @@ -148,11 +155,19 @@ impl StreamEventSender { result = send => match result { Ok(()) => Ok(()), Err(_) => { + eprintln!( + "skippy: stream receiver dropped for request {}; freeing the execution lane", + self.request_id + ); self.mark_receiver_unreachable(context); Err(OpenAiError::backend("stream receiver dropped")) } }, () = sleep => { + eprintln!( + "skippy: stream receiver stalled without draining for request {} after {:?}; freeing the execution lane", + self.request_id, self.stall_timeout + ); self.mark_receiver_unreachable(context); Err(OpenAiError::backend( "stream receiver stalled without draining", @@ -193,11 +208,19 @@ impl StreamEventSender { result = send => match result { Ok(()) => Ok(()), Err(_) => { + eprintln!( + "skippy: stream receiver dropped for request {} while delivering a terminal frame", + self.request_id + ); self.receiver_unreachable.store(true, Ordering::Release); Err(OpenAiError::backend("stream receiver dropped")) } }, () = sleep => { + eprintln!( + "skippy: stream receiver stalled without draining for request {} after {:?} while delivering a terminal frame", + self.request_id, self.stall_timeout + ); self.receiver_unreachable.store(true, Ordering::Release); Err(OpenAiError::backend( "stream receiver stalled without draining", @@ -959,6 +982,7 @@ impl StageOpenAiBackend { tx, tokio::runtime::Handle::current(), STREAM_SEND_STALL_TIMEOUT, + ids.request_id_string(), ); let mut chat_stream_parser = if let (true, Some(request), Some(metadata)) = (parse_chat_output, hook_request.clone(), chat_parse_metadata) diff --git a/crates/skippy-server/src/frontend/backend/tests.rs b/crates/skippy-server/src/frontend/backend/tests.rs index 3001fac8c6..51e08936ae 100644 --- a/crates/skippy-server/src/frontend/backend/tests.rs +++ b/crates/skippy-server/src/frontend/backend/tests.rs @@ -544,7 +544,12 @@ fn stalled_receiver_does_not_pin_the_generation_worker_forever() { .expect("channel has room for the first event"); let context = OpenAiRequestContext::new(); let rt = Runtime::new().expect("tokio runtime for stall test"); - let sender = StreamEventSender::new(tx, rt.handle().clone(), STREAM_SEND_STALL_TIMEOUT); + let sender = StreamEventSender::new( + tx, + rt.handle().clone(), + STREAM_SEND_STALL_TIMEOUT, + "test-request".to_owned(), + ); let sender_context = context.clone(); let (done_tx, done_rx) = std::sync::mpsc::channel(); @@ -585,7 +590,12 @@ fn stalled_receiver_self_cancels_after_the_stall_timeout_with_no_external_cancel .expect("channel has room for the first event"); let context = OpenAiRequestContext::new(); let rt = Runtime::new().expect("tokio runtime for stall test"); - let sender = StreamEventSender::new(tx, rt.handle().clone(), Duration::from_millis(50)); + let sender = StreamEventSender::new( + tx, + rt.handle().clone(), + Duration::from_millis(50), + "test-request".to_owned(), + ); let result = sender.send( Ok(GenerationStreamEvent::Delta("second".to_owned())), @@ -621,7 +631,12 @@ fn terminal_frames_are_delivered_after_the_request_is_cancelled() { let context = OpenAiRequestContext::new(); context.cancel(); let rt = Runtime::new().expect("tokio runtime for terminal-delivery test"); - let sender = StreamEventSender::new(tx, rt.handle().clone(), STREAM_SEND_STALL_TIMEOUT); + let sender = StreamEventSender::new( + tx, + rt.handle().clone(), + STREAM_SEND_STALL_TIMEOUT, + "test-request".to_owned(), + ); sender .send_terminal(Ok(GenerationStreamEvent::Done(FinishReason::Stop))) @@ -657,7 +672,18 @@ fn terminal_frames_are_dropped_once_the_receiver_is_proven_unreachable() { .expect("channel has room for the first event"); let context = OpenAiRequestContext::new(); let rt = Runtime::new().expect("tokio runtime for double-wait test"); - let sender = StreamEventSender::new(tx, rt.handle().clone(), Duration::from_millis(50)); + // Inject a generous stall timeout so the short-circuit assertion has a wide + // margin on a loaded CI runner: a terminal send that (wrongly) waited out + // the stall again would take at least `stall_timeout`, while the correct + // short-circuit is one atomic load. Deriving the bound from the timeout + // instead of a fixed wall-clock number keeps the two coupled. + let stall_timeout = Duration::from_millis(500); + let sender = StreamEventSender::new( + tx, + rt.handle().clone(), + stall_timeout, + "test-request".to_owned(), + ); let stalled = sender.send( Ok(GenerationStreamEvent::Delta("second".to_owned())), @@ -676,9 +702,11 @@ fn terminal_frames_are_dropped_once_the_receiver_is_proven_unreachable() { terminal.is_err(), "a proven-unreachable receiver must not be handed a terminal frame either" ); + // The short-circuit must complete in a small fraction of the injected + // stall timeout; a second wait would consume at least the whole timeout. assert!( - elapsed < Duration::from_millis(25), - "terminal send must short-circuit instead of waiting out the stall timeout again, took {elapsed:?}" + elapsed < stall_timeout / 5, + "terminal send must short-circuit instead of waiting out the stall timeout again, took {elapsed:?} (timeout {stall_timeout:?})" ); drop(rx); } From ecb2d61c9a3df6cc030213e51a36b5c6b047c465 Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Wed, 19 Aug 2026 10:45:44 +1000 Subject: [PATCH 5/5] fix(skippy): route SSE stall/drop diagnostics through Telemetry The stalled/dropped SSE consumer diagnostics used eprintln!, but skippy-server routes operator-facing signal through its structured Telemetry facility, not a logging facade. Replace the four eprintln! calls in StreamEventSender with a Telemetry::emit of stage.openai_stream_lane_freed, recording the receiver outcome (dropped vs stalled), the frame kind (in-flight vs terminal), the request id, and the stall timeout as structured attributes so a freed execution lane is correlated and observable without stderr scraping. Co-authored-by: Michael Neale Signed-off-by: Michael Neale --- crates/skippy-server/src/frontend/backend.rs | 46 ++++++++++++------- .../src/frontend/backend/tests.rs | 25 ++++++++++ 2 files changed, 55 insertions(+), 16 deletions(-) diff --git a/crates/skippy-server/src/frontend/backend.rs b/crates/skippy-server/src/frontend/backend.rs index 4d0e770014..5c5e3b9234 100644 --- a/crates/skippy-server/src/frontend/backend.rs +++ b/crates/skippy-server/src/frontend/backend.rs @@ -28,6 +28,7 @@ use crate::frontend::request::{ ensure_completion_runtime_features_supported, }; use crate::runtime_state::RuntimeSessionStats; +use crate::telemetry::Telemetry; use crate::telemetry::lifecycle_attrs; use crate::telemetry::now_unix_nanos; use async_trait::async_trait; @@ -103,6 +104,11 @@ struct StreamEventSender { /// timeout. Nothing further can reach the client, so later frames are /// dropped rather than waited on again. receiver_unreachable: AtomicBool, + /// Structured telemetry sink for the stall/drop diagnostics. `skippy-server` + /// routes operator-facing signal through `Telemetry`, not a logging facade, + /// so a freed execution lane is correlated by `request_id` and observable + /// without scraping stderr. + telemetry: Telemetry, } impl StreamEventSender { @@ -111,6 +117,7 @@ impl StreamEventSender { runtime: tokio::runtime::Handle, stall_timeout: Duration, request_id: String, + telemetry: Telemetry, ) -> Self { Self { tx, @@ -118,9 +125,27 @@ impl StreamEventSender { stall_timeout, request_id, receiver_unreachable: AtomicBool::new(false), + telemetry, } } + /// Emit a structured "execution lane freed" event when a consumer is found + /// gone or stalled. `outcome` names the failure (dropped vs stalled) and + /// `frame_kind` names which send path hit it (an in-flight event vs a + /// terminal frame), so an operator can tell a client cancellation apart + /// from a stalled consumer pinning a lane without log scraping. + fn emit_lane_freed(&self, outcome: &str, frame_kind: &str) { + let mut attrs = BTreeMap::new(); + attrs.insert(attr_key::REQUEST_ID.to_string(), json!(self.request_id)); + attrs.insert("skippy.stream.outcome".to_string(), json!(outcome)); + attrs.insert("skippy.stream.frame_kind".to_string(), json!(frame_kind)); + attrs.insert( + "skippy.stream.stall_timeout_ms".to_string(), + json!(self.stall_timeout.as_millis() as u64), + ); + self.telemetry.emit("stage.openai_stream_lane_freed", attrs); + } + /// Mark the receiver unreachable and free the request's execution lane. /// Called once nothing further sent on this channel could possibly reach /// the client: the receiver dropped, or it stayed full past @@ -155,19 +180,13 @@ impl StreamEventSender { result = send => match result { Ok(()) => Ok(()), Err(_) => { - eprintln!( - "skippy: stream receiver dropped for request {}; freeing the execution lane", - self.request_id - ); + self.emit_lane_freed("receiver_dropped", "in_flight"); self.mark_receiver_unreachable(context); Err(OpenAiError::backend("stream receiver dropped")) } }, () = sleep => { - eprintln!( - "skippy: stream receiver stalled without draining for request {} after {:?}; freeing the execution lane", - self.request_id, self.stall_timeout - ); + self.emit_lane_freed("receiver_stalled", "in_flight"); self.mark_receiver_unreachable(context); Err(OpenAiError::backend( "stream receiver stalled without draining", @@ -208,19 +227,13 @@ impl StreamEventSender { result = send => match result { Ok(()) => Ok(()), Err(_) => { - eprintln!( - "skippy: stream receiver dropped for request {} while delivering a terminal frame", - self.request_id - ); + self.emit_lane_freed("receiver_dropped", "terminal"); self.receiver_unreachable.store(true, Ordering::Release); Err(OpenAiError::backend("stream receiver dropped")) } }, () = sleep => { - eprintln!( - "skippy: stream receiver stalled without draining for request {} after {:?} while delivering a terminal frame", - self.request_id, self.stall_timeout - ); + self.emit_lane_freed("receiver_stalled", "terminal"); self.receiver_unreachable.store(true, Ordering::Release); Err(OpenAiError::backend( "stream receiver stalled without draining", @@ -983,6 +996,7 @@ impl StageOpenAiBackend { tokio::runtime::Handle::current(), STREAM_SEND_STALL_TIMEOUT, ids.request_id_string(), + self.telemetry.clone(), ); let mut chat_stream_parser = if let (true, Some(request), Some(metadata)) = (parse_chat_output, hook_request.clone(), chat_parse_metadata) diff --git a/crates/skippy-server/src/frontend/backend/tests.rs b/crates/skippy-server/src/frontend/backend/tests.rs index 51e08936ae..f1aea4f232 100644 --- a/crates/skippy-server/src/frontend/backend/tests.rs +++ b/crates/skippy-server/src/frontend/backend/tests.rs @@ -4,6 +4,27 @@ use openai_frontend::FinishReason; use serde_json::json; use tokio::runtime::Runtime; +/// A disabled telemetry sink for `StreamEventSender` construction in tests. +/// +/// `TelemetryLevel::Off` makes `emit` a no-op, so these tests exercise the +/// stall/drop control flow without needing a collector; the sink only has to +/// be a valid handle. +fn test_telemetry() -> crate::telemetry::Telemetry { + let config: skippy_protocol::StageConfig = serde_json::from_value(json!({ + "run_id": "run", + "topology_id": "topology", + "model_id": "org/model:Q4_K_M", + "stage_id": "stage-0", + "stage_index": 0, + "layer_start": 0, + "layer_end": 4, + "load_mode": "runtime-slice", + "bind_addr": "127.0.0.1:0", + })) + .expect("minimal stage config for telemetry"); + crate::telemetry::Telemetry::new(None, 1, config, crate::telemetry::TelemetryLevel::Off) +} + fn trusted_ids(session_id: &str) -> OpenAiGenerationIds { OpenAiGenerationIds::new_with_trust(OpenAiCacheHints::default(), Some(session_id), true) } @@ -549,6 +570,7 @@ fn stalled_receiver_does_not_pin_the_generation_worker_forever() { rt.handle().clone(), STREAM_SEND_STALL_TIMEOUT, "test-request".to_owned(), + test_telemetry(), ); let sender_context = context.clone(); @@ -595,6 +617,7 @@ fn stalled_receiver_self_cancels_after_the_stall_timeout_with_no_external_cancel rt.handle().clone(), Duration::from_millis(50), "test-request".to_owned(), + test_telemetry(), ); let result = sender.send( @@ -636,6 +659,7 @@ fn terminal_frames_are_delivered_after_the_request_is_cancelled() { rt.handle().clone(), STREAM_SEND_STALL_TIMEOUT, "test-request".to_owned(), + test_telemetry(), ); sender @@ -683,6 +707,7 @@ fn terminal_frames_are_dropped_once_the_receiver_is_proven_unreachable() { rt.handle().clone(), stall_timeout, "test-request".to_owned(), + test_telemetry(), ); let stalled = sender.send(