Enhancing /usage to support credits information (#3333)

* Extend /usage to support credits with sub commands to show only credits or context

Author: akhileshb

crates/chat-cli/src/api_client/error.rs

@@ -1,6 +1,7 @@
 use amzn_codewhisperer_client::operation::create_subscription_token::CreateSubscriptionTokenError;
 use amzn_codewhisperer_client::operation::generate_completions::GenerateCompletionsError;
 use amzn_codewhisperer_client::operation::get_profile::GetProfileError;
+use amzn_codewhisperer_client::operation::get_usage_limits::GetUsageLimitsError;
 use amzn_codewhisperer_client::operation::list_available_customizations::ListAvailableCustomizationsError;
 use amzn_codewhisperer_client::operation::list_available_models::ListAvailableModelsError;
 use amzn_codewhisperer_client::operation::list_available_profiles::ListAvailableProfilesError;
@@ -67,6 +68,9 @@ pub enum ApiClientError {
     #[error("{}", SdkErrorDisplay(.0))]
     CreateSubscriptionToken(#[from] SdkError<CreateSubscriptionTokenError, HttpResponse>),
 
+    #[error("{}", SdkErrorDisplay(.0))]
+    GetUsageLimitsError(#[from] SdkError<GetUsageLimitsError, HttpResponse>),
+
     /// Returned from the backend when the user input is too large to fit within the model context
     /// window.
     ///
@@ -130,6 +134,7 @@ impl ApiClientError {
             Self::ListAvailableModelsError(e) => sdk_status_code(e),
             Self::DefaultModelNotFound => None,
             Self::GetProfileError(e) => sdk_status_code(e),
+            Self::GetUsageLimitsError(e) => sdk_status_code(e),
         }
     }
 }
@@ -158,6 +163,7 @@ impl ReasonCode for ApiClientError {
             Self::ListAvailableModelsError(e) => sdk_error_code(e),
             Self::DefaultModelNotFound => "DefaultModelNotFound".to_string(),
             Self::GetProfileError(e) => sdk_error_code(e),
+            Self::GetUsageLimitsError(e) => sdk_error_code(e),
         }
     }
 }

crates/chat-cli/src/api_client/error_utils.rs

@@ -0,0 +1,30 @@
+use amzn_codewhisperer_client::types::AccessDeniedExceptionReason;
+use amzn_codewhisperer_client::operation::get_usage_limits::GetUsageLimitsError;
+
+use crate::api_client::ApiClientError;
+
+#[derive(Debug, PartialEq)]
+pub enum GetUsageLimitsErrorType {
+    FeatureNotSupported,
+    Other,
+}
+
+/// Classify GetUsageLimits API errors
+pub fn classify_get_usage_limits_error(api_error: &ApiClientError) -> GetUsageLimitsErrorType {
+    match api_error {
+        ApiClientError::GetUsageLimitsError(sdk_err) => {
+            match sdk_err.as_service_error() {
+                Some(GetUsageLimitsError::AccessDeniedError(access_denied)) => {
+                    match access_denied.reason() {
+                        Some(AccessDeniedExceptionReason::FeatureNotSupported) => {
+                            GetUsageLimitsErrorType::FeatureNotSupported
+                        },
+                        _ => GetUsageLimitsErrorType::Other,
+                    }
+                },
+                _ => GetUsageLimitsErrorType::Other,
+            }
+        },
+        _ => GetUsageLimitsErrorType::Other,
+    }
+}

crates/chat-cli/src/api_client/mod.rs

@@ -3,6 +3,7 @@ pub mod customization;
 mod delay_interceptor;
 mod endpoints;
 mod error;
+pub mod error_utils;
 pub mod model;
 mod opt_out;
 pub mod profile;
@@ -375,6 +376,15 @@ impl ApiClient {
             .map_err(ApiClientError::CreateSubscriptionToken)
     }
 
+    pub async fn get_usage_limits(&self) -> Result<amzn_codewhisperer_client::operation::get_usage_limits::GetUsageLimitsOutput, ApiClientError> {
+        self.client
+            .get_usage_limits()
+            .set_origin(Some(amzn_codewhisperer_client::types::Origin::from("KIRO_CLI")))
+            .send()
+            .await
+            .map_err(ApiClientError::GetUsageLimitsError)
+    }
+
     pub async fn send_message(&self, conversation: ConversationState) -> Result<SendMessageOutput, ApiClientError> {
         debug!("Sending conversation: {:#?}", conversation);
 

crates/chat-cli/src/cli/chat/cli/usage/mod.rs

@@ -1,4 +1,6 @@
 use clap::Args;
+use crossterm::style::Color;
+use crossterm::{execute, style};
 
 use crate::cli::chat::token_counter::TokenCount;
 use crate::cli::chat::{ChatError, ChatSession, ChatState};
@@ -19,19 +21,111 @@ pub struct DetailedUsageData {
     pub dropped_context_files: Vec<(String, String)>,
 }
 
+/// Billing usage data from API
+#[derive(Debug)]
+pub struct BillingUsageData {
+    pub status: BillingDataStatus,
+    pub plan_name: String,
+    pub overages_enabled: bool,
+    pub billing_cycle_reset: String,
+    pub usage_breakdowns: Vec<UsageBreakdownInfo>,
+    pub bonus_credits: Vec<BonusCredit>,
+}
+
+#[derive(Debug)]
+pub enum BillingDataStatus {
+    Available,
+    FeatureNotSupported,
+    BackendError(String),
+}
+
+/// Individual usage breakdown information
+#[derive(Debug)]
+pub struct UsageBreakdownInfo {
+    #[allow(dead_code)]
+    pub resource_type: String,
+    pub display_name: String,
+    pub used: f64,
+    pub limit: f64,
+    pub percentage: i32,
+}
+
+/// Individual bonus credit information
+#[derive(Debug)]
+pub struct BonusCredit {
+    pub name: String,
+    pub used: f64,
+    pub total: f64,
+    pub days_until_expiry: i64,
+}
+
 /// Arguments for the usage command that displays token usage statistics and context window
 /// information.
 ///
 /// This command shows how many tokens are being used by different components (context files, tools,
 /// assistant responses, and user prompts) within the current chat session's context window.
 #[deny(missing_docs)]
 #[derive(Debug, PartialEq, Args)]
-pub struct UsageArgs;
+pub struct UsageArgs {
+    /// Show only context window usage
+    #[arg(long)]
+    context: bool,
+    /// Show only credits and billing information
+    #[arg(long)]
+    credits: bool,
+}
 
 impl UsageArgs {
     pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+        match (self.context, self.credits) {
+            (true, false) => {
+                // Show only context window usage
+                self.show_context_usage(os, session).await
+            },
+            (false, true) => {
+                // Show only credits/billing information
+                self.show_credits_info(os, session).await
+            },
+            (false, false) => {
+                // Show both (default behavior)
+                self.show_full_usage(os, session).await
+            },
+            (true, true) => {
+                // Both flags specified - show error
+                execute!(
+                    session.stderr,
+                    style::SetForegroundColor(Color::Red),
+                    style::Print("Error: Cannot specify both --context and --credits flags\n"),
+                    style::SetForegroundColor(Color::Reset),
+                )?;
+                Ok(ChatState::PromptUser { skip_printing_tools: true })
+            }
+        }
+    }
+
+    async fn show_context_usage(&self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+        let usage_data = usage_data_provider::get_detailed_usage_data(session, os).await?;
+        usage_renderer::render_context_window(&usage_data, session).await?;
+        Ok(ChatState::PromptUser { skip_printing_tools: true })
+    }
+
+    async fn show_credits_info(&self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+        let billing_data = usage_data_provider::get_billing_usage_data(os).await?;
+        usage_renderer::render_billing_info(&billing_data, session, true).await?;
+        Ok(ChatState::PromptUser { skip_printing_tools: true })
+    }
+
+    async fn show_full_usage(&self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+        // Get both billing and context data
+        let billing_data = usage_data_provider::get_billing_usage_data(os).await?;
         let usage_data = usage_data_provider::get_detailed_usage_data(session, os).await?;
+
+        // Render billing information
+        usage_renderer::render_billing_info(&billing_data, session, false).await?;
+
+        // Render context window information
         usage_renderer::render_context_window(&usage_data, session).await?;
+        
         Ok(ChatState::PromptUser { skip_printing_tools: true })
     }
 }

crates/chat-cli/src/cli/chat/cli/usage/usage_data_provider.rs

@@ -1,3 +1,6 @@
+use chrono::{DateTime, Utc};
+
+use crate::api_client::error_utils::{classify_get_usage_limits_error, GetUsageLimitsErrorType};
 use crate::cli::chat::cli::model::context_window_tokens;
 use crate::cli::chat::token_counter::{CharCount, TokenCount};
 use crate::cli::chat::{ChatError, ChatSession};
@@ -34,6 +37,118 @@ pub(super) async fn get_detailed_usage_data(session: &mut ChatSession, os: &Os)
     })
 }
 
+/// Get billing usage data from API
+pub(super) async fn get_billing_usage_data(os: &Os) -> Result<super::BillingUsageData, ChatError> {
+    match os.client.get_usage_limits().await {
+        Ok(usage_limits) => {
+            let usage_breakdown = usage_limits.usage_breakdown_list();
+            
+            // Get plan info
+            let plan_name = usage_limits.subscription_info()
+                .map_or("Unknown", |si| si.subscription_title())
+                .to_string();
+
+            // Get overage status
+            let overages_enabled = usage_limits.overage_configuration()
+                .is_some_and(|config| config.overage_status().as_str() == "ENABLED");
+
+            // Get billing cycle reset date from main object
+            let billing_cycle_reset = usage_limits.next_date_reset()
+                .map_or_else(
+                    || "Billing cycle reset: Unknown".to_string(),
+                    |next_reset| {
+                        let reset_secs = next_reset.secs();
+                        let reset_date = DateTime::from_timestamp(reset_secs, 0).unwrap_or_else(Utc::now);
+                        format!("Billing cycle reset: {}", reset_date.format("%m/%d"))
+                    }
+                );
+
+            // Process all usage breakdowns
+            let mut usage_breakdowns = Vec::new();
+            let mut bonus_credits = Vec::new();
+
+            for item in usage_breakdown {
+                // Skip items without free trial info
+                if item.free_trial_info().is_none() {
+                    continue;
+                }
+
+                let resource_type = item.resource_type()
+                    .map_or("Unknown", |rt| rt.as_str())
+                    .to_string();
+                let display_name = item.display_name_plural()
+                    .or_else(|| item.display_name())
+                    .unwrap_or(&resource_type)
+                    .to_string();
+                let used = item.current_usage_with_precision().unwrap_or(0.0);
+                let limit = item.usage_limit_with_precision().unwrap_or(0.0);
+                let percentage = if limit > 0.0 { (used / limit * 100.0) as i32 } else { 0 };
+
+                usage_breakdowns.push(super::UsageBreakdownInfo {
+                    resource_type: resource_type.clone(),
+                    display_name: display_name.clone(),
+                    used,
+                    limit,
+                    percentage,
+                });
+
+                // Check for bonus credits in this item
+                if let Some(free_trial_info) = item.free_trial_info() {
+                    if free_trial_info.free_trial_status().map(|s| s.as_str()) == Some("ACTIVE") {
+                        let bonus_used = free_trial_info.current_usage_with_precision().unwrap_or(0.0);
+                        let bonus_total = free_trial_info.usage_limit_with_precision().unwrap_or(0.0);
+                        
+                        if let Some(expiry_timestamp) = free_trial_info.free_trial_expiry() {
+                            let expiry_secs = expiry_timestamp.secs();
+                            let expiry_date = DateTime::from_timestamp(expiry_secs, 0).unwrap_or_else(Utc::now);
+                            let now = Utc::now();
+                            let days_until_expiry = (expiry_date - now).num_days().max(0);
+                            
+                            bonus_credits.push(super::BonusCredit {
+                                name: display_name,
+                                used: bonus_used,
+                                total: bonus_total,
+                                days_until_expiry,
+                            });
+                        }
+                    }
+                }
+            }
+
+            Ok(super::BillingUsageData {
+                status: super::BillingDataStatus::Available,
+                plan_name,
+                overages_enabled,
+                billing_cycle_reset,
+                usage_breakdowns,
+                bonus_credits,
+            })
+        },
+        Err(err) => {
+            // Check if this is an AccessDeniedError with FEATURE_NOT_SUPPORTED reason
+            let is_feature_not_supported = matches!(
+                classify_get_usage_limits_error(&err), 
+                GetUsageLimitsErrorType::FeatureNotSupported
+            );
+            
+            let status = if is_feature_not_supported {
+                super::BillingDataStatus::FeatureNotSupported
+            } else {
+                super::BillingDataStatus::BackendError(err.to_string())
+            };
+
+            Ok(super::BillingUsageData {
+                status,
+                plan_name: "Unknown".to_string(),
+                overages_enabled: false,
+                billing_cycle_reset: "Unknown".to_string(),
+                usage_breakdowns: Vec::new(),
+                bonus_credits: Vec::new(),
+            })
+        },
+    }
+}
+
 /// Get total usage percentage (external API)
 pub async fn get_total_usage_percentage(session: &mut ChatSession, os: &Os) -> Result<f32, ChatError> {
     let data = get_detailed_usage_data(session, os).await?;