tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

fix (providers): handle error code for context length exceeded (#11283)

Browse Source
This commit is contained in:
Lifei Zhou
2026-08-17 07:56:02 +00:00
committed by GitHub
parent 9915ba7eb7
commit 85aac19404
3 changed files with 111 additions and 246 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-providers/src/http_status.rs
+111 -2
View File
@@ -107,7 +107,29 @@ fn parse_http_date(value: &str) -> Option<SystemTime> {
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<const N: usize>(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:?}"
);
}
}
}
crates/goose-providers/tests/context_length_byte_limit.rs
-58
View File
@@ -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}"
);
}
}
crates/goose/src/providers/utils.rs
-186
View File
@@ -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<Duration> {
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::<u64>().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<Value, ProviderError> {
let status = response.status();
let url = super::http_status::sanitize_url(response.url().as_str());
let payload: Option<Value> = 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))
);
}
}