diff --git a/README.md b/README.md index f6d8e4427..f6bb9b652 100644 --- a/README.md +++ b/README.md @@ -141,7 +141,7 @@ Use `aether-agent-core` as a Rust library to build your own agent in ~25 lines. #[tokio::main] async fn main() -> Result<(), Box> { // 1. Create a provider (reads ANTHROPIC_API_KEY from env) - let llm = AnthropicProvider::new(None)?; + let llm = AnthropicProvider::new(None); // 2. Spawn MCP tool servers from one or more mcp.json files let mut mcp_runtime = mcp(".").from_json_files(&["mcp.json"])?.spawn().await?; diff --git a/crates/aether-cli/src/provider_connection_args.rs b/crates/aether-cli/src/provider_connection_args.rs index ad96dcc97..bed77eade 100644 --- a/crates/aether-cli/src/provider_connection_args.rs +++ b/crates/aether-cli/src/provider_connection_args.rs @@ -1,4 +1,5 @@ use std::collections::BTreeMap; +use std::num::NonZeroU64; use std::str::FromStr; use llm::{ProviderAuthMode, ProviderConnectionOverride, ProviderConnectionOverrides}; @@ -7,7 +8,7 @@ use llm::{ProviderAuthMode, ProviderConnectionOverride, ProviderConnectionOverri pub struct ProviderConnectionArgs { #[arg( long = "provider", - value_name = "PROVIDER.url=URL|PROVIDER.auth=default|none|PROVIDER.request-model=MODEL|bedrock.inference-profile-arn=ARN" + value_name = "PROVIDER.url=URL|PROVIDER.auth=default|none|PROVIDER.request-model=MODEL|PROVIDER.idle-timeout-secs=SECS|bedrock.inference-profile-arn=ARN" )] pub providers: Vec, } @@ -35,7 +36,7 @@ impl FromStr for ProviderArg { let (key, setting) = split_key_value(value)?; let (provider, field) = key .split_once('.') - .ok_or_else(|| "provider override must be PROVIDER.url=URL, PROVIDER.auth=default|none, PROVIDER.request-model=MODEL, or bedrock.inference-profile-arn=ARN".to_string())?; + .ok_or_else(|| "provider override must be PROVIDER.url=URL, PROVIDER.auth=default|none, PROVIDER.request-model=MODEL, PROVIDER.idle-timeout-secs=SECS, or bedrock.inference-profile-arn=ARN".to_string())?; validate_provider(provider)?; if setting.trim().is_empty() { @@ -49,6 +50,7 @@ impl FromStr for ProviderArg { } "auth" => ProviderConnectionOverride::auth(parse_auth_mode(setting)?), "request-model" => ProviderConnectionOverride::request_model(setting), + "idle-timeout-secs" => ProviderConnectionOverride::idle_timeout_secs(parse_idle_timeout_secs(setting)?), "inference-profile-arn" => { if provider != "bedrock" { return Err("inference-profile-arn is only supported for the bedrock provider".to_string()); @@ -57,7 +59,8 @@ impl FromStr for ProviderArg { } _ => { return Err( - "provider override field must be url, auth, request-model, or inference-profile-arn".to_string() + "provider override field must be url, auth, request-model, idle-timeout-secs, or inference-profile-arn" + .to_string() ); } }; @@ -85,6 +88,10 @@ fn validate_url(url: &str) -> Result<(), String> { } } +fn parse_idle_timeout_secs(value: &str) -> Result { + value.parse().map_err(|_| "provider idle timeout must be a positive number of seconds".to_string()) +} + fn parse_auth_mode(value: &str) -> Result { match value { "default" => Ok(ProviderAuthMode::Default), @@ -122,6 +129,12 @@ mod tests { assert_eq!(arg.connection.request_model.as_deref(), Some("production-coding")); } + #[test] + fn parses_provider_idle_timeout() { + let arg: ProviderArg = "ollama.idle-timeout-secs=900".parse().unwrap(); + assert_eq!(arg.connection.idle_timeout_secs, NonZeroU64::new(900)); + } + #[test] fn combines_repeated_provider_overrides() { let args = ProviderConnectionArgs { @@ -171,6 +184,8 @@ mod tests { assert!("bedrock.url=".parse::().is_err()); assert!("bedrock.url=file:///tmp/proxy".parse::().is_err()); assert!("bedrock.auth=disabled".parse::().is_err()); + assert!("ollama.idle-timeout-secs=0".parse::().is_err()); + assert!("ollama.idle-timeout-secs=soon".parse::().is_err()); assert!("bedrock.region=us-west-2".parse::().is_err()); assert!("openai.inference-profile-arn=arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000".parse::().is_err()); } diff --git a/crates/llm/Cargo.toml b/crates/llm/Cargo.toml index 7edf8d4b9..bffe5f0e3 100644 --- a/crates/llm/Cargo.toml +++ b/crates/llm/Cargo.toml @@ -81,6 +81,7 @@ aether-doctest = { path = "../doctest" } axum = { workspace = true } http = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } tracing-subscriber = { workspace = true } [[example]] diff --git a/crates/llm/src/docs/llm_error.md b/crates/llm/src/docs/llm_error.md index b0b6b559e..b1f0a233a 100644 --- a/crates/llm/src/docs/llm_error.md +++ b/crates/llm/src/docs/llm_error.md @@ -6,9 +6,6 @@ Errors that can occur when interacting with LLM providers. - **`MissingApiKey`** -- A required environment variable (e.g. `ANTHROPIC_API_KEY`) is not set. - **`OAuthError`** -- OAuth authentication failed (when the `oauth` feature is enabled). -## Request/Response -- **`HttpClientCreation`** -- Failed to create the HTTP client (e.g. TLS configuration error). - ## Parsing - **`JsonParsing`** -- Failed to parse or serialize JSON (response body, tool arguments). - **`ToolParameterParsing`** -- Failed to parse tool parameters for a specific tool. diff --git a/crates/llm/src/docs/provider_connection_override.md b/crates/llm/src/docs/provider_connection_override.md index a403159d1..373ac7fd8 100644 --- a/crates/llm/src/docs/provider_connection_override.md +++ b/crates/llm/src/docs/provider_connection_override.md @@ -20,6 +20,12 @@ Request-target routing keeps the catalog model identity and its capabilities whi Use `url` for providers with resource-specific endpoints, such as Microsoft Foundry. A trailing slash is normalized before `/chat/completions` is appended. A provider with `requestModel` cannot appear more than once in an alloy because the target would be ambiguous. +Give a slow local model more time to produce its first token. A request fails with a retryable timeout once the provider has been silent for `idleTimeoutSecs` (default 300), before responding or between stream events: + +```json +{ "idleTimeoutSecs": 900 } +``` + Pin a Bedrock application inference profile: ```json diff --git a/crates/llm/src/error.rs b/crates/llm/src/error.rs index b972ebb60..1fa0519a6 100644 --- a/crates/llm/src/error.rs +++ b/crates/llm/src/error.rs @@ -12,9 +12,6 @@ pub enum LlmError { /// Environment variable not set or invalid #[error("{0} environment variable not set")] MissingApiKey(String), - /// HTTP client creation failed - #[error("Failed to create HTTP client: {0}")] - HttpClientCreation(String), /// Normalized provider-side failure carrying retry classification and /// support diagnostics (HTTP status, request ID, provider error code). #[error("{0}")] @@ -328,7 +325,6 @@ mod tests { assert!(!LlmError::from(ProviderError::api("x")).is_retryable()); assert!(!LlmError::from(ProviderError::authentication("x")).is_retryable()); assert!(!LlmError::MissingApiKey("x".into()).is_retryable()); - assert!(!LlmError::HttpClientCreation("x".into()).is_retryable()); assert!(!LlmError::IoError("x".into()).is_retryable()); assert!(!LlmError::JsonParsing("x".into()).is_retryable()); assert!(!LlmError::ToolParameterParsing { tool_name: "t".into(), error: "e".into() }.is_retryable()); diff --git a/crates/llm/src/parser.rs b/crates/llm/src/parser.rs index b2e247b2f..274c4385c 100644 --- a/crates/llm/src/parser.rs +++ b/crates/llm/src/parser.rs @@ -312,6 +312,7 @@ mod tests { auth_mode: Some(crate::ProviderAuthMode::None), request_model: Some("production-coding".to_string()), inference_profile_arn: None, + idle_timeout_secs: None, }, )]), )); diff --git a/crates/llm/src/provider.rs b/crates/llm/src/provider.rs index 6dcdfba99..8ada83f7e 100644 --- a/crates/llm/src/provider.rs +++ b/crates/llm/src/provider.rs @@ -5,7 +5,7 @@ use crate::Result as LlmResult; use crate::catalog::ReasoningEffortError; use std::future::Future; use std::pin::Pin; -use tokio_stream::{Stream, StreamExt}; +use tokio_stream::Stream; use utils::ReasoningEffort; use super::{Context, LlmResponse}; @@ -82,37 +82,6 @@ pub(crate) fn validate_reasoning(context: &Context, model: Option<&LlmModel>) -> Ok(()) } -/// Bridge a fallible request setup into an [`LlmResponseStream`]. -/// -/// `open` issues the request; `process` turns what it returns into a response -/// stream. A setup failure becomes the stream's single item, so providers never -/// hand-roll the yield-then-return dance and cannot drop an error on the way. -pub(crate) fn stream_from( - open: impl Future> + Send + 'static, - process: impl FnOnce(T) -> S + Send + 'static, -) -> LlmResponseStream -where - T: Send, - S: Stream> + Send + 'static, -{ - Box::pin(async_stream::stream! { - match open.await { - Ok(opened) => { - let mut stream = Box::pin(process(opened)); - while let Some(item) = stream.next().await { - yield item; - } - } - Err(error) => yield Err(error), - } - }) -} - -/// A response stream whose only item is `error`. -pub(crate) fn error_stream(error: LlmError) -> LlmResponseStream { - Box::pin(tokio_stream::once(Err(error))) -} - impl StreamingModelProvider for Box { fn stream_response(&self, context: &Context) -> LlmResponseStream { (**self).stream_response(context) diff --git a/crates/llm/src/provider_connection.rs b/crates/llm/src/provider_connection.rs index c265311e8..7cb1d4d92 100644 --- a/crates/llm/src/provider_connection.rs +++ b/crates/llm/src/provider_connection.rs @@ -1,4 +1,8 @@ use std::collections::BTreeMap; +use std::num::NonZeroU64; +use std::time::Duration; + +pub(crate) const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_mins(5); #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] #[serde(rename_all = "kebab-case")] @@ -8,12 +12,13 @@ pub enum ProviderAuthMode { None, } -#[derive(Clone, Debug, Default, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct ProviderConnectionConfig { pub base_url: Option, pub auth_mode: ProviderAuthMode, pub request_model: Option, pub inference_profile_arn: Option, + pub idle_timeout: Duration, } #[doc = include_str!("docs/provider_connection_override.md")] @@ -33,6 +38,8 @@ pub struct ProviderConnectionOverride { /// AWS Bedrock application inference profile ARN to route requests through. #[serde(default, skip_serializing_if = "Option::is_none")] pub inference_profile_arn: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub idle_timeout_secs: Option, } #[derive(Clone, Debug, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] @@ -48,6 +55,21 @@ impl ProviderConnectionConfig { auth_mode: value.auth_mode.unwrap_or_default(), request_model: value.request_model, inference_profile_arn: value.inference_profile_arn, + idle_timeout: value + .idle_timeout_secs + .map_or(DEFAULT_STREAM_IDLE_TIMEOUT, |secs| Duration::from_secs(secs.get())), + } + } +} + +impl Default for ProviderConnectionConfig { + fn default() -> Self { + Self { + base_url: None, + auth_mode: ProviderAuthMode::default(), + request_model: None, + inference_profile_arn: None, + idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, } } } @@ -69,6 +91,10 @@ impl ProviderConnectionOverride { Self { inference_profile_arn: Some(arn.into()), ..Self::default() } } + pub fn idle_timeout_secs(secs: NonZeroU64) -> Self { + Self { idle_timeout_secs: Some(secs), ..Self::default() } + } + pub fn merge(&mut self, override_value: Self) { if override_value.base_url.is_some() { self.base_url = override_value.base_url; @@ -82,6 +108,9 @@ impl ProviderConnectionOverride { if override_value.inference_profile_arn.is_some() { self.inference_profile_arn = override_value.inference_profile_arn; } + if override_value.idle_timeout_secs.is_some() { + self.idle_timeout_secs = override_value.idle_timeout_secs; + } } } @@ -142,6 +171,16 @@ mod tests { ); } + #[test] + fn idle_timeout_defaults_and_deserializes() { + let overrides: ProviderConnectionOverrides = + serde_json::from_str(r#"{"ollama":{"idleTimeoutSecs":900}}"#).unwrap(); + + assert_eq!(overrides.config_for("ollama").idle_timeout, Duration::from_mins(15)); + assert_eq!(overrides.config_for("anthropic").idle_timeout, Duration::from_mins(5)); + assert!(serde_json::from_str::(r#"{"ollama":{"idleTimeoutSecs":0}}"#).is_err()); + } + #[test] fn merge_replaces_inference_profile_arn() { let mut first = ProviderConnectionOverride::inference_profile_arn("arn:first"); diff --git a/crates/llm/src/providers/anthropic/provider.rs b/crates/llm/src/providers/anthropic/provider.rs index 620eee2dd..aca98dd3e 100644 --- a/crates/llm/src/providers/anthropic/provider.rs +++ b/crates/llm/src/providers/anthropic/provider.rs @@ -1,14 +1,13 @@ use super::mappers::{map_messages, map_tools}; -use super::streaming::process_anthropic_stream; +use super::streaming::decode_line; use super::types::{Request, Thinking}; use crate::provider::{ - LlmResponseStream, ProviderFactory, StreamingModelProvider, error_stream, get_context_window, validate_reasoning, + LlmResponseStream, ProviderFactory, StreamingModelProvider, get_context_window, validate_reasoning, }; -use crate::providers::http::{anthropic_code, rejected}; -use crate::{Context, LlmError, ProviderAuthMode, ProviderConnectionConfig, ProviderError, ReasoningEffort, Result}; -use async_stream; -use eventsource_stream::Eventsource; -use futures::StreamExt; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::http::{SseData, anthropic_code, http_client, open_sse}; +use crate::providers::response_stream::{OpenedStream, response_stream}; +use crate::{Context, LlmError, ProviderAuthMode, ProviderConnectionConfig, ReasoningEffort, Result}; use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; use reqwest::{Client, header}; use std::env; @@ -27,19 +26,19 @@ pub struct AnthropicProvider { base_url: Option, auth_mode: ProviderAuthMode, api_key: Option, + idle_timeout: Duration, } impl AnthropicProvider { - pub fn new(api_key: Option) -> Result { - let client = build_client()?; - - Ok(Self { - client, + pub fn new(api_key: Option) -> Self { + Self { + client: http_client(), model: "claude-sonnet-4-5-20250929".to_string(), base_url: Some("https://api.anthropic.com".to_string()), auth_mode: ProviderAuthMode::Default, api_key, - }) + idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, + } } pub fn with_model(mut self, model: &str) -> Self { @@ -57,10 +56,12 @@ impl AnthropicProvider { self.base_url = Some(base_url); } self.auth_mode = connection.auth_mode; + self.idle_timeout = connection.idle_timeout; self } pub(crate) fn build_request(&self, context: &Context) -> Result { + validate_reasoning(context, self.model().as_ref())?; let (system_prompt, messages) = map_messages(context.messages())?; let tools = if context.tools().is_empty() { None } else { Some(map_tools(context.tools())?) }; @@ -129,11 +130,7 @@ impl AnthropicProvider { Ok(headers) } - async fn send_request( - &self, - request: Request, - headers: header::HeaderMap, - ) -> Result>> { + async fn send_request(&self, request: Request, headers: header::HeaderMap) -> Result> { let base_url = self.base_url.as_deref().unwrap_or("https://api.anthropic.com"); let url = format!("{base_url}/v1/messages"); @@ -144,34 +141,17 @@ impl AnthropicProvider { ); debug!("Anthropic request headers: {}", format_headers(&headers)); - let response = self.client.post(&url).headers(headers).json(&request).send().await?; - - if !response.status().is_success() { - return Err(rejected(response, anthropic_code).await.into()); - } - - let event_stream = response.bytes_stream().eventsource(); - let processed_stream = event_stream.filter_map(|result| { - std::future::ready(match result { - Ok(event) => { - let data = event.data; - if data == "[DONE]" { None } else { Some(Ok(data)) } - } - Err(e) => Some(Err(ProviderError::stream_interrupted(e.to_string()).into())), - }) - }); - - Ok(processed_stream) + open_sse(self.client.post(&url).headers(headers).json(&request), anthropic_code).await } } impl ProviderFactory for AnthropicProvider { fn from_env() -> impl Future> + Send { - ready(Self::new(None)) + ready(Ok(Self::new(None))) } fn from_env_with_connection(connection: ProviderConnectionConfig) -> impl Future> + Send { - ready(Self::new(None).map(|provider| provider.with_connection(connection))) + ready(Ok(Self::new(None).with_connection(connection))) } fn with_model(self, model: &str) -> Self { @@ -188,43 +168,19 @@ impl StreamingModelProvider for AnthropicProvider { get_context_window("anthropic", &self.model) } - fn stream_response<'a>(&self, context: &Context) -> LlmResponseStream { - if let Err(error) = validate_reasoning(context, self.model().as_ref()) { - return error_stream(error); - } + fn stream_response(&self, context: &Context) -> LlmResponseStream { let provider = self.clone(); let context = context.clone(); - Box::pin(async_stream::stream! { - let headers = match provider.build_headers() { - Ok(result) => result, - Err(e) => { - yield Err(e); - return; - } - }; - - let request = match provider.build_request(&context) { - Ok(req) => req, - Err(e) => { - yield Err(e); - return; - } - }; - - let stream = match provider.send_request(request, headers).await { - Ok(stream) => stream, - Err(e) => { - yield Err(e); - return; - } - }; - - let mut anthropic_stream = Box::pin(process_anthropic_stream(stream)); - while let Some(result) = anthropic_stream.next().await { - yield result; - } - }) + response_stream( + async move { + let request = provider.build_request(&context)?; + let headers = provider.build_headers()?; + provider.send_request(request, headers).await + }, + |line, turn| decode_line(&line, turn), + self.idle_timeout, + ) } fn display_name(&self) -> String { @@ -232,10 +188,6 @@ impl StreamingModelProvider for AnthropicProvider { } } -fn build_client() -> Result { - Client::builder().timeout(Duration::from_mins(1)).build().map_err(|e| LlmError::HttpClientCreation(e.to_string())) -} - fn effort_to_budget_tokens(effort: ReasoningEffort) -> Option { Some(match effort { ReasoningEffort::Default | ReasoningEffort::Disabled => return None, @@ -273,16 +225,18 @@ mod tests { use crate::ToolDefinition; use crate::providers::anthropic::types::{SystemContent, SystemContentBlock}; + use crate::providers::test_capture_server::{CaptureServer, ResponseSpec, hello_context}; + use futures::StreamExt; use reqwest::header::AUTHORIZATION; + use std::time::Duration; fn create_test_provider() -> AnthropicProvider { - AnthropicProvider::new(Some("test-api-key".to_string())).unwrap().with_model("claude-sonnet-4-5-20250929") + AnthropicProvider::new(Some("test-api-key".to_string())).with_model("claude-sonnet-4-5-20250929") } #[tokio::test] async fn default_and_disabled_thinking_preserve_sampling_and_token_limit() { - use crate::providers::test_capture_server::CaptureServer; let mut server = CaptureServer::start_with_response(include_str!("../../../tests/fixtures/anthropic/01_minimal.sse")).await; let model = crate::LlmModel::all() @@ -293,7 +247,7 @@ mod tests { .unwrap(); let provider = create_test_provider().with_model(&model.model_id()).with_base_url(&server.base_url); for (effort, temperature) in [(ReasoningEffort::Default, 0.0), (ReasoningEffort::Disabled, 0.5)] { - let mut context = Context::new(vec![ChatMessage::user("Hello")], vec![]); + let mut context = hello_context(); context.set_reasoning_effort(effort); context.set_model_settings(crate::ModelSettings { temperature: Some(temperature), @@ -314,15 +268,22 @@ mod tests { } } - #[test] - fn test_provider_creation() { - let provider = AnthropicProvider::new(Some("test-api-key".to_string())); - assert!(provider.is_ok()); + #[tokio::test] + async fn stream_with_long_pauses_between_events_completes() { + let spec = ResponseSpec::sse(include_str!("../../../tests/fixtures/anthropic/01_minimal.sse")) + .paced(Duration::from_mins(2)); + let mut server = CaptureServer::start_with_spec(spec).await; + let provider = create_test_provider().with_base_url(&server.base_url); + + let responses = server.collect_on_paused_clock(provider.stream_response(&hello_context())).await; + + assert!(responses.iter().all(Result::is_ok), "{responses:?}"); + assert!(responses.iter().any(|response| matches!(response, Ok(crate::LlmResponse::Done { .. })))); } #[test] fn build_headers_uses_api_key() { - let provider = AnthropicProvider::new(Some("test-api-key".to_string())).unwrap(); + let provider = AnthropicProvider::new(Some("test-api-key".to_string())); let headers = provider.build_headers().expect("headers"); assert_eq!(headers.get("x-api-key").and_then(|value| value.to_str().ok()), Some("test-api-key")); assert!(headers.get(AUTHORIZATION).is_none()); @@ -332,7 +293,6 @@ mod tests { #[test] fn build_headers_skips_api_key_when_auth_is_none() { let provider = AnthropicProvider::new(None) - .unwrap() .with_connection(ProviderConnectionConfig { auth_mode: ProviderAuthMode::None, ..Default::default() }); let headers = provider.build_headers().expect("headers"); assert!(headers.get("x-api-key").is_none()); @@ -343,9 +303,7 @@ mod tests { fn test_build_request_simple() { let provider = create_test_provider(); - let context = Context::new(vec![ChatMessage::user("Hello")], vec![]); - - let request = provider.build_request(&context).unwrap(); + let request = provider.build_request(&hello_context()).unwrap(); assert_eq!(request.model, "claude-sonnet-4-5-20250929"); assert_eq!(request.max_tokens, DEFAULT_MAX_TOKENS); assert_eq!(request.messages.len(), 1); @@ -386,7 +344,7 @@ mod tests { #[test] fn test_build_request_with_caching() { - let provider = AnthropicProvider::new(Some("test-api-key".to_string())).unwrap(); // Caching is enabled by default + let provider = AnthropicProvider::new(Some("test-api-key".to_string())); // Caching is enabled by default let context = Context::new( vec![ChatMessage::system("Hello"), ChatMessage::user("Hello")], @@ -450,7 +408,7 @@ mod tests { #[test] fn test_build_request_thinking_bumps_max_tokens_if_needed() { - let provider = AnthropicProvider::new(Some("test-api-key".to_string())).unwrap(); + let provider = AnthropicProvider::new(Some("test-api-key".to_string())); let mut context = Context::new(vec![ChatMessage::user("Hi")], vec![]); context.set_model_settings(crate::ModelSettings { max_tokens: Some(500), ..Default::default() }); @@ -471,7 +429,7 @@ mod tests { #[test] fn test_anthropic_provider_display_name_default() { - let provider = AnthropicProvider::new(Some("test-api-key".to_string())).unwrap(); + let provider = AnthropicProvider::new(Some("test-api-key".to_string())); assert_eq!(provider.display_name(), "Anthropic (claude-sonnet-4-5-20250929)"); } diff --git a/crates/llm/src/providers/anthropic/streaming.rs b/crates/llm/src/providers/anthropic/streaming.rs index 455c96dd0..cb2db5c2c 100644 --- a/crates/llm/src/providers/anthropic/streaming.rs +++ b/crates/llm/src/providers/anthropic/streaming.rs @@ -1,16 +1,20 @@ use super::types::{ContentBlockDeltaData, ContentBlockStartData, StreamEvent}; -use crate::providers::stream_assembler::{StreamAssembler, assemble}; -use crate::{LlmError, LlmResponse, ProviderError, Result, StopReason}; -use futures::{Stream, StreamExt, stream}; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::response_stream::{OpenedStream, StreamAssembler, response_stream}; +use crate::{LlmError, LlmResponse, LlmResponseStream, ProviderError, Result, StopReason}; +use futures::Stream; +use std::future::ready; use tracing::debug; -pub fn process_anthropic_stream( - lines: impl Stream> + Send, -) -> impl Stream> + Send { - stream::iter([Ok(LlmResponse::Start)]).chain(assemble(lines, |line, turn| decode_line(&line, turn))) +pub fn process_anthropic_stream(lines: impl Stream> + Send + 'static) -> LlmResponseStream { + response_stream( + ready(Ok(OpenedStream::new(lines))), + |line, turn| decode_line(&line, turn), + DEFAULT_STREAM_IDLE_TIMEOUT, + ) } -fn decode_line(line: &str, turn: &mut StreamAssembler) -> Result> { +pub(super) fn decode_line(line: &str, turn: &mut StreamAssembler) -> Result> { match serde_json::from_str(line) { Ok(event) => decode_event(event, turn), Err(e) => { @@ -80,6 +84,7 @@ mod tests { use super::*; use crate::testing::llm_response; use crate::{ProviderErrorKind, TokenUsage}; + use futures::{StreamExt, stream}; use serde_json::{Value, json}; #[tokio::test] diff --git a/crates/llm/src/providers/bedrock/mantle/client.rs b/crates/llm/src/providers/bedrock/mantle/client.rs index d5250dabb..38d63a8e2 100644 --- a/crates/llm/src/providers/bedrock/mantle/client.rs +++ b/crates/llm/src/providers/bedrock/mantle/client.rs @@ -2,8 +2,10 @@ use tracing::debug; use super::auth::MantleAuth; use crate::catalog::transport::{ModelTransport, expand_api_template}; +use crate::providers::http::{SseData, http_client}; use crate::providers::openai_responses::mappers::{ResponsesRequestPolicy, build_wire_request}; -use crate::providers::openai_responses::transport::{ResponsesConnection, send}; +use crate::providers::openai_responses::transport::send; +use crate::providers::response_stream::OpenedStream; use crate::{Context, LlmError, Result}; /// Transport for Bedrock models served through the `OpenAI` Responses API. @@ -18,7 +20,7 @@ pub struct MantleClient { impl MantleClient { pub fn new(region: String, auth: MantleAuth, base_url_override: Option) -> Self { Self { - http: reqwest::Client::new(), + http: http_client(), region, auth, base_url_override: base_url_override.map(|url| url.trim_end_matches('/').to_string()), @@ -49,7 +51,7 @@ impl MantleClient { model: &str, transport: &ModelTransport, context: &Context, - ) -> Result { + ) -> Result> { let url = self.endpoint(transport)?; let body = build_wire_request(model, context, &ResponsesRequestPolicy::MANTLE)?; diff --git a/crates/llm/src/providers/bedrock/provider.rs b/crates/llm/src/providers/bedrock/provider.rs index 2d45d3f6c..636326d11 100644 --- a/crates/llm/src/providers/bedrock/provider.rs +++ b/crates/llm/src/providers/bedrock/provider.rs @@ -1,9 +1,12 @@ use super::mantle::{MantleAuth, MantleClient}; use super::mappers::{default_cache_point, map_messages, map_tools}; -use super::streaming::{converse_events, process_bedrock_stream}; +use super::streaming::{converse_events, decode_converse_event}; use crate::catalog::transport::ModelTransport; -use crate::provider::{LlmResponseStream, ProviderFactory, StreamingModelProvider, get_context_window, stream_from}; -use crate::providers::openai_responses::transport::process_connection; +use crate::provider::{ + LlmResponseStream, ProviderFactory, StreamingModelProvider, get_context_window, validate_reasoning, +}; +use crate::providers::openai_responses::streaming::decode_responses; +use crate::providers::response_stream::{OpenedStream, response_stream}; use crate::{Context, LlmError, ProviderAuthMode, ProviderConnectionConfig, ProviderError, Result}; use aws_config::Region; use aws_credential_types::provider::SharedCredentialsProvider; @@ -14,6 +17,7 @@ use aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver; use aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError; use aws_sdk_bedrockruntime::types::{ConverseStreamOutput, InferenceConfiguration}; use aws_sdk_bedrockruntime::{Client, Config}; +use std::time::Duration; use tracing::{error, info, warn}; const DEFAULT_MODEL: &str = "anthropic.claude-sonnet-4-5-20250929-v1:0"; @@ -34,6 +38,7 @@ pub struct BedrockProvider { mantle: MantleClient, model: String, inference_profile_arn: Option, + idle_timeout: Duration, } impl BedrockProvider { @@ -56,7 +61,7 @@ impl BedrockProvider { .or_else(region_from_env) .unwrap_or_else(|| DEFAULT_REGION.to_string()); let auth = mantle_auth(config.credentials_provider(), ®ion); - Self::assemble(Client::new(&config), region, auth, connection) + Self::from_parts(Client::new(&config), region, auth, connection) } /// Create a provider from explicit configuration without async credential discovery. @@ -71,16 +76,17 @@ impl BedrockProvider { ProviderAuthMode::None => MantleAuth::None, }; let client = build_client(credentials, ®ion, connection.base_url.as_deref(), connection.auth_mode); - Self::assemble(client, region, auth, connection) + Self::from_parts(client, region, auth, connection) } - fn assemble(client: Client, region: String, auth: MantleAuth, connection: ProviderConnectionConfig) -> Self { + fn from_parts(client: Client, region: String, auth: MantleAuth, connection: ProviderConnectionConfig) -> Self { let mantle = MantleClient::new(region, auth, connection.base_url.clone()); Self { client, mantle, model: DEFAULT_MODEL.to_string(), inference_profile_arn: connection.inference_profile_arn, + idle_timeout: connection.idle_timeout, } } @@ -115,6 +121,7 @@ impl BedrockProvider { &self, context: &Context, ) -> Result> { + validate_reasoning(context, self.model().as_ref())?; let cache_point = self.model().is_some_and(|m| m.supports_prompt_caching()).then(default_cache_point).transpose()?; let (system_blocks, messages) = map_messages(context.messages(), cache_point.as_ref())?; @@ -187,16 +194,15 @@ impl StreamingModelProvider for BedrockProvider { } fn stream_response(&self, context: &Context) -> LlmResponseStream { - if let Err(error) = crate::provider::validate_reasoning(context, self.model().as_ref()) { - return crate::provider::error_stream(error); - } let provider = self.clone(); let context = context.clone(); let Some(transport) = self.mantle_transport() else { - return stream_from(async move { provider.send_converse_stream(&context).await }, |receiver| { - process_bedrock_stream(converse_events(receiver)) - }); + return response_stream( + async move { Ok(OpenedStream::new(converse_events(provider.send_converse_stream(&context).await?))) }, + |event, turn| Ok(decode_converse_event(event, turn)), + self.idle_timeout, + ); }; if let Some(arn) = self.inference_profile_arn.as_deref() { @@ -207,9 +213,10 @@ impl StreamingModelProvider for BedrockProvider { ); } - stream_from( + response_stream( async move { provider.mantle.stream(&provider.model, &transport, &context).await }, - process_connection, + decode_responses(), + self.idle_timeout, ) } @@ -304,56 +311,12 @@ fn region_from_env() -> Option { mod tests { use super::*; use crate::catalog::Provider; - use crate::providers::test_capture_server::CaptureServer; + use crate::providers::test_capture_server::{CaptureServer, ResponseSpec, hello_context}; use crate::types::IsoString; use crate::{AssistantReasoning, ChatMessage, EncryptedReasoningContent, LlmModel, MessageId}; - use axum::Router; - use axum::body::Body; - use axum::extract::State; - use axum::http::{HeaderMap, Method, Request, StatusCode}; - use axum::response::IntoResponse; - use axum::routing::any; use futures::StreamExt; - use std::sync::Arc; - use tokio::net::TcpListener; - use tokio::sync::{Mutex, oneshot}; use utils::ReasoningEffort; - fn inference_profile_arn(model: &str) -> String { - format!("arn:aws:bedrock:us-west-2:000000000000:inference-profile/{model}") - } - - fn application_inference_profile_arn() -> &'static str { - "arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000" - } - - fn test_provider() -> BedrockProvider { - BedrockProvider::from_config(None, None, ProviderConnectionConfig::default()) - } - - /// A catalog model routed to the Responses transport, resolved from the - /// catalog so a models.dev sync that retires one model does not quietly - /// leave these tests exercising the Converse path instead. - fn mantle_model() -> String { - LlmModel::all() - .iter() - .find(|model| model.provider_enum() == Provider::Bedrock && model.transport().is_some()) - .expect("catalog must expose at least one Responses-transport Bedrock model") - .model_id() - .to_string() - } - - /// A provider talking to `server` over the Responses transport, unauthenticated. - async fn mantle_provider(server: &CaptureServer) -> BedrockProvider { - BedrockProvider::new(ProviderConnectionConfig { - base_url: Some(server.base_url.clone()), - auth_mode: ProviderAuthMode::None, - ..Default::default() - }) - .await - .with_model(&mantle_model()) - } - #[test] fn test_display_name() { assert_eq!(test_provider().display_name(), "Bedrock (anthropic.claude-sonnet-4-5-20250929-v1:0)"); @@ -373,19 +336,12 @@ mod tests { #[tokio::test] async fn auth_none_sends_unsigned_request_to_custom_endpoint() { - let endpoint = FakeBedrockEndpoint::start().await; - let provider = BedrockProvider::new(ProviderConnectionConfig { - base_url: Some(endpoint.url.clone()), - auth_mode: ProviderAuthMode::None, - ..Default::default() - }) - .await; + let mut server = CaptureServer::start_with_response("{}").await; + let provider = create_provider(&server.base_url, DEFAULT_MODEL).await; - let result = provider.send_converse_stream(&hello_context()).await; - let request = endpoint.request.await.expect("fake Bedrock endpoint received no request"); + let _ = provider.stream_response(&hello_context()).collect::>().await; + let request = server.captured().await; - assert!(result.is_err()); - assert_eq!(request.method, Method::POST); assert!(request.path.starts_with("/model/"), "{}", request.path); assert!(!request.headers.contains_key("authorization"), "request was signed: {:?}", request.headers); assert!( @@ -403,10 +359,6 @@ mod tests { } } - fn hello_context() -> Context { - Context::new(vec![ChatMessage::user("Hello")], vec![]) - } - #[tokio::test] async fn mantle_disabled_uses_none_without_summary() { let model = LlmModel::all() @@ -418,7 +370,7 @@ mod tests { }) .unwrap(); let mut server = CaptureServer::start_responses().await; - let provider = mantle_provider(&server).await.with_model(&model.model_id()); + let provider = create_provider(&server.base_url, &model.model_id()).await; let mut context = hello_context(); context.set_reasoning_effort(ReasoningEffort::Disabled); let responses = provider.stream_response(&context).collect::>().await; @@ -474,7 +426,6 @@ mod tests { #[tokio::test] async fn http_200_failed_server_error_is_retryable_with_diagnostics() { - use crate::providers::test_capture_server::ResponseSpec; let spec = ResponseSpec::sse(include_str!("../../../tests/fixtures/openai_responses/04_failed_server.sse")) .with_header("x-amzn-requestid", "amzn-req-123"); let mut server = CaptureServer::start_with_spec(spec).await; @@ -561,17 +512,11 @@ mod tests { #[tokio::test] async fn converse_shape_models_do_not_use_the_responses_endpoint() { - let endpoint = FakeBedrockEndpoint::start().await; - let provider = BedrockProvider::new(ProviderConnectionConfig { - base_url: Some(endpoint.url.clone()), - auth_mode: ProviderAuthMode::None, - ..Default::default() - }) - .await - .with_model(DEFAULT_MODEL); + let mut server = CaptureServer::start_with_response("{}").await; + let provider = create_provider(&server.base_url, DEFAULT_MODEL).await; let _ = provider.stream_response(&hello_context()).collect::>().await; - let request = endpoint.request.await.expect("fake Bedrock endpoint received no request"); + let request = server.captured().await; assert!(request.path.starts_with("/model/"), "{}", request.path); } @@ -686,20 +631,19 @@ mod tests { #[tokio::test] async fn separate_inference_profile_arn_is_used_as_request_model_id() { - let endpoint = FakeBedrockEndpoint::start().await; + let mut server = CaptureServer::start_with_response("{}").await; let provider = BedrockProvider::new(ProviderConnectionConfig { - base_url: Some(endpoint.url.clone()), + base_url: Some(server.base_url.clone()), auth_mode: ProviderAuthMode::None, - request_model: None, inference_profile_arn: Some(application_inference_profile_arn().to_string()), + ..Default::default() }) .await .with_model(DEFAULT_MODEL); - let result = provider.send_converse_stream(&hello_context()).await; - let request = endpoint.request.await.expect("fake Bedrock endpoint received no request"); + let _ = provider.stream_response(&hello_context()).collect::>().await; + let request = server.captured().await; - assert!(result.is_err()); assert!( request.path.contains("arn%3Aaws%3Abedrock%3Aus-west-2%3A000000000000%3Aapplication-inference-profile"), "{}", @@ -729,63 +673,38 @@ mod tests { assert!(!unknown_profile.model().unwrap().supports_prompt_caching()); } - struct FakeBedrockEndpoint { - url: String, - request: oneshot::Receiver, + fn inference_profile_arn(model: &str) -> String { + format!("arn:aws:bedrock:us-west-2:000000000000:inference-profile/{model}") } - struct CapturedRequest { - method: Method, - path: String, - headers: HeaderMap, + fn application_inference_profile_arn() -> &'static str { + "arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000" } - #[derive(Clone)] - struct FakeBedrockState { - request_tx: Arc>>>, - shutdown_tx: Arc>>>, + fn test_provider() -> BedrockProvider { + BedrockProvider::from_config(None, None, ProviderConnectionConfig::default()) } - impl FakeBedrockEndpoint { - async fn start() -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind fake Bedrock endpoint"); - let url = format!("http://{}", listener.local_addr().expect("fake Bedrock endpoint address")); - let (request_tx, request) = oneshot::channel(); - let (shutdown_tx, shutdown) = oneshot::channel(); - let state = FakeBedrockState { - request_tx: Arc::new(Mutex::new(Some(request_tx))), - shutdown_tx: Arc::new(Mutex::new(Some(shutdown_tx))), - }; - - let app = Router::new().fallback(any(capture_bedrock_request)).with_state(state); - tokio::spawn(async move { - axum::serve(listener, app) - .with_graceful_shutdown(async { - let _ = shutdown.await; - }) - .await - .expect("serve fake Bedrock endpoint"); - }); + fn mantle_model() -> String { + LlmModel::all() + .iter() + .find(|model| model.provider_enum() == Provider::Bedrock && model.transport().is_some()) + .expect("catalog must expose at least one Responses-transport Bedrock model") + .model_id() + .to_string() + } - Self { url, request } - } + async fn mantle_provider(server: &CaptureServer) -> BedrockProvider { + create_provider(&server.base_url, &mantle_model()).await } - async fn capture_bedrock_request( - State(state): State, - request: Request, - ) -> impl IntoResponse { - let (parts, _) = request.into_parts(); - if let Some(tx) = state.request_tx.lock().await.take() { - let _ = tx.send(CapturedRequest { - method: parts.method, - path: parts.uri.path().to_string(), - headers: parts.headers, - }); - } - if let Some(tx) = state.shutdown_tx.lock().await.take() { - let _ = tx.send(()); - } - (StatusCode::FORBIDDEN, "{}") + async fn create_provider(base_url: &str, model: &str) -> BedrockProvider { + BedrockProvider::new(ProviderConnectionConfig { + base_url: Some(base_url.to_string()), + auth_mode: ProviderAuthMode::None, + ..Default::default() + }) + .await + .with_model(model) } } diff --git a/crates/llm/src/providers/bedrock/streaming.rs b/crates/llm/src/providers/bedrock/streaming.rs index e1d5f0781..f2753ccb7 100644 --- a/crates/llm/src/providers/bedrock/streaming.rs +++ b/crates/llm/src/providers/bedrock/streaming.rs @@ -6,16 +6,23 @@ use aws_sdk_bedrockruntime::types::{ StopReason as BedrockStopReason, TokenUsage as BedrockTokenUsage, }; use aws_smithy_types::event_stream::RawMessage; -use futures::{Stream, StreamExt, stream}; +use futures::{Stream, stream}; +use std::future::ready; use tracing::{error, warn}; -use crate::providers::stream_assembler::{StreamAssembler, assemble}; -use crate::{LlmError, LlmResponse, ProviderError, StopReason, TokenUsage, Tokens}; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; + +use crate::providers::response_stream::{OpenedStream, StreamAssembler, response_stream}; +use crate::{LlmError, LlmResponse, LlmResponseStream, ProviderError, StopReason, TokenUsage, Tokens}; pub fn process_bedrock_stream( - events: impl Stream> + Send, -) -> impl Stream> + Send { - stream::iter([Ok(LlmResponse::Start)]).chain(assemble(events, |event, turn| Ok(decode_event(event, turn)))) + events: impl Stream> + Send + 'static, +) -> LlmResponseStream { + response_stream( + ready(Ok(OpenedStream::new(events))), + |event, turn| Ok(decode_converse_event(event, turn)), + DEFAULT_STREAM_IDLE_TIMEOUT, + ) } pub(crate) fn converse_events( @@ -69,7 +76,7 @@ impl From> for LlmError { } } -fn decode_event(event: ConverseStreamOutput, turn: &mut StreamAssembler) -> Vec { +pub(super) fn decode_converse_event(event: ConverseStreamOutput, turn: &mut StreamAssembler) -> Vec { let response = match event { ConverseStreamOutput::ContentBlockStart(event) => match event.start { Some(ContentBlockStart::ToolUse(tool)) => { @@ -123,6 +130,7 @@ mod tests { ContentBlockDeltaEvent, ContentBlockStartEvent, ContentBlockStopEvent, ConversationRole, ConverseStreamMetadataEvent, MessageStartEvent, MessageStopEvent, ToolUseBlockDelta, ToolUseBlockStart, }; + use futures::StreamExt; #[tokio::test] async fn test_text_stream() { diff --git a/crates/llm/src/providers/codex/provider.rs b/crates/llm/src/providers/codex/provider.rs index ae9ea51eb..653ac2c7f 100644 --- a/crates/llm/src/providers/codex/provider.rs +++ b/crates/llm/src/providers/codex/provider.rs @@ -1,11 +1,16 @@ use super::oauth::CodexTokenManager; -use crate::provider::{LlmResponseStream, StreamingModelProvider, get_context_window, stream_from}; +use crate::provider::{LlmResponseStream, StreamingModelProvider, get_context_window}; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::http::{SseData, http_client}; use crate::providers::openai_responses::mappers::{ResponsesRequestPolicy, build_wire_request}; -use crate::providers::openai_responses::transport::{ResponsesConnection, process_connection, send}; +use crate::providers::openai_responses::streaming::decode_responses; +use crate::providers::openai_responses::transport::send; +use crate::providers::response_stream::{OpenedStream, response_stream}; use crate::{Context, LlmError, Result}; use aether_auth::OAuthCredentialStorage; use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue}; use std::sync::Arc; +use std::time::Duration; use tracing::debug; const CODEX_API_BASE: &str = "https://chatgpt.com/backend-api/codex"; @@ -17,6 +22,7 @@ pub struct CodexProvider { client: reqwest::Client, model: String, token_manager: Arc, + idle_timeout: Duration, } impl CodexProvider { @@ -24,9 +30,10 @@ impl CodexProvider { let token_manager = CodexTokenManager::new(store, super::PROVIDER_ID); Self { base_url: CODEX_API_BASE.to_string(), - client: reqwest::Client::new(), + client: http_client(), model: "gpt-5.5".to_string(), token_manager: Arc::new(token_manager), + idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, } } @@ -34,6 +41,7 @@ impl CodexProvider { if let Some(base_url) = connection.base_url { self.base_url = base_url.trim_end_matches('/').to_string(); } + self.idle_timeout = connection.idle_timeout; self } @@ -70,7 +78,7 @@ impl CodexProvider { /// Uses manual SSE parsing because the Codex API does not return a /// `Content-Type: text/event-stream` header, which `reqwest_eventsource` /// (used by `async-openai`'s `create_stream`) requires. - async fn send_request(&self, request: serde_json::Value, headers: HeaderMap) -> Result { + async fn send_request(&self, request: serde_json::Value, headers: HeaderMap) -> Result> { let url = format!("{}/responses", self.base_url); debug!("Sending request to Codex API: {url}"); @@ -80,7 +88,7 @@ impl CodexProvider { ); match send(&self.client, &url, headers, request).await { - Ok(connection) => Ok(connection), + Ok(opened) => Ok(opened), Err(error) => { if error.provider().map(|provider| provider.kind) == Some(crate::ProviderErrorKind::Authentication) { self.token_manager.clear_cache().await; @@ -101,19 +109,17 @@ impl StreamingModelProvider for CodexProvider { } fn stream_response(&self, context: &Context) -> LlmResponseStream { - if let Err(error) = crate::provider::validate_reasoning(context, self.model().as_ref()) { - return crate::provider::error_stream(error); - } let provider = self.clone(); let context = context.clone(); - stream_from( + response_stream( async move { - let headers = provider.build_headers().await?; let request = provider.build_wire_request(&context)?; + let headers = provider.build_headers().await?; provider.send_request(request, headers).await }, - process_connection, + decode_responses(), + self.idle_timeout, ) } diff --git a/crates/llm/src/providers/gemini/provider.rs b/crates/llm/src/providers/gemini/provider.rs index 6699c2dde..4a09096af 100644 --- a/crates/llm/src/providers/gemini/provider.rs +++ b/crates/llm/src/providers/gemini/provider.rs @@ -1,14 +1,15 @@ -use crate::provider::get_context_window; -use crate::providers::http::openai_client; +use crate::provider::{get_context_window, validate_reasoning}; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::http::{http_client, openai_client}; use crate::providers::openai_compatible::{AetherOpenAiConfig, build_chat_request, create_custom_stream_generic}; +use crate::providers::response_stream::error_stream; use crate::{ Context, LlmError, LlmResponseStream, ProviderAuthMode, ProviderConnectionConfig, ProviderFactory, Result, StreamingModelProvider, }; -use async_stream::stream; -use futures::StreamExt; use std::env::var; use std::future::ready; +use std::time::Duration; pub const GEMINI_API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/openai/"; @@ -18,16 +19,26 @@ pub struct GeminiProvider { base_url: Option, auth_mode: ProviderAuthMode, model: String, + http: reqwest::Client, + idle_timeout: Duration, } impl GeminiProvider { pub fn new(api_key: Option) -> Self { - Self { api_key, base_url: None, auth_mode: ProviderAuthMode::Default, model: String::new() } + Self { + api_key, + base_url: None, + auth_mode: ProviderAuthMode::Default, + model: String::new(), + http: http_client(), + idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, + } } pub fn with_connection(mut self, connection: ProviderConnectionConfig) -> Self { self.base_url = connection.base_url; self.auth_mode = connection.auth_mode; + self.idle_timeout = connection.idle_timeout; self } @@ -51,7 +62,16 @@ impl GeminiProvider { fn build_openai_client(&self, api_key: &str) -> async_openai::Client { let api_base = self.base_url.as_deref().unwrap_or(GEMINI_API_BASE); let config = async_openai::config::OpenAIConfig::new().with_api_key(api_key).with_api_base(api_base); - openai_client(AetherOpenAiConfig::new(config, self.auth_mode), reqwest::Client::new()) + openai_client(AetherOpenAiConfig::new(config, self.auth_mode), self.http.clone()) + } + + fn try_stream_response(&self, context: &Context) -> Result { + validate_reasoning(context, self.model().as_ref())?; + let api_key = self.get_api_key()?; + let request = build_chat_request(&self.model, context, None)?; + + tracing::info!("Using Gemini API with API key (OpenAI-compatible endpoint)"); + Ok(create_custom_stream_generic(&self.build_openai_client(&api_key), request, self.idle_timeout)) } } @@ -80,37 +100,7 @@ impl StreamingModelProvider for GeminiProvider { } fn stream_response(&self, context: &Context) -> LlmResponseStream { - if let Err(error) = crate::provider::validate_reasoning(context, self.model().as_ref()) { - return crate::provider::error_stream(error); - } - let provider = self.clone(); - let context = context.clone(); - - Box::pin(stream! { - let api_key = match provider.get_api_key() { - Ok(key) => key, - Err(e) => { - yield Err(e); - return; - } - }; - - tracing::info!("Using Gemini API with API key (OpenAI-compatible endpoint)"); - let client = provider.build_openai_client(&api_key); - let request = match build_chat_request(&provider.model, &context, None) { - Ok(req) => req, - Err(e) => { - yield Err(e); - return; - } - }; - let mut inner_stream = - create_custom_stream_generic(&client, request); - - while let Some(result) = inner_stream.next().await { - yield result; - } - }) + self.try_stream_response(context).unwrap_or_else(error_stream) } fn display_name(&self) -> String { @@ -122,6 +112,7 @@ impl StreamingModelProvider for GeminiProvider { mod tests { use super::*; use async_openai::config::Config; + use futures::StreamExt; use reqwest::header::AUTHORIZATION; #[tokio::test] diff --git a/crates/llm/src/providers/generic.rs b/crates/llm/src/providers/generic.rs index 0526d5e65..4706dd035 100644 --- a/crates/llm/src/providers/generic.rs +++ b/crates/llm/src/providers/generic.rs @@ -2,15 +2,18 @@ use async_openai::Client; use async_openai::config::{Config, OpenAIConfig}; use reqwest::Url; use schemars::Schema; +use std::time::Duration; use crate::catalog::Provider; -use crate::provider::{error_stream, get_context_window, stream_from, validate_reasoning}; -use crate::providers::http::openai_client; +use crate::provider::{get_context_window, validate_reasoning}; +use crate::providers::http::{http_client, openai_client}; use crate::providers::openai_compatible::{ AetherOpenAiConfig, PromptCacheKeySource, build_chat_request, create_custom_stream_generic, }; use crate::providers::openai_responses::mappers::build_wire_request; -use crate::providers::openai_responses::transport::{process_connection, send}; +use crate::providers::openai_responses::streaming::decode_responses; +use crate::providers::openai_responses::transport::send; +use crate::providers::response_stream::{error_stream, response_stream}; use crate::tool_schema::normalize_for_moonshot; use crate::{ Context, LlmError, LlmModel, LlmResponseStream, ProviderAuthMode, ProviderConnectionConfig, Result, @@ -94,6 +97,7 @@ pub struct GenericProvider { chat_client: Client, model: String, request_model: Option, + idle_timeout: Duration, } impl GenericProvider { @@ -134,7 +138,7 @@ impl GenericProvider { connection.auth_mode, ); - let http = reqwest::Client::new(); + let http = http_client(); Ok(Self { config, chat_client: openai_client(openai_config.clone(), http.clone()), @@ -142,6 +146,7 @@ impl GenericProvider { http, model: config.default_model.to_string(), request_model: connection.request_model, + idle_timeout: connection.idle_timeout, }) } @@ -155,12 +160,13 @@ impl GenericProvider { impl StreamingModelProvider for GenericProvider { fn stream_response(&self, context: &Context) -> LlmResponseStream { - match &self.config.api { + let stream = match &self.config.api { Api::ChatCompletions { tool_schema_transform, prompt_cache_key } => { self.stream_chat_completions(context, *tool_schema_transform, *prompt_cache_key) } Api::Responses(policy) => self.stream_responses(context, policy), - } + }; + stream.unwrap_or_else(error_stream) } fn display_name(&self) -> String { @@ -182,41 +188,31 @@ impl GenericProvider { context: &Context, tool_schema_transform: Option, prompt_cache_key: PromptCacheKeySource, - ) -> LlmResponseStream { - if let Err(error) = validate_reasoning(context, self.model().as_ref()) { - return error_stream(error); - } - let mut request = match build_chat_request( - self.request_model.as_deref().unwrap_or(&self.model), - context, - tool_schema_transform, - ) { - Ok(request) => request, - Err(error) => return error_stream(error), - }; + ) -> Result { + validate_reasoning(context, self.model().as_ref())?; + let model = self.request_model.as_deref().unwrap_or(&self.model); + let mut request = build_chat_request(model, context, tool_schema_transform)?; request.prompt_cache_key = prompt_cache_key.resolve(context).map(String::from); - create_custom_stream_generic(&self.chat_client, request) + Ok(create_custom_stream_generic(&self.chat_client, request, self.idle_timeout)) } - fn stream_responses(&self, context: &Context, policy: &ResponsesRequestPolicy) -> LlmResponseStream { - let mut url = match Url::parse(&self.openai_config.url("/responses")) { - Ok(url) => url, - Err(error) => return error_stream(LlmError::ProviderRequest(error.to_string())), - }; - + fn stream_responses(&self, context: &Context, policy: &ResponsesRequestPolicy) -> Result { + let mut url = Url::parse(&self.openai_config.url("/responses")) + .map_err(|error| LlmError::ProviderRequest(error.to_string()))?; url.query_pairs_mut().extend_pairs(self.openai_config.query()); - let mut request = match build_wire_request(&self.model, context, policy) { - Ok(request) => request, - Err(error) => return error_stream(error), - }; + let mut request = build_wire_request(&self.model, context, policy)?; if let Some(model) = &self.request_model { request["model"] = model.clone().into(); } let http = self.http.clone(); let headers = self.openai_config.headers(); - stream_from(async move { send(&http, url.as_str(), headers, request).await }, process_connection) + Ok(response_stream( + async move { send(&http, url.as_str(), headers, request).await }, + decode_responses(), + self.idle_timeout, + )) } } diff --git a/crates/llm/src/providers/http.rs b/crates/llm/src/providers/http.rs index 5463c2af0..e9118f0e0 100644 --- a/crates/llm/src/providers/http.rs +++ b/crates/llm/src/providers/http.rs @@ -1,9 +1,41 @@ -use crate::ProviderError; +use crate::providers::response_stream::OpenedStream; +use crate::{LlmError, ProviderError, Result}; use async_openai::{Client, config::Config, error::OpenAIError, middleware::HttpRequestFactory}; -use reqwest::{Request, Response, header::HeaderMap}; +use eventsource_stream::Eventsource; +use futures::{Stream, StreamExt}; +use reqwest::{Request, RequestBuilder, Response, header::HeaderMap}; use serde_json::{Value, from_str}; +use std::future::ready; +use std::pin::Pin; +use std::time::Duration; use tower::{Service, ServiceExt, service_fn}; +pub(crate) type SseData = Pin> + Send>>; + +pub(crate) fn http_client() -> reqwest::Client { + reqwest::Client::builder().connect_timeout(Duration::from_secs(30)).build().expect("HTTP client should build") +} + +pub(crate) async fn open_sse( + request: RequestBuilder, + code: fn(&str) -> Option, +) -> Result> { + let response = request.send().await?; + if !response.status().is_success() { + return Err(rejected(response, code).await.into()); + } + + let http = HttpResponseMetadata::from(&response); + let data = response.bytes_stream().eventsource().filter_map(|event| { + ready(match event { + Ok(event) if event.data == "[DONE]" => None, + Ok(event) => Some(Ok(event.data)), + Err(error) => Some(Err(ProviderError::stream_interrupted(error.to_string()).into())), + }) + }); + Ok(OpenedStream::with_http(Box::pin(data), http)) +} + /// Preserve HTTP diagnostics before async-openai deserializes rejected responses. pub(crate) fn openai_client(config: T, service: U) -> Client where @@ -29,6 +61,17 @@ pub(crate) struct HttpResponseMetadata { pub(crate) request_id: Option, } +impl HttpResponseMetadata { + pub(crate) fn annotate(&self, error: LlmError) -> LlmError { + match error { + LlmError::Provider(provider) => { + provider.with_http_metadata(Some(self.status), self.request_id.clone()).into() + } + error => error, + } + } +} + impl From<&Response> for HttpResponseMetadata { fn from(response: &Response) -> Self { Self { status: response.status().as_u16(), request_id: extract_request_id(response.headers()) } diff --git a/crates/llm/src/providers/local/llama_cpp.rs b/crates/llm/src/providers/local/llama_cpp.rs index 8e7945f9c..943bb00ae 100644 --- a/crates/llm/src/providers/local/llama_cpp.rs +++ b/crates/llm/src/providers/local/llama_cpp.rs @@ -1,19 +1,25 @@ #![doc = include_str!(concat!(env!("OUT_DIR"), "/docs/llamacpp.md"))] use super::util::get_local_config; -use crate::providers::http::openai_client; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::http::{http_client, openai_client}; use crate::providers::openai::OpenAiChatProvider; use crate::{ProviderConnectionConfig, ProviderFactory, Result}; use async_openai::{Client, config::OpenAIConfig}; use std::future::ready; +use std::time::Duration; pub struct LlamaCppProvider { client: Client, + idle_timeout: Duration, } impl LlamaCppProvider { pub fn new(base_url: &str) -> Self { - Self { client: openai_client(get_local_config(base_url), reqwest::Client::new()) } + Self { + client: openai_client(get_local_config(base_url), http_client()), + idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, + } } } @@ -30,7 +36,7 @@ impl ProviderFactory for LlamaCppProvider { fn from_env_with_connection(connection: ProviderConnectionConfig) -> impl Future> + Send { let base_url = connection.base_url.as_deref().unwrap_or("http://localhost:8080/v1"); - ready(Ok(Self::new(base_url))) + ready(Ok(Self { idle_timeout: connection.idle_timeout, ..Self::new(base_url) })) } fn with_model(self, _model: &str) -> Self { @@ -53,4 +59,8 @@ impl OpenAiChatProvider for LlamaCppProvider { fn provider_name(&self) -> &'static str { "LlamaCpp" } + + fn idle_timeout(&self) -> Duration { + self.idle_timeout + } } diff --git a/crates/llm/src/providers/local/ollama.rs b/crates/llm/src/providers/local/ollama.rs index cbcf0c1ee..d464a1e28 100644 --- a/crates/llm/src/providers/local/ollama.rs +++ b/crates/llm/src/providers/local/ollama.rs @@ -1,20 +1,27 @@ #![doc = include_str!(concat!(env!("OUT_DIR"), "/docs/ollama.md"))] use super::util::get_local_config; -use crate::providers::http::openai_client; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::http::{http_client, openai_client}; use crate::providers::openai::OpenAiChatProvider; use crate::{ProviderConnectionConfig, ProviderFactory, Result}; use async_openai::{Client, config::OpenAIConfig}; use std::future::ready; +use std::time::Duration; pub struct OllamaProvider { model: String, client: Client, + idle_timeout: Duration, } impl OllamaProvider { pub fn new(model: &str, base_url: &str) -> Self { - Self { model: model.to_string(), client: openai_client(get_local_config(base_url), reqwest::Client::new()) } + Self { + model: model.to_string(), + client: openai_client(get_local_config(base_url), http_client()), + idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, + } } pub fn default(model: &str) -> Self { @@ -29,7 +36,7 @@ impl ProviderFactory for OllamaProvider { fn from_env_with_connection(connection: ProviderConnectionConfig) -> impl Future> + Send { let base_url = connection.base_url.as_deref().unwrap_or("http://localhost:11434/v1"); - ready(Ok(Self::new("", base_url))) + ready(Ok(Self { idle_timeout: connection.idle_timeout, ..Self::new("", base_url) })) } fn with_model(mut self, model: &str) -> Self { @@ -52,4 +59,8 @@ impl OpenAiChatProvider for OllamaProvider { fn provider_name(&self) -> &'static str { "Ollama" } + + fn idle_timeout(&self) -> Duration { + self.idle_timeout + } } diff --git a/crates/llm/src/providers/mod.rs b/crates/llm/src/providers/mod.rs index 2d949a3b9..7b7bd8f64 100644 --- a/crates/llm/src/providers/mod.rs +++ b/crates/llm/src/providers/mod.rs @@ -13,6 +13,6 @@ pub mod openai; pub mod openai_compatible; pub(crate) mod openai_responses; pub mod openrouter; -pub(crate) mod stream_assembler; +pub(crate) mod response_stream; #[cfg(test)] pub(crate) mod test_capture_server; diff --git a/crates/llm/src/providers/openai/provider.rs b/crates/llm/src/providers/openai/provider.rs index 6db4c5b69..b8b339cf5 100644 --- a/crates/llm/src/providers/openai/provider.rs +++ b/crates/llm/src/providers/openai/provider.rs @@ -2,9 +2,11 @@ use async_openai::{Client, config::Config, types::chat::CreateChatCompletionRequ use tracing::debug; use super::mappers::{map_messages, map_tools}; -use crate::provider::error_stream; +use crate::provider::validate_reasoning; use crate::providers::openai_compatible::create_custom_stream_generic; -use crate::{Context, LlmResponseStream, StreamingModelProvider}; +use crate::providers::response_stream::error_stream; +use crate::{Context, LlmResponseStream, Result, StreamingModelProvider}; +use std::time::Duration; /// A Provider that's compatible with `OpenAI`'s chat completion API /// Other providers (e.g. Ollama, Llama.cpp etc) that are "`OpenAI` compatible" should implement this trait @@ -14,31 +16,12 @@ pub trait OpenAiChatProvider { fn client(&self) -> &Client; fn model(&self) -> &str; fn provider_name(&self) -> &str; + fn idle_timeout(&self) -> Duration; } impl StreamingModelProvider for T { fn stream_response(&self, context: &Context) -> LlmResponseStream { - if let Err(error) = crate::provider::validate_reasoning(context, None) { - return crate::provider::error_stream(error); - } - let model = self.model().to_string(); - let messages = match map_messages(context.messages()) { - Ok(messages) => messages, - Err(e) => return error_stream(e), - }; - let message_count = messages.len(); - let tools = if context.tools().is_empty() { - None - } else { - match map_tools(context.tools(), None) { - Ok(t) => Some(t), - Err(e) => return error_stream(e), - } - }; - - debug!("Starting chat completion stream for model: {model} with {message_count} messages"); - let request = CreateChatCompletionRequest { model, messages, tools, stream: Some(true), ..Default::default() }; - create_custom_stream_generic(self.client(), request) + try_stream_response(self, context).unwrap_or_else(error_stream) } fn context_window(&self) -> Option { @@ -51,20 +34,50 @@ impl StreamingModelProvider for T { } } +fn try_stream_response(provider: &T, context: &Context) -> Result { + validate_reasoning(context, None)?; + let model = provider.model().to_string(); + let messages = map_messages(context.messages())?; + let message_count = messages.len(); + let tools = if context.tools().is_empty() { None } else { Some(map_tools(context.tools(), None)?) }; + + debug!("Starting chat completion stream for model: {model} with {message_count} messages"); + let request = CreateChatCompletionRequest { model, messages, tools, stream: Some(true), ..Default::default() }; + Ok(create_custom_stream_generic(provider.client(), request, provider.idle_timeout())) +} + #[cfg(test)] mod tests { use futures::StreamExt; use super::*; use crate::providers::local::ollama::OllamaProvider; - use crate::providers::test_capture_server::CaptureServer; - use crate::{ChatMessage, LlmResponse, Result}; + use crate::providers::test_capture_server::{CaptureServer, ResponseSpec, hello_context}; + use crate::{LlmError, LlmResponse, ProviderConnectionConfig, ProviderErrorKind, ProviderFactory}; + + #[tokio::test] + async fn local_provider_honors_configured_idle_timeout() { + let spec = ResponseSpec::sse(include_str!("../../../tests/fixtures/openai/01_minimal.sse")) + .paced(Duration::from_mins(2)); + let mut server = CaptureServer::start_with_spec(spec).await; + let connection = ProviderConnectionConfig { + base_url: Some(server.base_url.clone()), + idle_timeout: Duration::from_mins(1), + ..Default::default() + }; + let provider = OllamaProvider::from_env_with_connection(connection).await.unwrap().with_model("test-model"); + + let responses = server.collect_on_paused_clock(provider.stream_response(&hello_context())).await; + + let error = responses.last().and_then(|response| response.as_ref().err()).and_then(LlmError::provider); + assert_eq!(error.map(|error| error.kind), Some(ProviderErrorKind::Timeout), "{responses:?}"); + } #[tokio::test] async fn local_chat_providers_omit_cache_metadata() { let mut server = CaptureServer::start_chat_completions().await; let provider = OllamaProvider::new("test-model", &server.base_url); - let mut context = Context::new(vec![ChatMessage::user("Hello")], vec![]); + let mut context = hello_context(); context.set_prompt_cache_key(Some("prefix-abc".to_string())); context.set_session_affinity_key(Some("conversation-abc".to_string())); diff --git a/crates/llm/src/providers/openai_compatible/streaming.rs b/crates/llm/src/providers/openai_compatible/streaming.rs index ed816235b..35c41b26b 100644 --- a/crates/llm/src/providers/openai_compatible/streaming.rs +++ b/crates/llm/src/providers/openai_compatible/streaming.rs @@ -1,44 +1,47 @@ use super::types::{ChatCompletionStreamResponse, FinishReason, FunctionCallDelta, ToolCallDelta}; -use crate::provider::stream_from; -use crate::providers::stream_assembler::{StreamAssembler, assemble}; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::response_stream::{OpenedStream, StreamAssembler, response_stream}; use crate::{LlmError, LlmResponse, LlmResponseStream, ProviderError, Result, StopReason}; use async_openai::{Client, config::Config}; -use futures::{Stream, StreamExt, stream}; +use futures::{Stream, StreamExt}; use serde::Serialize; +use std::future::ready; +use std::time::Duration; use tracing::{debug, warn}; /// Generic streaming function that accepts any serializable request type. /// This enables providers to use custom request types while reusing the streaming logic. -pub fn create_custom_stream_generic(client: &Client, request: U) -> LlmResponseStream +pub fn create_custom_stream_generic(client: &Client, request: U, idle_timeout: Duration) -> LlmResponseStream where T: Config + Clone + 'static, U: Serialize + Send + 'static, { let client = client.clone(); - stream_from( + response_stream( async move { - client.chat().create_stream_byot::(request).await.map_err(|e| { - warn!("create_stream_byot failed: {e}"); - LlmError::from(e) - }) - }, - |stream| { - process_compatible_stream(stream.map(|result| { - result.map_err(|e| { + let chunks = + client.chat().create_stream_byot::(request).await.map_err(|e| { + warn!("create_stream_byot failed: {e}"); + LlmError::from(e) + })?; + Ok(OpenedStream::new(chunks.map(|chunk| { + chunk.map_err(|e| { warn!("Stream error from API: {e}"); - ProviderError::stream_interrupted(e.to_string()) + ProviderError::stream_interrupted(e.to_string()).into() }) - })) + }))) }, + decode_chunk, + idle_timeout, ) } -pub fn process_compatible_stream + Send>( - chunks: impl Stream> + Send, -) -> impl Stream> + Send { +pub fn process_compatible_stream + Send + 'static>( + chunks: impl Stream> + Send + 'static, +) -> LlmResponseStream { let chunks = chunks.map(|chunk| chunk.map_err(Into::into)); - stream::iter([Ok(LlmResponse::Start)]).chain(assemble(chunks, decode_chunk)) + response_stream(ready(Ok(OpenedStream::new(chunks))), decode_chunk, DEFAULT_STREAM_IDLE_TIMEOUT) } fn decode_chunk(mut chunk: ChatCompletionStreamResponse, turn: &mut StreamAssembler) -> Result> { @@ -104,6 +107,7 @@ mod tests { use crate::testing::{FakeHttpService, llm_response}; use crate::{ProviderErrorKind, TokenUsage}; use async_openai::config::OpenAIConfig; + use futures::stream; use reqwest::{Body, Method}; #[tokio::test] @@ -123,7 +127,9 @@ mod tests { .into() }); let client = openai_client(OpenAIConfig::new().with_api_key("test-key"), service.clone()); - let responses = create_custom_stream_generic(&client, request.clone()).collect::>().await; + let responses = create_custom_stream_generic(&client, request.clone(), DEFAULT_STREAM_IDLE_TIMEOUT) + .collect::>() + .await; assert_eq!(responses.len(), if status == 200 { 2 } else { 1 }, "{responses:?}"); if status == 200 { diff --git a/crates/llm/src/providers/openai_responses/streaming.rs b/crates/llm/src/providers/openai_responses/streaming.rs index 5b52c1814..f577b050a 100644 --- a/crates/llm/src/providers/openai_responses/streaming.rs +++ b/crates/llm/src/providers/openai_responses/streaming.rs @@ -1,8 +1,8 @@ use async_openai::types::responses::{OutputItem, ReasoningItem, ResponseUsage, Status}; -use futures::Stream; use serde::{Deserialize, Deserializer, de::Error as _}; +use tracing::debug; -use crate::providers::stream_assembler::{StreamAssembler, assemble}; +use crate::providers::response_stream::StreamAssembler; use crate::{LlmResponse, ProviderError, ProviderErrorKind, Result, StopReason, TokenUsage}; #[derive(Debug)] @@ -123,12 +123,13 @@ pub struct ResponsesErrorEvent { pub message: String, } -/// Process an `OpenAI` Responses event stream into `LlmResponse` items. -pub fn process_response_stream( - events: impl Stream> + Send, -) -> impl Stream> + Send { +pub(crate) fn decode_responses() -> impl FnMut(String, &mut StreamAssembler) -> Result> + Send { let mut started = false; - assemble(events, move |event, turn| { + move |data, turn| { + let event = serde_json::from_str::(&data).map_err(|error| { + debug!(data, %error, "Failed to decode Responses SSE event"); + ProviderError::stream_interrupted(format!("Invalid Responses SSE event: {error}")) + })?; if matches!(event, ResponsesStreamEvent::Created) { started = true; } else if !started && !event.may_precede_creation() { @@ -137,7 +138,7 @@ pub fn process_response_stream( ); } decode_event(event, turn) - }) + } } #[derive(Deserialize, Default)] @@ -165,7 +166,6 @@ fn decode_event(event: ResponsesStreamEvent, turn: &mut StreamAssembler) -> let incomplete = matches!(&event, ResponsesStreamEvent::Incomplete(_)); let responses = match event { - ResponsesStreamEvent::Created => vec![LlmResponse::Start], ResponsesStreamEvent::OutputTextDelta(e) if !e.delta.is_empty() => vec![LlmResponse::Text { chunk: e.delta }], ResponsesStreamEvent::OutputItemAdded(e) => match e.item { OutputItem::FunctionCall(call) => vec![turn.start_tool(e.output_index, call.call_id, call.name)], @@ -205,7 +205,8 @@ fn decode_event(event: ResponsesStreamEvent, turn: &mut StreamAssembler) -> let message = format!("Responses API error: {}", e.message); return Err(map_responses_error(e.code, message, ProviderErrorKind::Unknown).into()); } - ResponsesStreamEvent::Ignored + ResponsesStreamEvent::Created + | ResponsesStreamEvent::Ignored | ResponsesStreamEvent::OutputTextDelta(_) | ResponsesStreamEvent::ReasoningTextDelta(_) => vec![], }; @@ -216,10 +217,13 @@ fn decode_event(event: ResponsesStreamEvent, turn: &mut StreamAssembler) -> #[cfg(test)] mod tests { use super::*; - use crate::LlmError; + use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; + use crate::providers::response_stream::{OpenedStream, response_stream}; use crate::testing::llm_response; - use futures::{FutureExt, StreamExt, stream}; + use crate::{LlmError, LlmResponseStream}; + use futures::{FutureExt, Stream, StreamExt, stream}; use serde_json::{Value, json}; + use std::future::ready; #[tokio::test] async fn test_text_stream() { @@ -418,7 +422,7 @@ mod tests { (responses_stream().created().text(&["done"]).completed(), StopReason::EndTurn), (responses_stream().created().text(&["done"]).incomplete(), StopReason::Length), ] { - let events = events.build().into_iter().map(|event| serde_json::from_value(event).map_err(LlmError::from)); + let events = events.build().into_iter().map(|event| Ok(event.to_string())); let events = stream::iter(events).chain(stream::pending()); let responses = process_response_stream(events) .collect::>() @@ -460,7 +464,7 @@ mod tests { ] { let responses = process_events(responses_stream().error(code, "slow down").build()).await; - let err = responses[0].as_ref().expect_err("expected the error event to surface as Err"); + let err = responses[1].as_ref().expect_err("expected the error event to surface as Err"); assert_eq!(err.provider().map(|provider| provider.kind), Some(kind), "{code:?}: {err:?}"); assert_eq!(err.provider().and_then(|provider| provider.code.as_deref()), code); assert!(err.is_retryable(), "{code:?}: {err:?}"); @@ -479,7 +483,7 @@ mod tests { ] { let responses = process_events(responses_stream().failed(code, "model overloaded").build()).await; - let err = responses[0].as_ref().expect_err("expected the failure to surface as Err"); + let err = responses[1].as_ref().expect_err("expected the failure to surface as Err"); assert_eq!(err.provider().map(|provider| provider.kind), Some(kind), "{code:?}: {err:?}"); assert_eq!(err.provider().and_then(|provider| provider.code.as_deref()), code); assert_eq!(err.is_retryable(), retryable, "{code:?}: {err:?}"); @@ -505,18 +509,18 @@ mod tests { #[tokio::test] async fn test_stream_error_propagation_is_retryable() { - let events: Vec> = - vec![Err(ProviderError::stream_interrupted("connection lost").into())]; + let events: Vec> = vec![Err(ProviderError::stream_interrupted("connection lost").into())]; let responses: Vec<_> = process_response_stream(tokio_stream::iter(events)).collect().await; - let err = responses[0].as_ref().expect_err("expected upstream Err to surface as Err"); + assert!(matches!(responses[0], Ok(LlmResponse::Start))); + let err = responses[1].as_ref().expect_err("expected upstream Err to surface as Err"); assert_eq!( err.provider().map(|provider| provider.kind), Some(ProviderErrorKind::StreamInterrupted), "got {err:?}" ); - assert_eq!(responses.len(), 1); + assert_eq!(responses.len(), 2); assert!(err.is_retryable(), "mid-stream interrupts must be retryable"); } @@ -525,7 +529,7 @@ mod tests { let responses = process_events(responses_stream().text(&["leaked"]).build()).await; assert_eq!( - responses[0].as_ref().err().and_then(LlmError::provider).map(|provider| provider.kind), + responses[1].as_ref().err().and_then(LlmError::provider).map(|provider| provider.kind), Some(ProviderErrorKind::StreamInterrupted), "{responses:?}" ); @@ -655,18 +659,20 @@ mod tests { } async fn process_events(events: Vec) -> Vec> { - let events = events.into_iter().map(|event| serde_json::from_value(event).map_err(LlmError::from)); - process_response_stream(tokio_stream::iter(events)).collect().await + process_data(events.into_iter().map(|event| event.to_string()).collect()).await } async fn process_fixture(sse: &str) -> Vec> { - let events = sse - .lines() - .filter_map(|line| line.strip_prefix("data: ")) - .filter(|data| *data != "[DONE]") - .map(|data| serde_json::from_str(data).expect("fixture lines are JSON")) - .collect(); - process_events(events).await + let data = sse.lines().filter_map(|line| line.strip_prefix("data: ")).filter(|data| *data != "[DONE]"); + process_data(data.map(str::to_string).collect()).await + } + + async fn process_data(data: Vec) -> Vec> { + process_response_stream(tokio_stream::iter(data.into_iter().map(Ok))).collect().await + } + + fn process_response_stream(data: impl Stream> + Send + 'static) -> LlmResponseStream { + response_stream(ready(Ok(OpenedStream::new(data))), decode_responses(), DEFAULT_STREAM_IDLE_TIMEOUT) } fn find_usage(responses: &[Result]) -> Option { diff --git a/crates/llm/src/providers/openai_responses/transport.rs b/crates/llm/src/providers/openai_responses/transport.rs index 2ccdb0e01..c60bf7b57 100644 --- a/crates/llm/src/providers/openai_responses/transport.rs +++ b/crates/llm/src/providers/openai_responses/transport.rs @@ -1,60 +1,14 @@ -use std::pin::Pin; - -use super::streaming::{ResponsesStreamEvent, process_response_stream}; -use crate::providers::http::{HttpResponseMetadata, rejected, responses_code}; -use crate::{LlmResponse, ProviderError, Result}; -use eventsource_stream::Eventsource; -use futures::{Stream, StreamExt}; +use crate::Result; +use crate::providers::http::{SseData, open_sse, responses_code}; +use crate::providers::response_stream::OpenedStream; use reqwest::header::ACCEPT; use reqwest::{Client, header::HeaderMap}; -use tracing::debug; - -pub(crate) type ResponsesEventStream = Pin> + Send>>; - -pub(crate) struct ResponsesConnection { - pub(crate) events: ResponsesEventStream, - pub(crate) metadata: HttpResponseMetadata, -} pub(crate) async fn send( http: &Client, url: &str, headers: HeaderMap, body: serde_json::Value, -) -> Result { - let response = http.post(url).headers(headers).header(ACCEPT, "text/event-stream").json(&body).send().await?; - open_connection(response).await -} - -pub(crate) async fn open_connection(response: reqwest::Response) -> Result { - if !response.status().is_success() { - return Err(rejected(response, responses_code).await.into()); - } - let metadata = HttpResponseMetadata::from(&response); - Ok(ResponsesConnection { events: decode_response_sse(response), metadata }) -} - -pub(crate) fn process_connection( - connection: ResponsesConnection, -) -> Pin> + Send>> { - let ResponsesConnection { events, metadata } = connection; - Box::pin(process_response_stream(events).map(move |result| { - result.map_err(|error| match error.provider().cloned() { - Some(provider) => provider.with_http_metadata(Some(metadata.status), metadata.request_id.clone()).into(), - None => error, - }) - })) -} - -pub(crate) fn decode_response_sse(response: reqwest::Response) -> ResponsesEventStream { - Box::pin(response.bytes_stream().eventsource().filter_map(|result| { - std::future::ready(match result { - Ok(event) if event.data == "[DONE]" => None, - Ok(event) => Some(serde_json::from_str::(&event.data).map_err(|error| { - debug!(data = event.data, %error, "Failed to decode Responses SSE event"); - ProviderError::stream_interrupted(format!("Invalid Responses SSE event: {error}")).into() - })), - Err(error) => Some(Err(ProviderError::stream_interrupted(error.to_string()).into())), - }) - })) +) -> Result> { + open_sse(http.post(url).headers(headers).header(ACCEPT, "text/event-stream").json(&body), responses_code).await } diff --git a/crates/llm/src/providers/openrouter/provider.rs b/crates/llm/src/providers/openrouter/provider.rs index 20fcac269..67bf59e35 100644 --- a/crates/llm/src/providers/openrouter/provider.rs +++ b/crates/llm/src/providers/openrouter/provider.rs @@ -1,34 +1,44 @@ use super::types::OpenRouterChatRequest; -use crate::provider::{error_stream, get_context_window}; -use crate::providers::http::openai_client; +use crate::provider::{get_context_window, validate_reasoning}; +use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT; +use crate::providers::http::{http_client, openai_client}; use crate::providers::openai_compatible::{ AetherOpenAiConfig, build_chat_request, streaming::create_custom_stream_generic, }; +use crate::providers::response_stream::error_stream; use crate::{ Context, LlmError, LlmResponseStream, ProviderAuthMode, ProviderConnectionConfig, ProviderFactory, Result, StreamingModelProvider, }; use async_openai::{Client, config::OpenAIConfig}; use std::future::ready; +use std::time::Duration; pub struct OpenRouterProvider { client: Client, model: String, + idle_timeout: Duration, } impl OpenRouterProvider { - pub fn new(api_key: String, model: String) -> Result { + pub fn new(api_key: String, model: String) -> Self { let config = openai_config(Some(api_key), ProviderConnectionConfig::default()); - - let client = openai_client(config, reqwest::Client::new()); - Ok(Self { client, model }) + Self { client: openai_client(config, http_client()), model, idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT } } pub fn default(model: &str) -> Result { let api_key = std::env::var("OPENROUTER_API_KEY") .map_err(|_| LlmError::MissingApiKey("OPENROUTER_API_KEY".to_string()))?; - Self::new(api_key, model.to_string()) + Ok(Self::new(api_key, model.to_string())) + } + + fn try_stream_response(&self, context: &Context) -> Result { + validate_reasoning(context, self.model().as_ref())?; + let mut request = build_chat_request(&self.model, context, None)?; + request.prompt_cache_key = context.prompt_cache_key().map(String::from); + let request = OpenRouterChatRequest::from_compatible(request, context.session_affinity_key()); + Ok(create_custom_stream_generic(&self.client, request, self.idle_timeout)) } } @@ -57,17 +67,7 @@ impl StreamingModelProvider for OpenRouterProvider { } fn stream_response(&self, context: &Context) -> LlmResponseStream { - if let Err(error) = crate::provider::validate_reasoning(context, self.model().as_ref()) { - return crate::provider::error_stream(error); - } - let mut request = match build_chat_request(&self.model, context, None) { - Ok(request) => request, - Err(e) => return error_stream(e), - }; - request.prompt_cache_key = context.prompt_cache_key().map(String::from); - let request = OpenRouterChatRequest::from_compatible(request, context.session_affinity_key()); - - create_custom_stream_generic(&self.client, request) + self.try_stream_response(context).unwrap_or_else(error_stream) } fn display_name(&self) -> String { @@ -90,10 +90,11 @@ fn provider_from_connection(connection: ProviderConnectionConfig) -> Result None, }; + let idle_timeout = connection.idle_timeout; let config = openai_config(api_key, connection); - let client = openai_client(config, reqwest::Client::new()); + let client = openai_client(config, http_client()); - Ok(OpenRouterProvider { client, model: String::new() }) + Ok(OpenRouterProvider { client, model: String::new(), idle_timeout }) } #[cfg(test)] @@ -207,7 +208,7 @@ mod tests { fn provider_with_service(service: &FakeHttpService) -> OpenRouterProvider { let config = openai_config(Some("test-key".into()), ProviderConnectionConfig::default()); let client = openai_client(config, service.clone()); - OpenRouterProvider { client, model: "test-model".into() } + OpenRouterProvider { client, model: "test-model".into(), idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT } } fn response(status: u16, body: impl Into) -> reqwest::Response { diff --git a/crates/llm/src/providers/response_stream.rs b/crates/llm/src/providers/response_stream.rs new file mode 100644 index 000000000..95bacf6e1 --- /dev/null +++ b/crates/llm/src/providers/response_stream.rs @@ -0,0 +1,248 @@ +use crate::providers::http::HttpResponseMetadata; +use crate::{LlmError, LlmResponse, LlmResponseStream, ProviderError, Result, StopReason, ToolCallRequest}; +use async_stream::{stream, try_stream}; +use futures::{Stream, StreamExt}; +use std::collections::BTreeMap; +use std::fmt::Debug; +use std::future::Future; +use std::pin::pin; +use std::time::Duration; +use tracing::warn; + +pub(crate) fn response_stream( + open: impl Future>> + Send + 'static, + decode: impl FnMut(U, &mut StreamAssembler) -> Result> + Send + 'static, + idle_timeout: Duration, +) -> LlmResponseStream +where + T: Stream> + Send + 'static, + U: Send, + V: Ord + Debug + Send + 'static, +{ + Box::pin(stream! { + match within(idle_timeout, open).await.flatten() { + Ok(OpenedStream { events, http }) => { + let mut responses = pin!(decode_events(events, decode, idle_timeout)); + while let Some(response) = responses.next().await { + yield response.map_err(|error| match &http { + Some(http) => http.annotate(error), + None => error, + }); + } + } + Err(error) => yield Err(error), + } + }) +} + +pub(crate) fn error_stream(error: LlmError) -> LlmResponseStream { + Box::pin(futures::stream::iter([Err(error)])) +} + +pub(crate) struct OpenedStream { + events: T, + http: Option, +} + +impl OpenedStream { + pub fn new(events: T) -> Self { + Self { events, http: None } + } + + pub fn with_http(events: T, http: HttpResponseMetadata) -> Self { + Self { events, http: Some(http) } + } +} + +pub(crate) struct StreamAssembler { + tool_calls: BTreeMap, + stop_reason: Option, + termination: Termination, +} + +impl StreamAssembler { + pub fn new() -> Self { + Self { tool_calls: BTreeMap::new(), stop_reason: None, termination: Termination::Pending } + } + + pub fn start_tool(&mut self, index: T, id: String, name: String) -> LlmResponse { + let start = LlmResponse::ToolRequestStart { id: id.clone(), name: name.clone() }; + self.tool_calls.insert(index, ToolCallRequest { id, name, arguments: String::new() }); + start + } + + pub fn append_tool_args(&mut self, index: &T, chunk: String) -> Option { + if chunk.is_empty() { + return None; + } + + let Some(tool_call) = self.tool_calls.get_mut(index) else { + warn!("Received tool call arguments for unknown index {index:?}"); + return None; + }; + tool_call.arguments.push_str(&chunk); + Some(LlmResponse::ToolRequestArg { id: tool_call.id.clone(), chunk }) + } + + pub fn complete_tool(&mut self, index: &T) -> Option { + self.tool_calls.remove(index).map(|tool_call| LlmResponse::ToolRequestComplete { tool_call }) + } + + pub fn complete_tool_with(&mut self, index: &T, mut tool_call: ToolCallRequest) -> Vec { + let mut responses = Vec::new(); + match self.tool_calls.remove(index) { + Some(streamed) if tool_call.arguments.is_empty() => tool_call.arguments = streamed.arguments, + Some(_) => {} + None => { + responses + .push(LlmResponse::ToolRequestStart { id: tool_call.id.clone(), name: tool_call.name.clone() }); + } + } + responses.push(LlmResponse::ToolRequestComplete { tool_call }); + responses + } + + pub fn complete_all_tools(&mut self) -> Vec { + std::mem::take(&mut self.tool_calls) + .into_values() + .map(|tool_call| LlmResponse::ToolRequestComplete { tool_call }) + .collect() + } + + pub fn stop(&mut self, stop_reason: StopReason) { + self.stop_reason = Some(stop_reason); + } + + pub fn allow_eof(&mut self) { + self.termination = Termination::AtEof; + } + + pub fn finish_now(&mut self) { + self.termination = Termination::Immediate; + } + + fn finish(mut self) -> Result> { + if matches!(self.termination, Termination::Pending) { + return Err(ProviderError::stream_interrupted("stream ended before the provider's terminal event").into()); + } + + let mut responses = self.complete_all_tools(); + responses.push(LlmResponse::Done { stop_reason: self.stop_reason }); + Ok(responses) + } +} + +enum Termination { + Pending, + AtEof, + Immediate, +} + +fn decode_events( + events: impl Stream> + Send, + mut decode: impl FnMut(T, &mut StreamAssembler) -> Result> + Send, + idle_timeout: Duration, +) -> impl Stream> + Send +where + T: Send, + U: Ord + Debug + Send, +{ + try_stream! { + let mut events = pin!(events); + let mut turn = StreamAssembler::new(); + yield LlmResponse::Start; + + while let Some(event) = within(idle_timeout, events.next()).await? { + for response in decode(event?, &mut turn)? { + yield response; + } + + if matches!(turn.termination, Termination::Immediate) { + break; + } + } + + for response in turn.finish()? { + yield response; + } + } +} + +async fn within(idle_timeout: Duration, future: impl Future) -> Result { + tokio::time::timeout(idle_timeout, future) + .await + .map_err(|_| ProviderError::timeout(format!("Provider sent nothing for {}s", idle_timeout.as_secs())).into()) +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::stream; + use std::future::{pending, ready}; + + #[tokio::test(start_paused = true)] + async fn provider_that_never_responds_times_out() { + let responses = collect(text_stream(pending::>>>())).await; + + assert_eq!(responses.len(), 1, "{responses:?}"); + assert_timed_out(&responses); + } + + #[tokio::test(start_paused = true)] + async fn provider_that_goes_silent_mid_stream_times_out() { + let events = stream::iter([Ok("partial")]).chain(stream::pending()); + + let responses = collect(text_stream(ready(Ok(OpenedStream::new(events))))).await; + + assert!(matches!(&responses[..2], [Ok(LlmResponse::Start), Ok(LlmResponse::Text { .. })]), "{responses:?}"); + assert_timed_out(&responses); + } + + #[tokio::test(start_paused = true)] + async fn slow_stream_that_keeps_sending_events_completes() { + let events = stream::iter(["one", "two", "three"]).then(|text| async move { + tokio::time::sleep(Duration::from_secs(50)).await; + Ok(text) + }); + + let responses = collect(text_stream(ready(Ok(OpenedStream::new(events))))).await; + + assert!(responses.iter().all(Result::is_ok), "{responses:?}"); + assert!(matches!(responses.last(), Some(Ok(LlmResponse::Done { .. }))), "{responses:?}"); + } + + #[tokio::test] + async fn errors_after_opening_carry_the_http_diagnostics() { + let events = stream::iter([Err(ProviderError::stream_interrupted("connection reset").into())]); + let http = HttpResponseMetadata { status: 200, request_id: Some("req-1".to_string()) }; + + let responses = collect(text_stream(ready(Ok(OpenedStream::with_http(events, http))))).await; + + let error = responses.last().and_then(|response| response.as_ref().err()).and_then(LlmError::provider); + assert_eq!(error.and_then(|error| error.request_id.as_deref()), Some("req-1"), "{responses:?}"); + assert_eq!(error.and_then(|error| error.http_status), Some(200)); + } + + const IDLE_TIMEOUT: Duration = Duration::from_mins(1); + + fn text_stream(open: impl Future>> + Send + 'static) -> LlmResponseStream + where + T: Stream> + Send + 'static, + { + let decode = |text: &'static str, turn: &mut StreamAssembler| { + turn.allow_eof(); + Ok(vec![LlmResponse::Text { chunk: text.to_string() }]) + }; + response_stream(open, decode, IDLE_TIMEOUT) + } + + async fn collect(responses: LlmResponseStream) -> Vec> { + responses.collect().await + } + + fn assert_timed_out(responses: &[Result]) { + let kind = responses.last().and_then(|response| response.as_ref().err()).and_then(LlmError::provider); + assert_eq!(kind.map(|error| error.kind), Some(crate::ProviderErrorKind::Timeout), "{responses:?}"); + assert!(!responses.iter().any(|response| matches!(response, Ok(LlmResponse::Done { .. }))), "{responses:?}"); + } +} diff --git a/crates/llm/src/providers/stream_assembler.rs b/crates/llm/src/providers/stream_assembler.rs deleted file mode 100644 index b49a23cc5..000000000 --- a/crates/llm/src/providers/stream_assembler.rs +++ /dev/null @@ -1,118 +0,0 @@ -use std::collections::BTreeMap; -use std::fmt::Debug; - -use futures::{Stream, StreamExt}; -use tracing::warn; - -use crate::{LlmResponse, ProviderError, Result, StopReason, ToolCallRequest}; - -pub(crate) fn assemble( - events: impl Stream> + Send, - mut decode: impl FnMut(T, &mut StreamAssembler) -> Result> + Send, -) -> impl Stream> + Send -where - T: Send, - U: Ord + Debug + Send, -{ - async_stream::try_stream! { - let mut events = std::pin::pin!(events); - let mut turn = StreamAssembler::new(); - - while let Some(event) = events.next().await { - for response in decode(event?, &mut turn)? { - yield response; - } - if matches!(turn.termination, Termination::Immediate) { - break; - } - } - - for response in turn.finish()? { - yield response; - } - } -} - -pub(crate) struct StreamAssembler { - tool_calls: BTreeMap, - stop_reason: Option, - termination: Termination, -} - -impl StreamAssembler { - pub fn new() -> Self { - Self { tool_calls: BTreeMap::new(), stop_reason: None, termination: Termination::Pending } - } - - pub fn start_tool(&mut self, index: I, id: String, name: String) -> LlmResponse { - let start = LlmResponse::ToolRequestStart { id: id.clone(), name: name.clone() }; - self.tool_calls.insert(index, ToolCallRequest { id, name, arguments: String::new() }); - start - } - - pub fn append_tool_args(&mut self, index: &I, chunk: String) -> Option { - if chunk.is_empty() { - return None; - } - - let Some(tool_call) = self.tool_calls.get_mut(index) else { - warn!("Received tool call arguments for unknown index {index:?}"); - return None; - }; - tool_call.arguments.push_str(&chunk); - Some(LlmResponse::ToolRequestArg { id: tool_call.id.clone(), chunk }) - } - - pub fn complete_tool(&mut self, index: &I) -> Option { - self.tool_calls.remove(index).map(|tool_call| LlmResponse::ToolRequestComplete { tool_call }) - } - - pub fn complete_tool_with(&mut self, index: &I, mut tool_call: ToolCallRequest) -> Vec { - let mut responses = Vec::new(); - match self.tool_calls.remove(index) { - Some(streamed) if tool_call.arguments.is_empty() => tool_call.arguments = streamed.arguments, - Some(_) => {} - None => { - responses - .push(LlmResponse::ToolRequestStart { id: tool_call.id.clone(), name: tool_call.name.clone() }); - } - } - responses.push(LlmResponse::ToolRequestComplete { tool_call }); - responses - } - - pub fn complete_all_tools(&mut self) -> Vec { - std::mem::take(&mut self.tool_calls) - .into_values() - .map(|tool_call| LlmResponse::ToolRequestComplete { tool_call }) - .collect() - } - - pub fn stop(&mut self, stop_reason: StopReason) { - self.stop_reason = Some(stop_reason); - } - - pub fn allow_eof(&mut self) { - self.termination = Termination::AtEof; - } - - pub fn finish_now(&mut self) { - self.termination = Termination::Immediate; - } - - fn finish(mut self) -> Result> { - if matches!(self.termination, Termination::Pending) { - return Err(ProviderError::stream_interrupted("stream ended before the provider's terminal event").into()); - } - - let mut responses = self.complete_all_tools(); - responses.push(LlmResponse::Done { stop_reason: self.stop_reason }); - Ok(responses) - } -} - -enum Termination { - Pending, - AtEof, - Immediate, -} diff --git a/crates/llm/src/providers/test_capture_server.rs b/crates/llm/src/providers/test_capture_server.rs index 6620944d5..a60094226 100644 --- a/crates/llm/src/providers/test_capture_server.rs +++ b/crates/llm/src/providers/test_capture_server.rs @@ -1,12 +1,17 @@ +use crate::{ChatMessage, Context, LlmResponse, LlmResponseStream, Result}; +use axum::body::Body; use axum::extract::State; use axum::http::header::CONTENT_TYPE; use axum::http::{HeaderMap, StatusCode}; use axum::response::IntoResponse; use axum::routing::post; use axum::{Json, Router}; +use futures::{StreamExt, stream}; use serde_json::Value; use std::collections::HashMap; +use std::convert::Infallible; use std::sync::Arc; +use std::time::Duration; use tokio::net::TcpListener; use tokio::sync::mpsc; @@ -20,17 +25,38 @@ pub(crate) struct ResponseSpec { pub(crate) status: u16, pub(crate) body: String, pub(crate) headers: HashMap, + gap: Option, } impl ResponseSpec { pub(crate) fn sse(body: &str) -> Self { - Self { status: 200, body: body.to_string(), headers: HashMap::new() } + Self { status: 200, body: body.to_string(), headers: HashMap::new(), gap: None } } pub(crate) fn with_header(mut self, name: &str, value: &str) -> Self { self.headers.insert(name.to_string(), value.to_string()); self } + + pub(crate) fn paced(mut self, gap: Duration) -> Self { + self.gap = Some(gap); + self + } + + fn response_body(&self) -> Body { + let Some(gap) = self.gap else { + return Body::from(self.body.clone()); + }; + let events: Vec = self.body.split_inclusive("\n\n").map(str::to_string).collect(); + Body::from_stream(stream::iter(events).then(move |event| async move { + tokio::time::sleep(gap).await; + Ok::<_, Infallible>(event) + })) + } +} + +pub(crate) fn hello_context() -> Context { + Context::new(vec![ChatMessage::user("Hello")], vec![]) } struct CaptureState { @@ -64,6 +90,7 @@ impl CaptureServer { .route("/v1/messages", post(capture)) .route("/chat/completions", post(capture)) .route("/v1/chat/completions", post(capture)) + .route("/model/{model}/converse-stream", post(capture)) .with_state(Arc::new(CaptureState { sender, response: spec })); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); @@ -74,6 +101,14 @@ impl CaptureServer { pub(crate) async fn captured(&mut self) -> CapturedRequest { self.receiver.recv().await.expect("no request captured") } + + pub(crate) async fn collect_on_paused_clock(&mut self, responses: LlmResponseStream) -> Vec> { + let (responses, ()) = tokio::join!(responses.collect::>(), async { + self.captured().await; + tokio::time::pause(); + }); + responses + } } const RESPONSES_FIXTURE: &str = include_str!("../../tests/fixtures/openai_responses/01_minimal.sse"); @@ -94,5 +129,5 @@ async fn capture( response_headers.insert(name, value); } } - (status, response_headers, state.response.body.clone()) + (status, response_headers, state.response.response_body()) } diff --git a/packages/website/src/content/docs/aether/settings/llm-providers.mdx b/packages/website/src/content/docs/aether/settings/llm-providers.mdx index b91e28e51..18a9a9406 100644 --- a/packages/website/src/content/docs/aether/settings/llm-providers.mdx +++ b/packages/website/src/content/docs/aether/settings/llm-providers.mdx @@ -256,6 +256,7 @@ The top-level `providers` object (or per-agent `providers`) overrides how Aether | `auth` | Authentication mode: `default` (provider default) or `none` (skip auth). Use `none` only behind a trusted proxy that injects credentials. | | `requestModel` | Overrides the model id sent in requests without changing the catalog identity used for capabilities, context window, and display. | | `inferenceProfileArn` | Bedrock application inference profile ARN to route requests through. | +| `idleTimeoutSecs` | Seconds the provider may stay silent, before responding or between stream events, before the request fails with a retryable timeout. Defaults to 300; raise it for slow local models. | ```json title=".aether/settings.json" { @@ -271,7 +272,7 @@ The top-level `providers` object (or per-agent `providers`) overrides how Aether } ``` -The CLI equivalents are `--provider PROVIDER.url=...`, `--provider PROVIDER.auth=...`, `--provider PROVIDER.request-model=...`, and `--provider bedrock.inference-profile-arn=...`. +The CLI equivalents are `--provider PROVIDER.url=...`, `--provider PROVIDER.auth=...`, `--provider PROVIDER.request-model=...`, `--provider PROVIDER.idle-timeout-secs=...`, and `--provider bedrock.inference-profile-arn=...`. ## Reasoning effort