Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions crates/goose-cli/src/commands/configure.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use goose::config::{
use goose::conversation::message::Message;
use goose::model::ModelConfig;
use goose::providers::provider_test::test_provider_configuration;
use goose::providers::{create, providers};
use goose::providers::{create, providers, retry_operation, RetryConfig};
use goose::session::{SessionManager, SessionType};
use serde_json::Value;
use std::collections::HashMap;
Expand Down Expand Up @@ -570,13 +570,15 @@ pub async fn configure_provider_dialog() -> anyhow::Result<bool> {
}
}

// Attempt to fetch supported models for this provider
let spin = spinner();
spin.start("Attempting to fetch supported models...");
let models_res = {
let temp_model_config = ModelConfig::new(&provider_meta.default_model)?;
let temp_provider = create(provider_name, temp_model_config).await?;
temp_provider.fetch_recommended_models().await
retry_operation(&RetryConfig::default(), || async {
temp_provider.fetch_recommended_models().await
})
.await
};
spin.stop(style("Model fetch complete").green());

Expand Down
8 changes: 6 additions & 2 deletions crates/goose-server/src/routes/config_management.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,9 @@ use goose::providers::auto_detect::detect_provider_from_api_key;
use goose::providers::base::{ProviderMetadata, ProviderType};
use goose::providers::canonical::maybe_get_canonical_model;
use goose::providers::create_with_default_model;
use goose::providers::errors::ProviderError;
use goose::providers::providers as get_providers;
use goose::providers::{retry_operation, RetryConfig};
use goose::{
agents::execute_commands, agents::ExtensionConfig, config::permission::PermissionLevel,
slash_commands,
Expand Down Expand Up @@ -399,13 +401,15 @@ pub async fn get_provider_models(
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

let models_result = provider.fetch_recommended_models().await;
let models_result = retry_operation(&RetryConfig::default(), || async {
provider.fetch_recommended_models().await
})
.await;

match models_result {
Ok(Some(models)) => Ok(Json(models)),
Ok(None) => Ok(Json(Vec::new())),
Err(provider_error) => {
use goose::providers::errors::ProviderError;
let status_code = match provider_error {
// Permanent misconfigurations - client should fix configuration
ProviderError::Authentication(_) => StatusCode::BAD_REQUEST,
Expand Down
15 changes: 11 additions & 4 deletions crates/goose/src/providers/auto_detect.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::model::ModelConfig;
use crate::providers::retry::{retry_operation, RetryConfig};

pub async fn detect_provider_from_api_key(api_key: &str) -> Option<(String, Vec<String>)> {
let provider_tests = vec![
Expand All @@ -24,10 +25,16 @@ pub async fn detect_provider_from_api_key(api_key: &str) -> Option<(String, Vec<
)
.await
{
Ok(provider) => match provider.fetch_supported_models().await {
Ok(Some(models)) => Some((provider_name.to_string(), models)),
_ => None,
},
Ok(provider) => {
match retry_operation(&RetryConfig::default(), || async {
provider.fetch_supported_models().await
})
.await
{
Ok(Some(models)) => Some((provider_name.to_string(), models)),
_ => None,
}
}
Err(_) => None,
};

Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/providers/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,3 +42,4 @@ pub mod xai;
pub use factory::{
create, create_with_default_model, create_with_named_model, providers, refresh_custom_providers,
};
pub use retry::{retry_operation, RetryConfig};
37 changes: 13 additions & 24 deletions crates/goose/src/providers/openai.rs
Original file line number Diff line number Diff line change
Expand Up @@ -320,30 +320,19 @@ impl Provider for OpenAiProvider {

async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
let models_path = self.base_path.replace("v1/chat/completions", "v1/models");
let response = self
.with_retry(|| async {
let response = self.api_client.response_get(&models_path).await?;
let json = handle_response_openai_compat(response).await?;
if let Some(err_obj) = json.get("error") {
let msg = err_obj
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("unknown error");
return Err(ProviderError::Authentication(msg.to_string()));
}
Ok(json)
})
.await
.inspect_err(|e| {
tracing::warn!("Failed to fetch supported models from OpenAI: {:?}", e);
})?;

let data = response
.get("data")
.and_then(|v| v.as_array())
.ok_or_else(|| {
ProviderError::UsageError("Missing data field in JSON response".into())
})?;
let response = self.api_client.response_get(&models_path).await?;
let json = handle_response_openai_compat(response).await?;
if let Some(err_obj) = json.get("error") {
let msg = err_obj
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("unknown error");
return Err(ProviderError::Authentication(msg.to_string()));
}

let data = json.get("data").and_then(|v| v.as_array()).ok_or_else(|| {
ProviderError::UsageError("Missing data field in JSON response".into())
})?;
let mut models: Vec<String> = data
.iter()
.filter_map(|m| m.get("id").and_then(|v| v.as_str()).map(str::to_string))
Expand Down
62 changes: 55 additions & 7 deletions crates/goose/src/providers/retry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,10 @@ impl RetryConfig {
}
}

pub fn max_retries(&self) -> usize {
self.max_retries
}

pub fn delay_for_attempt(&self, attempt: usize) -> Duration {
if attempt == 0 {
return Duration::from_millis(0);
Expand All @@ -67,6 +71,56 @@ impl RetryConfig {
}
}

pub fn should_retry(error: &ProviderError) -> bool {
matches!(
error,
ProviderError::RateLimitExceeded { .. }
| ProviderError::ServerError(_)
| ProviderError::RequestFailed(_)
)
}

pub async fn retry_operation<F, Fut, T>(
config: &RetryConfig,
operation: F,
) -> Result<T, ProviderError>
where
F: Fn() -> Fut + Send,
Fut: Future<Output = Result<T, ProviderError>> + Send,
T: Send,
{
let mut attempts = 0;

loop {
match operation().await {
Ok(result) => return Ok(result),
Err(error) => {
if should_retry(&error) && attempts < config.max_retries {
attempts += 1;
tracing::warn!(
"Request failed, retrying ({}/{}): {:?}",
attempts,
config.max_retries,
error
);

let delay = match &error {
ProviderError::RateLimitExceeded {
retry_delay: Some(d),
..
} => *d,
_ => config.delay_for_attempt(attempts),
};

sleep(delay).await;
continue;
}
return Err(error);
}
}
}
}

/// Trait for retry functionality to keep Provider dyn-compatible
#[async_trait]
pub trait ProviderRetry {
Expand All @@ -87,12 +141,7 @@ pub trait ProviderRetry {
return match operation().await {
Ok(result) => Ok(result),
Err(error) => {
let should_retry = matches!(
error,
ProviderError::RateLimitExceeded { .. } | ProviderError::ServerError(_)
);

if should_retry && attempts < config.max_retries {
if should_retry(&error) && attempts < config.max_retries {
attempts += 1;
tracing::warn!(
"Request failed, retrying ({}/{}): {:?}",
Expand Down Expand Up @@ -130,7 +179,6 @@ pub trait ProviderRetry {
}
}

// Let specific providers define their retry config if desired
impl<P: Provider> ProviderRetry for P {
fn retry_config(&self) -> RetryConfig {
Provider::retry_config(self)
Expand Down