From d25184f800bcc353badb59cdd75d921501f4af10 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Thu, 27 Aug 2026 18:31:38 -0400 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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/5] 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();