fix: broken state for telemetryClientId, user agent versioning

Author: brandonskiser

crates/chat-cli/src/aws_common/user_agent_override_interceptor.rs

@@ -18,13 +18,16 @@ use http::header::{
     InvalidHeaderValue,
     USER_AGENT,
 };
-use tracing::warn;
+use tracing::{
+    trace,
+    warn,
+};
 
 /// The environment variable name of additional user agent metadata we include in the user agent
 /// string. This is used in AWS CloudShell where they want to track usage by version.
 const AWS_TOOLING_USER_AGENT: &str = "AWS_TOOLING_USER_AGENT";
 
-const VERSION_HEADER: &str = "Version";
+const VERSION_HEADER: &str = "appVersion";
 const VERSION_VALUE: &str = env!("CARGO_PKG_VERSION");
 
 #[derive(Debug)]
@@ -91,47 +94,56 @@ impl Intercept for UserAgentOverrideInterceptor {
         // Allow for overriding the user agent by an earlier interceptor (so, for example,
         // tests can use `AwsUserAgent::for_tests()`) by attempting to grab one out of the
         // config bag before creating one.
-        let ua: Cow<'_, AwsUserAgent> = cfg.load::<AwsUserAgent>().map(Cow::Borrowed).map_or_else(
-            || {
+        let ua: Cow<'_, AwsUserAgent> = match cfg.get_mut::<AwsUserAgent>() {
+            Some(ua) => {
+                apply_additional_metadata(&self.env, ua);
+                Cow::Borrowed(ua)
+            },
+            None => {
                 let api_metadata = cfg
                     .load::<ApiMetadata>()
                     .ok_or(UserAgentOverrideInterceptorError::MissingApiMetadata)?;
 
-                let aws_tooling_user_agent = env.get(AWS_TOOLING_USER_AGENT);
-                let mut ua = AwsUserAgent::new_from_environment(env, api_metadata.clone());
-
-                let ver = format!("{VERSION_HEADER}/{VERSION_VALUE}");
-                match AdditionalMetadata::new(clean_metadata(&ver)) {
-                    Ok(md) => {
-                        ua.add_additional_metadata(md);
-                    },
-                    Err(err) => panic!("Failed to parse version: {err}"),
-                };
+                let mut ua = AwsUserAgent::new_from_environment(self.env.clone(), api_metadata.clone());
 
                 let maybe_app_name = cfg.load::<AppName>();
                 if let Some(app_name) = maybe_app_name {
                     ua.set_app_name(app_name.clone());
                 }
-                if let Ok(val) = aws_tooling_user_agent {
-                    match AdditionalMetadata::new(clean_metadata(&val)) {
-                        Ok(md) => {
-                            ua.add_additional_metadata(md);
-                        },
-                        Err(err) => warn!(%err, %val, "Failed to parse {AWS_TOOLING_USER_AGENT}"),
-                    };
-                }
 
-                Ok(Cow::Owned(ua))
+                apply_additional_metadata(&env, &mut ua);
+
+                Cow::Owned(ua)
             },
-            Result::<_, UserAgentOverrideInterceptorError>::Ok,
-        )?;
+        };
+
+        trace!(?ua, "setting user agent");
 
         let headers = context.request_mut().headers_mut();
         headers.insert(USER_AGENT.as_str(), ua.aws_ua_header());
         Ok(())
     }
 }
 
+fn apply_additional_metadata(env: &Env, ua: &mut AwsUserAgent) {
+    let ver = format!("{VERSION_HEADER}/{VERSION_VALUE}");
+    match AdditionalMetadata::new(clean_metadata(&ver)) {
+        Ok(md) => {
+            ua.add_additional_metadata(md);
+        },
+        Err(err) => panic!("Failed to parse version: {err}"),
+    };
+
+    if let Ok(val) = env.get(AWS_TOOLING_USER_AGENT) {
+        match AdditionalMetadata::new(clean_metadata(&val)) {
+            Ok(md) => {
+                ua.add_additional_metadata(md);
+            },
+            Err(err) => warn!(%err, %val, "Failed to parse {AWS_TOOLING_USER_AGENT}"),
+        };
+    }
+}
+
 fn clean_metadata(s: &str) -> String {
     let valid_character = |c: char| -> bool {
         match c {

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

@@ -266,13 +266,13 @@ impl Database {
     /// Get the client ID used for telemetry requests.
     pub fn get_client_id(&mut self) -> Result<Option<Uuid>, DatabaseError> {
         Ok(self
-            .get_entry::<String>(Table::State, CLIENT_ID_KEY)?
+            .get_json_entry::<String>(Table::State, CLIENT_ID_KEY)?
             .and_then(|s| Uuid::from_str(&s).ok()))
     }
 
     /// Set the client ID used for telemetry requests.
     pub fn set_client_id(&mut self, client_id: Uuid) -> Result<usize, DatabaseError> {
-        self.set_entry(Table::State, CLIENT_ID_KEY, client_id.to_string())
+        self.set_json_entry(Table::State, CLIENT_ID_KEY, client_id.to_string())
     }
 
     /// Get the start URL used for IdC login.
@@ -333,6 +333,19 @@ impl Database {
         self.set_json_entry(Table::Conversations, path, state)
     }
 
+    pub async fn get_secret(&self, key: &str) -> Result<Option<Secret>, DatabaseError> {
+        Ok(self.get_entry::<String>(Table::Auth, key)?.map(Into::into))
+    }
+
+    pub async fn set_secret(&self, key: &str, value: &str) -> Result<(), DatabaseError> {
+        self.set_entry(Table::Auth, key, value)?;
+        Ok(())
+    }
+
+    pub async fn delete_secret(&self, key: &str) -> Result<(), DatabaseError> {
+        self.delete_entry(Table::Auth, key)
+    }
+
     // Private functions. Do not expose.
 
     fn migrate(self) -> Result<Self, DatabaseError> {
@@ -428,19 +441,6 @@ impl Database {
 
         Ok(map)
     }
-
-    pub async fn get_secret(&self, key: &str) -> Result<Option<Secret>, DatabaseError> {
-        Ok(self.get_entry::<String>(Table::Auth, key)?.map(Into::into))
-    }
-
-    pub async fn set_secret(&self, key: &str, value: &str) -> Result<(), DatabaseError> {
-        self.set_entry(Table::Auth, key, value)?;
-        Ok(())
-    }
-
-    pub async fn delete_secret(&self, key: &str) -> Result<(), DatabaseError> {
-        self.delete_entry(Table::Auth, key)
-    }
 }
 
 fn max_migration<C: Deref<Target = Connection>>(conn: &C) -> Option<i64> {

crates/fig_desktop_api/src/requests/profile.rs

@@ -8,6 +8,7 @@ use fig_telemetry_core::{
     QProfileSwitchIntent,
     TelemetryResult,
 };
+use tracing::debug;
 
 use super::{
     RequestResult,
@@ -59,17 +60,15 @@ pub async fn set_profile(request: SetProfileRequest) -> RequestResult {
 }
 
 pub async fn list_available_profiles(_request: ListAvailableProfilesRequest) -> RequestResult {
-    Ok(
-        ServerOriginatedSubMessage::ListAvailableProfilesResponse(ListAvailableProfilesResponse {
-            profiles: fig_api_client::profile::list_available_profiles()
-                .await
-                .iter()
-                .map(|profile| fig_proto::fig::Profile {
-                    arn: profile.arn.clone(),
-                    profile_name: profile.profile_name.clone(),
-                })
-                .collect(),
+    debug!("listing available profiles");
+    let profiles = fig_api_client::profile::list_available_profiles()
+        .await
+        .iter()
+        .map(|profile| fig_proto::fig::Profile {
+            arn: profile.arn.clone(),
+            profile_name: profile.profile_name.clone(),
         })
-        .into(),
-    )
+        .collect();
+    debug!(?profiles, "fetched profiles");
+    Ok(ServerOriginatedSubMessage::ListAvailableProfilesResponse(ListAvailableProfilesResponse { profiles }).into())
 }