From 85aac194044aadbb58cfb62b1b927e919be89652 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Mon, 17 Aug 2026 07:56:02 +0000 Subject: [PATCH] fix (providers): handle error code for context length exceeded (#11283) --- crates/goose-providers/src/http_status.rs | 113 ++++++++++- .../tests/context_length_byte_limit.rs | 58 ------ crates/goose/src/providers/utils.rs | 186 ------------------ 3 files changed, 111 insertions(+), 246 deletions(-) delete mode 100644 crates/goose-providers/tests/context_length_byte_limit.rs diff --git a/crates/goose-providers/src/http_status.rs b/crates/goose-providers/src/http_status.rs index 98d6e608e..44ed287d2 100644 --- a/crates/goose-providers/src/http_status.rs +++ b/crates/goose-providers/src/http_status.rs @@ -107,7 +107,29 @@ fn parse_http_date(value: &str) -> Option { None } -pub fn is_context_length_exceeded_message(text: &str) -> bool { +fn is_context_length_exceeded(payload: Option<&Value>, message: &str) -> bool { + let payload_exceeded = payload + .and_then(|payload| payload.get("error")) + .is_some_and(|error| { + error + .get("code") + .and_then(Value::as_str) + .is_some_and(|code| code.eq_ignore_ascii_case("context_length_exceeded")) + || match ( + error.get("n_prompt_tokens").and_then(Value::as_f64), + error.get("n_ctx").and_then(Value::as_f64), + ) { + (Some(prompt_tokens), Some(context_limit)) => { + context_limit > 0.0 && prompt_tokens > context_limit + } + _ => false, + } + }); + + payload_exceeded || is_context_length_exceeded_message(message) +} + +fn is_context_length_exceeded_message(text: &str) -> bool { let text_lower = text.to_lowercase(); let direct_context_phrases = [ @@ -247,7 +269,7 @@ pub fn map_http_error_to_provider_error( StatusCode::PAYLOAD_TOO_LARGE => ProviderError::ContextLengthExceeded(extract_message()), StatusCode::BAD_REQUEST => { let payload_str = extract_message(); - if is_context_length_exceeded_message(&payload_str) { + if is_context_length_exceeded(payload.as_ref(), &payload_str) { ProviderError::ContextLengthExceeded(payload_str) } else { ProviderError::RequestFailed(format!("Bad request (400): {}", payload_str)) @@ -356,6 +378,15 @@ mod tests { h } + fn error_payload(fields: [(&str, Value); N]) -> Value { + let mut error = json!({ "message": "invalid request" }); + let error = error.as_object_mut().unwrap(); + for (key, value) in fields { + error.insert(key.to_string(), value); + } + json!({ "error": error }) + } + #[test] fn retry_after_prefers_body_seconds_over_header() { let payload = json!({ @@ -468,6 +499,11 @@ mod tests { "Input token count exceeds the maximum number of tokens allowed", "Please reduce the length of the messages", "prompt is too long for this model", + "Server received a request which exceeds maximum allowed content length. RequestSize(bytes): 34021227, Limit(bytes): 33554432.", + "Request body size exceeds the maximum allowed limit", + "Request body is too large", + "Request payload too large", + "Content-Length exceeds the maximum allowed request size", ]; for message in messages { @@ -487,6 +523,9 @@ mod tests { "temperature exceeds maximum allowed value", "schema is too long", "metadata length exceeds maximum allowed", + "Invalid request body: temperature exceeds maximum allowed value", + "tools[0].description content length exceeds maximum allowed", + "response content length exceeds maximum allowed", ]; for message in messages { @@ -496,4 +535,74 @@ mod tests { ); } } + + #[test] + fn structured_context_length_classifier_handles_supported_and_invalid_payloads() { + let cases = [ + ( + "context overflow code", + error_payload([("code", json!("Context_Length_Exceeded"))]), + true, + ), + ( + "prompt exceeds context limit", + error_payload([ + ("code", json!(400)), + ("n_prompt_tokens", json!(49202)), + ("n_ctx", json!(49152)), + ]), + true, + ), + ( + "generic output token limit", + error_payload([( + "message", + json!("max_tokens must be less than or equal to 4096"), + )]), + false, + ), + ( + "prompt count only", + error_payload([("n_prompt_tokens", json!(49202))]), + false, + ), + ( + "null token counts", + error_payload([("n_prompt_tokens", Value::Null), ("n_ctx", Value::Null)]), + false, + ), + ( + "zero context limit", + error_payload([("n_prompt_tokens", json!(1)), ("n_ctx", json!(0))]), + false, + ), + ( + "equal token counts", + error_payload([("n_prompt_tokens", json!(49152)), ("n_ctx", json!(49152))]), + false, + ), + ( + "top-level token counts", + json!({ + "error": { "message": "invalid request" }, + "n_prompt_tokens": 49202, + "n_ctx": 49152 + }), + false, + ), + ]; + + for (case, payload, expected_context_overflow) in cases { + let error = map_http_error_to_provider_error( + StatusCode::BAD_REQUEST, + Some(payload), + "http://test/endpoint", + ); + assert_eq!( + matches!(&error, ProviderError::ContextLengthExceeded(_)), + expected_context_overflow, + "unexpected classification for {case}: {error:?}" + ); + } + } } diff --git a/crates/goose-providers/tests/context_length_byte_limit.rs b/crates/goose-providers/tests/context_length_byte_limit.rs deleted file mode 100644 index cbe429b03..000000000 --- a/crates/goose-providers/tests/context_length_byte_limit.rs +++ /dev/null @@ -1,58 +0,0 @@ -use goose_providers::errors::ProviderError; -use goose_providers::http_status::{ - is_context_length_exceeded_message, map_http_error_to_provider_error, -}; -use reqwest::StatusCode; -use serde_json::json; - -#[test] -fn byte_size_limit_messages_classify_as_context_length_exceeded() { - let messages = [ - "Server received a request which exceeds maximum allowed content length. RequestSize(bytes): 34021227, Limit(bytes): 33554432.", - "Request body size exceeds the maximum allowed limit", - "Request body is too large", - "Request payload too large", - "Content-Length exceeds the maximum allowed request size", - ]; - - for message in messages { - assert!( - is_context_length_exceeded_message(message), - "expected context-length match for: {message}" - ); - } -} - -#[test] -fn byte_size_limit_bad_request_maps_to_context_length_exceeded() { - let message = "Server received a request which exceeds maximum allowed content length. RequestSize(bytes): 34021227, Limit(bytes): 33554432."; - let error = map_http_error_to_provider_error( - StatusCode::BAD_REQUEST, - Some(json!({ "error": { "message": message } })), - "https://example.com/v1/messages", - ); - - assert_eq!( - error, - ProviderError::ContextLengthExceeded(message.to_string()) - ); -} - -#[test] -fn generic_length_errors_are_not_context_length_exceeded() { - let messages = [ - "metadata length exceeds maximum allowed", - "temperature exceeds maximum allowed value", - "Invalid request body: temperature exceeds maximum allowed value", - "tools[0].description content length exceeds maximum allowed", - "response content length exceeds maximum allowed", - "max_tokens must be less than or equal to 4096", - ]; - - for message in messages { - assert!( - !is_context_length_exceeded_message(message), - "expected generic bad request for: {message}" - ); - } -} diff --git a/crates/goose/src/providers/utils.rs b/crates/goose/src/providers/utils.rs index 32636e343..562d2a27e 100644 --- a/crates/goose/src/providers/utils.rs +++ b/crates/goose/src/providers/utils.rs @@ -1,15 +1,12 @@ use crate::config::paths::Paths; use anyhow::{anyhow, Result}; use fs_err::File; -use goose_providers::errors::{GoogleErrorCode, ProviderError}; use goose_providers::request_log::{install_logger, RequestLogHandle, RequestLogger}; -use reqwest::{Response, StatusCode}; use serde_json::Value; use std::error::Error; use std::io::{BufWriter, Write}; use std::path::PathBuf; use std::sync::OnceLock; -use std::time::Duration; use uuid::Uuid; pub fn filter_extensions_from_system_prompt(system: &str) -> String { @@ -37,16 +34,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") @@ -56,114 +43,6 @@ pub fn is_google_model(payload: &Value) -> bool { .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" - && goose_providers::http_status::is_context_length_exceeded_message(&error_msg) - { - 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") { @@ -380,69 +259,4 @@ mod tests { assert_eq!(is_google_model(&payload), expected_result); } } - - #[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_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)) - ); - } }