feat: q builder id pro subscription flow (#143)

Author: hayemaxi

crates/chat-cli/src/api_client/clients/client.rs

@@ -1,6 +1,8 @@
 use amzn_codewhisperer_client::Client as CodewhispererClient;
+use amzn_codewhisperer_client::operation::create_subscription_token::CreateSubscriptionTokenOutput;
 use amzn_codewhisperer_client::types::{
     OptOutPreference,
+    SubscriptionStatus,
     TelemetryEvent,
     UserContext,
 };
@@ -123,6 +125,21 @@ impl Client {
             ]),
         }
     }
+
+    pub async fn create_subscription_token(&self) -> Result<CreateSubscriptionTokenOutput, ApiClientError> {
+        match &self.inner {
+            inner::Inner::Codewhisperer(client) => client
+                .create_subscription_token()
+                .send()
+                .await
+                .map_err(ApiClientError::CreateSubscriptionToken),
+            inner::Inner::Mock => Ok(CreateSubscriptionTokenOutput::builder()
+                .set_encoded_verification_url(Some("test/url".to_string()))
+                .set_status(Some(SubscriptionStatus::Inactive))
+                .set_token(Some("test-token".to_string()))
+                .build()?),
+        }
+    }
 }
 
 #[cfg(test)]

crates/chat-cli/src/api_client/clients/streaming_client.rs

@@ -147,6 +147,15 @@ impl StreamingClient {
                                 err.meta().message()
                                 == Some("Encountered unexpectedly high load when processing the request, please try again.")
                             });
+                        let is_monthly_limit_err = e
+                            .raw_response()
+                            .and_then(|resp| resp.body().bytes())
+                            .and_then(|bytes| match String::from_utf8(bytes.to_vec()) {
+                                Ok(s) => Some(s.contains("MONTHLY_REQUEST_COUNT")),
+                                Err(_) => None,
+                            })
+                            .unwrap_or(false);
+
                         if is_quota_breach {
                             Err(ApiClientError::QuotaBreach("quota has reached its limit"))
                         } else if is_context_window_overflow {
@@ -157,6 +166,8 @@ impl StreamingClient {
                                 .and_then(|err| err.meta().request_id())
                                 .map(|s| s.to_string());
                             Err(ApiClientError::ModelOverloadedError { request_id })
+                        } else if is_monthly_limit_err {
+                            Err(ApiClientError::MonthlyLimitReached)
                         } else {
                             Err(e.into())
                         }

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

@@ -1,3 +1,4 @@
+use amzn_codewhisperer_client::operation::create_subscription_token::CreateSubscriptionTokenError;
 use amzn_codewhisperer_client::operation::generate_completions::GenerateCompletionsError;
 use amzn_codewhisperer_client::operation::list_available_customizations::ListAvailableCustomizationsError;
 use amzn_codewhisperer_client::operation::list_available_profiles::ListAvailableProfilesError;
@@ -52,6 +53,13 @@ pub enum ApiClientError {
     #[error("quota has reached its limit")]
     QuotaBreach(&'static str),
 
+    // Separate from quota breach (somehow)
+    #[error("monthly query limit reached")]
+    MonthlyLimitReached,
+
+    #[error("{}", SdkErrorDisplay(.0))]
+    CreateSubscriptionToken(#[from] SdkError<CreateSubscriptionTokenError, HttpResponse>),
+
     /// Returned from the backend when the user input is too large to fit within the model context
     /// window.
     ///
@@ -88,11 +96,13 @@ impl ReasonCode for ApiClientError {
             ApiClientError::QDeveloperChatResponseStream(e) => sdk_error_code(e),
             ApiClientError::ListAvailableProfilesError(e) => sdk_error_code(e),
             ApiClientError::SendTelemetryEvent(e) => sdk_error_code(e),
+            ApiClientError::CreateSubscriptionToken(e) => sdk_error_code(e),
             ApiClientError::QuotaBreach(_) => "QuotaBreachError".to_string(),
             ApiClientError::ContextWindowOverflow => "ContextWindowOverflow".to_string(),
             ApiClientError::SmithyBuild(_) => "SmithyBuildError".to_string(),
             ApiClientError::AuthError(_) => "AuthError".to_string(),
             ApiClientError::ModelOverloadedError { .. } => "ModelOverloadedError".to_string(),
+            ApiClientError::MonthlyLimitReached => "MonthlyLimitReached".to_string(),
         }
     }
 }
@@ -150,6 +160,10 @@ mod tests {
                 QDeveloperSendMessageError::unhandled("<unhandled>"),
                 response(),
             )),
+            ApiClientError::CreateSubscriptionToken(SdkError::service_error(
+                CreateSubscriptionTokenError::unhandled("<unhandled>"),
+                response(),
+            )),
             ApiClientError::CodewhispererChatResponseStream(SdkError::service_error(
                 CodewhispererChatResponseStreamError::unhandled("<unhandled>"),
                 raw_message(),

crates/chat-cli/src/auth/builder_id.rs

@@ -41,6 +41,10 @@ use aws_smithy_runtime_api::client::identity::{
 };
 use aws_smithy_types::error::display::DisplayErrorContext;
 use aws_types::region::Region;
+use eyre::{
+    Result,
+    eyre,
+};
 use time::OffsetDateTime;
 use tracing::{
     debug,
@@ -568,6 +572,17 @@ impl ResolveIdentity for BearerResolver {
     }
 }
 
+pub async fn is_idc_user(database: &Database) -> Result<bool> {
+    if cfg!(test) {
+        return Ok(false);
+    }
+    if let Ok(Some(token)) = BuilderIdToken::load(database).await {
+        Ok(token.token_type() == TokenType::IamIdentityCenter)
+    } else {
+        Err(eyre!("No auth token found - is the user signed in?"))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/chat-cli/src/cli/chat/command.rs

@@ -60,6 +60,9 @@ pub enum Command {
     },
     Mcp,
     Model,
+    Subscribe {
+        manage: bool,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -841,6 +844,10 @@ impl Command {
                 },
                 "mcp" => Self::Mcp,
                 "model" => Self::Model,
+                "subscribe" => {
+                    let manage = parts.contains(&"--manage");
+                    Self::Subscribe { manage }
+                },
                 unknown_command => {
                     let looks_like_path = {
                         let after_slash_command_str = parts[1..].join(" ");