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 5541704a8..56eb37c5f 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -460,14 +460,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()); @@ -2931,6 +2944,7 @@ mod tests { staking_farm: config::StakingFarmConfig::default(), aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), + stream_watchdog: config::StreamWatchdogConfig::default(), credit_allocation: config::CreditAllocationConfig::default(), ita: config::ItaAttestationConfig::default(), }; @@ -3048,6 +3062,7 @@ mod tests { staking_farm: config::StakingFarmConfig::default(), aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), + stream_watchdog: config::StreamWatchdogConfig::default(), credit_allocation: config::CreditAllocationConfig::default(), ita: config::ItaAttestationConfig::default(), }; diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index e87c3fd30..ac983c882 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -329,6 +329,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. /// @@ -2114,10 +2118,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" ); } @@ -2900,10 +2907,13 @@ async fn completions_inner( Err(e) => { stream_error_count .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + 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 5942246b9..5726894a7 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -151,6 +151,7 @@ pub fn test_config() -> ApiConfig { enabled: true, ..config::UsageReportingConfig::default() }, + stream_watchdog: config::StreamWatchdogConfig::default(), credit_allocation: config::CreditAllocationConfig::default(), ita: config::ItaAttestationConfig::default(), } diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index be8c56370..edd07def7 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, /// Posting-time credit allocation policy. The order is persisted with /// every attributed usage charge, so changing it never rewrites history. pub credit_allocation: CreditAllocationConfig, @@ -82,6 +83,7 @@ impl ApiConfig { aml: AmlConfig::from_env()?, ita: ItaAttestationConfig::from_env()?, usage_reporting: UsageReportingConfig::from_env()?, + stream_watchdog: StreamWatchdogConfig::from_env()?, credit_allocation: CreditAllocationConfig::from_env()?, }) } @@ -373,6 +375,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 @@ -1526,6 +1579,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] #[serial] fn credit_allocation_defaults_and_policy_version_boundaries() { diff --git a/crates/inference_providers/src/attested/nearai/mod.rs b/crates/inference_providers/src/attested/nearai/mod.rs index 7fcf5ec4c..1e8c32268 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/inference_providers/src/mock.rs b/crates/inference_providers/src/mock.rs index 9fb572895..96e579bdd 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,8 +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. @@ -731,7 +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, })), @@ -757,7 +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, })), @@ -781,7 +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, })), @@ -899,6 +913,27 @@ 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; + } + + /// 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) { @@ -1076,21 +1111,29 @@ 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)); } - config + if config.stream_stall_override { + return Ok(Box::pin(stream::pending())); + } + 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) @@ -1183,7 +1226,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/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 1bb776d70..8d40ad32a 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -114,6 +114,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>, @@ -255,6 +258,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!( @@ -268,6 +278,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, @@ -294,36 +319,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; @@ -339,7 +377,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; } }; @@ -525,6 +563,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, @@ -536,6 +594,8 @@ where if matches!(&self.state, StreamState::Streaming) { match Pin::new(&mut self.inner).poll_next(cx) { Poll::Ready(Some(Ok(event))) => { + self.idle_armed = false; + if event.is_done_marker() { if event.raw_bytes.len() > MAX_PROVIDER_DONE_EVENT_BYTES { let error = inference_providers::CompletionError::CompletionError( @@ -617,10 +677,46 @@ where Poll::Ready(None) => self.begin_finalizing(), Poll::Ready(Some(Err(err))) => { // Capture error for stop_reason in usage recording (handled in Drop). + self.idle_armed = false; self.last_error = Some(err.clone()); 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; + } + 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.begin_finalizing(); + return Poll::Ready(Some(Err(timeout))); + } } continue; } @@ -659,7 +755,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; } @@ -771,6 +891,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) @@ -847,7 +968,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( @@ -889,9 +1010,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. /// @@ -1610,6 +1737,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, @@ -1779,15 +1909,18 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { }; // Get the LLM stream - let attributed_stream = match self + let provider_result = self .inference_provider_pool .chat_completion_stream_with_attribution( chat_params, request.body_hash.clone(), routing_hints, + self.stream_idle_timeouts + .map(|timeouts| timeouts.first_token), ) - .await - { + .await; + + let attributed_stream = match provider_result { Ok(pair) => pair, Err(e) => { // Guard will decrement counter on drop @@ -2459,6 +2592,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, @@ -2676,6 +2812,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, @@ -2849,6 +2988,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, @@ -3003,6 +3145,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, @@ -3131,6 +3276,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, @@ -3342,6 +3490,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, @@ -3371,6 +3522,669 @@ 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>>); + + struct CaptureWriter; + + impl std::io::Write for CaptureWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + 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()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CaptureWriter { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + CaptureWriter + } + } + + 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( + 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 events = capture_logs_for(request_id, || { + 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, + saw_upstream_done_marker: 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, + }; + }); + + 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] + 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"] + ); + } + + #[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, + { + 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, + saw_upstream_done_marker: 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_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"), + raw_passthrough: true, + chunk: None, + } + } + + #[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( + 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 cd3724150..e9bf7388b 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))); @@ -158,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, @@ -194,6 +232,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). @@ -2569,24 +2609,52 @@ 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://...) - 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(); @@ -3576,18 +3644,24 @@ 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` 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( &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); // Extract model_pub_key from params.extra for routing let model_pub_key_str = params @@ -3630,34 +3704,56 @@ impl InferenceProviderPool { |provider| { let params = params_for_provider.clone(); let request_hash = request_hash.clone(); + 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 _; while leading_control.len() < MAX_LEADING_CONTROL_EVENTS - && matches!(peekable.peek().await, Some(Ok(event)) if event.chunk.is_none()) + && peek_is_control(&mut peekable, first_event_timeout) + .await + .map_err(fail_prefill)? { if let Some(event) = peekable.next().await { leading_control.push(event); } } - if let Some(Err(error)) = peekable.peek().await { - if Self::classify_retry_decision(error).starts_with("retryable_") { - let error = error.clone(); - provider.pin_chat_connection(&request_hash, ""); - provider.unpin_chat_connection(""); + 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_") { + 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)) } }, ) @@ -3704,52 +3800,48 @@ impl InferenceProviderPool { // cap we return the stream without pinning a sticky-routing mapping. let mut leading_control: Vec> = Vec::new(); - { + let peek_first_event = async { use futures::StreamExt as _; while leading_control.len() < MAX_LEADING_CONTROL_EVENTS - && matches!(peekable.peek().await, Some(Ok(event)) if event.chunk.is_none()) + && peek_is_control(&mut peekable, None).await == Ok(true) { 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, - "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; + 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 } - None + Some(Err(error)) => Some(error.clone()), + None => None, } - Some(Err(error)) => Some(error.clone()), - None => None, }; + + 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, }) @@ -3786,6 +3878,7 @@ impl InferenceProviderPool { mut hints: ChatRoutingHints, ) -> Result { let model_id = params.model.clone(); + let forwarded_request_id = forwarded_request_id(¶ms.extra); // Non-streaming requests may carry policy hints such as // `fallback_disabled`, but do not carry the stream-side prefix hash or // token estimate. Multi-capacity models still get context routing @@ -3848,6 +3941,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; @@ -9057,6 +9151,200 @@ 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" + ); + } + + #[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(), @@ -10737,6 +11025,7 @@ mod tests { fallback_disabled: true, ..Default::default() }, + None, ) .await; assert!(disabled.is_err(), "the primary stream error must surface"); @@ -10750,6 +11039,7 @@ mod tests { fallback_params(&model_id), "enabled".to_string(), ChatRoutingHints::default(), + None, ) .await .expect("default policy should allow the registered fallback stream"); diff --git a/env.example b/env.example index d944dbfa6..48b40c5b6 100644 --- a/env.example +++ b/env.example @@ -166,6 +166,18 @@ 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 + # Credit allocation # ============================================================================= # Posting-time funding priority. Use each API/database credit type exactly once.