diff --git a/Cargo.lock b/Cargo.lock index ccc796526..e79fe7c84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6562,6 +6562,7 @@ dependencies = [ "tokio-stream", "tokio-test", "tracing", + "tracing-subscriber", "url", "urlencoding", "utoipa", diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index b800ef9db..a3fbc1bf2 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -440,14 +440,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()); @@ -2824,6 +2837,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(), }; @@ -2938,6 +2952,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 f1aaa7759..f5c7a5e12 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -320,6 +320,10 @@ fn completion_stream_error_category(e: &inference_providers::CompletionError) -> } } +fn sanitized_stream_error(e: &inference_providers::CompletionError) -> String { + services::inference_provider_pool::InferenceProviderPool::safe_error_detail(e) +} + /// 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,10 +1852,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, "Completion stream error" ); } @@ -2492,10 +2499,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/api/tests/common/mod.rs b/crates/api/tests/common/mod.rs index 561a52b64..0e70c5c1b 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -136,6 +136,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 8d0cfa260..54d2c2077 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -35,6 +35,7 @@ pub struct ApiConfig { pub staking_farm: StakingFarmConfig, pub aml: AmlConfig, pub usage_reporting: UsageReportingConfig, + pub stream_watchdog: StreamWatchdogConfig, pub ita: ItaAttestationConfig, } @@ -79,6 +80,7 @@ impl ApiConfig { aml: AmlConfig::from_env()?, ita: ItaAttestationConfig::from_env()?, usage_reporting: UsageReportingConfig::from_env()?, + stream_watchdog: StreamWatchdogConfig::from_env()?, }) } } @@ -304,6 +306,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 @@ -1255,6 +1308,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/inference_providers/src/attested/nearai/mod.rs b/crates/inference_providers/src/attested/nearai/mod.rs index 22a4a6c72..98f6a0f1c 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/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..432e93902 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>, @@ -187,6 +190,13 @@ 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| { + super::inference_provider_pool::InferenceProviderPool::safe_error_detail(error) + }); // Create span with context BEFORE any early returns so all error logs have context let _span = tracing::error_span!( @@ -200,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, @@ -226,36 +251,49 @@ 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!(%organization_id, %model_id, model = %self.model_name, - stream_completed = self.stream_completed, - stream_error = self.last_error.is_some(), - "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, - "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!(%chat_id, %organization_id, %model_id, model = %self.model_name, - stream_completed = self.stream_completed, - stream_error = self.last_error.is_some(), - "Stream interrupted before usage stats received (client disconnect or provider error)"); - } else { - tracing::error!(%chat_id, %organization_id, %model_id, model = %self.model_name, - "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; } (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; @@ -271,7 +309,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; } }; @@ -463,6 +501,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, @@ -479,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()))); } @@ -546,7 +606,42 @@ 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()); + self.state = StreamState::Done; + return Poll::Ready(Some(Err(timeout))); + } } } StreamState::Finalizing(ref mut future) => match future.as_mut().poll(cx) { @@ -619,6 +714,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) @@ -692,7 +788,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( @@ -734,9 +830,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. /// @@ -1454,6 +1556,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, @@ -1621,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 @@ -2353,6 +2469,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, @@ -2525,6 +2644,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, @@ -2678,6 +2800,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, @@ -2805,6 +2930,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, @@ -3015,6 +3143,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, @@ -3043,6 +3174,438 @@ 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}")) + } + } + + fn interrupted_stream_event( + request_id: Uuid, + last_token_time: Option, + last_chat_id: Option, + last_error: Option, + ) -> serde_json::Value { + 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" + }; + 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, + idle_timeouts: None, + idle_timer: None, + idle_armed: false, + metric_tags: vec![], + concurrent_counter: None, + last_usage_stats, + 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"] + ); + } + + #[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()); + } + + #[tokio::test] + async fn an_interrupted_stream_reports_the_error_it_holds() { + 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 = + interrupted_stream_event(Uuid::new_v4(), None, chat_id, Some(failure.clone())); + + 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}" + ); + } + + 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" + ); + } + + #[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"] + ); + } + + 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" + ); + } + + 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()); + 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/crates/services/src/inference_provider_pool/mod.rs b/crates/services/src/inference_provider_pool/mod.rs index b82650554..b537bda17 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))); @@ -194,6 +203,8 @@ fn record_backend_key_divergence( /// 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). @@ -2422,8 +2433,32 @@ 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 - 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://...) @@ -3443,6 +3478,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 @@ -3548,6 +3584,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 @@ -3599,6 +3636,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 @@ -3663,6 +3701,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; diff --git a/env.example b/env.example index 3ffc5fde7..f200200f5 100644 --- a/env.example +++ b/env.example @@ -144,6 +144,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) # =============================================================================