diff --git a/crates/goose-providers/src/google_response.rs b/crates/goose-providers/src/google_response.rs new file mode 100644 index 000000000000..db4ef0248303 --- /dev/null +++ b/crates/goose-providers/src/google_response.rs @@ -0,0 +1,217 @@ +use crate::errors::{GoogleErrorCode, ProviderError}; +use crate::http_status::sanitize_url; +use reqwest::{Response, StatusCode}; +use serde_json::Value; +use std::time::Duration; + +fn format_server_error_message(status_code: StatusCode, payload: Option<&Value>) -> String { + match payload { + Some(Value::Null) | None => format!( + "HTTP {}: No response body received from server", + status_code.as_u16() + ), + Some(payload) => format!("HTTP {}: {}", status_code.as_u16(), payload), + } +} + +pub fn is_google_model(payload: &Value) -> bool { + payload + .get("model") + .and_then(|model| model.as_str()) + .unwrap_or("") + .to_lowercase() + .contains("google") +} + +fn get_google_final_status(status: StatusCode, payload: Option<&Value>) -> StatusCode { + if status.is_success() { + if let Some(payload) = payload { + if let Some(error) = payload.get("error") { + if let Some(code) = error.get("code").and_then(|code| code.as_u64()) { + if let Some(google_error) = GoogleErrorCode::from_code(code) { + return google_error.to_status_code(); + } + } + } + } + } + status +} + +fn parse_google_retry_delay(payload: &Value) -> Option { + payload + .get("error") + .and_then(|error| error.get("details")) + .and_then(|details| details.as_array()) + .and_then(|details| { + details.iter().find_map(|detail| { + if detail + .get("@type") + .and_then(|type_name| type_name.as_str()) + .is_some_and(|type_name| type_name.ends_with("RetryInfo")) + { + detail + .get("retryDelay") + .and_then(|delay| delay.as_str()) + .and_then(|delay| delay.strip_suffix('s')) + .and_then(|seconds| seconds.parse::().ok()) + .map(Duration::from_secs) + } else { + None + } + }) + }) +} + +/// Handle response from Google Gemini API-compatible endpoints. +pub async fn handle_response_google_compat(response: Response) -> Result { + let status = response.status(); + let url = sanitize_url(response.url().as_str()); + let payload: Option = response.json().await.ok(); + let final_status = get_google_final_status(status, payload.as_ref()); + + match final_status { + StatusCode::OK => payload + .ok_or_else(|| ProviderError::RequestFailed("Response body is not valid JSON".to_string())), + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => Err(ProviderError::Authentication( + format!( + "Authentication failed for {url}. Please ensure your API keys are valid and have the required permissions. Status: {}. Response: {:?}", + final_status, payload + ), + )), + StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND => { + let mut error_msg = "Unknown error".to_string(); + if let Some(payload) = &payload { + if let Some(error) = payload.get("error") { + error_msg = error + .get("message") + .and_then(|message| message.as_str()) + .unwrap_or("Unknown error") + .to_string(); + let error_status = error + .get("status") + .and_then(|status| status.as_str()) + .unwrap_or("Unknown status"); + if error_status == "INVALID_ARGUMENT" + && error_msg.to_lowercase().contains("exceeds") + { + return Err(ProviderError::ContextLengthExceeded(error_msg)); + } + } + } + tracing::debug!( + "{}", + format!( + "Provider request failed with status: {}. Payload: {:?}", + final_status, payload + ) + ); + Err(ProviderError::RequestFailed(format!( + "Request failed with status {final_status} at {url}. Message: {error_msg}" + ))) + } + StatusCode::TOO_MANY_REQUESTS => { + let retry_delay = payload.as_ref().and_then(parse_google_retry_delay); + Err(ProviderError::RateLimitExceeded { + details: format!("{:?}", payload), + retry_delay, + }) + } + _ if final_status.is_server_error() => Err(ProviderError::ServerError(format!( + "Server error ({}) at {url}: {}", + final_status, + format_server_error_message(final_status, payload.as_ref()) + ))), + _ => { + tracing::debug!( + "{}", + format!( + "Provider request failed with status: {}. Payload: {:?}", + final_status, payload + ) + ); + Err(ProviderError::RequestFailed(format!( + "Request failed with status {final_status} at {url}" + ))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn detects_google_model_names() { + for (payload, expected) in [ + (json!({ "model": "google_gemini" }), true), + (json!({ "model": "microsoft_bing" }), false), + (json!({ "model": "" }), false), + (json!({}), false), + (json!({ "model": "Google_XYZ" }), true), + (json!({ "model": "google_abc" }), true), + ] { + assert_eq!(is_google_model(&payload), expected); + } + } + + #[test] + fn google_final_status_preserves_success_without_error_payload() { + let result = get_google_final_status(StatusCode::OK, Some(&json!({}))); + assert_eq!(result, StatusCode::OK); + } + + #[test] + fn google_final_status_uses_payload_error_code_for_success_statuses() { + for (error_code, status, expected_status) in [ + (200, None, StatusCode::OK), + (429, Some(StatusCode::OK), StatusCode::TOO_MANY_REQUESTS), + (400, Some(StatusCode::OK), StatusCode::BAD_REQUEST), + (401, Some(StatusCode::OK), StatusCode::UNAUTHORIZED), + (403, Some(StatusCode::OK), StatusCode::FORBIDDEN), + (404, Some(StatusCode::OK), StatusCode::NOT_FOUND), + (500, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR), + (503, Some(StatusCode::OK), StatusCode::SERVICE_UNAVAILABLE), + (999, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR), + (500, Some(StatusCode::BAD_REQUEST), StatusCode::BAD_REQUEST), + ( + 404, + Some(StatusCode::INTERNAL_SERVER_ERROR), + StatusCode::INTERNAL_SERVER_ERROR, + ), + ] { + let payload = if status.is_some() { + json!({ + "error": { + "code": error_code, + "message": "Error message" + } + }) + } else { + json!({}) + }; + + let result = get_google_final_status(status.unwrap_or(StatusCode::OK), Some(&payload)); + assert_eq!(result, expected_status); + } + } + + #[test] + fn parses_google_retry_delay() { + let payload = json!({ + "error": { + "details": [ + { + "@type": "type.googleapis.com/google.rpc.RetryInfo", + "retryDelay": "42s" + } + ] + } + }); + assert_eq!( + parse_google_retry_delay(&payload), + Some(Duration::from_secs(42)) + ); + } +} diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index a4939b6fd1ea..e41eb0563e5a 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -3,6 +3,7 @@ pub mod avian; pub mod canonical; pub mod errors; pub mod function_name; +pub mod google_response; pub mod http_status; pub mod json; pub mod metadata; diff --git a/crates/goose/src/providers/utils.rs b/crates/goose/src/providers/utils.rs index 82cbcd832614..0ca99a50c28c 100644 --- a/crates/goose/src/providers/utils.rs +++ b/crates/goose/src/providers/utils.rs @@ -1,10 +1,10 @@ use super::base::Usage; -use super::errors::GoogleErrorCode; use crate::config::paths::Paths; use crate::providers::errors::ProviderError; use anyhow::Result; use base64::Engine; pub use goose_providers::function_name::{is_valid_function_name, sanitize_function_name}; +pub use goose_providers::google_response::{handle_response_google_compat, is_google_model}; pub use goose_providers::json::{ json_escape_control_chars_in_string, safely_parse_json, unescape_json_values, }; @@ -12,13 +12,11 @@ use goose_types::ModelConfig; pub use goose_types::{ extract_reasoning_effort, is_openai_responses_model, openai_reasoning_effort_for_thinking, }; -use reqwest::{Response, StatusCode}; use rmcp::model::{AnnotateAble, ImageContent, RawImageContent}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::io::Read; use std::path::Path; -use std::time::Duration; #[derive(Debug, Copy, Clone, Serialize, Deserialize)] pub enum ImageFormat { @@ -71,131 +69,6 @@ pub fn filter_extensions_from_system_prompt(system: &str) -> String { } } -fn format_server_error_message(status_code: StatusCode, payload: Option<&Value>) -> String { - match payload { - Some(Value::Null) | None => format!( - "HTTP {}: No response body received from server", - status_code.as_u16() - ), - Some(p) => format!("HTTP {}: {}", status_code.as_u16(), p), - } -} - -pub fn is_google_model(payload: &Value) -> bool { - payload - .get("model") - .and_then(|m| m.as_str()) - .unwrap_or("") - .to_lowercase() - .contains("google") -} - -/// Extracts `StatusCode` from response status or payload error code. -/// This function first checks the status code of the response. If the status is successful (2xx), -/// it then checks the payload for any error codes and maps them to appropriate `StatusCode`. -/// If the status is not successful (e.g., 4xx or 5xx), the original status code is returned. -fn get_google_final_status(status: StatusCode, payload: Option<&Value>) -> StatusCode { - // If the status is successful, check for an error in the payload - if status.is_success() { - if let Some(payload) = payload { - if let Some(error) = payload.get("error") { - if let Some(code) = error.get("code").and_then(|c| c.as_u64()) { - if let Some(google_error) = GoogleErrorCode::from_code(code) { - return google_error.to_status_code(); - } - } - } - } - } - status -} - -fn parse_google_retry_delay(payload: &Value) -> Option { - payload - .get("error") - .and_then(|error| error.get("details")) - .and_then(|details| details.as_array()) - .and_then(|details_array| { - details_array.iter().find_map(|detail| { - if detail - .get("@type") - .and_then(|t| t.as_str()) - .is_some_and(|s| s.ends_with("RetryInfo")) - { - detail - .get("retryDelay") - .and_then(|delay| delay.as_str()) - .and_then(|s| s.strip_suffix('s')) - .and_then(|num| num.parse::().ok()) - .map(Duration::from_secs) - } else { - None - } - }) - }) -} - -/// Handle response from Google Gemini API-compatible endpoints. -/// -/// Processes HTTP responses, handling specific statuses and parsing the payload -/// for error messages. Logs the response payload for debugging purposes. -/// -/// ### References -/// - Error Codes: https://ai.google.dev/gemini-api/docs/troubleshooting?lang=python -/// -/// ### Arguments -/// - `response`: The HTTP response to process. -/// -/// ### Returns -/// - `Ok(Value)`: Parsed JSON on success. -/// - `Err(ProviderError)`: Describes the failure reason. -pub async fn handle_response_google_compat(response: Response) -> Result { - let status = response.status(); - let url = super::http_status::sanitize_url(response.url().as_str()); - let payload: Option = response.json().await.ok(); - let final_status = get_google_final_status(status, payload.as_ref()); - - match final_status { - StatusCode::OK => payload.ok_or_else( || ProviderError::RequestFailed("Response body is not valid JSON".to_string()) ), - StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => { - Err(ProviderError::Authentication(format!("Authentication failed for {url}. Please ensure your API keys are valid and have the required permissions. \ - Status: {}. Response: {:?}", final_status, payload ))) - } - StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND => { - let mut error_msg = "Unknown error".to_string(); - if let Some(payload) = &payload { - if let Some(error) = payload.get("error") { - error_msg = error.get("message").and_then(|m| m.as_str()).unwrap_or("Unknown error").to_string(); - let error_status = error.get("status").and_then(|s| s.as_str()).unwrap_or("Unknown status"); - if error_status == "INVALID_ARGUMENT" && error_msg.to_lowercase().contains("exceeds") { - return Err(ProviderError::ContextLengthExceeded(error_msg.to_string())); - } - } - } - tracing::debug!( - "{}", format!("Provider request failed with status: {}. Payload: {:?}", final_status, payload) - ); - Err(ProviderError::RequestFailed(format!("Request failed with status {} at {url}. Message: {}", final_status, error_msg))) - } - StatusCode::TOO_MANY_REQUESTS => { - let retry_delay = payload.as_ref().and_then(parse_google_retry_delay); - Err(ProviderError::RateLimitExceeded { - details: format!("{:?}", payload), - retry_delay, - }) - } - _ if final_status.is_server_error() => Err(ProviderError::ServerError( - format!("Server error ({}) at {url}: {}", final_status, format_server_error_message(final_status, payload.as_ref())), - )), - _ => { - tracing::debug!( - "{}", format!("Provider request failed with status: {}. Payload: {:?}", final_status, payload) - ); - Err(ProviderError::RequestFailed(format!("Request failed with status {} at {url}", final_status))) - } - } -} - /// Extract the model name from a JSON object. Common with most providers to have this top level attribute. pub fn get_model(data: &Value) -> String { if let Some(model) = data.get("model") { @@ -528,53 +401,6 @@ mod tests { } } - #[test] - fn test_get_google_final_status_success() { - let status = StatusCode::OK; - let payload = json!({}); - let result = get_google_final_status(status, Some(&payload)); - assert_eq!(result, StatusCode::OK); - } - - #[test] - fn test_get_google_final_status_with_error_code() { - // Test error code mappings for different payload error codes - let test_cases = vec![ - // (error code, status, expected status code) - (200, None, StatusCode::OK), - (429, Some(StatusCode::OK), StatusCode::TOO_MANY_REQUESTS), - (400, Some(StatusCode::OK), StatusCode::BAD_REQUEST), - (401, Some(StatusCode::OK), StatusCode::UNAUTHORIZED), - (403, Some(StatusCode::OK), StatusCode::FORBIDDEN), - (404, Some(StatusCode::OK), StatusCode::NOT_FOUND), - (500, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR), - (503, Some(StatusCode::OK), StatusCode::SERVICE_UNAVAILABLE), - (999, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR), - (500, Some(StatusCode::BAD_REQUEST), StatusCode::BAD_REQUEST), - ( - 404, - Some(StatusCode::INTERNAL_SERVER_ERROR), - StatusCode::INTERNAL_SERVER_ERROR, - ), - ]; - - for (error_code, status, expected_status) in test_cases { - let payload = if let Some(_status) = status { - json!({ - "error": { - "code": error_code, - "message": "Error message" - } - }) - } else { - json!({}) - }; - - let result = get_google_final_status(status.unwrap_or(StatusCode::OK), Some(&payload)); - assert_eq!(result, expected_status); - } - } - #[test] fn test_safely_parse_json() { // Test valid JSON that should parse without escaping (contains proper escape sequence) @@ -667,24 +493,6 @@ mod tests { ); } - #[test] - fn test_parse_google_retry_delay() { - let payload = json!({ - "error": { - "details": [ - { - "@type": "type.googleapis.com/google.rpc.RetryInfo", - "retryDelay": "42s" - } - ] - } - }); - assert_eq!( - parse_google_retry_delay(&payload), - Some(Duration::from_secs(42)) - ); - } - #[test] fn test_is_openai_responses_model_matches_o_and_gpt5_families() { for model in [