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

fix: detect low balance and prompt for top up (#7166)

Browse Source 
Signed-off-by: raj-subhankar <subhankar.rj@gmail.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: raj-subhankar <subhankar.rj@gmail.com>
This commit is contained in:
Michael Neale
2026-02-19 13:20:16 +11:00
committed by GitHub
parent 3600c84e4b
commit 629108d0fc
13 changed files with 461 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/mod.rs
+39 -1
View File
@@ -45,7 +45,8 @@ use goose::conversation::message::{ActionRequiredData, Message, MessageContent};
use rustyline::EditMode;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::io::IsTerminal;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
@@ -967,6 +968,7 @@ impl CliSession {
let mut progress_bars = output::McpSpinners::new();
let cancel_token_clone = cancel_token.clone();
let mut markdown_buffer = streaming_buffer::MarkdownBuffer::new();
let mut prompted_credits_urls: HashSet<String> = HashSet::new();
let mut thinking_header_shown = false;
use futures::StreamExt;
@@ -1041,6 +1043,11 @@ impl CliSession {
emit_stream_event(&StreamEvent::Message { message: message.clone() });
} else if !is_json_mode {
output::render_message_streaming(&message, &mut markdown_buffer, &mut thinking_header_shown, self.debug);
maybe_open_credits_top_up_url(
&message,
interactive,
&mut prompted_credits_urls,
);
}
}
}
@@ -1452,6 +1459,37 @@ impl CliSession {
}
}
fn maybe_open_credits_top_up_url(
message: &Message,
interactive: bool,
prompted_credits_urls: &mut HashSet<String>,
) {
if !interactive || !std::io::stdout().is_terminal() {
return;
}
let Some(url) = output::get_credits_top_up_url(message) else {
return;
};
if !prompted_credits_urls.insert(url.clone()) {
return;
}
let should_open = cliclack::confirm("Open the top-up URL in your browser?")
.initial_value(false)
.interact()
.unwrap_or(false);
if should_open && webbrowser::open(&url).is_err() {
output::render_text(
"Could not open browser automatically. Visit the URL above.",
Some(Color::Yellow),
true,
);
}
}
fn emit_stream_event(event: &StreamEvent) {
if let Ok(json) = serde_json::to_string(event) {
println!("{}", json);
crates/goose-cli/src/session/output.rs
+67 -5
View File
@@ -3,7 +3,8 @@ use bat::WrappingMode;
use console::{measure_text_width, style, Color, Term};
use goose::config::Config;
use goose::conversation::message::{
ActionRequiredData, Message, MessageContent, ToolRequest, ToolResponse,
ActionRequiredData, Message, MessageContent, SystemNotificationContent, SystemNotificationType,
ToolRequest, ToolResponse,
};
use goose::providers::canonical::maybe_get_canonical_model;
#[cfg(target_os = "windows")]
@@ -245,8 +246,6 @@ pub fn render_message(message: &Message, debug: bool) {
print_markdown("Thinking was redacted", theme);
}
MessageContent::SystemNotification(notification) => {
use goose::conversation::message::SystemNotificationType;
match notification.notification_type {
SystemNotificationType::ThinkingMessage => {
show_thinking();
@@ -256,6 +255,9 @@ pub fn render_message(message: &Message, debug: bool) {
hide_thinking();
println!("\n{}", style(&notification.msg).yellow());
}
SystemNotificationType::CreditsExhausted => {
render_credits_exhausted_notification(notification);
}
}
}
_ => {
@@ -329,8 +331,6 @@ pub fn render_message_streaming(
print_markdown("Thinking was redacted", theme);
}
MessageContent::SystemNotification(notification) => {
use goose::conversation::message::SystemNotificationType;
match notification.notification_type {
SystemNotificationType::ThinkingMessage => {
show_thinking();
@@ -341,6 +341,10 @@ pub fn render_message_streaming(
hide_thinking();
println!("\n{}", style(&notification.msg).yellow());
}
SystemNotificationType::CreditsExhausted => {
flush_markdown_buffer(buffer, theme);
render_credits_exhausted_notification(notification);
}
}
}
_ => {
@@ -353,6 +357,40 @@ pub fn render_message_streaming(
let _ = std::io::stdout().flush();
}
fn render_credits_exhausted_notification(notification: &SystemNotificationContent) {
hide_thinking();
println!("\n{}", style(&notification.msg).yellow());
if let Some(url) = notification
.data
.as_ref()
.and_then(|d| d.get("top_up_url"))
.and_then(|v| v.as_str())
{
println!(
"{}",
style(format!("Visit this URL to top up credits: {url}")).yellow()
);
}
}
pub fn get_credits_top_up_url(message: &Message) -> Option<String> {
message.content.iter().find_map(|content| {
let MessageContent::SystemNotification(notification) = content else {
return None;
};
if notification.notification_type != SystemNotificationType::CreditsExhausted {
return None;
}
notification
.data
.as_ref()
.and_then(|d| d.get("top_up_url"))
.and_then(|v| v.as_str())
.map(str::to_string)
})
}
pub fn flush_markdown_buffer(buffer: &mut MarkdownBuffer, theme: Theme) {
let remaining = buffer.flush();
if !remaining.is_empty() {
@@ -1434,6 +1472,7 @@ impl McpSpinners {
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::env;
#[test]
@@ -1501,4 +1540,27 @@ mod tests {
"/v/l/p/w/m/components/file.txt"
);
}
#[test]
fn test_get_credits_top_up_url_from_credits_notification() {
let message = Message::assistant().with_system_notification_with_data(
SystemNotificationType::CreditsExhausted,
"Insufficient credits",
json!({"top_up_url": "https://router.tetrate.ai/billing"}),
);
assert_eq!(
get_credits_top_up_url(&message).as_deref(),
Some("https://router.tetrate.ai/billing")
);
}
#[test]
fn test_get_credits_top_up_url_ignores_non_credits_notification() {
let message = Message::assistant().with_system_notification_with_data(
SystemNotificationType::InlineMessage,
"hello",
json!({"top_up_url": "https://router.tetrate.ai/billing"}),
);
assert_eq!(get_credits_top_up_url(&message), None);
}
}
crates/goose/src/agents/agent.rs
+23
View File
@@ -1500,6 +1500,29 @@ impl Agent {
}
}
}
Err(ref provider_err @ ProviderError::CreditsExhausted { details: _, ref top_up_url }) => {
crate::posthog::emit_error(provider_err.telemetry_type(), &provider_err.to_string());
error!("Error: {}", provider_err);
let user_msg = if top_up_url.is_some() {
"Please add credits to your account, then resend your message to continue.".to_string()
} else {
"Please check your account with your provider to add more credits, then resend your message to continue.".to_string()
};
let notification_data = serde_json::json!({
"top_up_url": top_up_url,
});
yield AgentEvent::Message(
Message::assistant().with_system_notification_with_data(
SystemNotificationType::CreditsExhausted,
user_msg,
notification_data,
)
);
break;
}
Err(ref provider_err) => {
crate::posthog::emit_error(provider_err.telemetry_type(), &provider_err.to_string());
error!("Error: {}", provider_err);
crates/goose/src/conversation/message.rs
+1
View File
@@ -164,6 +164,7 @@ pub struct FrontendToolRequest {
pub enum SystemNotificationType {
ThinkingMessage,
InlineMessage,
CreditsExhausted,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
crates/goose/src/providers/errors.rs
+7
View File
@@ -30,6 +30,12 @@ pub enum ProviderError {
#[error("Unsupported operation: {0}")]
NotImplemented(String),
#[error("Credits exhausted: {details}")]
CreditsExhausted {
details: String,
top_up_url: Option<String>,
},
}
impl ProviderError {
@@ -43,6 +49,7 @@ impl ProviderError {
ProviderError::ExecutionError(_) => "execution",
ProviderError::UsageError(_) => "usage",
ProviderError::NotImplemented(_) => "not_implemented",
ProviderError::CreditsExhausted { .. } => "credits_exhausted",
}
}
}
crates/goose/src/providers/openai_compatible.rs
+68
View File
@@ -176,6 +176,10 @@ pub fn map_http_error_to_provider_error(
StatusCode::NOT_FOUND => {
ProviderError::RequestFailed(format!("Resource not found (404): {}", extract_message()))
}
StatusCode::PAYMENT_REQUIRED => ProviderError::CreditsExhausted {
details: extract_message(),
top_up_url: None,
},
StatusCode::PAYLOAD_TOO_LARGE => ProviderError::ContextLengthExceeded(extract_message()),
StatusCode::BAD_REQUEST => {
let payload_str = extract_message();
@@ -251,3 +255,67 @@ pub fn stream_openai_compat(
}
}))
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use test_case::test_case;
#[test_case(
StatusCode::PAYMENT_REQUIRED,
Some(json!({"error": {"message": "Insufficient credits to complete this request"}})),
"CreditsExhausted"
; "402 with payload"
)]
#[test_case(
StatusCode::PAYMENT_REQUIRED,
None,
"CreditsExhausted"
; "402 without payload"
)]
#[test_case(
StatusCode::TOO_MANY_REQUESTS,
Some(json!({"error": {"message": "Rate limit exceeded"}})),
"RateLimitExceeded"
; "429 rate limit"
)]
#[test_case(
StatusCode::UNAUTHORIZED,
None,
"Authentication"
; "401 unauthorized"
)]
#[test_case(
StatusCode::BAD_REQUEST,
Some(json!({"error": {"message": "This request exceeds the maximum context length"}})),
"ContextLengthExceeded"
; "400 context length"
)]
#[test_case(
StatusCode::INTERNAL_SERVER_ERROR,
None,
"ServerError"
; "500 server error"
)]
fn http_status_maps_to_expected_error(
status: StatusCode,
payload: Option<Value>,
expected_variant: &str,
) {
let err = map_http_error_to_provider_error(status, payload);
let actual = err.telemetry_type();
let expected_telemetry = match expected_variant {
"CreditsExhausted" => "credits_exhausted",
"RateLimitExceeded" => "rate_limit",
"Authentication" => "auth",
"ContextLengthExceeded" => "context_length",
"ServerError" => "server",
other => panic!("Unknown variant: {other}"),
};
assert_eq!(
actual, expected_telemetry,
"Expected {expected_variant}, got error: {err:?}"
);
}
}
crates/goose/src/providers/tetrate.rs
+137 -58
View File
@@ -1,7 +1,10 @@
use super::api_client::{ApiClient, AuthMethod};
use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata};
use super::errors::ProviderError;
use super::openai_compatible::{handle_status_openai_compat, stream_openai_compat};
use super::openai_compatible::{
handle_response_openai_compat, handle_status_openai_compat, map_http_error_to_provider_error,
stream_openai_compat,
};
use super::retry::ProviderRetry;
use super::utils::RequestLog;
use crate::config::signup_tetrate::TETRATE_DEFAULT_MODEL;
@@ -13,9 +16,12 @@ use futures::future::BoxFuture;
use crate::model::ModelConfig;
use crate::providers::formats::openai::create_request;
use rmcp::model::Tool;
use serde_json::Value;
const TETRATE_PROVIDER_NAME: &str = "tetrate";
// Tetrate Agent Router Service can run many models, we suggest the default
pub const TETRATE_DOC_URL: &str = "https://router.tetrate.ai";
pub const TETRATE_BILLING_URL: &str = "https://router.tetrate.ai/billing";
pub const TETRATE_KNOWN_MODELS: &[&str] = &[
"claude-opus-4-1",
"claude-3-7-sonnet-latest",
@@ -28,7 +34,6 @@ pub const TETRATE_KNOWN_MODELS: &[&str] = &[
"gpt-5-nano",
"gpt-4.1",
];
pub const TETRATE_DOC_URL: &str = "https://router.tetrate.ai";
#[derive(serde::Serialize)]
pub struct TetrateProvider {
@@ -44,7 +49,6 @@ impl TetrateProvider {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let config = crate::config::Config::global();
let api_key: String = config.get_secret("TETRATE_API_KEY")?;
// API host for LLM endpoints (/v1/chat/completions, /v1/models)
let host: String = config
.get_param("TETRATE_HOST")
.unwrap_or_else(|_| "https://api.router.tetrate.ai".to_string());
@@ -61,6 +65,27 @@ impl TetrateProvider {
name: TETRATE_PROVIDER_NAME.to_string(),
})
}
fn enrich_credits_error(err: ProviderError) -> ProviderError {
match err {
ProviderError::CreditsExhausted { details, .. } => ProviderError::CreditsExhausted {
details,
top_up_url: Some(TETRATE_BILLING_URL.to_string()),
},
other => other,
}
}
fn error_from_tetrate_error_payload(payload: Value) -> ProviderError {
let code = payload
.get("error")
.and_then(|e| e.get("code"))
.and_then(|c| c.as_u64())
.unwrap_or(500) as u16;
let status = reqwest::StatusCode::from_u16(code)
.unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR);
Self::enrich_credits_error(map_http_error_to_provider_error(status, Some(payload)))
}
}
impl ProviderDef for TetrateProvider {
@@ -130,7 +155,34 @@ impl Provider for TetrateProvider {
.api_client
.response_post(Some(session_id), "v1/chat/completions", &payload)
.await?;
handle_status_openai_compat(resp).await
let resp = handle_status_openai_compat(resp)
.await
.map_err(Self::enrich_credits_error)?;
let is_json = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.map(|v| v.to_ascii_lowercase())
.is_some_and(|v| v.contains("json"));
if is_json {
// Streaming responses should be SSE; when we get JSON instead, parse it to map
// explicit error payloads and otherwise fail as a protocol mismatch.
let body = handle_response_openai_compat(resp)
.await
.map_err(Self::enrich_credits_error)?;
if body.get("error").is_some() {
return Err(Self::error_from_tetrate_error_payload(body));
}
return Err(ProviderError::ExecutionError(
"Expected streaming response but received non-streaming payload"
.to_string(),
));
}
Ok(resp)
})
.await
.inspect_err(|e| {
@@ -140,69 +192,96 @@ impl Provider for TetrateProvider {
stream_openai_compat(response, log)
}
/// Fetch supported models from Tetrate Agent Router Service API (only models with tool support)
/// Fetch supported models from Tetrate Agent Router Service API
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
// Use the existing api_client which already has authentication configured
let response = match self
let response = self
.api_client
.request(None, "v1/models")
.response_get()
.response_get(None, "v1/models")
.await
{
Ok(response) => response,
Err(e) => {
return Err(ProviderError::ExecutionError(format!(
"Failed to fetch models from Tetrate API: {}. Please check your API key and account at {}",
e, TETRATE_DOC_URL
)));
}
};
.map_err(|e| ProviderError::RequestFailed(e.to_string()))?;
let json = handle_response_openai_compat(response).await?;
let json: serde_json::Value = response.json().await.map_err(|e| {
ProviderError::ExecutionError(format!(
"Failed to parse Tetrate API response: {}. Please check your API key and account at {}",
e, TETRATE_DOC_URL
))
})?;
// Check for error in response
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::ExecutionError(format!(
"Tetrate API error: {}. Please check your API key and account at {}",
msg, TETRATE_DOC_URL
)));
// Tetrate can return errors in 200 OK responses, so check explicitly
if json.get("error").is_some() {
return Err(Self::error_from_tetrate_error_payload(json));
}
// The response format from /v1/models is expected to be OpenAI-compatible
// It should have a "data" field with an array of model objects
let data = json.get("data").and_then(|v| v.as_array()).ok_or_else(|| {
ProviderError::ExecutionError(format!(
"Tetrate API response missing 'data' field. Please check your API key and account at {}",
TETRATE_DOC_URL
))
let arr = json.get("data").and_then(|v| v.as_array()).ok_or_else(|| {
ProviderError::RequestFailed("Missing 'data' array in models response".to_string())
})?;
let mut models: Vec<String> = data
let mut models: Vec<String> = arr
.iter()
.filter_map(|model| {
let id = model.get("id").and_then(|v| v.as_str())?;
let supports_computer_use = model
.get("supports_computer_use")
.and_then(|v| v.as_bool())
.unwrap_or(false);
if supports_computer_use {
Some(id.to_string())
} else {
None
}
})
.filter_map(|m| m.get("id").and_then(|v| v.as_str()).map(str::to_string))
.collect();
models.sort();
Ok(models)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn enrich_adds_dashboard_url() {
let err = ProviderError::CreditsExhausted {
details: "out of credits".to_string(),
top_up_url: None,
};
match TetrateProvider::enrich_credits_error(err) {
ProviderError::CreditsExhausted { top_up_url, .. } => {
assert_eq!(
top_up_url.as_deref(),
Some("https://router.tetrate.ai/billing")
);
}
_ => panic!("Expected CreditsExhausted variant"),
}
}
#[test]
fn enrich_passes_through_other_errors() {
let err = ProviderError::ServerError("boom".to_string());
assert!(matches!(
TetrateProvider::enrich_credits_error(err),
ProviderError::ServerError(_)
));
}
#[test]
fn error_payload_maps_credits_and_adds_billing_url() {
let payload = json!({
"error": {
"code": 402,
"message": "Insufficient credits"
}
});
match TetrateProvider::error_from_tetrate_error_payload(payload) {
ProviderError::CreditsExhausted {
details,
top_up_url,
} => {
assert!(details.contains("Insufficient credits"));
assert_eq!(top_up_url.as_deref(), Some(TETRATE_BILLING_URL));
}
other => panic!("Expected CreditsExhausted, got {other:?}"),
}
}
#[test]
fn error_payload_maps_authentication() {
let payload = json!({
"error": {
"code": 401,
"message": "Invalid API key"
}
});
match TetrateProvider::error_from_tetrate_error_payload(payload) {
ProviderError::Authentication(msg) => {
assert!(msg.contains("Invalid API key"));
}
other => panic!("Expected Authentication, got {other:?}"),
}
}
}