From 269e82a4f6289f2fd79d874f016fa64bd0d442d2 Mon Sep 17 00:00:00 2001 From: Lucas Kim Date: Mon, 11 May 2026 14:34:45 -0400 Subject: [PATCH] feat: add GOOSE_PROVIDER_TIMEOUT config for provider API timeouts --- crates/goose/src/providers/api_client.rs | 96 ++++++++++++++++++++++-- crates/goose/src/providers/databricks.rs | 23 ++---- crates/goose/src/providers/litellm.rs | 10 +-- crates/goose/src/providers/ollama.rs | 11 ++- crates/goose/src/providers/openai.rs | 25 +++--- 5 files changed, 116 insertions(+), 49 deletions(-) diff --git a/crates/goose/src/providers/api_client.rs b/crates/goose/src/providers/api_client.rs index 86109073ad..cf3deebce3 100644 --- a/crates/goose/src/providers/api_client.rs +++ b/crates/goose/src/providers/api_client.rs @@ -12,6 +12,8 @@ use std::fs::read_to_string; use std::path::PathBuf; use std::time::Duration; +const GOOSE_PROVIDER_TIMEOUT: &str = "GOOSE_PROVIDER_TIMEOUT"; + pub struct ApiClient { client: Client, host: String, @@ -279,11 +281,7 @@ pub struct ApiRequestBuilder<'a> { impl ApiClient { pub fn new(host: String, auth: AuthMethod) -> Result { - Self::with_timeout( - host, - auth, - Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS), - ) + Self::with_timeout(host, auth, resolve_provider_timeout(None)) } pub fn with_timeout(host: String, auth: AuthMethod, timeout: Duration) -> Result { @@ -433,6 +431,22 @@ impl ApiClient { } } +pub(crate) fn resolve_provider_timeout(provider_timeout_key: Option<&str>) -> Duration { + let config = crate::config::Config::global(); + let timeout_secs = provider_timeout_key + .and_then(|key| config.get_param::(key).ok()) + .filter(|seconds| *seconds > 0) + .or_else(|| { + config + .get_param::(GOOSE_PROVIDER_TIMEOUT) + .ok() + .filter(|seconds| *seconds > 0) + }) + .unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS); + + Duration::from_secs(timeout_secs) +} + impl<'a> ApiRequestBuilder<'a> { pub fn header(mut self, key: &str, value: &str) -> Result { let header_name = HeaderName::from_bytes(key.as_bytes())?; @@ -645,6 +659,78 @@ mod tests { use super::*; use test_case::test_case; + #[test] + fn test_api_client_default_timeout() { + let _guard = env_lock::lock_env([(GOOSE_PROVIDER_TIMEOUT, None::<&str>)]); + let client = ApiClient::new("http://localhost:8080".to_string(), AuthMethod::NoAuth) + .expect("client should build"); + + assert_eq!( + client.timeout, + Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS) + ); + } + + #[test] + fn test_api_client_uses_goose_provider_timeout() { + let _guard = env_lock::lock_env([(GOOSE_PROVIDER_TIMEOUT, Some("123"))]); + let client = ApiClient::new("http://localhost:8080".to_string(), AuthMethod::NoAuth) + .expect("client should build"); + + assert_eq!(client.timeout, Duration::from_secs(123)); + } + + #[test] + fn test_api_client_ignores_zero_goose_provider_timeout() { + let _guard = env_lock::lock_env([(GOOSE_PROVIDER_TIMEOUT, Some("0"))]); + let client = ApiClient::new("http://localhost:8080".to_string(), AuthMethod::NoAuth) + .expect("client should build"); + + assert_eq!( + client.timeout, + Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS) + ); + } + + #[test] + fn test_api_client_with_timeout_overrides_goose_provider_timeout() { + let _guard = env_lock::lock_env([(GOOSE_PROVIDER_TIMEOUT, Some("123"))]); + let client = ApiClient::with_timeout( + "http://localhost:8080".to_string(), + AuthMethod::NoAuth, + Duration::from_secs(42), + ) + .expect("client should build"); + + assert_eq!(client.timeout, Duration::from_secs(42)); + } + + #[test] + fn test_resolve_provider_timeout_prefers_provider_specific_timeout() { + let _guard = env_lock::lock_env([ + ("OPENAI_TIMEOUT", Some("321")), + (GOOSE_PROVIDER_TIMEOUT, Some("123")), + ]); + + assert_eq!( + resolve_provider_timeout(Some("OPENAI_TIMEOUT")), + Duration::from_secs(321) + ); + } + + #[test] + fn test_resolve_provider_timeout_uses_goose_timeout_when_provider_specific_unset() { + let _guard = env_lock::lock_env([ + ("OPENAI_TIMEOUT", None::<&str>), + (GOOSE_PROVIDER_TIMEOUT, Some("123")), + ]); + + assert_eq!( + resolve_provider_timeout(Some("OPENAI_TIMEOUT")), + Duration::from_secs(123) + ); + } + #[test_case(Some("test-session_id-456"), None, Some("test-session_id-456"); "header set")] #[test_case(Some("new-session"), Some(("Agent-Session-Id", "old-session")), Some("new-session"); "replaces existing")] #[test_case(None, Some(("Agent-Session-Id", "old-session")), None; "removes existing on none")] diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 2695871beb..a6216f35c1 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -4,13 +4,9 @@ use futures::future::BoxFuture; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::sync::{Arc, Mutex}; -use std::time::Duration; -use super::api_client::{ApiClient, AuthMethod, AuthProvider}; -use super::base::{ - ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata, - DEFAULT_PROVIDER_TIMEOUT_SECS, -}; +use super::api_client::{resolve_provider_timeout, ApiClient, AuthMethod, AuthProvider}; +use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; use super::embedding::EmbeddingCapable; use super::errors::ProviderError; use super::formats::databricks::create_request; @@ -36,7 +32,6 @@ use serde_json::json; const DEFAULT_CLIENT_ID: &str = "databricks-cli"; const DEFAULT_REDIRECT_URL: &str = "http://localhost"; const DEFAULT_SCOPES: &[&str] = &["all-apis", "offline_access"]; - const DATABRICKS_PROVIDER_NAME: &str = "databricks"; pub const DATABRICKS_DEFAULT_MODEL: &str = "databricks-claude-sonnet-4"; const DATABRICKS_DEFAULT_FAST_MODEL: &str = "databricks-claude-haiku-4-5"; @@ -171,11 +166,8 @@ impl DatabricksProvider { token_cache: token_cache.clone(), })); - let api_client = ApiClient::with_timeout( - host, - auth_method, - Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS), - )?; + let api_client = + ApiClient::with_timeout(host, auth_method, resolve_provider_timeout(None))?; let mut provider = Self { api_client, @@ -239,11 +231,8 @@ impl DatabricksProvider { token_cache: token_cache.clone(), })); - let api_client = ApiClient::with_timeout( - host, - auth_method, - Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS), - )?; + let api_client = + ApiClient::with_timeout(host, auth_method, resolve_provider_timeout(None))?; Ok(Self { api_client, diff --git a/crates/goose/src/providers/litellm.rs b/crates/goose/src/providers/litellm.rs index 8e6f08f518..0d12fa514e 100644 --- a/crates/goose/src/providers/litellm.rs +++ b/crates/goose/src/providers/litellm.rs @@ -4,10 +4,9 @@ use futures::future::BoxFuture; use serde_json::{json, Value}; use std::collections::HashMap; -use super::api_client::{ApiClient, AuthMethod}; +use super::api_client::{resolve_provider_timeout, ApiClient, AuthMethod}; use super::base::{ ConfigKey, MessageStream, ModelInfo, Provider, ProviderDef, ProviderMetadata, ProviderUsage, - DEFAULT_PROVIDER_TIMEOUT_SECS, }; use super::embedding::EmbeddingCapable; use super::errors::ProviderError; @@ -49,9 +48,7 @@ impl LiteLLMProvider { .get("LITELLM_CUSTOM_HEADERS") .cloned() .map(parse_custom_headers); - let timeout_secs: u64 = config - .get_param("LITELLM_TIMEOUT") - .unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS); + let timeout = resolve_provider_timeout(Some("LITELLM_TIMEOUT")); let auth = if api_key.is_empty() { AuthMethod::NoAuth @@ -59,8 +56,7 @@ impl LiteLLMProvider { AuthMethod::BearerToken(api_key) }; - let mut api_client = - ApiClient::with_timeout(host, auth, std::time::Duration::from_secs(timeout_secs))?; + let mut api_client = ApiClient::with_timeout(host, auth, timeout)?; if let Some(headers) = custom_headers { let mut header_map = reqwest::header::HeaderMap::new(); diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index 049e6a24ee..6ebfdb9e90 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -1,4 +1,4 @@ -use super::api_client::{ApiClient, AuthMethod}; +use super::api_client::{resolve_provider_timeout, ApiClient, AuthMethod}; use super::base::{ ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata, DEFAULT_PROVIDER_TIMEOUT_SECS, @@ -136,8 +136,7 @@ impl OllamaProvider { .get_param("OLLAMA_HOST") .unwrap_or_else(|_| OLLAMA_HOST.to_string()); - let timeout: Duration = - Duration::from_secs(config.get_param("OLLAMA_TIMEOUT").unwrap_or(OLLAMA_TIMEOUT)); + let timeout = resolve_provider_timeout(Some("OLLAMA_TIMEOUT")); let base = if host.starts_with("http://") || host.starts_with("https://") { host.clone() @@ -174,7 +173,11 @@ impl OllamaProvider { model: ModelConfig, config: DeclarativeProviderConfig, ) -> Result { - let timeout = Duration::from_secs(config.timeout_seconds.unwrap_or(OLLAMA_TIMEOUT)); + let timeout = config + .timeout_seconds + .filter(|seconds| *seconds > 0) + .map(Duration::from_secs) + .unwrap_or_else(|| resolve_provider_timeout(None)); let base = if config.base_url.starts_with("http://") || config.base_url.starts_with("https://") { diff --git a/crates/goose/src/providers/openai.rs b/crates/goose/src/providers/openai.rs index 379b9594ff..45519728d0 100644 --- a/crates/goose/src/providers/openai.rs +++ b/crates/goose/src/providers/openai.rs @@ -1,7 +1,5 @@ -use super::api_client::{ApiClient, AuthMethod}; -use super::base::{ - ConfigKey, ModelInfo, Provider, ProviderDef, ProviderMetadata, DEFAULT_PROVIDER_TIMEOUT_SECS, -}; +use super::api_client::{resolve_provider_timeout, ApiClient, AuthMethod}; +use super::base::{ConfigKey, ModelInfo, Provider, ProviderDef, ProviderMetadata}; use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse}; use super::errors::ProviderError; use super::formats::openai::{create_request, get_usage, response_to_message}; @@ -225,19 +223,13 @@ impl OpenAiProvider { let organization: Option = config.get_param("OPENAI_ORGANIZATION").ok(); let project: Option = config.get_param("OPENAI_PROJECT").ok(); - let timeout_secs: u64 = config - .get_param("OPENAI_TIMEOUT") - .unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS); + let timeout = resolve_provider_timeout(Some("OPENAI_TIMEOUT")); let auth = match api_key { Some(key) if !key.is_empty() => AuthMethod::BearerToken(key), _ => AuthMethod::NoAuth, }; - let mut api_client = ApiClient::with_timeout( - parsed.host, - auth, - std::time::Duration::from_secs(timeout_secs), - )?; + let mut api_client = ApiClient::with_timeout(parsed.host, auth, timeout)?; if !parsed.query_params.is_empty() { api_client = api_client.with_query(parsed.query_params); @@ -354,16 +346,17 @@ impl OpenAiProvider { Self::derive_base_path(url.path()) }; - let timeout_secs = config + let timeout = config .timeout_seconds - .unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS); + .filter(|seconds| *seconds > 0) + .map(std::time::Duration::from_secs) + .unwrap_or_else(|| resolve_provider_timeout(None)); let auth = match api_key { Some(key) if !key.is_empty() => AuthMethod::BearerToken(key), _ => AuthMethod::NoAuth, }; - let mut api_client = - ApiClient::with_timeout(host, auth, std::time::Duration::from_secs(timeout_secs))?; + let mut api_client = ApiClient::with_timeout(host, auth, timeout)?; // Add custom headers if present if let Some(headers) = &config.headers {