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

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>

crates/goose-cli/src/session/mod.rs

@@ -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

@@ -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);
+    }
 }