ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: add GOOSE_PROVIDER_TIMEOUT config for provider API timeouts

Browse Source
This commit is contained in:
Lucas Kim
2026-05-11 14:34:45 -04:00
parent 401f8e86ba
commit 269e82a4f6
5 changed files with 116 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/providers/api_client.rs
+91 -5
View File
@@ -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> {
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<Self> {
@@ -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::<u64>(key).ok())
.filter(|seconds| *seconds > 0)
.or_else(|| {
config
.get_param::<u64>(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<Self> {
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")]
crates/goose/src/providers/databricks.rs
+6 -17
View File
@@ -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,
crates/goose/src/providers/litellm.rs
+3 -7
View File
@@ -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();
crates/goose/src/providers/ollama.rs
+7 -4
View File
@@ -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<Self> {
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://") {
crates/goose/src/providers/openai.rs
+9 -16
View File
@@ -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<String> = config.get_param("OPENAI_ORGANIZATION").ok();
let project: Option<String> = 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 {