diff --git a/crates/chat-cli/src/api_client/error.rs b/crates/chat-cli/src/api_client/error.rs index 4ac80f329c..ab14c7f0b2 100644 --- a/crates/chat-cli/src/api_client/error.rs +++ b/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; @@ -96,6 +97,10 @@ pub enum ApiClientError { #[error("failed to load credentials: {}", .0)] Credentials(CredentialsError), + // Get usgae limits error + #[error("{}", SdkErrorDisplay(.0))] + GetUsageLimits(#[from] SdkError), + #[error(transparent)] ListAvailableModelsError(#[from] SdkError), @@ -127,6 +132,7 @@ impl ApiClientError { Self::ModelOverloadedError { status_code, .. } => *status_code, Self::MonthlyLimitReached { status_code } => *status_code, Self::Credentials(_e) => None, + Self::GetUsageLimits(e) => sdk_status_code(e), Self::ListAvailableModelsError(e) => sdk_status_code(e), Self::DefaultModelNotFound => None, Self::GetProfileError(e) => sdk_status_code(e), @@ -155,6 +161,7 @@ impl ReasonCode for ApiClientError { Self::ModelOverloadedError { .. } => "ModelOverloadedError".to_string(), Self::MonthlyLimitReached { .. } => "MonthlyLimitReached".to_string(), Self::Credentials(_) => "CredentialsError".to_string(), + Self::GetUsageLimits(e) => sdk_error_code(e), Self::ListAvailableModelsError(e) => sdk_error_code(e), Self::DefaultModelNotFound => "DefaultModelNotFound".to_string(), Self::GetProfileError(e) => sdk_error_code(e), diff --git a/crates/chat-cli/src/api_client/mod.rs b/crates/chat-cli/src/api_client/mod.rs index f21b448b77..89c1286d85 100644 --- a/crates/chat-cli/src/api_client/mod.rs +++ b/crates/chat-cli/src/api_client/mod.rs @@ -13,13 +13,18 @@ use std::time::Duration; use amzn_codewhisperer_client::Client as CodewhispererClient; use amzn_codewhisperer_client::operation::create_subscription_token::CreateSubscriptionTokenOutput; +use amzn_codewhisperer_client::operation::get_usage_limits::GetUsageLimitsOutput; use amzn_codewhisperer_client::types::Origin::Cli; use amzn_codewhisperer_client::types::{ Model, OptInFeatureToggle, OptOutPreference, + ResourceType, SubscriptionStatus, TelemetryEvent, + UsageBreakdown, + UsageLimitList, + UsageLimitType, UserContext, }; use amzn_codewhisperer_streaming_client::Client as CodewhispererStreamingClient; @@ -582,6 +587,47 @@ impl ApiClient { } } + pub async fn get_usage_limits( + &self, + resource_type: Option, + ) -> Result { + if cfg!(test) { + use std::time::{ + Duration as StdDuration, + SystemTime, + }; + let mock_limits = UsageLimitList::builder() + .r#type(UsageLimitType::AgenticRequest) + .current_usage(1000) + .total_usage_limit(1100) + .build() + .unwrap(); + use aws_smithy_types::DateTime as SmithyDateTime; + let next_reset = SmithyDateTime::from(SystemTime::now() + StdDuration::from_secs(14 * 24 * 3600)); + let usage_breakdown = UsageBreakdown::builder() + .current_usage(1_234) + .current_overages(34) + .usage_limit(1_000) + .overage_charges(12.34) + .next_date_reset(next_reset) + .build()?; + + return Ok(GetUsageLimitsOutput::builder() + .limits(mock_limits) + .usage_breakdown_list(usage_breakdown.clone()) + .build()); + } + + // Currently for QDev we should only use AgenticRequest for this API + self.client + .get_usage_limits() + .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone())) + .set_resource_type(resource_type) + .send() + .await + .map_err(ApiClientError::GetUsageLimits) + } + /// Only meant for testing. Do not use outside of testing responses. pub fn set_mock_output(&mut self, json: serde_json::Value) { let mut mock = Vec::new(); diff --git a/crates/chat-cli/src/cli/chat/cli/usage.rs b/crates/chat-cli/src/cli/chat/cli/usage.rs index 6e6fe0c961..18366221fe 100644 --- a/crates/chat-cli/src/cli/chat/cli/usage.rs +++ b/crates/chat-cli/src/cli/chat/cli/usage.rs @@ -1,3 +1,12 @@ +use std::convert::TryFrom; +use std::time::SystemTime; + +use amzn_codewhisperer_client::types::{ + OverageStatus, + ResourceType, + SubscriptionType, + UsageBreakdown, +}; use clap::Args; use crossterm::style::{ Attribute, @@ -190,6 +199,82 @@ impl UsageArgs { )), )?; + match os.client.get_usage_limits(None).await { + Ok(resp) => { + tracing::debug!(?resp, "Raw get_usage_limits response"); + // Subscription tier + if let Some(sub) = resp.subscription_info() { + let tier_str = match sub.r#type() { + SubscriptionType::QDeveloperStandaloneFree => "Free tier", + SubscriptionType::QDeveloperStandalone => "Pro tier", + SubscriptionType::QDeveloperStandaloneProPlus => "Pro Plus tier", + _ => "", + }; + queue!( + session.stderr, + style::Print("\n"), + style::SetAttribute(Attribute::Bold), + style::Print(format!("📊 {} Usage limits\n", tier_str)), + style::SetAttribute(Attribute::Reset), + )?; + } + + // Usage breakdown + let list: &[UsageBreakdown] = resp.usage_breakdown_list(); + if list.is_empty() { + queue!(session.stderr, style::Print("\nUsage information unavailable\n\n"),)?; + } else { + let ub = list + .iter() + .find(|b| matches!(b.resource_type(), Some(ResourceType::AgenticRequest))) + .unwrap_or_else(|| list.first().expect("UsageBreakdown list is not null")); + + let current = ub.current_usage(); + let limit = ub.usage_limit(); + let overage_charges = ub.overage_charges(); + let reset_local = match ub.next_date_reset() { + Some(dt) => { + // DateTime → SystemTime + match SystemTime::try_from(*dt) { + Ok(st) => { + let local: chrono::DateTime = st.into(); + local.format("%m/%d/%Y at %H:%M:%S").to_string() + }, + Err(_) => "1st of next month 12:00:00 GMT".to_string(), + } + }, + None => "1st of next month 12:00:00 GMT".to_string(), + }; + + // Overage status + let overage_msg = match resp.overage_configuration().map(|c| c.overage_status()) { + Some(OverageStatus::Enabled) => format!("${:.2} incurred in overages", overage_charges), + Some(OverageStatus::Disabled) => "Overage disabled by admin".to_string(), + _ => String::new(), + }; + + queue!( + session.stderr, + // Line 1: queries used + style::Print(format!("• {} of {} queries used\n", current, limit)), + // Line 2: overage info + style::Print(format!("• {}\n", overage_msg)), + // Line 3: reset time + style::Print(format!("• Limits reset on {}\n\n", reset_local)), + )?; + } + }, + Err(e) => { + tracing::error!(error = ?e, "Failed to load usage limits with full error"); + queue!( + session.stderr, + style::SetForegroundColor(Color::Red), + style::Print(format!("\nFailed to load usage limits: {}\n\n", e)), + style::SetForegroundColor(Color::Reset), + )?; + }, + } + queue!( session.stderr, style::SetAttribute(Attribute::Bold), diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index e155099450..18bf890021 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -171,9 +171,10 @@ use crate::util::{ ui, }; -const LIMIT_REACHED_TEXT: &str = color_print::cstr! { "You've used all your free requests for this month. You have two options: -1. Upgrade to a paid subscription for increased limits. See our Pricing page for what's included> https://aws.amazon.com/q/developer/pricing/ -2. Wait until next month when your limit automatically resets." }; +const LIMIT_REACHED_TEXT: &str = color_print::cstr! { "You've used all your free requests for this month. You have three options: +1. Upgrade your subscription tier for increased limits. See our Pricing page for what's included in each tier> https://aws.amazon.com/q/developer/pricing/ +2. Enable overages in the Q Developer console so that you can continue to work beyond your monthly limit. Learn more: https://docs.aws.amazon.com/console/amazonq/subscription +3. Wait until next month when your limit automatically resets." }; pub const EXTRA_HELP: &str = color_print::cstr! {" MCP: @@ -1030,7 +1031,7 @@ impl ChatSession { execute!( self.stderr, style::SetForegroundColor(Color::Yellow), - style::Print("Monthly request limit reached"), + style::Print("Monthly request limit reached\n"), style::SetForegroundColor(Color::Reset), )?; @@ -1056,12 +1057,18 @@ impl ChatSession { } else { execute!( self.stderr, - style::SetForegroundColor(Color::Yellow), - style::Print(format!(" - {limits_text}\n\n")), + style::Print("\n"), + style::Print( + "To increase your capacity, ask your account administrator to upgrade your subscription tier or enable overages.\n" + ), + style::Print("Learn more: "), + style::SetForegroundColor(Color::Blue), + style::Print("https://docs.aws.amazon.com/console/amazonq/subscription\n\n"), style::SetForegroundColor(Color::Reset), )?; } + self.conversation.reset_next_user_message(); self.inner = Some(ChatState::PromptUser { skip_printing_tools: false, });