From fde246f7077694429daa7411b1ae8e55d691f788 Mon Sep 17 00:00:00 2001 From: wsp1911 Date: Wed, 24 Jun 2026 18:21:08 +0800 Subject: [PATCH] fix(ai): align stream TTFT with first effective output - apply TTFT to first effective streamed output instead of HTTP 200 - allow null TTFT config to wait indefinitely - remove reasoning-specific TTFT override - clarify stream timeout UI copy and next-turn activation --- src/crates/adapters/ai-adapters/src/client.rs | 12 +- .../adapters/ai-adapters/src/client/sse.rs | 24 +++- src/crates/adapters/ai-adapters/src/lib.rs | 2 +- .../src/providers/anthropic/request.rs | 3 +- .../src/providers/gemini/code_assist.rs | 10 +- .../src/providers/gemini/request.rs | 10 +- .../ai-adapters/src/providers/openai/chat.rs | 3 +- .../src/providers/openai/codex_chatgpt.rs | 10 +- .../src/providers/openai/responses.rs | 10 +- .../src/stream/stream_handler/anthropic.rs | 24 +++- .../src/stream/stream_handler/gemini.rs | 20 ++- .../src/stream/stream_handler/mod.rs | 114 +++++++++++++++++- .../src/stream/stream_handler/openai.rs | 20 ++- .../src/stream/stream_handler/responses.rs | 107 +++++++++++++--- .../tests/common/sse_fixture_server.rs | 5 + .../tests/common/stream_test_harness.rs | 16 ++- .../ai-adapters/tests/stream_test_harness.rs | 29 +++++ .../core/src/infrastructure/ai/mod.rs | 44 ++++--- .../assembly/core/src/service/config/types.rs | 22 +++- .../config/components/AIModelConfig.tsx | 38 +++++- .../src/locales/en-US/settings/ai-model.json | 3 +- .../src/locales/zh-CN/settings/ai-model.json | 3 +- .../src/locales/zh-TW/settings/ai-model.json | 3 +- 23 files changed, 445 insertions(+), 87 deletions(-) diff --git a/src/crates/adapters/ai-adapters/src/client.rs b/src/crates/adapters/ai-adapters/src/client.rs index e36537d774..1d56482420 100644 --- a/src/crates/adapters/ai-adapters/src/client.rs +++ b/src/crates/adapters/ai-adapters/src/client.rs @@ -37,22 +37,20 @@ pub struct StreamResponse { pub trace_handle: Option, } -/// Default time to wait for the first response headers / stream body to start. +/// Default time to wait for the first effective streamed output after a request starts. pub const DEFAULT_STREAM_TTFT_TIMEOUT_SECS: u64 = 30; /// Default idle time between streamed chunks once the stream has started. pub const DEFAULT_STREAM_IDLE_TIMEOUT_SECS: u64 = 45; -/// Minimum TTFT for models with explicit reasoning enabled. -pub const REASONING_STREAM_TTFT_TIMEOUT_SECS: u64 = 45; - /// Runtime stream behavior shared across provider implementations. #[derive(Debug, Clone, Default)] pub struct StreamOptions { /// Maximum idle time between streamed chunks. `None` means wait indefinitely. pub idle_timeout: Option, - /// Maximum time to wait for HTTP response headers when opening a stream. - /// `None` means wait indefinitely. + /// Maximum time to wait for the first effective streamed output (text, + /// reasoning, or tool-call data) after a request starts. `None` means wait + /// indefinitely. pub ttft_timeout: Option, } @@ -100,7 +98,7 @@ impl AIClient { self.stream_options.idle_timeout } - /// Returns the configured time-to-first-token timeout for opening a stream, if any. + /// Returns the configured timeout for the first effective streamed output, if any. pub fn stream_ttft_timeout(&self) -> Option { self.stream_options.ttft_timeout } diff --git a/src/crates/adapters/ai-adapters/src/client/sse.rs b/src/crates/adapters/ai-adapters/src/client/sse.rs index c10e223079..11a7f05ac3 100644 --- a/src/crates/adapters/ai-adapters/src/client/sse.rs +++ b/src/crates/adapters/ai-adapters/src/client/sse.rs @@ -55,11 +55,18 @@ where fn format_ttft_timeout_error(label: &str, ttft_timeout: Option) -> String { let timeout_secs = ttft_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); format!( - "{} TTFT timeout after {}s waiting for response headers", + "{} TTFT timeout after {}s waiting for first effective stream output", label, timeout_secs ) } +fn remaining_ttft_timeout( + started_at: std::time::Instant, + ttft_timeout: Option, +) -> Option { + ttft_timeout.map(|timeout| timeout.saturating_sub(started_at.elapsed())) +} + fn format_transport_error(label: &str, error: &reqwest::Error) -> String { format!("{} connection failed: {}", label, error) } @@ -116,6 +123,7 @@ where reqwest::Response, mpsc::UnboundedSender>, Option>, + Option, ), { let mut last_error = None; @@ -279,7 +287,8 @@ where let (tx, rx) = mpsc::unbounded_channel(); let (tx_raw, rx_raw) = mpsc::unbounded_channel(); - spawn_handler(response, tx, Some(tx_raw)); + let remaining_ttft_timeout = remaining_ttft_timeout(request_start_time, ttft_timeout); + spawn_handler(response, tx, Some(tx_raw), remaining_ttft_timeout); return Ok(StreamResponse { stream: Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(rx)), @@ -311,6 +320,17 @@ mod tests { ); assert!(message.contains("TTFT timeout after 30s")); + assert!(message.contains("first effective stream output")); + } + + #[test] + fn remaining_ttft_timeout_subtracts_elapsed_request_time() { + let start = std::time::Instant::now() - Duration::from_secs(2); + let remaining = remaining_ttft_timeout(start, Some(Duration::from_secs(5))); + + let remaining = remaining.expect("remaining timeout"); + assert!(remaining <= Duration::from_secs(3)); + assert!(remaining > Duration::from_secs(2)); } #[test] diff --git a/src/crates/adapters/ai-adapters/src/lib.rs b/src/crates/adapters/ai-adapters/src/lib.rs index b9ffa1d5a9..a095b624d7 100644 --- a/src/crates/adapters/ai-adapters/src/lib.rs +++ b/src/crates/adapters/ai-adapters/src/lib.rs @@ -11,7 +11,7 @@ pub mod types; pub use client::{ AIClient, StreamOptions, StreamResponse, DEFAULT_STREAM_IDLE_TIMEOUT_SECS, - DEFAULT_STREAM_TTFT_TIMEOUT_SECS, REASONING_STREAM_TTFT_TIMEOUT_SECS, + DEFAULT_STREAM_TTFT_TIMEOUT_SECS, }; pub use model_selector::{ classify_model_selector, resolve_cache_model_selector, resolve_required_model_selector, diff --git a/src/crates/adapters/ai-adapters/src/providers/anthropic/request.rs b/src/crates/adapters/ai-adapters/src/providers/anthropic/request.rs index baed0fad17..888a14f740 100644 --- a/src/crates/adapters/ai-adapters/src/providers/anthropic/request.rs +++ b/src/crates/adapters/ai-adapters/src/providers/anthropic/request.rs @@ -350,12 +350,13 @@ pub(crate) async fn send_stream( ttft_timeout, trace, || apply_headers(client, client.client.post(&url), &url), - move |response, tx, tx_raw| { + move |response, tx, tx_raw, remaining_ttft_timeout| { tokio::spawn(handle_anthropic_stream( response, tx, tx_raw, inline_think_in_text, + remaining_ttft_timeout, idle_timeout, )); }, diff --git a/src/crates/adapters/ai-adapters/src/providers/gemini/code_assist.rs b/src/crates/adapters/ai-adapters/src/providers/gemini/code_assist.rs index 0c00f7312a..9cad79192f 100644 --- a/src/crates/adapters/ai-adapters/src/providers/gemini/code_assist.rs +++ b/src/crates/adapters/ai-adapters/src/providers/gemini/code_assist.rs @@ -182,8 +182,14 @@ pub(crate) async fn send_stream( ttft_timeout, trace, || apply_headers(client, client.client.post(&url)), - move |response, tx, tx_raw| { - tokio::spawn(handle_gemini_stream(response, tx, tx_raw, idle_timeout)); + move |response, tx, tx_raw, remaining_ttft_timeout| { + tokio::spawn(handle_gemini_stream( + response, + tx, + tx_raw, + remaining_ttft_timeout, + idle_timeout, + )); }, ) .await diff --git a/src/crates/adapters/ai-adapters/src/providers/gemini/request.rs b/src/crates/adapters/ai-adapters/src/providers/gemini/request.rs index f865c2829b..8df79808cf 100644 --- a/src/crates/adapters/ai-adapters/src/providers/gemini/request.rs +++ b/src/crates/adapters/ai-adapters/src/providers/gemini/request.rs @@ -346,8 +346,14 @@ pub(crate) async fn send_stream( ttft_timeout, trace, || apply_headers(client, client.client.post(&url)), - move |response, tx, tx_raw| { - tokio::spawn(handle_gemini_stream(response, tx, tx_raw, idle_timeout)); + move |response, tx, tx_raw, remaining_ttft_timeout| { + tokio::spawn(handle_gemini_stream( + response, + tx, + tx_raw, + remaining_ttft_timeout, + idle_timeout, + )); }, ) .await diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/chat.rs b/src/crates/adapters/ai-adapters/src/providers/openai/chat.rs index e45704bf45..be8a08c0ba 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/chat.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/chat.rs @@ -100,12 +100,13 @@ pub(crate) async fn send_stream( ttft_timeout, trace, || common::apply_headers(client, client.client.post(&url)), - move |response, tx, tx_raw| { + move |response, tx, tx_raw, remaining_ttft_timeout| { tokio::spawn(handle_openai_stream( response, tx, tx_raw, inline_think_in_text, + remaining_ttft_timeout, idle_timeout, )); }, diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs index decd573626..509e68c6c6 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs @@ -170,8 +170,14 @@ pub(crate) async fn send_stream( ttft_timeout, trace, || common::apply_headers(client, client.client.post(&url)), - move |response, tx, tx_raw| { - tokio::spawn(handle_responses_stream(response, tx, tx_raw, idle_timeout)); + move |response, tx, tx_raw, remaining_ttft_timeout| { + tokio::spawn(handle_responses_stream( + response, + tx, + tx_raw, + remaining_ttft_timeout, + idle_timeout, + )); }, ) .await diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs index 359038b7b4..b5fc6527c5 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs @@ -134,8 +134,14 @@ pub(crate) async fn send_stream( ttft_timeout, trace, || common::apply_headers(client, client.client.post(&url)), - move |response, tx, tx_raw| { - tokio::spawn(handle_responses_stream(response, tx, tx_raw, idle_timeout)); + move |response, tx, tx_raw, remaining_ttft_timeout| { + tokio::spawn(handle_responses_stream( + response, + tx, + tx_raw, + remaining_ttft_timeout, + idle_timeout, + )); }, ) .await diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs index eb6498cff9..49e2595edd 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs @@ -1,6 +1,6 @@ use super::inline_think::InlineThinkParser; use super::stream_stats::StreamStats; -use super::{next_stream_item, TimedStreamItem}; +use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, TimedStreamItem}; use crate::stream::types::anthropic::{ AnthropicSSEError, ContentBlock, ContentBlockDelta, ContentBlockStart, MessageDelta, MessageStart, Usage, @@ -26,16 +26,18 @@ pub async fn handle_anthropic_stream( tx_event: mpsc::UnboundedSender>, tx_raw_sse: Option>, inline_think_in_text: bool, + ttft_timeout: Option, idle_timeout: Option, ) { let mut stream = response.bytes_stream().eventsource(); let mut usage = Usage::default(); let mut stats = StreamStats::new("Anthropic"); let mut inline_think_parser = InlineThinkParser::new(inline_think_in_text); + let mut timeout_controller = StreamTimeoutController::new(ttft_timeout, idle_timeout); let mut received_finish_reason = false; loop { - let sse = match next_stream_item(&mut stream, idle_timeout).await { + let sse = match next_stream_item(&mut stream, &timeout_controller).await { TimedStreamItem::Item(Ok(sse)) => sse, TimedStreamItem::End => { if received_finish_reason { @@ -60,7 +62,18 @@ pub async fn handle_anthropic_stream( let _ = tx_event.send(Err(anyhow!(error_msg))); return; } - TimedStreamItem::TimedOut => { + TimedStreamItem::TimedOut(StreamTimeoutStage::Ttft) => { + let timeout_secs = ttft_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); + let error_msg = format!( + "Anthropic stream TTFT timeout after {}s waiting for first effective output", + timeout_secs + ); + stats.log_summary("ttft_timeout"); + error!("{}", error_msg); + let _ = tx_event.send(Err(anyhow!(error_msg))); + return; + } + TimedStreamItem::TimedOut(StreamTimeoutStage::Idle) => { let timeout_secs = idle_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); let error_msg = format!( "SSE Timeout: idle timeout waiting for SSE after {}s", @@ -133,6 +146,7 @@ pub async fn handle_anthropic_stream( ) { emit_normalized_response( &mut inline_think_parser, + &mut timeout_controller, &tx_event, &mut stats, UnifiedResponse::from(content_block_start), @@ -154,6 +168,7 @@ pub async fn handle_anthropic_stream( match UnifiedResponse::try_from(content_block_delta) { Ok(unified_response) => emit_normalized_response( &mut inline_think_parser, + &mut timeout_controller, &tx_event, &mut stats, unified_response, @@ -188,6 +203,7 @@ pub async fn handle_anthropic_stream( } emit_normalized_response( &mut inline_think_parser, + &mut timeout_controller, &tx_event, &mut stats, unified_response, @@ -344,11 +360,13 @@ fn trace_unified_response_if_useful(response: &UnifiedResponse) { fn emit_normalized_response( inline_think_parser: &mut InlineThinkParser, + timeout_controller: &mut StreamTimeoutController, tx_event: &mpsc::UnboundedSender>, stats: &mut StreamStats, unified_response: UnifiedResponse, ) { for normalized_response in inline_think_parser.normalize_response(unified_response) { + timeout_controller.observe_unified_response(&normalized_response); trace_unified_response_if_useful(&normalized_response); stats.record_unified_response(&normalized_response); let _ = tx_event.send(Ok(normalized_response)); diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs index 49b0e1a3cd..a274db0505 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs @@ -1,5 +1,5 @@ use super::stream_stats::StreamStats; -use super::{next_stream_item, TimedStreamItem}; +use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, TimedStreamItem}; use crate::stream::types::gemini::GeminiSSEData; use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; @@ -80,15 +80,17 @@ pub async fn handle_gemini_stream( response: Response, tx_event: mpsc::UnboundedSender>, tx_raw_sse: Option>, + ttft_timeout: Option, idle_timeout: Option, ) { let mut stream = response.bytes_stream().eventsource(); let mut received_finish_reason = false; let mut tool_call_state = GeminiToolCallState::new(); let mut stats = StreamStats::new("Gemini"); + let mut timeout_controller = StreamTimeoutController::new(ttft_timeout, idle_timeout); loop { - let sse = match next_stream_item(&mut stream, idle_timeout).await { + let sse = match next_stream_item(&mut stream, &timeout_controller).await { TimedStreamItem::Item(Ok(sse)) => sse, TimedStreamItem::End => { if received_finish_reason { @@ -108,7 +110,18 @@ pub async fn handle_gemini_stream( let _ = tx_event.send(Err(anyhow!(error_msg))); return; } - TimedStreamItem::TimedOut => { + TimedStreamItem::TimedOut(StreamTimeoutStage::Ttft) => { + let timeout_secs = ttft_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); + let error_msg = format!( + "Gemini stream TTFT timeout after {}s waiting for first effective output", + timeout_secs + ); + stats.log_summary("ttft_timeout"); + error!("{}", error_msg); + let _ = tx_event.send(Err(anyhow!(error_msg))); + return; + } + TimedStreamItem::TimedOut(StreamTimeoutStage::Idle) => { let timeout_secs = idle_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); let error_msg = format!("Gemini SSE stream timeout after {}s", timeout_secs); stats.log_summary("sse_stream_timeout"); @@ -189,6 +202,7 @@ pub async fn handle_gemini_stream( ); for unified_response in unified_responses { + timeout_controller.observe_unified_response(&unified_response); stats.record_unified_response(&unified_response); let _ = tx_event.send(Ok(unified_response)); } diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/mod.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/mod.rs index 5691d99c4c..a759a39866 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/mod.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/mod.rs @@ -6,31 +6,93 @@ mod responses; mod stream_stats; use futures::{Stream, StreamExt}; -use std::time::Duration; +use std::time::{Duration, Instant}; + +use crate::stream::types::unified::UnifiedResponse; pub use anthropic::handle_anthropic_stream; pub use gemini::handle_gemini_stream; pub use openai::handle_openai_stream; pub use responses::handle_responses_stream; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StreamTimeoutStage { + Ttft, + Idle, +} + pub(super) enum TimedStreamItem { Item(T), End, - TimedOut, + TimedOut(StreamTimeoutStage), +} + +pub(super) struct StreamTimeoutController { + first_effective_output_deadline: Option, + idle_timeout: Option, + first_effective_output_seen: bool, +} + +impl StreamTimeoutController { + pub(super) fn new( + ttft_timeout: Option, + idle_timeout: Option, + ) -> Self { + Self { + first_effective_output_deadline: ttft_timeout.map(|timeout| Instant::now() + timeout), + idle_timeout, + first_effective_output_seen: false, + } + } + + pub(super) fn observe_unified_response(&mut self, response: &UnifiedResponse) { + if is_effective_stream_output(response) { + self.first_effective_output_seen = true; + } + } + + pub(super) fn timeout_for_wait(&self) -> (Option, StreamTimeoutStage) { + if !self.first_effective_output_seen { + return ( + self.first_effective_output_deadline + .map(|deadline| deadline.saturating_duration_since(Instant::now())), + StreamTimeoutStage::Ttft, + ); + } + + (self.idle_timeout, StreamTimeoutStage::Idle) + } +} + +fn is_effective_stream_output(response: &UnifiedResponse) -> bool { + response.text.as_ref().is_some_and(|text| !text.is_empty()) + || response + .reasoning_content + .as_ref() + .is_some_and(|reasoning| !reasoning.is_empty()) + || response.tool_call.as_ref().is_some_and(|tool_call| { + tool_call.id.is_some() + || tool_call.name.is_some() + || tool_call + .arguments + .as_ref() + .is_some_and(|arguments| !arguments.is_empty()) + }) } pub(super) async fn next_stream_item( stream: &mut S, - idle_timeout: Option, + timeout_controller: &StreamTimeoutController, ) -> TimedStreamItem where S: Stream + Unpin, { - match idle_timeout { - Some(idle_timeout) => match tokio::time::timeout(idle_timeout, stream.next()).await { + let (timeout, stage) = timeout_controller.timeout_for_wait(); + match timeout { + Some(timeout) => match tokio::time::timeout(timeout, stream.next()).await { Ok(Some(item)) => TimedStreamItem::Item(item), Ok(None) => TimedStreamItem::End, - Err(_) => TimedStreamItem::TimedOut, + Err(_) => TimedStreamItem::TimedOut(stage), }, None => match stream.next().await { Some(item) => TimedStreamItem::Item(item), @@ -38,3 +100,43 @@ where }, } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::stream::types::unified::UnifiedToolCall; + + #[test] + fn effective_output_includes_text_reasoning_and_tool_calls() { + assert!(is_effective_stream_output(&UnifiedResponse { + text: Some("hello".to_string()), + ..Default::default() + })); + assert!(is_effective_stream_output(&UnifiedResponse { + reasoning_content: Some("thinking".to_string()), + ..Default::default() + })); + assert!(is_effective_stream_output(&UnifiedResponse { + tool_call: Some(UnifiedToolCall { + tool_call_index: Some(0), + id: Some("call_1".to_string()), + name: Some("search".to_string()), + arguments: None, + arguments_is_snapshot: false, + }), + ..Default::default() + })); + } + + #[test] + fn empty_control_only_response_is_not_effective_output() { + assert!(!is_effective_stream_output(&UnifiedResponse { + finish_reason: Some("stop".to_string()), + ..Default::default() + })); + assert!(!is_effective_stream_output(&UnifiedResponse { + provider_metadata: Some(serde_json::json!({ "status": "ok" })), + ..Default::default() + })); + } +} diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs index 5ec482f0ff..cc4ab202bd 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs @@ -1,6 +1,6 @@ use super::inline_think::InlineThinkParser; use super::stream_stats::StreamStats; -use super::{next_stream_item, TimedStreamItem}; +use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, TimedStreamItem}; use crate::stream::types::openai::OpenAISSEData; use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; @@ -73,10 +73,12 @@ pub async fn handle_openai_stream( tx_event: mpsc::UnboundedSender>, tx_raw_sse: Option>, inline_think_in_text: bool, + ttft_timeout: Option, idle_timeout: Option, ) { let mut stream = response.bytes_stream().eventsource(); let mut stats = StreamStats::new("OpenAI"); + let mut timeout_controller = StreamTimeoutController::new(ttft_timeout, idle_timeout); // Track whether a chunk with `finish_reason` was received. // Some providers (e.g. MiniMax) close the stream after the final chunk // without sending `[DONE]`, so we treat `Ok(None)` as a normal termination @@ -85,7 +87,7 @@ pub async fn handle_openai_stream( let mut normalizer = OpenAIResponseNormalizer::new(inline_think_in_text); loop { - let sse = match next_stream_item(&mut stream, idle_timeout).await { + let sse = match next_stream_item(&mut stream, &timeout_controller).await { TimedStreamItem::Item(Ok(sse)) => sse, TimedStreamItem::End => { if received_finish_reason { @@ -109,7 +111,18 @@ pub async fn handle_openai_stream( let _ = tx_event.send(Err(anyhow!(error_msg))); return; } - TimedStreamItem::TimedOut => { + TimedStreamItem::TimedOut(StreamTimeoutStage::Ttft) => { + let timeout_secs = ttft_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); + let error_msg = format!( + "OpenAI stream TTFT timeout after {}s waiting for first effective output", + timeout_secs + ); + stats.log_summary("ttft_timeout"); + error!("{}", error_msg); + let _ = tx_event.send(Err(anyhow!(error_msg))); + return; + } + TimedStreamItem::TimedOut(StreamTimeoutStage::Idle) => { let timeout_secs = idle_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); let error_msg = format!("SSE stream timeout after {}s", timeout_secs); stats.log_summary("sse_stream_timeout"); @@ -224,6 +237,7 @@ pub async fn handle_openai_stream( } for normalized_response in normalized_responses { + timeout_controller.observe_unified_response(&normalized_response); if normalized_response.finish_reason.is_some() { received_finish_reason = true; } diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs index ffb4616b70..25fbf05d33 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs @@ -1,5 +1,5 @@ use super::stream_stats::StreamStats; -use super::{next_stream_item, TimedStreamItem}; +use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, TimedStreamItem}; use crate::stream::types::responses::{ parse_responses_output_item, ResponsesCompleted, ResponsesDone, ResponsesStreamEvent, }; @@ -46,10 +46,12 @@ impl InProgressToolCall { } fn emit_unified_response( + timeout_controller: &mut StreamTimeoutController, tx_event: &mpsc::UnboundedSender>, stats: &mut StreamStats, unified_response: UnifiedResponse, ) { + timeout_controller.observe_unified_response(&unified_response); trace!( target: AI_STREAM_RESPONSE_TARGET, "Responses unified response: {:?}", @@ -60,6 +62,7 @@ fn emit_unified_response( } fn emit_tool_call_item( + timeout_controller: &mut StreamTimeoutController, tx_event: &mpsc::UnboundedSender>, stats: &mut StreamStats, output_index: Option, @@ -67,7 +70,7 @@ fn emit_tool_call_item( ) { if let Some(unified_response) = parse_responses_output_item(item_value, output_index) { if unified_response.tool_call.is_some() { - emit_unified_response(tx_event, stats, unified_response); + emit_unified_response(timeout_controller, tx_event, stats, unified_response); } } } @@ -85,6 +88,7 @@ fn cleanup_tool_call_tracking( } fn handle_function_call_arguments_delta( + timeout_controller: &mut StreamTimeoutController, tx_event: &mpsc::UnboundedSender>, stats: &mut StreamStats, output_index: Option, @@ -128,11 +132,12 @@ fn handle_function_call_arguments_delta( }), ..Default::default() }; - emit_unified_response(tx_event, stats, unified_response); + emit_unified_response(timeout_controller, tx_event, stats, unified_response); Ok(()) } fn handle_function_call_output_item_done( + timeout_controller: &mut StreamTimeoutController, tx_event: &mpsc::UnboundedSender>, stats: &mut StreamStats, event_output_index: Option, @@ -149,14 +154,14 @@ fn handle_function_call_output_item_done( }); let Some(output_index) = output_index else { - emit_tool_call_item(tx_event, stats, event_output_index, item_value); + emit_tool_call_item(timeout_controller, tx_event, stats, event_output_index, item_value); return; }; let Some(tc) = tool_calls_by_output_index.get_mut(&output_index) else { // The provider may send `output_item.done` with an output_index even when the // earlier `output_item.added` event was omitted or missed. Fall back to the full item. - emit_tool_call_item(tx_event, stats, Some(output_index), item_value); + emit_tool_call_item(timeout_controller, tx_event, stats, Some(output_index), item_value); return; }; @@ -194,7 +199,7 @@ fn handle_function_call_output_item_done( }), ..Default::default() }; - emit_unified_response(tx_event, stats, unified_response); + emit_unified_response(timeout_controller, tx_event, stats, unified_response); } } @@ -227,6 +232,7 @@ pub async fn handle_responses_stream( response: Response, tx_event: mpsc::UnboundedSender>, tx_raw_sse: Option>, + ttft_timeout: Option, idle_timeout: Option, ) { let mut stream = response.bytes_stream().eventsource(); @@ -236,9 +242,10 @@ pub async fn handle_responses_stream( let mut tool_calls_by_output_index: HashMap = HashMap::new(); let mut tool_call_index_by_id: HashMap = HashMap::new(); let mut stats = StreamStats::new("Responses"); + let mut timeout_controller = StreamTimeoutController::new(ttft_timeout, idle_timeout); loop { - let sse = match next_stream_item(&mut stream, idle_timeout).await { + let sse = match next_stream_item(&mut stream, &timeout_controller).await { TimedStreamItem::Item(Ok(sse)) => sse, TimedStreamItem::End => { if received_finish_reason { @@ -258,7 +265,18 @@ pub async fn handle_responses_stream( let _ = tx_event.send(Err(anyhow!(error_msg))); return; } - TimedStreamItem::TimedOut => { + TimedStreamItem::TimedOut(StreamTimeoutStage::Ttft) => { + let timeout_secs = ttft_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); + let error_msg = format!( + "Responses stream TTFT timeout after {}s waiting for first effective output", + timeout_secs + ); + stats.log_summary("ttft_timeout"); + error!("{}", error_msg); + let _ = tx_event.send(Err(anyhow!(error_msg))); + return; + } + TimedStreamItem::TimedOut(StreamTimeoutStage::Idle) => { let timeout_secs = idle_timeout.map(|timeout| timeout.as_secs()).unwrap_or(0); let error_msg = format!("Responses SSE stream timeout after {}s", timeout_secs); stats.log_summary("sse_stream_timeout"); @@ -345,7 +363,12 @@ pub async fn handle_responses_stream( text: Some(delta), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); } } "response.reasoning_text.delta" | "response.reasoning_summary_text.delta" => { @@ -354,11 +377,17 @@ pub async fn handle_responses_stream( reasoning_content: Some(delta), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); } } "response.function_call_arguments.delta" => { if let Err(err) = handle_function_call_arguments_delta( + &mut timeout_controller, &tx_event, &mut stats, event.output_index, @@ -381,6 +410,7 @@ pub async fn handle_responses_stream( // For tool calls, prefer streaming deltas and only use item.done as a tail-filler / fallback. if item_value.get("type").and_then(Value::as_str) == Some("function_call") { handle_function_call_output_item_done( + &mut timeout_controller, &tx_event, &mut stats, event.output_index, @@ -398,7 +428,12 @@ pub async fn handle_responses_stream( unified_response.text = None; } if unified_response.text.is_some() || unified_response.tool_call.is_some() { - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); } } } @@ -444,7 +479,12 @@ pub async fn handle_responses_stream( ), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); } } } @@ -461,7 +501,12 @@ pub async fn handle_responses_stream( finish_reason: Some("stop".to_string()), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); continue; } Some(Err(e)) => { @@ -479,7 +524,12 @@ pub async fn handle_responses_stream( finish_reason: Some("stop".to_string()), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); continue; } } @@ -496,7 +546,12 @@ pub async fn handle_responses_stream( finish_reason: Some("stop".to_string()), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); continue; } Some(Err(e)) => { @@ -513,7 +568,12 @@ pub async fn handle_responses_stream( finish_reason: Some("stop".to_string()), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); continue; } } @@ -565,7 +625,12 @@ pub async fn handle_responses_stream( finish_reason: Some(finish_reason), ..Default::default() }; - emit_unified_response(&tx_event, &mut stats, unified_response); + emit_unified_response( + &mut timeout_controller, + &tx_event, + &mut stats, + unified_response, + ); continue; } _ => {} @@ -578,7 +643,7 @@ mod tests { use super::{ super::stream_stats::StreamStats, extract_api_error_message, handle_function_call_arguments_delta, handle_function_call_output_item_done, - InProgressToolCall, + InProgressToolCall, StreamTimeoutController, }; use serde_json::json; use std::collections::HashMap; @@ -632,8 +697,10 @@ mod tests { let mut tool_calls_by_output_index: HashMap = HashMap::new(); let mut tool_call_index_by_id: HashMap = HashMap::new(); let mut stats = StreamStats::new("Responses"); + let mut timeout_controller = StreamTimeoutController::new(None, None); handle_function_call_output_item_done( + &mut timeout_controller, &tx_event, &mut stats, Some(3), @@ -666,8 +733,10 @@ mod tests { let (tx_event, _rx_event) = mpsc::unbounded_channel(); let mut tool_calls_by_output_index: HashMap = HashMap::new(); let mut stats = StreamStats::new("Responses"); + let mut timeout_controller = StreamTimeoutController::new(None, None); let err = handle_function_call_arguments_delta( + &mut timeout_controller, &tx_event, &mut stats, None, @@ -684,8 +753,10 @@ mod tests { let (tx_event, _rx_event) = mpsc::unbounded_channel(); let mut tool_calls_by_output_index: HashMap = HashMap::new(); let mut stats = StreamStats::new("Responses"); + let mut timeout_controller = StreamTimeoutController::new(None, None); let err = handle_function_call_arguments_delta( + &mut timeout_controller, &tx_event, &mut stats, Some(2), diff --git a/src/crates/adapters/ai-adapters/tests/common/sse_fixture_server.rs b/src/crates/adapters/ai-adapters/tests/common/sse_fixture_server.rs index 047348effc..16d6be26cb 100644 --- a/src/crates/adapters/ai-adapters/tests/common/sse_fixture_server.rs +++ b/src/crates/adapters/ai-adapters/tests/common/sse_fixture_server.rs @@ -18,6 +18,7 @@ use tokio_stream::StreamExt; pub struct FixtureSseServerOptions { pub chunk_size: usize, pub chunk_delay: Duration, + pub initial_delay: Duration, } impl Default for FixtureSseServerOptions { @@ -25,6 +26,7 @@ impl Default for FixtureSseServerOptions { Self { chunk_size: 23, chunk_delay: Duration::from_millis(1), + initial_delay: Duration::ZERO, } } } @@ -81,6 +83,9 @@ async fn stream_fixture_handler(State(state): State) -> impl In let (tx, rx) = mpsc::channel::(8); tokio::spawn(async move { + if !state.options.initial_delay.is_zero() { + tokio::time::sleep(state.options.initial_delay).await; + } let chunk_size = state.options.chunk_size.max(1); for chunk in state.payload.chunks(chunk_size) { if tx.send(Bytes::copy_from_slice(chunk)).await.is_err() { diff --git a/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs b/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs index 8aa01bd003..23937c97e4 100644 --- a/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs +++ b/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs @@ -52,6 +52,8 @@ pub struct StreamFixtureRunOptions { pub server_options: FixtureSseServerOptions, pub request_timeout: Duration, pub process_timeout: Duration, + pub ttft_timeout: Option, + pub idle_timeout: Option, pub openai_inline_think_in_text: bool, pub anthropic_inline_think_in_text: bool, pub log_raw_sse: bool, @@ -63,6 +65,8 @@ impl Default for StreamFixtureRunOptions { server_options: FixtureSseServerOptions::default(), request_timeout: Duration::from_secs(5), process_timeout: Duration::from_secs(5), + ttft_timeout: None, + idle_timeout: None, openai_inline_think_in_text: false, anthropic_inline_think_in_text: false, log_raw_sse: false, @@ -137,7 +141,8 @@ pub async fn run_stream_fixture_with_options( tx_event, Some(tx_raw_sse), options.openai_inline_think_in_text, - None, + options.ttft_timeout, + options.idle_timeout, )); } StreamFixtureProvider::Anthropic => { @@ -146,7 +151,8 @@ pub async fn run_stream_fixture_with_options( tx_event, Some(tx_raw_sse), options.anthropic_inline_think_in_text, - None, + options.ttft_timeout, + options.idle_timeout, )); } StreamFixtureProvider::Gemini => { @@ -154,7 +160,8 @@ pub async fn run_stream_fixture_with_options( response, tx_event, Some(tx_raw_sse), - None, + options.ttft_timeout, + options.idle_timeout, )); } StreamFixtureProvider::Responses => { @@ -162,7 +169,8 @@ pub async fn run_stream_fixture_with_options( response, tx_event, Some(tx_raw_sse), - None, + options.ttft_timeout, + options.idle_timeout, )); } } diff --git a/src/crates/adapters/ai-adapters/tests/stream_test_harness.rs b/src/crates/adapters/ai-adapters/tests/stream_test_harness.rs index 98c703c9be..c8578a77af 100644 --- a/src/crates/adapters/ai-adapters/tests/stream_test_harness.rs +++ b/src/crates/adapters/ai-adapters/tests/stream_test_harness.rs @@ -16,6 +16,7 @@ async fn stream_test_harness_fails_fast_when_fixture_processing_stalls() { server_options: FixtureSseServerOptions { chunk_size: 1, chunk_delay: Duration::from_millis(50), + ..Default::default() }, process_timeout: Duration::from_millis(20), ..Default::default() @@ -23,3 +24,31 @@ async fn stream_test_harness_fails_fast_when_fixture_processing_stalls() { ) .await; } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ttft_timeout_waits_for_first_effective_stream_output_not_http_200() { + let output = run_stream_fixture_with_options( + StreamFixtureProvider::OpenAi, + "stream/openai/tool_args_split_with_usage.sse", + StreamFixtureRunOptions { + server_options: FixtureSseServerOptions { + initial_delay: Duration::from_millis(60), + ..Default::default() + }, + ttft_timeout: Some(Duration::from_millis(20)), + process_timeout: Duration::from_secs(1), + ..Default::default() + }, + ) + .await; + + let error = output.result.expect_err("fixture should fail with TTFT timeout"); + assert!( + error + .error + .to_string() + .contains("TTFT timeout after 0s waiting for first effective output"), + "unexpected error: {}", + error.error + ); +} diff --git a/src/crates/assembly/core/src/infrastructure/ai/mod.rs b/src/crates/assembly/core/src/infrastructure/ai/mod.rs index 3fba730dae..c394bea000 100644 --- a/src/crates/assembly/core/src/infrastructure/ai/mod.rs +++ b/src/crates/assembly/core/src/infrastructure/ai/mod.rs @@ -12,13 +12,13 @@ pub use bitfun_ai_adapters::stream as ai_stream_handlers; pub use bitfun_ai_adapters::{ AIClient, StreamOptions, StreamResponse, DEFAULT_STREAM_IDLE_TIMEOUT_SECS, - DEFAULT_STREAM_TTFT_TIMEOUT_SECS, REASONING_STREAM_TTFT_TIMEOUT_SECS, + DEFAULT_STREAM_TTFT_TIMEOUT_SECS, }; pub use client_factory::{ get_global_ai_client_factory, initialize_global_ai_client_factory, AIClientFactory, }; -use crate::service::config::types::{AIConfig, AIModelConfig, ReasoningMode}; +use crate::service::config::types::{AIConfig, AIModelConfig}; pub fn build_stream_options(config: &AIConfig) -> StreamOptions { build_stream_options_for_model(config, None) @@ -26,29 +26,13 @@ pub fn build_stream_options(config: &AIConfig) -> StreamOptions { pub fn build_stream_options_for_model( config: &AIConfig, - model_config: Option<&AIModelConfig>, + _model_config: Option<&AIModelConfig>, ) -> StreamOptions { let idle_timeout = config.stream_idle_timeout_secs.map(Duration::from_secs); - let base_ttft_secs = config - .stream_ttft_timeout_secs - .or(Some(DEFAULT_STREAM_TTFT_TIMEOUT_SECS)); - - let ttft_secs = match (base_ttft_secs, model_config) { - (Some(secs), Some(model)) - if matches!( - model.effective_reasoning_mode(), - ReasoningMode::Enabled | ReasoningMode::Adaptive - ) => - { - Some(secs.max(REASONING_STREAM_TTFT_TIMEOUT_SECS)) - } - (secs, _) => secs, - }; - StreamOptions { idle_timeout, - ttft_timeout: ttft_secs.map(Duration::from_secs), + ttft_timeout: config.stream_ttft_timeout_secs.map(Duration::from_secs), } } @@ -58,17 +42,31 @@ mod tests { use crate::service::config::types::AIModelConfig; #[test] - fn reasoning_models_use_extended_ttft_timeout() { + fn model_reasoning_mode_does_not_override_ttft_timeout() { let config = AIConfig::default(); let mut model = AIModelConfig::default(); - model.reasoning_mode = Some(ReasoningMode::Enabled); + model.reasoning_mode = Some(crate::service::config::types::ReasoningMode::Enabled); let options = build_stream_options_for_model(&config, Some(&model)); assert_eq!( options.ttft_timeout, - Some(Duration::from_secs(REASONING_STREAM_TTFT_TIMEOUT_SECS)) + Some(Duration::from_secs(DEFAULT_STREAM_TTFT_TIMEOUT_SECS)) + ); + assert_eq!( + options.idle_timeout, + Some(Duration::from_secs(DEFAULT_STREAM_IDLE_TIMEOUT_SECS)) ); + } + + #[test] + fn explicit_none_ttft_timeout_means_wait_indefinitely() { + let mut config = AIConfig::default(); + config.stream_ttft_timeout_secs = None; + + let options = build_stream_options_for_model(&config, None); + + assert_eq!(options.ttft_timeout, None); assert_eq!( options.idle_timeout, Some(Duration::from_secs(DEFAULT_STREAM_IDLE_TIMEOUT_SECS)) diff --git a/src/crates/assembly/core/src/service/config/types.rs b/src/crates/assembly/core/src/service/config/types.rs index 9e5ab803f1..624605f8c5 100644 --- a/src/crates/assembly/core/src/service/config/types.rs +++ b/src/crates/assembly/core/src/service/config/types.rs @@ -770,7 +770,7 @@ fn default_stream_idle_timeout() -> Option { Some(45) } -/// Default time-to-first-token timeout while opening a stream. +/// Default timeout while waiting for the first effective streamed output. fn default_stream_ttft_timeout() -> Option { Some(30) } @@ -2146,6 +2146,26 @@ mod tests { assert!(config.review_teams.contains_key("default")); } + #[test] + fn deserializes_explicit_null_stream_ttft_timeout_as_none() { + let config: AIConfig = serde_json::from_value(serde_json::json!({ + "models": [], + "agent_models": {}, + "func_agent_models": {}, + "default_models": {}, + "agent_profiles": {}, + "proxy": { + "enabled": false, + "url": "" + }, + "stream_ttft_timeout_secs": null + })) + .expect("config with explicit null stream_ttft_timeout_secs should deserialize"); + + assert_eq!(config.stream_ttft_timeout_secs, None); + assert_eq!(config.stream_idle_timeout_secs, Some(45)); + } + #[test] fn app_logging_defaults_to_sensitive_diagnostics_enabled() { let config: AppLoggingConfig = serde_json::from_value(serde_json::json!({ diff --git a/src/web-ui/src/infrastructure/config/components/AIModelConfig.tsx b/src/web-ui/src/infrastructure/config/components/AIModelConfig.tsx index 727999311d..62b0d55c66 100644 --- a/src/web-ui/src/infrastructure/config/components/AIModelConfig.tsx +++ b/src/web-ui/src/infrastructure/config/components/AIModelConfig.tsx @@ -2446,6 +2446,38 @@ const AIModelConfig: React.FC = () => { ); }; + const streamTtftTimeoutLabel = ( + + {t('streamTtftTimeout.label')} + + + + + + + ); + + const streamIdleTimeoutLabel = ( + + {t('streamIdleTimeout.label')} + + + + + + + ); + return ( @@ -2587,7 +2619,7 @@ const AIModelConfig: React.FC = () => { { )} > { />