fix: save device registration when launching chat

Author: brandonskiser

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)