From d25184f800bcc353badb59cdd75d921501f4af10 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Thu, 27 Aug 2026 18:31:38 -0400 Subject: [PATCH 1/9] Log request id, duration and error detail on stream failures The stream-error log discarded the error text it already held, and neither path recorded a request id or a duration, so none of the roughly 500 records a day could be joined to another log line or say what actually failed. Both paths now carry request_id and error_detail, and the interrupted-stream path also carries total_duration_ms and ms_since_last_token. --- Cargo.lock | 1 + crates/api/src/routes/completions.rs | 2 + crates/services/Cargo.toml | 1 + crates/services/src/completions/mod.rs | 206 ++++++++++++++++++++++++- 4 files changed, 208 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0f21dca81..7f2333d94 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6560,6 +6560,7 @@ dependencies = [ "tokio-stream", "tokio-test", "tracing", + "tracing-subscriber", "url", "urlencoding", "utoipa", diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index f1aaa7759..11534d5ef 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -1849,9 +1849,11 @@ async fn chat_completions_inner( .fetch_add(1, std::sync::atomic::Ordering::Relaxed); if count == 0 { tracing::error!( + %request_id, %organization_id, model = %model_for_err, error_type = %completion_stream_error_category(&e), + error_detail = %e, "Completion stream error" ); } diff --git a/crates/services/Cargo.toml b/crates/services/Cargo.toml index d54cd1a9d..25d48cb9e 100644 --- a/crates/services/Cargo.toml +++ b/crates/services/Cargo.toml @@ -72,3 +72,4 @@ tokio-test = "0.4" async-trait = "0.1" futures = "0.3" mockall = "0.14" +tracing-subscriber = { version = "0.3", features = ["json"] } diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index cc5e39e46..9e63aa5f5 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -187,6 +187,11 @@ where let api_key_id = self.api_key_id; let model_id = self.model_id; let inference_type = self.inference_type; + let total_duration_ms = self.service_start_time.elapsed().as_millis() as u64; + let ms_since_last_token = self + .last_token_time + .map(|last| last.elapsed().as_millis() as u64); + let error_detail = self.last_error.as_ref().map(|error| error.to_string()); // Create span with context BEFORE any early returns so all error logs have context let _span = tracing::error_span!( @@ -227,9 +232,13 @@ where // (stream_completed == true) after e.g. a backend queue abort before // the first token. That is a provider error, not a mystery. if !self.stream_completed || self.last_error.is_some() { - tracing::warn!(%organization_id, %model_id, model = %self.model_name, + tracing::warn!(%request_id, %organization_id, %model_id, + model = %self.model_name, stream_completed = self.stream_completed, stream_error = self.last_error.is_some(), + error_detail = error_detail.as_deref(), + total_duration_ms, + ms_since_last_token, "Stream interrupted before usage stats or chat_id received (client disconnect or provider error)"); } else { tracing::error!(%organization_id, %model_id, model = %self.model_name, @@ -239,9 +248,13 @@ where } (None, Some(chat_id)) => { if !self.stream_completed || self.last_error.is_some() { - tracing::warn!(%chat_id, %organization_id, %model_id, model = %self.model_name, + tracing::warn!(%request_id, %chat_id, %organization_id, %model_id, + model = %self.model_name, stream_completed = self.stream_completed, stream_error = self.last_error.is_some(), + error_detail = error_detail.as_deref(), + total_duration_ms, + ms_since_last_token, "Stream interrupted before usage stats received (client disconnect or provider error)"); } else { tracing::error!(%chat_id, %organization_id, %model_id, model = %self.model_name, @@ -3043,6 +3056,195 @@ mod tests { ); } + #[derive(Clone, Default)] + struct CapturedLogs(Arc>>); + + impl std::io::Write for CapturedLogs { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CapturedLogs { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + + impl CapturedLogs { + fn event_containing(&self, needle: &str) -> serde_json::Value { + let raw = self.0.lock().unwrap().clone(); + String::from_utf8(raw) + .expect("log output is utf8") + .lines() + .map(|line| serde_json::from_str::(line).expect("json log line")) + .find(|event| { + event["fields"]["message"] + .as_str() + .is_some_and(|message| message.contains(needle)) + }) + .unwrap_or_else(|| panic!("no log event containing {needle}")) + } + } + + /// Mirrors the production JSON layer, which sets `with_current_span(false)` and + /// `with_span_list(false)` (`crates/api/src/main.rs`). Under those settings a + /// `request_id` carried only by the enclosing span is discarded, so these + /// assertions fail unless the fields are on the event itself. + fn interrupted_stream_event( + request_id: Uuid, + last_token_time: Option, + last_chat_id: Option, + last_error: Option, + ) -> serde_json::Value { + let needle = if last_chat_id.is_some() { + "Stream interrupted before usage stats received" + } else { + "Stream interrupted before usage stats or chat_id received" + }; + let logs = CapturedLogs::default(); + let subscriber = tracing_subscriber::fmt() + .json() + .with_current_span(false) + .with_span_list(false) + .with_max_level(tracing::Level::WARN) + .with_writer(logs.clone()) + .finish(); + + tracing::subscriber::with_default(subscriber, || { + let _interrupted = InterceptStream { + inner: stream::iter::>>( + vec![], + ), + attestation_service: Arc::new(MockAttestationService), + usage_service: Arc::new(MockUsageService), + metrics_service: Arc::new(CapturingMetricsService::new()), + request_id, + organization_id: Uuid::new_v4(), + workspace_id: Uuid::new_v4(), + api_key_id: Uuid::new_v4(), + model_id: Uuid::new_v4(), + model_name: "test-model".to_string(), + inference_type: crate::usage::ports::InferenceType::ChatCompletionStream, + service_start_time: Instant::now(), + provider_start_time: Instant::now(), + first_token_received: last_token_time.is_some(), + first_token_time: last_token_time, + ttft_ms: None, + token_count: 0, + last_token_time, + total_itl_ms: 0.0, + metric_tags: vec![], + concurrent_counter: None, + last_usage_stats: None, + last_chat_id, + stream_completed: false, + response_id: None, + last_finish_reason: None, + last_error, + state: StreamState::Streaming, + attestation_supported: true, + store_provider_chat_signature: true, + provider_attribution: crate::usage::ProviderAttribution::default(), + cache_write_cost_per_token: None, + requested_service_tier: None, + provider_service_tier: None, + latency_reporter: None, + }; + }); + + logs.event_containing(needle) + } + + #[tokio::test] + async fn interrupted_stream_logs_request_id_and_total_duration() { + let request_id = Uuid::new_v4(); + let event = interrupted_stream_event(request_id, None, None, None); + + assert_eq!( + event["fields"]["request_id"], + serde_json::Value::String(request_id.to_string()), + "without this field the record cannot be joined to any other log line" + ); + assert!( + event["fields"]["total_duration_ms"].is_u64(), + "duration must be numeric, got {}", + event["fields"]["total_duration_ms"] + ); + } + + /// The arm reached once a chat_id has arrived logs separately from the one that + /// runs before it, so both carry the fields or half the records stay unjoinable. + #[tokio::test] + async fn an_interrupted_stream_holding_a_chat_id_logs_the_same_fields() { + let request_id = Uuid::new_v4(); + let event = interrupted_stream_event(request_id, None, Some("chat-abc".to_string()), None); + + assert_eq!( + event["fields"]["request_id"], + serde_json::Value::String(request_id.to_string()) + ); + assert!(event["fields"]["total_duration_ms"].is_u64()); + } + + /// `stream_error` only says an error existed. Both arms hold the error itself, so + /// both must report it or the record still cannot say what went wrong. + #[tokio::test] + async fn an_interrupted_stream_reports_the_error_it_holds() { + let failure = + inference_providers::CompletionError::CompletionError("upstream gone".to_string()); + + for chat_id in [None, Some("chat-abc".to_string())] { + let event = + interrupted_stream_event(Uuid::new_v4(), None, chat_id, Some(failure.clone())); + + assert_eq!( + event["fields"]["stream_error"], + serde_json::Value::Bool(true) + ); + assert_eq!( + event["fields"]["error_detail"], + serde_json::Value::String(failure.to_string()), + "the error is in scope at this site and must not be reduced to a boolean" + ); + } + + let clean = interrupted_stream_event(Uuid::new_v4(), None, None, None); + assert_eq!( + clean["fields"]["error_detail"], + serde_json::Value::Null, + "a client disconnect carries no error, so the field must be absent" + ); + } + + /// A stream that died before the first token must stay distinguishable from one + /// that died after it, so the field is omitted rather than zeroed, which also + /// keeps it numeric and therefore comparable in a query. + #[tokio::test] + async fn interrupted_stream_separates_no_token_from_a_measured_gap() { + let before_any_token = interrupted_stream_event(Uuid::new_v4(), None, None, None); + assert_eq!( + before_any_token["fields"]["ms_since_last_token"], + serde_json::Value::Null, + "no token arrived, so the field must be absent rather than zero or a string" + ); + + let after_a_token = + interrupted_stream_event(Uuid::new_v4(), Some(Instant::now()), None, None); + assert!( + after_a_token["fields"]["ms_since_last_token"].is_u64(), + "a delivered token must produce a numeric gap, got {}", + after_a_token["fields"]["ms_since_last_token"] + ); + } + // ============================================ // vLLM error mapping tests (is_external: false) // ============================================ From 42dd809b373f739dd88c0022f0a8a2d4c31b1c03 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Fri, 28 Aug 2026 00:40:27 -0400 Subject: [PATCH 2/9] Carry request id on the remaining stream outcomes The other outcome logs in record_usage_and_metrics reported without a request id, leaving half the stream outcomes unjoinable. They now carry it, and the three that report a completed stream also carry total_duration_ms. --- crates/services/src/completions/mod.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 9e63aa5f5..a01310a44 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -241,7 +241,9 @@ where ms_since_last_token, "Stream interrupted before usage stats or chat_id received (client disconnect or provider error)"); } else { - tracing::error!(%organization_id, %model_id, model = %self.model_name, + tracing::error!(%request_id, %organization_id, %model_id, + model = %self.model_name, + total_duration_ms, "Stream completed but no usage stats and no chat_id available"); } return; @@ -257,18 +259,22 @@ where ms_since_last_token, "Stream interrupted before usage stats received (client disconnect or provider error)"); } else { - tracing::error!(%chat_id, %organization_id, %model_id, model = %self.model_name, + tracing::error!(%request_id, %chat_id, %organization_id, %model_id, + model = %self.model_name, + total_duration_ms, "Stream completed but no usage stats available"); } return; } (Some(usage), None) => { tracing::error!( + %request_id, prompt_tokens = usage.prompt_tokens, completion_tokens = usage.completion_tokens, %organization_id, %model_id, model = %self.model_name, + total_duration_ms, "Stream ended but no chat_id available" ); return; @@ -284,7 +290,7 @@ where let handle = match tokio::runtime::Handle::try_current() { Ok(h) => h, Err(_) => { - tracing::error!("Cannot record usage: no Tokio runtime available"); + tracing::error!(%request_id, "Cannot record usage: no Tokio runtime available"); return; } }; From 5ec7fefc9c445f60a6b4a270b16c2c566666dbf3 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Wed, 2 Sep 2026 11:12:12 -0400 Subject: [PATCH 3/9] Redact stream error text and join the chat id to its request A provider copies its upstream message verbatim into HttpError, so an in-stream failure could put a client URL into an error log, and no line carried both the chat id and the request id, so a failed signature lookup could not be traced to its request. Error text is now redacted at every stream-failure site, the text-completion site gains the fields it lacked, and both chat mapping sites log the forwarded request id. --- crates/api/src/routes/completions.rs | 11 ++++- .../src/attested/nearai/mod.rs | 2 +- crates/services/src/completions/mod.rs | 49 ++++++++++++------- .../src/inference_provider_pool/mod.rs | 15 +++++- 4 files changed, 55 insertions(+), 22 deletions(-) diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index 11534d5ef..e573ddae1 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -320,6 +320,11 @@ fn completion_stream_error_category(e: &inference_providers::CompletionError) -> } } +/// Upstream messages are copied verbatim into `HttpError` and can carry a client URL. +fn sanitized_stream_error(e: &inference_providers::CompletionError) -> String { + services::inference_provider_pool::InferenceProviderPool::sanitize_error_message(&e.to_string()) +} + /// Returns an OpenAI-compatible `error.type` for a stream-level completion error. /// Used in the `data: {"error":{...}}` SSE frame so clients can branch on the type. /// @@ -1848,12 +1853,13 @@ async fn chat_completions_inner( let count = error_count_inner .fetch_add(1, std::sync::atomic::Ordering::Relaxed); if count == 0 { + let error_detail = sanitized_stream_error(&e); tracing::error!( %request_id, %organization_id, model = %model_for_err, error_type = %completion_stream_error_category(&e), - error_detail = %e, + %error_detail, "Completion stream error" ); } @@ -2494,10 +2500,13 @@ async fn completions_inner( ))) }), Err(e) => { + let error_detail = sanitized_stream_error(&e); tracing::error!( + %request_id, %organization_id, model = %model_for_err, error_type = %completion_stream_error_category(&e), + %error_detail, "Text completion stream error" ); Some(Ok::(sse_error_frame(&e))) diff --git a/crates/inference_providers/src/attested/nearai/mod.rs b/crates/inference_providers/src/attested/nearai/mod.rs index 7fac16809..ac8a98d0f 100644 --- a/crates/inference_providers/src/attested/nearai/mod.rs +++ b/crates/inference_providers/src/attested/nearai/mod.rs @@ -95,7 +95,7 @@ fn format_error_chain(e: &E) -> String { /// strips; the corresponding HTTP header names are `X-Request-Id`, `X-Org-Id`, /// and `X-Workspace-Id`. Exposed as `pub(crate)` so `external/mod.rs` can use /// the same constants instead of hardcoding the strings. -pub(crate) mod tracing_headers { +pub mod tracing_headers { /// UUIDv4 generated per request by cloud-api. Join key across all hops. pub const REQUEST_ID: &str = "x_request_id"; /// Organization UUID of the authenticated API key owner. diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index a01310a44..bca9a6cb8 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -191,7 +191,11 @@ where let ms_since_last_token = self .last_token_time .map(|last| last.elapsed().as_millis() as u64); - let error_detail = self.last_error.as_ref().map(|error| error.to_string()); + let error_detail = self.last_error.as_ref().map(|error| { + super::inference_provider_pool::InferenceProviderPool::sanitize_error_message( + &error.to_string(), + ) + }); // Create span with context BEFORE any early returns so all error logs have context let _span = tracing::error_span!( @@ -711,7 +715,7 @@ impl CompletionServiceImpl { workspace_id: Uuid, ) { extra.insert( - "x_request_id".to_string(), + inference_providers::attested::nearai::tracing_headers::REQUEST_ID.to_string(), serde_json::Value::String(request_id.to_string()), ); extra.insert( @@ -3100,10 +3104,8 @@ mod tests { } } - /// Mirrors the production JSON layer, which sets `with_current_span(false)` and - /// `with_span_list(false)` (`crates/api/src/main.rs`). Under those settings a - /// `request_id` carried only by the enclosing span is discarded, so these - /// assertions fail unless the fields are on the event itself. + /// Mirrors production's `with_current_span(false)` / `with_span_list(false)` + /// (`crates/api/src/main.rs`), which discard anything carried only by a span. fn interrupted_stream_event( request_id: Uuid, last_token_time: Option, @@ -3186,8 +3188,7 @@ mod tests { ); } - /// The arm reached once a chat_id has arrived logs separately from the one that - /// runs before it, so both carry the fields or half the records stay unjoinable. + /// A second arm runs once a chat_id has arrived; both must carry the fields. #[tokio::test] async fn an_interrupted_stream_holding_a_chat_id_logs_the_same_fields() { let request_id = Uuid::new_v4(); @@ -3200,12 +3201,13 @@ mod tests { assert!(event["fields"]["total_duration_ms"].is_u64()); } - /// `stream_error` only says an error existed. Both arms hold the error itself, so - /// both must report it or the record still cannot say what went wrong. + /// `stream_error` only says an error existed, and the upstream text it carries + /// can hold a client URL, so both arms must report it and both must redact. #[tokio::test] async fn an_interrupted_stream_reports_the_error_it_holds() { - let failure = - inference_providers::CompletionError::CompletionError("upstream gone".to_string()); + let failure = inference_providers::CompletionError::CompletionError( + "Error fetching image https://records.example.com/scan.png?sig=abc: 403".to_string(), + ); for chat_id in [None, Some("chat-abc".to_string())] { let event = @@ -3215,10 +3217,20 @@ mod tests { event["fields"]["stream_error"], serde_json::Value::Bool(true) ); - assert_eq!( - event["fields"]["error_detail"], - serde_json::Value::String(failure.to_string()), - "the error is in scope at this site and must not be reduced to a boolean" + let detail = event["fields"]["error_detail"] + .as_str() + .expect("the error is in scope here and must not be reduced to a boolean"); + assert!( + detail.contains("[URL_REDACTED]"), + "a client-supplied URL must not reach the logs, got {detail}" + ); + assert!( + !detail.contains("records.example.com"), + "redaction must remove the host as well, got {detail}" + ); + assert!( + detail.contains("Error fetching image"), + "redaction must keep the diagnostic text, got {detail}" ); } @@ -3230,9 +3242,8 @@ mod tests { ); } - /// A stream that died before the first token must stay distinguishable from one - /// that died after it, so the field is omitted rather than zeroed, which also - /// keeps it numeric and therefore comparable in a query. + /// Omitted rather than zeroed, so "died before the first token" stays distinct + /// and the field stays numeric for queries. #[tokio::test] async fn interrupted_stream_separates_no_token_from_a_measured_gap() { let before_any_token = interrupted_stream_event(Uuid::new_v4(), None, None, None); diff --git a/crates/services/src/inference_provider_pool/mod.rs b/crates/services/src/inference_provider_pool/mod.rs index 132725e63..e369f9785 100644 --- a/crates/services/src/inference_provider_pool/mod.rs +++ b/crates/services/src/inference_provider_pool/mod.rs @@ -58,6 +58,15 @@ impl BackendModelMetadata { } } +fn forwarded_request_id( + extra: &std::collections::HashMap, +) -> Option { + extra + .get(inference_providers::attested::nearai::tracing_headers::REQUEST_ID) + .and_then(|value| value.as_str()) + .map(str::to_string) +} + fn merge_positive_max(stored: &mut Option, candidate: Option) { if let Some(candidate) = candidate.filter(|value| *value > 0) { *stored = Some(stored.map_or(candidate, |stored| stored.max(candidate))); @@ -2185,7 +2194,7 @@ impl InferenceProviderPool { } /// Sanitize error message by removing sensitive information like IP addresses, URLs, and internal details - fn sanitize_error_message(error: &str) -> String { + pub fn sanitize_error_message(error: &str) -> String { let mut sanitized = error.to_string(); // Remove URLs (http://..., https://...) @@ -3204,6 +3213,7 @@ impl InferenceProviderPool { mut hints: ChatRoutingHints, ) -> Result { let model_id = params.model.clone(); + let forwarded_request_id = forwarded_request_id(¶ms.extra); // Extract model_pub_key from params.extra for routing let model_pub_key_str = params @@ -3307,6 +3317,7 @@ impl InferenceProviderPool { let chat_id = chat_chunk.id.clone(); tracing::info!( chat_id = %chat_id, + request_id = forwarded_request_id.as_deref(), "Storing chat_id mapping for streaming completion" ); // Pin the dedicated TLS connection so signature fetches @@ -3351,6 +3362,7 @@ impl InferenceProviderPool { request_hash: String, ) -> Result { let model_id = params.model.clone(); + let forwarded_request_id = forwarded_request_id(¶ms.extra); // Non-streaming requests carry no service-side routing hints (that // path predates PR #838's estimator and stays byte-identical for // single-capacity models); multi-tier models still get context @@ -3413,6 +3425,7 @@ impl InferenceProviderPool { let chat_id = response.response.id.clone(); tracing::info!( chat_id = %chat_id, + request_id = forwarded_request_id.as_deref(), "Storing chat_id mapping for non-streaming completion" ); self.store_chat_id_mapping(chat_id.clone(), provider).await; From f92608398cf66b4eaeed900c5b933c5267af139b Mon Sep 17 00:00:00 2001 From: neo-sky Date: Wed, 2 Sep 2026 12:39:20 -0400 Subject: [PATCH 4/9] Fail a stalled stream with a typed timeout A silent upstream closed the stream with no application error, so a truncated answer looked identical to a complete one. InterceptStream now fails a stream that produces nothing for its idle budget, using a longer bound before the first token because a large context can prefill for minutes. The watchdog is off unless STREAM_WATCHDOG_ENABLED is set. --- crates/api/src/lib.rs | 19 +- crates/api/src/routes/completions.rs | 1 - crates/api/tests/common/mod.rs | 1 + crates/config/src/types.rs | 134 ++++++++++++++ crates/services/src/completions/mod.rs | 230 ++++++++++++++++++++++++- env.example | 13 ++ 6 files changed, 387 insertions(+), 11 deletions(-) diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 0844a4932..05c7fd90c 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -424,14 +424,27 @@ pub async fn init_domain_services_with_pool( as Arc; // Create completion service with usage tracking (needs usage_service) - let completion_service = Arc::new(services::CompletionServiceImpl::new( + let mut completion_service_impl = services::CompletionServiceImpl::new( inference_provider_pool.clone(), attestation_service.clone(), usage_service.clone(), metrics_service.clone(), models_repo.clone() as Arc, org_limit_repository, - )); + ); + if config.stream_watchdog.enabled { + completion_service_impl = completion_service_impl.with_stream_idle_timeouts( + services::completions::StreamIdleTimeouts { + first_token: std::time::Duration::from_secs( + config.stream_watchdog.first_token_seconds, + ), + between_tokens: std::time::Duration::from_secs( + config.stream_watchdog.between_tokens_seconds, + ), + }, + ); + } + let completion_service = Arc::new(completion_service_impl); let brave_search_provider = Arc::new(services::responses::tools::brave::BraveWebSearchProvider::new()); @@ -2759,6 +2772,7 @@ mod tests { staking_farm: config::StakingFarmConfig::default(), aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), + stream_watchdog: config::StreamWatchdogConfig::default(), ita: config::ItaAttestationConfig::default(), }; @@ -2869,6 +2883,7 @@ mod tests { staking_farm: config::StakingFarmConfig::default(), aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), + stream_watchdog: config::StreamWatchdogConfig::default(), ita: config::ItaAttestationConfig::default(), }; diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index e573ddae1..43dc957ae 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -320,7 +320,6 @@ fn completion_stream_error_category(e: &inference_providers::CompletionError) -> } } -/// Upstream messages are copied verbatim into `HttpError` and can carry a client URL. fn sanitized_stream_error(e: &inference_providers::CompletionError) -> String { services::inference_provider_pool::InferenceProviderPool::sanitize_error_message(&e.to_string()) } diff --git a/crates/api/tests/common/mod.rs b/crates/api/tests/common/mod.rs index bf7f657d7..77ebec377 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -130,6 +130,7 @@ pub fn test_config() -> ApiConfig { enabled: true, ..config::UsageReportingConfig::default() }, + stream_watchdog: config::StreamWatchdogConfig::default(), ita: config::ItaAttestationConfig::default(), } } diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index cd07ed2c0..29f4d5a73 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -27,6 +27,7 @@ pub struct ApiConfig { pub staking_farm: StakingFarmConfig, pub aml: AmlConfig, pub usage_reporting: UsageReportingConfig, + pub stream_watchdog: StreamWatchdogConfig, pub ita: ItaAttestationConfig, } @@ -61,6 +62,7 @@ impl ApiConfig { aml: AmlConfig::from_env()?, ita: ItaAttestationConfig::from_env()?, usage_reporting: UsageReportingConfig::from_env()?, + stream_watchdog: StreamWatchdogConfig::from_env()?, }) } } @@ -286,6 +288,57 @@ fn parse_optional_i32_env(key: &str, default: Option) -> Result } } +/// Idle bounds that fail a stream producing no data, off unless an operator +/// enables it. The first-token bound is separate because a large context can +/// prefill for minutes before any token appears. +#[derive(Debug, Clone)] +pub struct StreamWatchdogConfig { + pub enabled: bool, + pub first_token_seconds: u64, + pub between_tokens_seconds: u64, +} + +impl Default for StreamWatchdogConfig { + fn default() -> Self { + Self { + enabled: false, + first_token_seconds: 300, + between_tokens_seconds: 90, + } + } +} + +impl StreamWatchdogConfig { + pub fn from_env() -> Result { + let defaults = Self::default(); + let config = Self { + enabled: parse_bool_env("STREAM_WATCHDOG_ENABLED", defaults.enabled)?, + first_token_seconds: parse_u64_env( + "STREAM_WATCHDOG_FIRST_TOKEN_SECONDS", + defaults.first_token_seconds, + )?, + between_tokens_seconds: parse_u64_env( + "STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS", + defaults.between_tokens_seconds, + )?, + }; + + if config.first_token_seconds == 0 || config.between_tokens_seconds == 0 { + return Err("stream watchdog timeouts must be greater than zero".to_string()); + } + if config.first_token_seconds > 3_600 || config.between_tokens_seconds > 3_600 { + return Err("stream watchdog timeouts must not exceed 3600".to_string()); + } + if config.first_token_seconds < config.between_tokens_seconds { + return Err("STREAM_WATCHDOG_FIRST_TOKEN_SECONDS must not be below \ + STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS" + .to_string()); + } + + Ok(config) + } +} + /// Operational limits for the programmatic usage-reporting API. /// /// Reporting is disabled by default because its production indexes are built @@ -1232,6 +1285,87 @@ mod tests { ); } + struct StreamWatchdogEnvGuard { + values: [(&'static str, Option); 3], + } + + impl StreamWatchdogEnvGuard { + const KEYS: [&'static str; 3] = [ + "STREAM_WATCHDOG_ENABLED", + "STREAM_WATCHDOG_FIRST_TOKEN_SECONDS", + "STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS", + ]; + + fn cleared() -> Self { + let guard = Self { + values: Self::KEYS.map(|key| (key, std::env::var_os(key))), + }; + for key in Self::KEYS { + std::env::remove_var(key); + } + guard + } + } + + impl Drop for StreamWatchdogEnvGuard { + fn drop(&mut self) { + for (key, value) in &mut self.values { + match value.take() { + Some(value) => std::env::set_var(*key, value), + None => std::env::remove_var(*key), + } + } + } + } + + #[test] + #[serial] + fn stream_watchdog_is_off_by_default_with_a_longer_first_token_bound() { + let _env = StreamWatchdogEnvGuard::cleared(); + + let config = StreamWatchdogConfig::from_env().unwrap(); + + assert!(!config.enabled); + assert_eq!(config.first_token_seconds, 300); + assert_eq!(config.between_tokens_seconds, 90); + assert!(config.first_token_seconds >= config.between_tokens_seconds); + } + + #[test] + #[serial] + fn stream_watchdog_rejects_a_first_token_bound_below_the_between_token_bound() { + let _env = StreamWatchdogEnvGuard::cleared(); + std::env::set_var("STREAM_WATCHDOG_FIRST_TOKEN_SECONDS", "30"); + std::env::set_var("STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS", "90"); + + let error = StreamWatchdogConfig::from_env().unwrap_err(); + + assert!(error.contains("STREAM_WATCHDOG_FIRST_TOKEN_SECONDS")); + assert!(error.contains("STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS")); + } + + #[test] + #[serial] + fn stream_watchdog_rejects_a_zero_bound_that_would_fail_every_stream() { + let _env = StreamWatchdogEnvGuard::cleared(); + std::env::set_var("STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS", "0"); + + let error = StreamWatchdogConfig::from_env().unwrap_err(); + + assert!(error.contains("greater than zero")); + } + + #[test] + #[serial] + fn stream_watchdog_rejects_a_bound_beyond_an_hour() { + let _env = StreamWatchdogEnvGuard::cleared(); + std::env::set_var("STREAM_WATCHDOG_FIRST_TOKEN_SECONDS", "3601"); + + let error = StreamWatchdogConfig::from_env().unwrap_err(); + + assert!(error.contains("3600")); + } + #[test] fn test_is_admin_email() { let config = AuthConfig { diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index bca9a6cb8..0715315b5 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -96,6 +96,9 @@ where last_token_time: Option, /// Accumulated inter-token latency for average calculation total_itl_ms: f64, + idle_timeouts: Option, + idle_timer: Option>>, + idle_armed: bool, // Pre-allocated low-cardinality metric tags (for Datadog/OTLP) metric_tags: Vec, concurrent_counter: Option>, @@ -486,6 +489,26 @@ where } } +#[derive(Debug, Clone, Copy)] +pub struct StreamIdleTimeouts { + pub first_token: Duration, + pub between_tokens: Duration, +} + +impl InterceptStream +where + S: Stream> + Unpin, +{ + fn idle_budget(&self) -> Option { + let timeouts = self.idle_timeouts?; + Some(if self.first_token_received { + timeouts.between_tokens + } else { + timeouts.first_token + }) + } +} + impl Stream for InterceptStream where S: Stream> + Unpin, @@ -555,6 +578,7 @@ where } } } + self.idle_armed = false; return Poll::Ready(Some(Ok(event.clone()))); } Poll::Ready(None) => { @@ -569,7 +593,41 @@ where self.last_error = Some(err.clone()); return Poll::Ready(Some(Err(err.clone()))); } - Poll::Pending => return Poll::Pending, + Poll::Pending => { + let Some(budget) = self.idle_budget() else { + return Poll::Pending; + }; + if !self.idle_armed { + let deadline = tokio::time::Instant::now() + budget; + match &mut self.idle_timer { + Some(timer) => timer.as_mut().reset(deadline), + None => { + self.idle_timer = + Some(Box::pin(tokio::time::sleep_until(deadline))); + } + } + self.idle_armed = true; + } + let timer = self + .idle_timer + .as_mut() + .expect("idle timer is created when the stream arms"); + if timer.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + let stalled_during = if self.first_token_received { + "generation" + } else { + "prefill" + }; + let timeout = inference_providers::CompletionError::Timeout { + operation: stalled_during.to_string(), + timeout_seconds: budget.as_secs(), + }; + self.idle_armed = false; + self.last_error = Some(timeout.clone()); + return Poll::Ready(Some(Err(timeout))); + } } } StreamState::Finalizing(ref mut future) => match future.as_mut().poll(cx) { @@ -642,6 +700,7 @@ pub struct CompletionServiceImpl { org_concurrent_limits: Cache, /// Repository for fetching organization concurrent limits organization_limit_repository: Arc, + stream_idle_timeouts: Option, } /// TTL for organization concurrent limit cache (5 minutes) @@ -757,9 +816,15 @@ impl CompletionServiceImpl { concurrent_limit: DEFAULT_CONCURRENT_LIMIT, org_concurrent_limits, organization_limit_repository, + stream_idle_timeouts: None, } } + pub fn with_stream_idle_timeouts(mut self, timeouts: StreamIdleTimeouts) -> Self { + self.stream_idle_timeouts = Some(timeouts); + self + } + /// Extract tools and tool_choice from the extra HashMap if present and /// parseable as the typed `ToolDefinition` / `ToolChoice` shapes. /// @@ -1477,6 +1542,9 @@ impl CompletionServiceImpl { ttft_ms: None, token_count: 0, last_token_time: None, + idle_timeouts: self.stream_idle_timeouts, + idle_timer: None, + idle_armed: false, total_itl_ms: 0.0, metric_tags, concurrent_counter, @@ -2376,6 +2444,9 @@ mod tests { token_count: 0, last_token_time: None, total_itl_ms: 0.0, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, metric_tags, concurrent_counter: None, last_usage_stats: None, @@ -2548,6 +2619,9 @@ mod tests { token_count: 0, last_token_time: None, total_itl_ms: 0.0, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, metric_tags: CompletionServiceImpl::create_metric_tags("test-model"), concurrent_counter: None, last_usage_stats: None, @@ -2701,6 +2775,9 @@ mod tests { token_count: 0, last_token_time: None, total_itl_ms: 0.0, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, metric_tags, concurrent_counter: None, last_usage_stats: None, @@ -2828,6 +2905,9 @@ mod tests { token_count: 0, last_token_time: None, total_itl_ms: 0.0, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, metric_tags, concurrent_counter: None, last_usage_stats: None, @@ -3038,6 +3118,9 @@ mod tests { token_count: 0, last_token_time: None, total_itl_ms: 0.0, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, metric_tags: vec![], concurrent_counter: Some(counter.clone()), last_usage_stats: None, @@ -3104,8 +3187,6 @@ mod tests { } } - /// Mirrors production's `with_current_span(false)` / `with_span_list(false)` - /// (`crates/api/src/main.rs`), which discard anything carried only by a span. fn interrupted_stream_event( request_id: Uuid, last_token_time: Option, @@ -3149,6 +3230,9 @@ mod tests { token_count: 0, last_token_time, total_itl_ms: 0.0, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, metric_tags: vec![], concurrent_counter: None, last_usage_stats: None, @@ -3188,7 +3272,6 @@ mod tests { ); } - /// A second arm runs once a chat_id has arrived; both must carry the fields. #[tokio::test] async fn an_interrupted_stream_holding_a_chat_id_logs_the_same_fields() { let request_id = Uuid::new_v4(); @@ -3201,8 +3284,6 @@ mod tests { assert!(event["fields"]["total_duration_ms"].is_u64()); } - /// `stream_error` only says an error existed, and the upstream text it carries - /// can hold a client URL, so both arms must report it and both must redact. #[tokio::test] async fn an_interrupted_stream_reports_the_error_it_holds() { let failure = inference_providers::CompletionError::CompletionError( @@ -3242,8 +3323,6 @@ mod tests { ); } - /// Omitted rather than zeroed, so "died before the first token" stays distinct - /// and the field stays numeric for queries. #[tokio::test] async fn interrupted_stream_separates_no_token_from_a_measured_gap() { let before_any_token = interrupted_stream_event(Uuid::new_v4(), None, None, None); @@ -3262,6 +3341,141 @@ mod tests { ); } + fn watched_stream(inner: S, timeouts: Option) -> InterceptStream + where + S: Stream> + Unpin, + { + InterceptStream { + inner, + attestation_service: Arc::new(MockAttestationService), + usage_service: Arc::new(MockUsageService), + metrics_service: Arc::new(CapturingMetricsService::new()), + request_id: Uuid::new_v4(), + organization_id: Uuid::new_v4(), + workspace_id: Uuid::new_v4(), + api_key_id: Uuid::new_v4(), + model_id: Uuid::new_v4(), + model_name: "test-model".to_string(), + inference_type: crate::usage::ports::InferenceType::ChatCompletionStream, + service_start_time: Instant::now(), + provider_start_time: Instant::now(), + first_token_received: false, + first_token_time: None, + ttft_ms: None, + token_count: 0, + last_token_time: None, + total_itl_ms: 0.0, + idle_timeouts: timeouts, + idle_timer: None, + idle_armed: false, + metric_tags: vec![], + concurrent_counter: None, + last_usage_stats: None, + last_chat_id: None, + stream_completed: false, + response_id: None, + last_finish_reason: None, + last_error: None, + state: StreamState::Streaming, + attestation_supported: true, + store_provider_chat_signature: true, + provider_attribution: crate::usage::ProviderAttribution::default(), + cache_write_cost_per_token: None, + requested_service_tier: None, + provider_service_tier: None, + latency_reporter: None, + } + } + + fn token_event() -> SSEEvent { + SSEEvent { + raw_bytes: Bytes::from("data: ..."), + raw_passthrough: true, + chunk: Some(StreamChunk::Chat(ChatCompletionChunk { + id: "chat-watchdog".to_string(), + object: "chat.completion.chunk".to_string(), + created: 1234567890, + model: "test-model".to_string(), + choices: vec![], + usage: None, + service_tier: None, + prompt_token_ids: None, + system_fingerprint: None, + modality: None, + extra: Default::default(), + })), + } + } + + #[tokio::test(start_paused = true)] + async fn a_stalled_stream_fails_with_a_typed_timeout() { + let inner = stream::iter(vec![Ok(token_event())]).chain(stream::pending()); + let mut watched = Box::pin(watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + )); + + assert!( + matches!(watched.next().await, Some(Ok(_))), + "the first token must reach the client before the watchdog is relevant" + ); + + match watched.next().await { + Some(Err(inference_providers::CompletionError::Timeout { + operation, + timeout_seconds, + })) => { + assert_eq!( + operation, "generation", + "a stall after the first token is a generation stall, not a prefill one" + ); + assert_eq!( + timeout_seconds, 90, + "the between-token budget applies once a token has arrived" + ); + } + other => panic!("a silent upstream must surface as a typed timeout, got {other:?}"), + } + } + + #[tokio::test(start_paused = true)] + async fn a_slow_prefill_is_not_mistaken_for_a_stall() { + let inner = stream::once(Box::pin(async { + tokio::time::sleep(Duration::from_secs(200)).await; + Ok(token_event()) + })); + let mut watched = Box::pin(watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + )); + + assert!( + matches!(watched.next().await, Some(Ok(_))), + "a 200s prefill is inside the 300s first-token budget and must survive; \ + applying the 90s between-token budget before the first token would kill it" + ); + } + + #[tokio::test(start_paused = true)] + async fn an_unwatched_stream_is_never_timed_out() { + let inner = stream::iter(vec![Ok(token_event())]).chain(stream::pending()); + let mut watched = Box::pin(watched_stream(inner, None)); + + assert!(matches!(watched.next().await, Some(Ok(_)))); + + let parked = tokio::time::timeout(Duration::from_secs(7_200), watched.next()).await; + assert!( + parked.is_err(), + "with no thresholds configured the stream must stay parked rather than fail" + ); + } + // ============================================ // vLLM error mapping tests (is_external: false) // ============================================ diff --git a/env.example b/env.example index a35615a7b..67ca053c3 100644 --- a/env.example +++ b/env.example @@ -140,6 +140,19 @@ USAGE_REPORTING_MAX_CONCURRENT_REQUESTS=4 USAGE_REPORTING_TOKEN_MAX_CONCURRENT_REQUESTS=2 USAGE_REPORTING_REQUEST_TIMEOUT_SECONDS=15 +# ============================================================================= +# Streaming Stall Watchdog +# ============================================================================= +# Fails a stream that produces no data for its idle budget, so a silent upstream +# surfaces as an error instead of a truncated answer. Disabled until the bounds +# below are confirmed against production inter-token gaps; the first-token bound +# is longer because a large context can prefill for minutes. Startup refuses a +# zero bound, a bound above 3600, or a first-token bound below the between-token +# bound. +STREAM_WATCHDOG_ENABLED=false +STREAM_WATCHDOG_FIRST_TOKEN_SECONDS=300 +STREAM_WATCHDOG_BETWEEN_TOKENS_SECONDS=90 + # ============================================================================= # AWS S3 Configuration (for file uploads) # ============================================================================= From 6f9f86b4b2f7eac7dd2ac65c8992186ef18c7ff3 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Wed, 2 Sep 2026 13:59:17 -0400 Subject: [PATCH 5/9] Keep upstream error text out of stream logs HTTP failures now report only their status, since the SSE parser copies an upstream message verbatim and it can echo customer input. The watchdog also ends the stream after one timeout instead of re-arming, deadlines the provider call, reports duration on failures that reach the billing path, and lets keepalives clear the idle timer. --- crates/api/src/routes/completions.rs | 2 +- crates/services/src/completions/mod.rs | 214 ++++++++++++++---- .../src/inference_provider_pool/mod.rs | 26 +++ 3 files changed, 199 insertions(+), 43 deletions(-) diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index 43dc957ae..f5c7a5e12 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -321,7 +321,7 @@ fn completion_stream_error_category(e: &inference_providers::CompletionError) -> } fn sanitized_stream_error(e: &inference_providers::CompletionError) -> String { - services::inference_provider_pool::InferenceProviderPool::sanitize_error_message(&e.to_string()) + services::inference_provider_pool::InferenceProviderPool::safe_error_detail(e) } /// Returns an OpenAI-compatible `error.type` for a stream-level completion error. diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 0715315b5..432e93902 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -195,9 +195,7 @@ where .last_token_time .map(|last| last.elapsed().as_millis() as u64); let error_detail = self.last_error.as_ref().map(|error| { - super::inference_provider_pool::InferenceProviderPool::sanitize_error_message( - &error.to_string(), - ) + super::inference_provider_pool::InferenceProviderPool::safe_error_detail(error) }); // Create span with context BEFORE any early returns so all error logs have context @@ -212,6 +210,21 @@ where ) .entered(); + if error_detail.is_some() { + tracing::warn!( + %request_id, + %organization_id, + %model_id, + model = %self.model_name, + chat_id = self.last_chat_id.as_deref(), + error_detail = error_detail.as_deref(), + stream_completed = self.stream_completed, + total_duration_ms, + ms_since_last_token, + "Stream failed" + ); + } + let ( input_tokens, output_tokens, @@ -238,38 +251,37 @@ where // keeps polling, so the stream still ends "normally" // (stream_completed == true) after e.g. a backend queue abort before // the first token. That is a provider error, not a mystery. - if !self.stream_completed || self.last_error.is_some() { - tracing::warn!(%request_id, %organization_id, %model_id, - model = %self.model_name, - stream_completed = self.stream_completed, - stream_error = self.last_error.is_some(), - error_detail = error_detail.as_deref(), - total_duration_ms, - ms_since_last_token, - "Stream interrupted before usage stats or chat_id received (client disconnect or provider error)"); - } else { - tracing::error!(%request_id, %organization_id, %model_id, - model = %self.model_name, - total_duration_ms, - "Stream completed but no usage stats and no chat_id available"); + if self.last_error.is_none() { + if !self.stream_completed { + tracing::warn!(%request_id, %organization_id, %model_id, + model = %self.model_name, + total_duration_ms, + ms_since_last_token, + "Stream interrupted before usage stats or chat_id received \ + (client disconnect)"); + } else { + tracing::error!(%request_id, %organization_id, %model_id, + model = %self.model_name, + total_duration_ms, + "Stream completed but no usage stats and no chat_id available"); + } } return; } (None, Some(chat_id)) => { - if !self.stream_completed || self.last_error.is_some() { - tracing::warn!(%request_id, %chat_id, %organization_id, %model_id, - model = %self.model_name, - stream_completed = self.stream_completed, - stream_error = self.last_error.is_some(), - error_detail = error_detail.as_deref(), - total_duration_ms, - ms_since_last_token, - "Stream interrupted before usage stats received (client disconnect or provider error)"); - } else { - tracing::error!(%request_id, %chat_id, %organization_id, %model_id, - model = %self.model_name, - total_duration_ms, - "Stream completed but no usage stats available"); + if self.last_error.is_none() { + if !self.stream_completed { + tracing::warn!(%request_id, %chat_id, %organization_id, %model_id, + model = %self.model_name, + total_duration_ms, + ms_since_last_token, + "Stream interrupted before usage stats received (client disconnect)"); + } else { + tracing::error!(%request_id, %chat_id, %organization_id, %model_id, + model = %self.model_name, + total_duration_ms, + "Stream completed but no usage stats available"); + } } return; } @@ -525,6 +537,8 @@ where // carry no tokens: pass them through untouched so // the route can forward their raw bytes, but keep // them out of TTFT/ITL metrics and chat tracking. + self.idle_armed = false; + if event.chunk.is_none() { return Poll::Ready(Some(Ok(event.clone()))); } @@ -578,7 +592,6 @@ where } } } - self.idle_armed = false; return Poll::Ready(Some(Ok(event.clone()))); } Poll::Ready(None) => { @@ -626,6 +639,7 @@ where }; self.idle_armed = false; self.last_error = Some(timeout.clone()); + self.state = StreamState::Done; return Poll::Ready(Some(Err(timeout))); } } @@ -1712,15 +1726,26 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { }; // Get the LLM stream - let attributed_stream = match self + let provider_call = self .inference_provider_pool .chat_completion_stream_with_attribution( chat_params, request.body_hash.clone(), routing_hints, - ) - .await - { + ); + let provider_result = match self.stream_idle_timeouts { + Some(timeouts) => tokio::time::timeout(timeouts.first_token, provider_call) + .await + .unwrap_or_else(|_| { + Err(inference_providers::CompletionError::Timeout { + operation: "prefill".to_string(), + timeout_seconds: timeouts.first_token.as_secs(), + }) + }), + None => provider_call.await, + }; + + let attributed_stream = match provider_result { Ok(pair) => pair, Err(e) => { // Guard will decrement counter on drop @@ -3193,7 +3218,19 @@ mod tests { last_chat_id: Option, last_error: Option, ) -> serde_json::Value { - let needle = if last_chat_id.is_some() { + stream_drop_event(request_id, last_token_time, last_chat_id, last_error, None) + } + + fn stream_drop_event( + request_id: Uuid, + last_token_time: Option, + last_chat_id: Option, + last_error: Option, + last_usage_stats: Option, + ) -> serde_json::Value { + let needle = if last_error.is_some() { + "Stream failed" + } else if last_chat_id.is_some() { "Stream interrupted before usage stats received" } else { "Stream interrupted before usage stats or chat_id received" @@ -3235,7 +3272,7 @@ mod tests { idle_armed: false, metric_tags: vec![], concurrent_counter: None, - last_usage_stats: None, + last_usage_stats, last_chat_id, stream_completed: false, response_id: None, @@ -3294,10 +3331,6 @@ mod tests { let event = interrupted_stream_event(Uuid::new_v4(), None, chat_id, Some(failure.clone())); - assert_eq!( - event["fields"]["stream_error"], - serde_json::Value::Bool(true) - ); let detail = event["fields"]["error_detail"] .as_str() .expect("the error is in scope here and must not be reduced to a boolean"); @@ -3462,6 +3495,103 @@ mod tests { ); } + fn control_event() -> SSEEvent { + SSEEvent { + raw_bytes: Bytes::from(": keepalive\n\n"), + raw_passthrough: true, + chunk: None, + } + } + + #[tokio::test] + async fn a_failure_after_usage_arrives_still_reports_duration() { + let event = stream_drop_event( + Uuid::new_v4(), + Some(Instant::now()), + Some("chat-billing".to_string()), + Some(inference_providers::CompletionError::Timeout { + operation: "generation".to_string(), + timeout_seconds: 90, + }), + Some(inference_providers::TokenUsage::new(12, 34)), + ); + + assert!( + event["fields"]["total_duration_ms"].is_u64(), + "providers that report usage continuously set it on the first chunk, so a later \ + failure reaches the billing arm and would otherwise record no duration at all" + ); + assert!(event["fields"]["error_detail"].is_string()); + } + + #[test] + fn an_upstream_http_error_never_carries_its_message_into_logs() { + let detail = + super::super::inference_provider_pool::InferenceProviderPool::safe_error_detail( + &inference_providers::CompletionError::HttpError { + status_code: 400, + message: "Invalid content in message: my private prompt".to_string(), + is_external: true, + }, + ); + + assert!( + !detail.contains("my private prompt"), + "the SSE parser copies an upstream error.message verbatim, so it can echo \ + customer input and must never reach a log line" + ); + assert!(detail.contains("400")); + } + + #[tokio::test(start_paused = true)] + async fn a_timed_out_stream_ends_instead_of_firing_again() { + let inner = stream::iter(vec![Ok(token_event())]).chain(stream::pending()); + let mut watched = Box::pin(watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + )); + + assert!(matches!(watched.next().await, Some(Ok(_)))); + assert!(matches!(watched.next().await, Some(Err(_)))); + + assert!( + watched.next().await.is_none(), + "the route keeps polling after an error, so a synthesized timeout must end \ + the stream rather than re-arm and fire every budget forever" + ); + } + + #[tokio::test(start_paused = true)] + async fn keepalives_keep_a_live_stream_alive() { + let inner = Box::pin(stream::iter(vec![Ok(token_event())]).chain(stream::unfold( + (), + |()| async { + tokio::time::sleep(Duration::from_secs(60)).await; + Some((Ok(control_event()), ())) + }, + ))); + let mut watched = Box::pin(watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + )); + + assert!(matches!(watched.next().await, Some(Ok(_)))); + + for _ in 0..4 { + assert!( + matches!(watched.next().await, Some(Ok(_))), + "a control frame arriving inside the budget proves the upstream is alive \ + and must clear the idle deadline" + ); + } + } + #[tokio::test(start_paused = true)] async fn an_unwatched_stream_is_never_timed_out() { let inner = stream::iter(vec![Ok(token_event())]).chain(stream::pending()); diff --git a/crates/services/src/inference_provider_pool/mod.rs b/crates/services/src/inference_provider_pool/mod.rs index e369f9785..c853f5bc5 100644 --- a/crates/services/src/inference_provider_pool/mod.rs +++ b/crates/services/src/inference_provider_pool/mod.rs @@ -160,6 +160,8 @@ fn record_provider_attempt( /// stream return or growing the stash unbounded (issue #701). const MAX_LEADING_CONTROL_EVENTS: usize = 32; +const MAX_LOGGED_ERROR_DETAIL: usize = 200; + /// EMA α for TTFT during warmup (first TTFT_WARMUP_SAMPLES observations). const TTFT_EWMA_ALPHA_WARMUP: f64 = 0.5; /// EMA α for TTFT after warmup (stable tracking). @@ -2193,6 +2195,30 @@ impl InferenceProviderPool { } } + pub fn safe_error_detail(error: &inference_providers::CompletionError) -> String { + use inference_providers::CompletionError as E; + + match error { + E::HttpError { + status_code, + is_external, + .. + } => format!("upstream http {status_code} (external={is_external})"), + E::Timeout { + operation, + timeout_seconds, + } => format!("timed out after {timeout_seconds}s during {operation}"), + other => { + let mut detail = Self::sanitize_error_message(&other.to_string()); + if let Some((cut, _)) = detail.char_indices().nth(MAX_LOGGED_ERROR_DETAIL) { + detail.truncate(cut); + detail.push_str("...[truncated]"); + } + detail + } + } + } + /// Sanitize error message by removing sensitive information like IP addresses, URLs, and internal details pub fn sanitize_error_message(error: &str) -> String { let mut sanitized = error.to_string(); From 1fad52934662b04b10c3e0f6e5ce30e15e8f9d1e Mon Sep 17 00:00:00 2001 From: neo-sky Date: Fri, 4 Sep 2026 17:41:17 -0400 Subject: [PATCH 6/9] Fix stream watchdog error and timeout handling --- crates/inference_providers/src/mock.rs | 16 ++ crates/services/src/completions/mod.rs | 217 ++++++++++++++++-- .../src/inference_provider_pool/mod.rs | 167 +++++++++++--- 3 files changed, 359 insertions(+), 41 deletions(-) diff --git a/crates/inference_providers/src/mock.rs b/crates/inference_providers/src/mock.rs index 0aaf165dd..03351565b 100644 --- a/crates/inference_providers/src/mock.rs +++ b/crates/inference_providers/src/mock.rs @@ -653,6 +653,9 @@ struct MockConfig { error_override: Option, /// When set, all chat completion streams are created successfully and then yield this error. stream_error_override: Option, + /// When set, all chat completion streams are created successfully and then never + /// yield an event, modelling an upstream that returns headers and goes silent. + stream_stall_override: bool, /// When set, all embeddings calls return this error instead of generating a response embedding_error_override: Option, /// When set, all audio transcription calls return this error instead of a response. @@ -728,6 +731,7 @@ impl MockProvider { default_response: ResponseTemplate::new("1. 2. 3."), error_override: None, stream_error_override: None, + stream_stall_override: false, embedding_error_override: None, audio_transcription_error_override: None, })), @@ -753,6 +757,7 @@ impl MockProvider { default_response: ResponseTemplate::new("1. 2. 3."), error_override: None, stream_error_override: None, + stream_stall_override: false, embedding_error_override: None, audio_transcription_error_override: None, })), @@ -776,6 +781,7 @@ impl MockProvider { default_response: ResponseTemplate::new("1. 2. 3."), error_override: None, stream_error_override: None, + stream_stall_override: false, embedding_error_override: None, audio_transcription_error_override: None, })), @@ -884,6 +890,13 @@ impl MockProvider { config.stream_error_override = error; } + /// Set a stream stall override — when set, chat completion stream creation succeeds + /// and the returned stream never yields an event. + pub async fn set_stream_stall_override(&self, stalled: bool) { + let mut config = self.config.lock().await; + config.stream_stall_override = stalled; + } + /// Override the embeddings response with an error (useful for testing error paths). /// Pass `None` to clear the override. pub async fn set_embedding_error_override(&self, error: Option) { @@ -1070,6 +1083,9 @@ impl crate::InferenceProvider for MockProvider { let stream = stream::iter(vec![Err(error.clone())]); return Ok(Box::pin(stream)); } + if config.stream_stall_override { + return Ok(Box::pin(stream::pending())); + } config .expectations .iter() diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 432e93902..26b293b1a 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -182,6 +182,24 @@ where }) } + fn create_pin_release_future(&self) -> FinalizeFuture { + if !self.attestation_supported || !self.store_provider_chat_signature { + return Box::pin(async {}); + } + + let chat_id = match &self.last_chat_id { + Some(id) => id.clone(), + None => return Box::pin(async {}), + }; + + let attestation_service = self.attestation_service.clone(); + Box::pin(async move { + attestation_service + .release_chat_signature_pin(&chat_id) + .await; + }) + } + /// Record usage and metrics. Called from Drop to ensure it always runs. fn record_usage_and_metrics(&self) { let request_id = self.request_id; @@ -603,6 +621,7 @@ where // Capture error for stop_reason in usage recording (handled in Drop) // Note: We intentionally skip Finalizing state (attestation) for errors // because partial completions cannot be verified by clients + self.idle_armed = false; self.last_error = Some(err.clone()); return Poll::Ready(Some(Err(err.clone()))); } @@ -639,7 +658,8 @@ where }; self.idle_armed = false; self.last_error = Some(timeout.clone()); - self.state = StreamState::Done; + let release_future = self.create_pin_release_future(); + self.state = StreamState::Finalizing(release_future); return Poll::Ready(Some(Err(timeout))); } } @@ -1726,24 +1746,16 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { }; // Get the LLM stream - let provider_call = self + let provider_result = self .inference_provider_pool .chat_completion_stream_with_attribution( chat_params, request.body_hash.clone(), routing_hints, - ); - let provider_result = match self.stream_idle_timeouts { - Some(timeouts) => tokio::time::timeout(timeouts.first_token, provider_call) - .await - .unwrap_or_else(|_| { - Err(inference_providers::CompletionError::Timeout { - operation: "prefill".to_string(), - timeout_seconds: timeouts.first_token.as_secs(), - }) - }), - None => provider_call.await, - }; + self.stream_idle_timeouts + .map(|timeouts| timeouts.first_token), + ) + .await; let attributed_stream = match provider_result { Ok(pair) => pair, @@ -3374,6 +3386,103 @@ mod tests { ); } + #[derive(Default)] + struct PinRecordingAttestationService { + released: std::sync::Mutex>, + } + + #[async_trait::async_trait] + impl crate::attestation::ports::AttestationServiceTrait for PinRecordingAttestationService { + async fn get_chat_signature( + &self, + _chat_id: &str, + _signing_algo: Option, + ) -> Result< + crate::attestation::models::SignatureLookupResult, + crate::attestation::AttestationError, + > { + Err(crate::attestation::AttestationError::InternalError( + "not used".to_string(), + )) + } + + async fn store_chat_signature_from_provider( + &self, + _chat_id: &str, + ) -> Result<(), crate::attestation::AttestationError> { + Ok(()) + } + + async fn store_chat_signature( + &self, + _chat_id: &str, + _request_hash: String, + _response_hash: String, + ) -> Result<(), crate::attestation::AttestationError> { + Ok(()) + } + + async fn release_chat_signature_pin(&self, chat_id: &str) { + self.released + .lock() + .expect("released mutex should not poison") + .push(chat_id.to_string()); + } + + async fn store_response_signature( + &self, + _response_id: &str, + _request_hash: String, + _response_hash: String, + ) -> Result<(), crate::attestation::AttestationError> { + Ok(()) + } + + async fn get_attestation_report( + &self, + model: Option, + signing_algo: Option, + nonce: Option, + signing_address: Option, + include_tls_fingerprint: bool, + provider_filter: Option, + ) -> Result< + crate::attestation::models::AttestationReport, + crate::attestation::AttestationError, + > { + MockAttestationService + .get_attestation_report( + model, + signing_algo, + nonce, + signing_address, + include_tls_fingerprint, + provider_filter, + ) + .await + } + + async fn get_ita_attestation_token( + &self, + query: crate::attestation::ita::ItaTokenQuery, + ) -> Result + { + MockAttestationService + .get_ita_attestation_token(query) + .await + } + + async fn verify_vpc_signature( + &self, + timestamp: i64, + signature: String, + ) -> Result { + MockAttestationService + .verify_vpc_signature(timestamp, signature) + .await + } + } + fn watched_stream(inner: S, timeouts: Option) -> InterceptStream where S: Stream> + Unpin, @@ -3495,6 +3604,86 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn an_upstream_error_restarts_the_idle_budget() { + let inner = stream::once(Box::pin(async { Ok(token_event()) })) + .chain(stream::once(Box::pin(async { + tokio::time::sleep(Duration::from_secs(60)).await; + Err(inference_providers::CompletionError::CompletionError( + "upstream failed".to_string(), + )) + }))) + .chain(stream::pending()); + let mut watched = Box::pin(watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + )); + + let started = tokio::time::Instant::now(); + assert!(matches!(watched.next().await, Some(Ok(_)))); + assert!(matches!( + watched.next().await, + Some(Err(inference_providers::CompletionError::CompletionError( + _ + ))) + )); + + match watched.next().await { + Some(Err(inference_providers::CompletionError::Timeout { .. })) => { + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(150), + "the budget must restart when the error arrives at 60s, giving a timeout at \ + 150s; a deadline still measured from the last token fires at 90s, cutting \ + the budget to the 30s that happened to remain. Fired at {elapsed:?}" + ); + } + other => panic!( + "expected the watchdog to fire on the silence after the error, got {other:?}" + ), + } + } + + #[tokio::test(start_paused = true)] + async fn a_watchdog_timeout_releases_the_signature_routing_pin() { + let attestation = Arc::new(PinRecordingAttestationService::default()); + let inner = stream::iter(vec![Ok(token_event())]).chain(stream::pending()); + let mut watched = watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + ); + watched.attestation_service = attestation.clone(); + let mut watched = Box::pin(watched); + + assert!(matches!(watched.next().await, Some(Ok(_)))); + assert!(matches!( + watched.next().await, + Some(Err(inference_providers::CompletionError::Timeout { .. })) + )); + assert!(watched.next().await.is_none()); + + std::mem::forget(watched); + + let released = attestation + .released + .lock() + .expect("released mutex should not poison") + .clone(); + assert_eq!( + released, + vec!["chat-watchdog".to_string()], + "a stream the watchdog tears down has no signature to fetch, so the routing pin must \ + be released explicitly; leaving it held grows the provider's chat_id map for the \ + life of the process" + ); + } + fn control_event() -> SSEEvent { SSEEvent { raw_bytes: Bytes::from(": keepalive\n\n"), diff --git a/crates/services/src/inference_provider_pool/mod.rs b/crates/services/src/inference_provider_pool/mod.rs index b537bda17..8b5d5d4da 100644 --- a/crates/services/src/inference_provider_pool/mod.rs +++ b/crates/services/src/inference_provider_pool/mod.rs @@ -2462,19 +2462,23 @@ impl InferenceProviderPool { let mut sanitized = error.to_string(); // Remove URLs (http://..., https://...) - let url_regex = Regex::new(r"https?://[^\s)]+").unwrap(); + static URL: std::sync::OnceLock = std::sync::OnceLock::new(); + let url_regex = URL.get_or_init(|| Regex::new(r"https?://[^\s)]+").unwrap()); sanitized = url_regex .replace_all(&sanitized, "[URL_REDACTED]") .to_string(); // Remove standalone IP addresses with ports (e.g., 192.168.0.1:8000) - let ip_port_regex = Regex::new(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}:\d+\b").unwrap(); + static IP_PORT: std::sync::OnceLock = std::sync::OnceLock::new(); + let ip_port_regex = + IP_PORT.get_or_init(|| Regex::new(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}:\d+\b").unwrap()); sanitized = ip_port_regex .replace_all(&sanitized, "[IP_REDACTED]") .to_string(); // Remove standalone IP addresses (e.g., 192.168.0.1) - let ip_regex = Regex::new(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b").unwrap(); + static IP: std::sync::OnceLock = std::sync::OnceLock::new(); + let ip_regex = IP.get_or_init(|| Regex::new(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b").unwrap()); sanitized = ip_regex .replace_all(&sanitized, "[IP_REDACTED]") .to_string(); @@ -3466,16 +3470,20 @@ impl InferenceProviderPool { hints: ChatRoutingHints, ) -> Result { Ok(self - .chat_completion_stream_with_attribution(params, request_hash, hints) + .chat_completion_stream_with_attribution(params, request_hash, hints, None) .await? .stream) } + /// `first_event_timeout` bounds the wait for the first stream event only. + /// Widening it to span provider selection caps the retry loop and starves + /// fallback, which the per-attempt TTFB guards already cover. pub async fn chat_completion_stream_with_attribution( &self, mut params: ChatCompletionParams, request_hash: String, mut hints: ChatRoutingHints, + first_event_timeout: Option, ) -> Result { let model_id = params.model.clone(); let forwarded_request_id = forwarded_request_id(¶ms.extra); @@ -3567,36 +3575,59 @@ impl InferenceProviderPool { // cap we return the stream without pinning a sticky-routing mapping. let mut leading_control: Vec> = Vec::new(); - { - use futures::StreamExt as _; - while leading_control.len() < MAX_LEADING_CONTROL_EVENTS - && matches!(peekable.peek().await, Some(Ok(event)) if event.chunk.is_none()) + let peek_first_event = async { { - if let Some(ev) = peekable.next().await { - leading_control.push(ev); + use futures::StreamExt as _; + while leading_control.len() < MAX_LEADING_CONTROL_EVENTS + && matches!(peekable.peek().await, Some(Ok(event)) if event.chunk.is_none()) + { + if let Some(ev) = peekable.next().await { + leading_control.push(ev); + } } } - } - let first_error = match peekable.peek().await { - Some(Ok(event)) => { - if let Some(inference_providers::StreamChunk::Chat(chat_chunk)) = &event.chunk { - let chat_id = chat_chunk.id.clone(); - tracing::info!( - chat_id = %chat_id, + match peekable.peek().await { + Some(Ok(event)) => { + if let Some(inference_providers::StreamChunk::Chat(chat_chunk)) = &event.chunk { + let chat_id = chat_chunk.id.clone(); + tracing::info!( + chat_id = %chat_id, + request_id = forwarded_request_id.as_deref(), + "Storing chat_id mapping for streaming completion" + ); + // Pin the dedicated TLS connection so signature fetches + // reuse the same connection that served this completion. + provider.pin_chat_connection(&request_hash, &chat_id); + pinned = true; + self.store_chat_id_mapping(chat_id, provider.clone()).await; + } + None + } + Some(Err(error)) => Some(error.clone()), + None => None, + } + }; + + let first_error = match first_event_timeout { + Some(budget) => match tokio::time::timeout(budget, peek_first_event).await { + Ok(first_error) => first_error, + Err(_elapsed) => { + provider.pin_chat_connection(&request_hash, ""); + provider.unpin_chat_connection(""); + tracing::warn!( + model = %model_id, request_id = forwarded_request_id.as_deref(), - "Storing chat_id mapping for streaming completion" + timeout_seconds = budget.as_secs(), + "Upstream sent no stream event before the first-event deadline" ); - // Pin the dedicated TLS connection so signature fetches - // reuse the same connection that served this completion. - provider.pin_chat_connection(&request_hash, &chat_id); - pinned = true; - self.store_chat_id_mapping(chat_id, provider.clone()).await; + return Err(CompletionError::Timeout { + operation: "prefill".to_string(), + timeout_seconds: budget.as_secs(), + }); } - None - } - Some(Err(error)) => Some(error.clone()), - None => None, + }, + None => peek_first_event.await, }; if !pinned { // Clean up orphaned pending client when peek fails or yields no chat_id @@ -8771,6 +8802,88 @@ mod tests { // the live request fail. The fallback only triggers on a genuine // request-level failure, which is deterministic to inject with a mock. + #[tokio::test(start_paused = true)] + async fn a_silent_upstream_is_bounded_by_the_first_event_deadline() { + use inference_providers::mock::MockProvider; + + let pool = InferenceProviderPool::new(None, ExternalProvidersConfig::default()); + let model_id = "stalling/model".to_string(); + let provider = Arc::new(MockProvider::new_accept_all()); + provider.set_stream_stall_override(true).await; + pool.register_provider(model_id.clone(), provider).await; + + let started = tokio::time::Instant::now(); + let result = pool + .chat_completion_stream_with_attribution( + fallback_params(&model_id), + "h".to_string(), + ChatRoutingHints::default(), + Some(Duration::from_secs(300)), + ) + .await; + + match result { + Err(CompletionError::Timeout { + operation, + timeout_seconds, + }) => { + assert_eq!(operation, "prefill"); + assert_eq!(timeout_seconds, 300); + assert!( + started.elapsed() >= Duration::from_secs(300), + "the deadline must actually bound the wait for the first event" + ); + } + Err(other) => panic!( + "an upstream that returns headers and then goes silent must surface as a typed \ + prefill timeout rather than hanging, got {other:?}" + ), + Ok(_) => { + panic!("an upstream that never yields an event must not produce a usable stream") + } + } + } + + #[tokio::test(start_paused = true)] + async fn the_first_event_deadline_does_not_cap_provider_selection() { + use inference_providers::mock::{MockProvider, RequestMatcher, ResponseTemplate}; + + let pool = InferenceProviderPool::new(None, ExternalProvidersConfig::default()); + let model_id = "slow-selection/model".to_string(); + + let failing = Arc::new(MockProvider::new_accept_all()); + failing + .set_error_override(Some(CompletionError::HttpError { + status_code: 503, + message: "busy".to_string(), + is_external: false, + })) + .await; + let healthy = Arc::new(MockProvider::new_accept_all()); + healthy + .when(RequestMatcher::Any) + .respond_with(ResponseTemplate::new("served after a retry")) + .await; + + pool.register_provider(model_id.clone(), failing).await; + pool.register_provider(model_id.clone(), healthy).await; + + let served = pool + .chat_completion_stream_with_attribution( + fallback_params(&model_id), + "h".to_string(), + ChatRoutingHints::default(), + Some(Duration::from_secs(300)), + ) + .await; + + assert!( + served.is_ok(), + "the deadline covers the wait for the first event, not provider selection; capping \ + the whole call lets a slow retry round consume it and starves fallback" + ); + } + fn fallback_params(model: &str) -> inference_providers::ChatCompletionParams { inference_providers::ChatCompletionParams { model: model.to_string(), From 2a7a765c5ca476e186e412192ebcbb0a57195325 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Sat, 5 Sep 2026 14:49:41 -0400 Subject: [PATCH 7/9] Bound the wait for the blank line after [DONE] An upstream that sends the marker and then holds the connection open pinned the stream until the L4 reaper. The bound finalizes rather than erroring, since the answer is already complete once the marker lands. --- crates/services/src/completions/mod.rs | 55 +++++++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 8c0b5c612..5cf6aa682 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -761,7 +761,31 @@ where self.begin_finalizing(); return Poll::Ready(Some(Err(err))); } - Poll::Pending => return Poll::Pending, + Poll::Pending => { + let Some(budget) = self.idle_budget() else { + return Poll::Pending; + }; + if !self.idle_armed { + let deadline = tokio::time::Instant::now() + budget; + match &mut self.idle_timer { + Some(timer) => timer.as_mut().reset(deadline), + None => { + self.idle_timer = + Some(Box::pin(tokio::time::sleep_until(deadline))); + } + } + self.idle_armed = true; + } + let timer = self + .idle_timer + .as_mut() + .expect("idle timer is created when the stream arms"); + if timer.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + self.idle_armed = false; + self.begin_finalizing(); + } } continue; } @@ -4008,6 +4032,35 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn a_stall_after_the_done_marker_ends_the_stream_cleanly() { + let inner = stream::iter(vec![ + Ok(token_event()), + Ok(terminal_control_event(b"data: [DONE]\n")), + ]) + .chain(stream::pending()); + let mut watched = Box::pin(watched_stream( + inner, + Some(StreamIdleTimeouts { + first_token: Duration::from_secs(300), + between_tokens: Duration::from_secs(90), + }), + )); + + assert!(matches!(watched.next().await, Some(Ok(_)))); + assert!(matches!(watched.next().await, Some(Ok(_)))); + assert!( + watched.next().await.is_none(), + "an upstream that sends [DONE] and then holds the connection open must not \ + pin the stream until the L4 reaper collects it" + ); + assert!( + watched.last_error.is_none(), + "the answer is complete once the marker lands, so waiting out the optional \ + blank line must not be reported as a failed stream" + ); + } + #[tokio::test] async fn a_failure_after_usage_arrives_still_reports_duration() { let event = stream_drop_event( From 7e7f6d70c1606ce21bfd8a1b086c10fe03d205d7 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Mon, 7 Sep 2026 16:53:42 -0400 Subject: [PATCH 8/9] Apply the prefill deadline per read and before the provider call --- crates/inference_providers/src/mock.rs | 52 +++- .../src/inference_provider_pool/mod.rs | 249 +++++++++++++----- 2 files changed, 234 insertions(+), 67 deletions(-) diff --git a/crates/inference_providers/src/mock.rs b/crates/inference_providers/src/mock.rs index 7164571cb..2120acea4 100644 --- a/crates/inference_providers/src/mock.rs +++ b/crates/inference_providers/src/mock.rs @@ -21,7 +21,7 @@ use bytes::Bytes; use futures_util::stream; use sha2::{Digest, Sha256}; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::{Mutex, RwLock}; /// Lightweight PII detector used only by [`MockProvider::privacy_classify_raw`] @@ -651,11 +651,13 @@ struct MockConfig { default_response: ResponseTemplate, /// When set, all chat completion calls return this error instead of generating a response error_override: Option, + call_stall_override: bool, /// When set, all chat completion streams are created successfully and then yield this error. stream_error_override: Option, /// When set, all chat completion streams are created successfully and then never /// yield an event, modelling an upstream that returns headers and goes silent. stream_stall_override: bool, + control_prelude_override: Option<(usize, Duration)>, /// When set, all embeddings calls return this error instead of generating a response embedding_error_override: Option, /// When set, all audio transcription calls return this error instead of a response. @@ -734,8 +736,10 @@ impl MockProvider { expectations: Vec::new(), default_response: ResponseTemplate::new("1. 2. 3."), error_override: None, + call_stall_override: false, stream_error_override: None, stream_stall_override: false, + control_prelude_override: None, embedding_error_override: None, audio_transcription_error_override: None, })), @@ -761,8 +765,10 @@ impl MockProvider { expectations: Vec::new(), default_response: ResponseTemplate::new("1. 2. 3."), error_override: None, + call_stall_override: false, stream_error_override: None, stream_stall_override: false, + control_prelude_override: None, embedding_error_override: None, audio_transcription_error_override: None, })), @@ -786,8 +792,10 @@ impl MockProvider { expectations: Vec::new(), default_response: ResponseTemplate::new("1. 2. 3."), error_override: None, + call_stall_override: false, stream_error_override: None, stream_stall_override: false, + control_prelude_override: None, embedding_error_override: None, audio_transcription_error_override: None, })), @@ -912,6 +920,20 @@ impl MockProvider { config.stream_stall_override = stalled; } + /// Set a call stall override: chat completion stream creation never returns, + /// as on a provider that peeks the first upstream payload before returning. + pub async fn set_call_stall_override(&self, stalled: bool) { + let mut config = self.config.lock().await; + config.call_stall_override = stalled; + } + + /// Emit `count` keepalive control events, one per `interval`, ahead of the + /// response chunks. + pub async fn set_control_prelude_override(&self, prelude: Option<(usize, Duration)>) { + let mut config = self.config.lock().await; + config.control_prelude_override = prelude; + } + /// Override the embeddings response with an error (useful for testing error paths). /// Pass `None` to clear the override. pub async fn set_embedding_error_override(&self, error: Option) { @@ -1093,11 +1115,15 @@ impl crate::InferenceProvider for MockProvider { } // Check for matching expectation (and error override) - let response_template = { + let (response_template, control_prelude) = { let config = self.config.lock().await; if let Some(ref error) = config.error_override { return Err(error.clone()); } + if config.call_stall_override { + drop(config); + return std::future::pending().await; + } if let Some(ref error) = config.stream_error_override { let stream = stream::iter(vec![Err(error.clone())]); return Ok(Box::pin(stream)); @@ -1105,12 +1131,13 @@ impl crate::InferenceProvider for MockProvider { if config.stream_stall_override { return Ok(Box::pin(stream::pending())); } - config + let template = config .expectations .iter() .find(|exp| exp.matcher.matches(¶ms)) .map(|exp| exp.response.clone()) - .unwrap_or_else(|| config.default_response.clone()) + .unwrap_or_else(|| config.default_response.clone()); + (template, config.control_prelude_override) }; // Calculate input tokens from messages (rough estimate: 1 word ≈ 1 token) @@ -1203,7 +1230,22 @@ impl crate::InferenceProvider for MockProvider { .chain(stream_error.into_iter().map(Err)), ); - Ok(Box::pin(stream)) + let Some((count, interval)) = control_prelude else { + return Ok(Box::pin(stream)); + }; + let prelude = stream::unfold(count, move |remaining| async move { + if remaining == 0 { + return None; + } + tokio::time::sleep(interval).await; + let event = SSEEvent { + raw_bytes: Bytes::from_static(b": keepalive\n\n"), + chunk: None, + raw_passthrough: true, + }; + Some((Ok(event), remaining - 1)) + }); + Ok(Box::pin(futures_util::StreamExt::chain(prelude, stream))) } async fn chat_completion( diff --git a/crates/services/src/inference_provider_pool/mod.rs b/crates/services/src/inference_provider_pool/mod.rs index 4403d98f5..85b647c79 100644 --- a/crates/services/src/inference_provider_pool/mod.rs +++ b/crates/services/src/inference_provider_pool/mod.rs @@ -167,6 +167,35 @@ fn reinsert_pubkey_pin(params: &mut ChatCompletionParams, pub_key: Option<&str>) } } +fn release_pending_route(provider: &Arc, request_hash: &str) { + provider.pin_chat_connection(request_hash, ""); + provider.unpin_chat_connection(""); +} + +async fn peek_is_control( + peekable: &mut inference_providers::PeekableStreamingResult, + idle_budget: Option, +) -> Result { + let peeked = match idle_budget { + Some(budget) => tokio::time::timeout(budget, peekable.peek()) + .await + .map_err(|_elapsed| budget)?, + None => peekable.peek().await, + }; + Ok(matches!(peeked, Some(Ok(event)) if event.chunk.is_none())) +} + +fn reattach_leading_control( + leading_control: Vec>, + peekable: inference_providers::PeekableStreamingResult, +) -> StreamingResult { + if leading_control.is_empty() { + return Box::pin(peekable); + } + use futures::StreamExt as _; + Box::pin(futures::stream::iter(leading_control).chain(peekable)) +} + fn record_backend_key_divergence( metrics: Option<&dyn crate::metrics::MetricsServiceTrait>, model_name: &str, @@ -3592,7 +3621,8 @@ impl InferenceProviderPool { .stream) } - /// `first_event_timeout` bounds the wait for the first stream event only. + /// `first_event_timeout` is a per-attempt idle budget: it bounds the wait for + /// the provider call and for each subsequent read up to the first stream event. /// Widening it to span provider selection caps the retry loop and starves /// fallback, which the per-attempt TTFB guards already cover. pub async fn chat_completion_stream_with_attribution( @@ -3649,61 +3679,53 @@ impl InferenceProviderPool { let forwarded_request_id = forwarded_request_id.clone(); let model_id = model_id.clone(); async move { - let stream = provider - .chat_completion_stream(params, request_hash.clone()) - .await?; + let fail_prefill = |budget: Duration| { + release_pending_route(&provider, &request_hash); + tracing::warn!( + model = %model_id, + request_id = forwarded_request_id.as_deref(), + timeout_seconds = budget.as_secs(), + "Upstream accepted the request and then produced no \ + stream event before the prefill deadline" + ); + CompletionError::Timeout { + operation: "prefill".to_string(), + timeout_seconds: budget.as_secs(), + } + }; + // The NEAR AI index route peeks the first payload before returning, + // so headers-then-silence hangs here, not in the drain below. + let open_stream = + provider.chat_completion_stream(params, request_hash.clone()); + let stream = match first_event_timeout { + Some(budget) => tokio::time::timeout(budget, open_stream) + .await + .map_err(|_elapsed| fail_prefill(budget))?, + None => open_stream.await, + }?; let mut peekable = StreamingResultExt::peekable(stream); let mut leading_control = Vec::new(); use futures::StreamExt as _; - let await_first_event = async { - while leading_control.len() < MAX_LEADING_CONTROL_EVENTS - && matches!(peekable.peek().await, Some(Ok(event)) if event.chunk.is_none()) - { - if let Some(event) = peekable.next().await { - leading_control.push(event); - } - } - peekable - .peek() + while leading_control.len() < MAX_LEADING_CONTROL_EVENTS + && peek_is_control(&mut peekable, first_event_timeout) .await - .and_then(|first| first.as_ref().err().cloned()) - }; - let first_error = match first_event_timeout { - Some(budget) => { - match tokio::time::timeout(budget, await_first_event).await { - Ok(first_error) => first_error, - Err(_elapsed) => { - provider.pin_chat_connection(&request_hash, ""); - provider.unpin_chat_connection(""); - tracing::warn!( - model = %model_id, - request_id = forwarded_request_id.as_deref(), - timeout_seconds = budget.as_secs(), - "Upstream accepted the request and then produced no \ - stream event before the prefill deadline" - ); - return Err(CompletionError::Timeout { - operation: "prefill".to_string(), - timeout_seconds: budget.as_secs(), - }); - } - } + .map_err(fail_prefill)? + { + if let Some(event) = peekable.next().await { + leading_control.push(event); } - None => await_first_event.await, - }; + } + let first_error = peekable + .peek() + .await + .and_then(|first| first.as_ref().err().cloned()); if let Some(error) = first_error { if Self::classify_retry_decision(&error).starts_with("retryable_") { - provider.pin_chat_connection(&request_hash, ""); - provider.unpin_chat_connection(""); + release_pending_route(&provider, &request_hash); return Err(error); } } - let primed: StreamingResult = if leading_control.is_empty() { - Box::pin(peekable) - } else { - Box::pin(futures::stream::iter(leading_control).chain(peekable)) - }; - Ok(primed) + Ok(reattach_leading_control(leading_control, peekable)) } }, ) @@ -3751,14 +3773,12 @@ impl InferenceProviderPool { let mut leading_control: Vec> = Vec::new(); let peek_first_event = async { + use futures::StreamExt as _; + while leading_control.len() < MAX_LEADING_CONTROL_EVENTS + && peek_is_control(&mut peekable, None).await == Ok(true) { - use futures::StreamExt as _; - while leading_control.len() < MAX_LEADING_CONTROL_EVENTS - && matches!(peekable.peek().await, Some(Ok(event)) if event.chunk.is_none()) - { - if let Some(ev) = peekable.next().await { - leading_control.push(ev); - } + if let Some(ev) = peekable.next().await { + leading_control.push(ev); } } @@ -3787,20 +3807,13 @@ impl InferenceProviderPool { let first_error = peek_first_event.await; if !pinned { // Clean up orphaned pending client when peek fails or yields no chat_id - provider.pin_chat_connection(&request_hash, ""); - provider.unpin_chat_connection(""); + release_pending_route(&provider, &request_hash); } if let Some(error) = first_error { return Err(error); } - let stream: StreamingResult = if leading_control.is_empty() { - Box::pin(peekable) - } else { - use futures::StreamExt as _; - Box::pin(futures::stream::iter(leading_control).chain(peekable)) - }; Ok(AttributedChatCompletionStream { - stream, + stream: reattach_leading_control(leading_control, peekable), provider_attribution, latency_reporter, }) @@ -9051,6 +9064,118 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn a_provider_that_peeks_before_returning_is_still_bounded() { + use inference_providers::mock::MockProvider; + + let pool = InferenceProviderPool::new(None, ExternalProvidersConfig::default()); + let model_id = "peeking/model".to_string(); + let provider = Arc::new(MockProvider::new_accept_all()); + provider.set_call_stall_override(true).await; + pool.register_provider(model_id.clone(), provider).await; + + let started = tokio::time::Instant::now(); + let result = pool + .chat_completion_stream_with_attribution( + fallback_params(&model_id), + "h".to_string(), + ChatRoutingHints::default(), + Some(Duration::from_secs(300)), + ) + .await; + + match result { + Err(CompletionError::Timeout { + operation, + timeout_seconds, + }) => { + assert_eq!(operation, "prefill"); + assert_eq!(timeout_seconds, 300); + assert!(started.elapsed() >= Duration::from_secs(300)); + } + Err(other) => panic!( + "a provider that waits for the first upstream payload inside its own call must \ + still hit the prefill deadline, got {other:?}" + ), + Ok(_) => panic!("a provider call that never returns must not produce a usable stream"), + } + } + + #[tokio::test(start_paused = true)] + async fn upstream_keepalives_extend_the_first_event_deadline() { + use inference_providers::mock::{MockProvider, RequestMatcher, ResponseTemplate}; + + let pool = InferenceProviderPool::new(None, ExternalProvidersConfig::default()); + let model_id = "keepalive/model".to_string(); + let provider = Arc::new(MockProvider::new_accept_all()); + provider + .when(RequestMatcher::Any) + .respond_with(ResponseTemplate::new("answered after a long prefill")) + .await; + provider + .set_control_prelude_override(Some((5, Duration::from_secs(200)))) + .await; + pool.register_provider(model_id.clone(), provider).await; + + let served = pool + .chat_completion_stream_with_attribution( + fallback_params(&model_id), + "h".to_string(), + ChatRoutingHints::default(), + Some(Duration::from_secs(300)), + ) + .await; + + assert!( + served.is_ok(), + "an upstream sending keepalives every 200s must survive a 300s deadline; the budget \ + is idle time between reads, not total time to the first data chunk" + ); + } + + #[tokio::test(start_paused = true)] + async fn a_keepalive_gap_past_the_deadline_still_times_out() { + use inference_providers::mock::{MockProvider, RequestMatcher, ResponseTemplate}; + + let pool = InferenceProviderPool::new(None, ExternalProvidersConfig::default()); + let model_id = "slow-keepalive/model".to_string(); + let provider = Arc::new(MockProvider::new_accept_all()); + provider + .when(RequestMatcher::Any) + .respond_with(ResponseTemplate::new("answered too late")) + .await; + provider + .set_control_prelude_override(Some((1, Duration::from_secs(400)))) + .await; + pool.register_provider(model_id.clone(), provider).await; + + let result = pool + .chat_completion_stream_with_attribution( + fallback_params(&model_id), + "h".to_string(), + ChatRoutingHints::default(), + Some(Duration::from_secs(300)), + ) + .await; + + match result { + Err(CompletionError::Timeout { + operation, + timeout_seconds, + }) => { + assert_eq!(operation, "prefill"); + assert_eq!(timeout_seconds, 300); + } + Err(other) => panic!( + "a single 400s gap exceeds the 300s idle budget and must surface as a prefill \ + timeout, got {other:?}" + ), + Ok(_) => panic!( + "resetting the deadline on each read must not remove the bound on any one read" + ), + } + } + fn fallback_params(model: &str) -> inference_providers::ChatCompletionParams { inference_providers::ChatCompletionParams { model: model.to_string(), From 7b5a2a98bac297d4396985e3d7a8341f18da624f Mon Sep 17 00:00:00 2001 From: neo-sky Date: Mon, 7 Sep 2026 17:19:23 -0400 Subject: [PATCH 9/9] Capture drop-path logs through the global subscriber --- crates/services/src/completions/mod.rs | 81 +++++++++++++++++--------- 1 file changed, 52 insertions(+), 29 deletions(-) diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 5cf6aa682..2790454bd 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -3524,12 +3524,24 @@ mod tests { ); } + // Callsite interest is cached process-wide, so a warn first reached by a + // drop test with no subscriber stays dead for every later thread-local one. + static ACTIVE_CAPTURE: std::sync::Mutex> = std::sync::Mutex::new(None); + #[derive(Clone, Default)] struct CapturedLogs(Arc>>); - impl std::io::Write for CapturedLogs { + struct CaptureWriter; + + impl std::io::Write for CaptureWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.0.lock().unwrap().extend_from_slice(buf); + let active = ACTIVE_CAPTURE.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(logs) = active.as_ref() { + logs.0 + .lock() + .unwrap_or_else(|e| e.into_inner()) + .extend_from_slice(buf); + } Ok(buf.len()) } @@ -3538,28 +3550,41 @@ mod tests { } } - impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CapturedLogs { + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CaptureWriter { type Writer = Self; fn make_writer(&'a self) -> Self::Writer { - self.clone() + CaptureWriter } } - impl CapturedLogs { - fn event_containing(&self, needle: &str) -> serde_json::Value { - let raw = self.0.lock().unwrap().clone(); - String::from_utf8(raw) - .expect("log output is utf8") - .lines() - .map(|line| serde_json::from_str::(line).expect("json log line")) - .find(|event| { - event["fields"]["message"] - .as_str() - .is_some_and(|message| message.contains(needle)) - }) - .unwrap_or_else(|| panic!("no log event containing {needle}")) - } + fn capture_logs_for(request_id: Uuid, emit: impl FnOnce()) -> Vec { + static CAPTURING: std::sync::Mutex<()> = std::sync::Mutex::new(()); + static INSTALLED: std::sync::Once = std::sync::Once::new(); + INSTALLED.call_once(|| { + let subscriber = tracing_subscriber::fmt() + .json() + .with_current_span(false) + .with_span_list(false) + .with_max_level(tracing::Level::WARN) + .with_writer(CaptureWriter) + .finish(); + let _ = tracing::subscriber::set_global_default(subscriber); + }); + + let _serialized = CAPTURING.lock().unwrap_or_else(|e| e.into_inner()); + let logs = CapturedLogs::default(); + *ACTIVE_CAPTURE.lock().unwrap_or_else(|e| e.into_inner()) = Some(logs.clone()); + emit(); + *ACTIVE_CAPTURE.lock().unwrap_or_else(|e| e.into_inner()) = None; + + let raw = logs.0.lock().unwrap_or_else(|e| e.into_inner()).clone(); + String::from_utf8(raw) + .expect("log output is utf8") + .lines() + .map(|line| serde_json::from_str::(line).expect("json log line")) + .filter(|event| event["fields"]["request_id"] == request_id.to_string()) + .collect() } fn interrupted_stream_event( @@ -3585,16 +3610,7 @@ mod tests { } else { "Stream interrupted before usage stats or chat_id received" }; - let logs = CapturedLogs::default(); - let subscriber = tracing_subscriber::fmt() - .json() - .with_current_span(false) - .with_span_list(false) - .with_max_level(tracing::Level::WARN) - .with_writer(logs.clone()) - .finish(); - - tracing::subscriber::with_default(subscriber, || { + let events = capture_logs_for(request_id, || { let _interrupted = InterceptStream { inner: stream::iter::>>( vec![], @@ -3640,7 +3656,14 @@ mod tests { }; }); - logs.event_containing(needle) + events + .into_iter() + .find(|event| { + event["fields"]["message"] + .as_str() + .is_some_and(|message| message.contains(needle)) + }) + .unwrap_or_else(|| panic!("no log event containing {needle}")) } #[tokio::test]