bskiser/fix token refresh (#1920)

* chore: add logs for refreshing token

* fix: save device registration when launching chat

Author: brandonskiser

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

@@ -45,6 +45,7 @@ use time::OffsetDateTime;
 use tracing::{
     debug,
     error,
+    trace,
     warn,
 };
 
@@ -123,16 +124,26 @@ impl DeviceRegistration {
 
     /// Loads the OIDC registered client from the secret store, deleting it if it is expired.
     async fn load_from_secret_store(database: &Database, region: &Region) -> Result<Option<Self>, AuthError> {
+        trace!(?region, "loading device registration from secret store");
         let device_registration = database.get_secret(Self::SECRET_KEY).await?;
 
         if let Some(device_registration) = device_registration {
             // check that the data is not expired, assume it is invalid if not present
             let device_registration: Self = serde_json::from_str(&device_registration.0)?;
 
             if let Some(client_secret_expires_at) = device_registration.client_secret_expires_at {
-                if !is_expired(&client_secret_expires_at) && device_registration.region == region.as_ref() {
+                let is_expired = is_expired(&client_secret_expires_at);
+                let registration_region_is_valid = device_registration.region == region.as_ref();
+                trace!(
+                    ?is_expired,
+                    ?registration_region_is_valid,
+                    "checking if device registration is valid"
+                );
+                if !is_expired && registration_region_is_valid {
                     return Ok(Some(device_registration));
                 }
+            } else {
+                warn!("no expiration time found for the client secret");
             }
         }
 
@@ -291,19 +302,25 @@ impl BuilderIdToken {
                 match token {
                     Some(token) => {
                         let region = token.region.clone().map_or(OIDC_BUILDER_ID_REGION, Region::new);
-
                         let client = client(region.clone());
-                        // if token is expired try to refresh
+
                         if token.is_expired() {
+                            trace!("token is expired, refreshing");
                             token.refresh_token(&client, database, &region).await
                         } else {
                             Ok(Some(token))
                         }
                     },
-                    None => Ok(None),
+                    None => {
+                        debug!("secret stored in the database was empty");
+                        Ok(None)
+                    },
                 }
             },
-            Ok(None) => Ok(None),
+            Ok(None) => {
+                debug!("no secret found in the database");
+                Ok(None)
+            },
             Err(err) => {
                 error!(%err, "Error getting builder id token from keychain");
                 Err(err)?

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

@@ -27,7 +27,10 @@ use serde_json::{
 };
 use settings::Settings;
 use thiserror::Error;
-use tracing::info;
+use tracing::{
+    info,
+    trace,
+};
 use uuid::Uuid;
 
 use crate::cli::ConversationState;
@@ -334,15 +337,18 @@ impl Database {
     }
 
     pub async fn get_secret(&self, key: &str) -> Result<Option<Secret>, DatabaseError> {
+        trace!(key, "getting secret");
         Ok(self.get_entry::<String>(Table::Auth, key)?.map(Into::into))
     }
 
     pub async fn set_secret(&self, key: &str, value: &str) -> Result<(), DatabaseError> {
+        trace!(key, "setting secret");
         self.set_entry(Table::Auth, key, value)?;
         Ok(())
     }
 
     pub async fn delete_secret(&self, key: &str) -> Result<(), DatabaseError> {
+        trace!(key, "deleting secret");
         self.delete_entry(Table::Auth, key)
     }
 

crates/fig_auth/src/consts.rs

@@ -2,7 +2,7 @@ use aws_types::region::Region;
 
 pub(crate) const CLIENT_NAME: &str = "Amazon Q Developer for command line";
 
-pub(crate) const OIDC_BUILDER_ID_REGION: Region = Region::from_static("us-east-1");
+pub const OIDC_BUILDER_ID_REGION: Region = Region::from_static("us-east-1");
 
 /// The scopes requested for OIDC
 ///

crates/fig_auth/src/lib.rs

@@ -1,5 +1,5 @@
 pub mod builder_id;
-mod consts;
+pub mod consts;
 mod error;
 pub mod pkce;
 mod scope;

crates/q_cli/src/cli/mod.rs

@@ -46,8 +46,13 @@ use eyre::{
     bail,
 };
 use feed::Feed;
-use fig_auth::builder_id::BuilderIdToken;
+use fig_auth::builder_id::{
+    BuilderIdToken,
+    DeviceRegistration,
+};
+use fig_auth::consts::OIDC_BUILDER_ID_REGION;
 use fig_auth::is_logged_in;
+use fig_auth::pkce::Region;
 use fig_auth::secret_store::SecretStore;
 use fig_ipc::local::open_ui_element;
 use fig_log::{
@@ -71,6 +76,7 @@ use tracing::{
     Level,
     debug,
     error,
+    warn,
 };
 
 use self::integrations::IntegrationsSubcommands;
@@ -380,10 +386,35 @@ impl Cli {
             assert_logged_in().await?;
         }
 
+        // Save credentials from the macOS keychain to sqlite.
+        // On Linux, this essentially just rewrites to the database.
         let secret_store = SecretStore::new().await.ok();
         if let Some(secret_store) = secret_store {
             if let Ok(database) = database().map_err(|err| error!(?err, "failed to open database")) {
                 if let Ok(token) = BuilderIdToken::load(&secret_store, false).await {
+                    // Save the device registration. This is required for token refresh to succeed.
+                    if let Some(token) = token.as_ref() {
+                        let region = token.region.clone().map_or(OIDC_BUILDER_ID_REGION, Region::new);
+                        match DeviceRegistration::load_from_secret_store(&secret_store, &region).await {
+                            Ok(Some(reg)) => match serde_json::to_string(&reg) {
+                                Ok(reg) => {
+                                    database
+                                        .set_auth_value("codewhisperer:odic:device-registration", reg)
+                                        .map_err(|err| error!(?err, "failed to write device registration to auth db"))
+                                        .ok();
+                                },
+                                Err(err) => error!(?err, "failed to serialize the device registration"),
+                            },
+                            Ok(None) => {
+                                warn!(?token, "no device registration found for token");
+                            },
+                            Err(err) => {
+                                error!(?err, "failed to load device registration");
+                            },
+                        }
+                    }
+
+                    // Next, save the token.
                     if let Ok(token) = serde_json::to_string(&token) {
                         database
                             .set_auth_value("codewhisperer:odic:token", token)