fix: telemetry credential provider for cognito (#131)

Author: brandonskiser

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

@@ -81,7 +81,7 @@ impl Client {
     ) -> Result<(), ApiClientError> {
         match &self.inner {
             inner::Inner::Codewhisperer(client) => {
-                let _ = client
+                client
                     .send_telemetry_event()
                     .telemetry_event(telemetry_event)
                     .user_context(user_context)
@@ -91,7 +91,7 @@ impl Client {
                     })
                     .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone()))
                     .send()
-                    .await;
+                    .await?;
                 Ok(())
             },
             inner::Inner::Mock => Ok(()),

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

@@ -1,6 +1,7 @@
 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;
+use amzn_codewhisperer_client::operation::send_telemetry_event::SendTelemetryEventError;
 pub use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError;
 use amzn_codewhisperer_streaming_client::types::error::ChatResponseStreamError as CodewhispererChatResponseStreamError;
 use amzn_consolas_client::operation::generate_recommendations::GenerateRecommendationsError;
@@ -29,6 +30,10 @@ pub enum ApiClientError {
     #[error("{}", SdkErrorDisplay(.0))]
     ListAvailableServices(#[from] SdkError<ListCustomizationsError, HttpResponse>),
 
+    // Telemetry client error
+    #[error("{}", SdkErrorDisplay(.0))]
+    SendTelemetryEvent(#[from] SdkError<SendTelemetryEventError, HttpResponse>),
+
     // Send message errors
     #[error("{}", SdkErrorDisplay(.0))]
     CodewhispererGenerateAssistantResponse(#[from] SdkError<GenerateAssistantResponseError, HttpResponse>),

crates/chat-cli/src/telemetry/cognito.rs

@@ -1,3 +1,5 @@
+use std::time::SystemTime;
+
 use amzn_toolkit_telemetry_client::config::BehaviorVersion;
 use aws_credential_types::provider::error::CredentialsError;
 use aws_credential_types::{
@@ -8,6 +10,10 @@ use aws_sdk_cognitoidentity::primitives::{
     DateTime,
     DateTimeFormat,
 };
+use tracing::{
+    trace,
+    warn,
+};
 
 use crate::aws_common::app_name;
 use crate::database::{
@@ -20,6 +26,8 @@ pub async fn get_cognito_credentials_send(
     database: &mut Database,
     telemetry_stage: &TelemetryStage,
 ) -> Result<Credentials, CredentialsError> {
+    trace!("Creating new cognito credentials");
+
     let conf = aws_sdk_cognitoidentity::Config::builder()
         .behavior_version(BehaviorVersion::v2025_01_17())
         .region(telemetry_stage.region.clone())
@@ -80,6 +88,10 @@ pub async fn get_cognito_credentials(
             session_token,
             expiration,
         }) => {
+            if is_expired(expiration.as_ref()) {
+                return get_cognito_credentials_send(database, telemetry_stage).await;
+            }
+
             let Some(access_key_id) = access_key_id else {
                 return get_cognito_credentials_send(database, telemetry_stage).await;
             };
@@ -104,12 +116,12 @@ pub async fn get_cognito_credentials(
 
 #[derive(Debug)]
 pub struct CognitoProvider {
-    credentials: Credentials,
+    telemetry_stage: TelemetryStage,
 }
 
 impl CognitoProvider {
-    pub fn new(credentials: Credentials) -> CognitoProvider {
-        CognitoProvider { credentials }
+    pub fn new(telemetry_stage: TelemetryStage) -> CognitoProvider {
+        CognitoProvider { telemetry_stage }
     }
 }
 
@@ -118,7 +130,36 @@ impl provider::ProvideCredentials for CognitoProvider {
     where
         Self: 'a,
     {
-        provider::future::ProvideCredentials::new(async { Ok(self.credentials.clone()) })
+        provider::future::ProvideCredentials::new(async {
+            match Database::new().await {
+                Ok(mut db) => get_cognito_credentials(&mut db, &self.telemetry_stage).await,
+                Err(err) => Err(CredentialsError::provider_error(format!(
+                    "failed to get database: {:?}",
+                    err
+                ))),
+            }
+        })
+    }
+}
+
+fn is_expired(expiration: Option<&String>) -> bool {
+    let expiration = if let Some(v) = expiration {
+        v
+    } else {
+        warn!("no cognito expiration was saved");
+        return true;
+    };
+
+    match DateTime::from_str(expiration, DateTimeFormat::DateTime) {
+        Ok(expiration) => {
+            // Check if the expiration is at least after five minutes after the current time.
+            let curr: DateTime = (SystemTime::now() + std::time::Duration::from_secs(60 * 5)).into();
+            expiration < curr
+        },
+        Err(err) => {
+            warn!(?err, "invalid cognito expiration was saved");
+            true
+        },
     }
 }
 

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

@@ -25,10 +25,7 @@ use amzn_toolkit_telemetry_client::{
     Config,
 };
 use aws_credential_types::provider::SharedCredentialsProvider;
-use cognito::{
-    CognitoProvider,
-    get_cognito_credentials,
-};
+use cognito::CognitoProvider;
 use endpoint::StaticEndpoint;
 pub use install_method::{
     InstallMethod,
@@ -149,7 +146,7 @@ impl TelemetryThread {
         let (tx, mut rx) = mpsc::unbounded_channel();
         let handle = tokio::spawn(async move {
             while let Some(event) = rx.recv().await {
-                trace!("Sending telemetry event: {:?}", event);
+                trace!("TelemetryThread received new telemetry event: {:?}", event);
                 telemetry_client.send_event(event).await;
             }
         });
@@ -319,24 +316,21 @@ impl TelemetryClient {
             && database.settings.get_bool(Setting::TelemetryEnabled).unwrap_or(true);
 
         // If telemetry is disabled we do not emit using toolkit_telemetry
-        let toolkit_telemetry_client = match telemetry_enabled {
-            true => match get_cognito_credentials(database, &TelemetryStage::EXTERNAL_PROD).await {
-                Ok(credentials) => Some(ToolkitTelemetryClient::from_conf(
-                    Config::builder()
-                        .http_client(crate::aws_common::http_client::client())
-                        .behavior_version(BehaviorVersion::v2025_01_17())
-                        .endpoint_resolver(StaticEndpoint(TelemetryStage::EXTERNAL_PROD.endpoint))
-                        .app_name(app_name())
-                        .region(TelemetryStage::EXTERNAL_PROD.region.clone())
-                        .credentials_provider(SharedCredentialsProvider::new(CognitoProvider::new(credentials)))
-                        .build(),
-                )),
-                Err(err) => {
-                    error!("Failed to acquire cognito credentials: {err}");
-                    None
-                },
-            },
-            false => None,
+        let toolkit_telemetry_client = if telemetry_enabled {
+            Some(ToolkitTelemetryClient::from_conf(
+                Config::builder()
+                    .http_client(crate::aws_common::http_client::client())
+                    .behavior_version(BehaviorVersion::v2025_01_17())
+                    .endpoint_resolver(StaticEndpoint(TelemetryStage::EXTERNAL_PROD.endpoint))
+                    .app_name(app_name())
+                    .region(TelemetryStage::EXTERNAL_PROD.region.clone())
+                    .credentials_provider(SharedCredentialsProvider::new(CognitoProvider::new(
+                        TelemetryStage::EXTERNAL_PROD,
+                    )))
+                    .build(),
+            ))
+        } else {
+            None
         };
 
         fn client_id(env: &Env, database: &mut Database, telemetry_enabled: bool) -> Result<Uuid, TelemetryError> {
@@ -400,38 +394,43 @@ impl TelemetryClient {
             {
                 Ok(event) => event,
                 Err(err) => {
-                    error!(err =% DisplayErrorContext(err), "Failed to send telemetry event");
+                    error!(err =% DisplayErrorContext(err), "Failed to send cw telemetry event");
                     return;
                 },
             };
 
+            let event = TelemetryEvent::ChatAddMessageEvent(chat_add_message_event);
+            debug!(
+                ?event,
+                ?user_context,
+                telemetry_enabled = self.telemetry_enabled,
+                "Sending cw telemetry event"
+            );
             if let Err(err) = self
                 .codewhisperer_client
-                .send_telemetry_event(
-                    TelemetryEvent::ChatAddMessageEvent(chat_add_message_event),
-                    user_context,
-                    self.telemetry_enabled,
-                )
+                .send_telemetry_event(event, user_context, self.telemetry_enabled)
                 .await
             {
-                error!(err =% DisplayErrorContext(err), "Failed to send telemetry event");
+                error!(err =% DisplayErrorContext(err), "Failed to send cw telemetry event");
             }
         }
     }
 
     async fn send_telemetry_toolkit_metric(&self, event: Event) {
         let Some(toolkit_telemetry_client) = self.toolkit_telemetry_client.clone() else {
+            trace!("not sending toolkit metric - client does not exist");
             return;
         };
         let client_id = self.client_id;
         let Some(metric_datum) = event.into_metric_datum() else {
+            trace!("not sending toolkit metric - metric datum does not exist");
             return;
         };
 
         let product = AwsProduct::CodewhispererTerminal;
         let metric_name = metric_datum.metric_name().to_owned();
 
-        debug!(?product, ?metric_datum, "Posting metrics");
+        debug!(?client_id, ?product, ?metric_datum, "Sending toolkit telemetry event");
         if let Err(err) = toolkit_telemetry_client
             .post_metrics()
             .aws_product(product)
@@ -445,7 +444,7 @@ impl TelemetryClient {
             .await
             .map_err(DisplayErrorContext)
         {
-            error!(%err, ?metric_name, "Failed to post metric");
+            error!(%err, ?metric_name, "Failed to post toolkit metric");
         }
     }