Add auth login failure telemetry for cli (#3317)

abhraina-aws · web-flow · commit 1d352621a7df · 2025-10-30T10:35:23.000-07:00
* add auth login failure telemetry
diff --git a/crates/chat-cli/src/auth/builder_id.rs b/crates/chat-cli/src/auth/builder_id.rs
@@ -273,6 +273,16 @@ pub enum TokenType {
     IamIdentityCenter,
 }
 
+impl From<Option<&str>> for TokenType {
+    fn from(start_url: Option<&str>) -> Self {
+        match start_url {
+            Some(url) if url == START_URL => TokenType::BuilderId,
+            None => TokenType::BuilderId,
+            Some(_) => TokenType::IamIdentityCenter,
+        }
+    }
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct BuilderIdToken {
     pub access_token: Secret,
@@ -302,7 +312,10 @@ impl BuilderIdToken {
     }
 
     /// Load the token from the keychain, refresh the token if it is expired and return it
-    pub async fn load(database: &Database) -> Result<Option<Self>, AuthError> {
+    pub async fn load(
+        database: &Database,
+        telemetry: Option<&crate::telemetry::TelemetryThread>,
+    ) -> Result<Option<Self>, AuthError> {
         // Can't use #[cfg(test)] without breaking lints, and we don't want to require
         // authentication in order to run ChatSession tests. Hence, adding this here with cfg!(test)
         if cfg!(test) {
@@ -328,7 +341,7 @@ impl BuilderIdToken {
 
                         if token.is_expired() {
                             trace!("token is expired, refreshing");
-                            token.refresh_token(&client, database, &region).await
+                            token.refresh_token(&client, database, &region, telemetry).await
                         } else {
                             trace!(?token, "found a valid token");
                             Ok(Some(token))
@@ -357,6 +370,7 @@ impl BuilderIdToken {
         client: &Client,
         database: &Database,
         region: &Region,
+        telemetry: Option<&crate::telemetry::TelemetryThread>,
     ) -> Result<Option<Self>, AuthError> {
         let Some(refresh_token) = &self.refresh_token else {
             warn!("no refresh token was found");
@@ -416,6 +430,25 @@ impl BuilderIdToken {
                 let display_err = DisplayErrorContext(&err);
                 error!("Failed to refresh builder id access token: {}", display_err);
 
+                // Send telemetry for refresh failure
+                if let Some(telemetry) = telemetry {
+                    let auth_method = match self.token_type() {
+                        TokenType::BuilderId => "BuilderId",
+                        TokenType::IamIdentityCenter => "IdentityCenter",
+                    };
+                    let oauth_flow = match self.oauth_flow {
+                        OAuthFlow::DeviceCode => "DeviceCode",
+                        OAuthFlow::Pkce => "PKCE",
+                    };
+                    let error_code = match &err {
+                        SdkError::ServiceError(service_err) => service_err.err().meta().code().map(|s| s.to_string()),
+                        _ => None,
+                    };
+                    telemetry
+                        .send_auth_failed(auth_method, oauth_flow, "TokenRefresh", error_code)
+                        .ok();
+                }
+
                 // if the error is the client's fault, clear the token
                 if let SdkError::ServiceError(service_err) = &err {
                     if !service_err.err().is_slow_down_exception() {
@@ -471,11 +504,7 @@ impl BuilderIdToken {
     }
 
     pub fn token_type(&self) -> TokenType {
-        match &self.start_url {
-            Some(url) if url == START_URL => TokenType::BuilderId,
-            None => TokenType::BuilderId,
-            Some(_) => TokenType::IamIdentityCenter,
-        }
+        TokenType::from(self.start_url.as_deref())
     }
 
     /// Check if the token is for the internal amzn start URL (`https://amzn.awsapps.com/start`),
@@ -498,6 +527,7 @@ pub async fn poll_create_token(
     device_code: String,
     start_url: Option<String>,
     region: Option<String>,
+    telemetry: &crate::telemetry::TelemetryThread,
 ) -> PollCreateToken {
     let region = region.clone().map_or(OIDC_BUILDER_ID_REGION, Region::new);
     let client = client(region.clone());
@@ -538,6 +568,20 @@ pub async fn poll_create_token(
         },
         Err(err) => {
             error!(?err, "Failed to poll for builder id token");
+
+            // Send telemetry for device code failure
+            let auth_method = match TokenType::from(start_url.as_deref()) {
+                TokenType::BuilderId => "BuilderId",
+                TokenType::IamIdentityCenter => "IdentityCenter",
+            };
+            let error_code = match &err {
+                SdkError::ServiceError(service_err) => service_err.err().meta().code().map(|s| s.to_string()),
+                _ => None,
+            };
+            telemetry
+                .send_auth_failed(auth_method, "DeviceCode", "NewLogin", error_code)
+                .ok();
+
             PollCreateToken::Error(err.into())
         },
     }
@@ -550,7 +594,7 @@ pub async fn is_logged_in(database: &mut Database) -> bool {
         return true;
     }
 
-    match BuilderIdToken::load(database).await {
+    match BuilderIdToken::load(database, None).await {
         Ok(Some(_)) => true,
         Ok(None) => {
             info!("not logged in - no valid token found");
@@ -585,7 +629,7 @@ pub async fn logout(database: &mut Database) -> Result<(), AuthError> {
 pub async fn get_start_url_and_region(database: &Database) -> (Option<String>, Option<String>) {
     // NOTE: Database provides direct methods to access the start_url and region, but they are not
     // guaranteed to be up to date in the chat session. Example: login is changed mid-chat session.
-    let token = BuilderIdToken::load(database).await;
+    let token = BuilderIdToken::load(database, None).await;
     match token {
         Ok(Some(t)) => (t.start_url, t.region),
         _ => (None, None),
@@ -603,7 +647,7 @@ impl ResolveIdentity for BearerResolver {
     ) -> IdentityFuture<'a> {
         IdentityFuture::new_boxed(Box::pin(async {
             let database = Database::new().await?;
-            match BuilderIdToken::load(&database).await? {
+            match BuilderIdToken::load(&database, None).await? {
                 Some(token) => Ok(Identity::new(
                     Token::new(token.access_token.0.clone(), Some(token.expires_at.into())),
                     Some(token.expires_at.into()),
@@ -618,7 +662,7 @@ 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 {
+    if let Ok(Some(token)) = BuilderIdToken::load(database, None).await {
         Ok(token.token_type() == TokenType::IamIdentityCenter)
     } else {
         Err(eyre!("No auth token found - is the user signed in?"))
diff --git a/crates/chat-cli/src/cli/user.rs b/crates/chat-cli/src/cli/user.rs
@@ -143,7 +143,16 @@ impl LoginArgs {
                             ]);
                             let ctrl_c_stream = ctrl_c();
                             tokio::select! {
-                                res = registration.finish(&client, Some(&mut os.database)) => res?,
+                                res = registration.finish(&client, Some(&mut os.database)) => {
+                                    if let Err(err) = res {
+                                        let auth_method = match crate::auth::builder_id::TokenType::from(start_url.as_deref()) {
+                                            crate::auth::builder_id::TokenType::BuilderId => "BuilderId",
+                                            crate::auth::builder_id::TokenType::IamIdentityCenter => "IdentityCenter",
+                                        };
+                                        os.telemetry.send_auth_failed(auth_method, "PKCE", "NewLogin", None).ok();
+                                        return Err(err.into());
+                                    }
+                                },
                                 Ok(_) = ctrl_c_stream => {
                                     #[allow(clippy::exit)]
                                     exit(1);
@@ -194,7 +203,7 @@ pub struct WhoamiArgs {
 
 impl WhoamiArgs {
     pub async fn execute(self, os: &mut Os) -> Result<ExitCode> {
-        let builder_id = BuilderIdToken::load(&os.database).await;
+        let builder_id = BuilderIdToken::load(&os.database, Some(&os.telemetry)).await;
 
         match builder_id {
             Ok(Some(token)) => {
@@ -245,7 +254,7 @@ pub enum LicenseType {
 }
 
 pub async fn profile(os: &mut Os) -> Result<ExitCode> {
-    if let Ok(Some(token)) = BuilderIdToken::load(&os.database).await {
+    if let Ok(Some(token)) = BuilderIdToken::load(&os.database, Some(&os.telemetry)).await {
         if matches!(token.token_type(), TokenType::BuilderId) {
             bail!("This command is only available for Pro users");
         }
@@ -314,6 +323,7 @@ async fn try_device_authorization(os: &mut Os, start_url: Option<String>, region
             device_auth.device_code.clone(),
             start_url.clone(),
             region.clone(),
+            &os.telemetry,
         )
         .await
         {
diff --git a/crates/chat-cli/src/telemetry/core.rs b/crates/chat-cli/src/telemetry/core.rs
@@ -23,6 +23,7 @@ use crate::telemetry::definitions::metrics::{
     CodewhispererterminalAddChatMessage,
     CodewhispererterminalAgentConfigInit,
     CodewhispererterminalAgentContribution,
+    CodewhispererterminalAuthFailed,
     CodewhispererterminalChatSlashCommandExecuted,
     CodewhispererterminalCliSubcommandExecuted,
     CodewhispererterminalMcpServerInit,
@@ -500,6 +501,24 @@ impl Event {
                 }
                 .into_metric_datum(),
             ),
+            EventType::AuthFailed {
+                auth_method,
+                oauth_flow,
+                error_type,
+                error_code,
+            } => Some(
+                CodewhispererterminalAuthFailed {
+                    create_time: self.created_time,
+                    value: None,
+                    credential_start_url: self.credential_start_url.map(Into::into),
+                    codewhispererterminal_in_cloudshell: None,
+                    codewhispererterminal_auth_method: Some(auth_method.into()),
+                    oauth_flow: Some(oauth_flow.into()),
+                    codewhispererterminal_error_type: Some(error_type.into()),
+                    codewhispererterminal_error_code: error_code.map(Into::into),
+                }
+                .into_metric_datum(),
+            ),
             EventType::DailyHeartbeat {} => Some(
                 AmazonqcliDailyHeartbeat {
                     create_time: self.created_time,
@@ -594,6 +613,12 @@ pub struct AgentConfigInitArgs {
 #[serde(tag = "type")]
 pub enum EventType {
     UserLoggedIn {},
+    AuthFailed {
+        auth_method: String,
+        oauth_flow: String,
+        error_type: String,
+        error_code: Option<String>,
+    },
     RefreshCredentials {
         request_id: String,
         result: TelemetryResult,
diff --git a/crates/chat-cli/src/telemetry/mod.rs b/crates/chat-cli/src/telemetry/mod.rs
@@ -235,6 +235,21 @@ impl TelemetryThread {
         Ok(self.tx.send(Event::new(EventType::UserLoggedIn {}))?)
     }
 
+    pub fn send_auth_failed(
+        &self,
+        auth_method: &str,
+        oauth_flow: &str,
+        error_type: &str,
+        error_code: Option<String>,
+    ) -> Result<(), TelemetryError> {
+        Ok(self.tx.send(Event::new(EventType::AuthFailed {
+            auth_method: auth_method.to_string(),
+            oauth_flow: oauth_flow.to_string(),
+            error_type: error_type.to_string(),
+            error_code,
+        }))?)
+    }
+
     pub fn send_daily_heartbeat(&self) -> Result<(), TelemetryError> {
         Ok(self.tx.send(Event::new(EventType::DailyHeartbeat {}))?)
     }
diff --git a/crates/chat-cli/telemetry_definitions.json b/crates/chat-cli/telemetry_definitions.json
@@ -65,6 +65,21 @@
       "type": "string",
       "description": "The oauth authentication flow executed by the user, e.g. device code or PKCE"
     },
+    {
+      "name": "codewhispererterminal_authMethod",
+      "type": "string",
+      "description": "The authentication method used, e.g. BuilderId or IdentityCenter"
+    },
+    {
+      "name": "codewhispererterminal_errorType",
+      "type": "string",
+      "description": "The type of authentication error, e.g. TokenRefresh or NewLogin"
+    },
+    {
+      "name": "codewhispererterminal_errorCode",
+      "type": "string",
+      "description": "The specific error code from the authentication service"
+    },
     {
       "name": "result",
       "type": "string",
@@ -369,6 +384,19 @@
         { "type": "codewhispererterminal_inCloudshell" }
       ]
     },
+    {
+      "name": "codewhispererterminal_authFailed",
+      "description": "Emitted when authentication fails",
+      "passive": false,
+      "metadata": [
+        { "type": "credentialStartUrl" },
+        { "type": "codewhispererterminal_inCloudshell" },
+        { "type": "codewhispererterminal_authMethod" },
+        { "type": "oauthFlow" },
+        { "type": "codewhispererterminal_errorType" },
+        { "type": "codewhispererterminal_errorCode", "required": false }
+      ]
+    },
     {
       "name": "codewhispererterminal_refreshCredentials",
       "description": "Emitted when users refresh their credentials",
