Merge branch 'main' into background-server-load

Author: dingfeli

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

@@ -52,49 +52,19 @@ use crate::auth::AuthError;
 use crate::auth::consts::*;
 use crate::auth::scope::is_scopes;
 use crate::aws_common::app_name;
-use crate::database::Database;
-use crate::database::secret_store::{
+use crate::database::{
+    Database,
     Secret,
-    SecretStore,
 };
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Deserialize)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub enum OAuthFlow {
     DeviceCode,
     // This must remain backwards compatible
     #[serde(alias = "PKCE")]
     Pkce,
 }
 
-// Implement Serialize manually to ensure proper serialization
-impl serde::Serialize for OAuthFlow {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: serde::Serializer,
-    {
-        match *self {
-            OAuthFlow::DeviceCode => serializer.serialize_str("DeviceCode"),
-            OAuthFlow::Pkce => serialize_pkce(serializer),
-        }
-    }
-}
-
-fn serialize_pkce<S>(serializer: S) -> Result<S::Ok, S::Error>
-where
-    S: serde::Serializer,
-{
-    serializer.serialize_str("PKCE")
-}
-
-impl std::fmt::Display for OAuthFlow {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match *self {
-            OAuthFlow::DeviceCode => write!(f, "DeviceCode"),
-            OAuthFlow::Pkce => write!(f, "PKCE"),
-        }
-    }
-}
-
 /// Indicates if an expiration time has passed, there is a small 1 min window that is removed
 /// so the token will not expire in transit
 fn is_expired(expiration_time: &OffsetDateTime) -> bool {
@@ -152,8 +122,8 @@ impl DeviceRegistration {
     }
 
     /// Loads the OIDC registered client from the secret store, deleting it if it is expired.
-    async fn load_from_secret_store(secret_store: &SecretStore, region: &Region) -> Result<Option<Self>, AuthError> {
-        let device_registration = secret_store.get(Self::SECRET_KEY).await?;
+    async fn load_from_secret_store(database: &Database, region: &Region) -> Result<Option<Self>, AuthError> {
+        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
@@ -167,7 +137,7 @@ impl DeviceRegistration {
         }
 
         // delete the data if its expired or invalid
-        if let Err(err) = secret_store.delete(Self::SECRET_KEY).await {
+        if let Err(err) = database.delete_secret(Self::SECRET_KEY).await {
             error!(?err, "Failed to delete device registration from keychain");
         }
 
@@ -181,7 +151,7 @@ impl DeviceRegistration {
         client: &Client,
         region: &Region,
     ) -> Result<Self, AuthError> {
-        match Self::load_from_secret_store(&database.secret_store, region).await {
+        match Self::load_from_secret_store(database, region).await {
             Ok(Some(registration)) if registration.oauth_flow == OAuthFlow::DeviceCode => match &registration.scopes {
                 Some(scopes) if is_scopes(scopes) => return Ok(registration),
                 _ => warn!("Invalid scopes in device registration, ignoring"),
@@ -210,17 +180,17 @@ impl DeviceRegistration {
             SCOPES.iter().map(|s| (*s).to_owned()).collect(),
         );
 
-        if let Err(err) = device_registration.save(&database.secret_store).await {
+        if let Err(err) = device_registration.save(database).await {
             error!(?err, "Failed to write device registration to keychain");
         }
 
         Ok(device_registration)
     }
 
     /// Saves to the passed secret store.
-    pub async fn save(&self, secret_store: &SecretStore) -> Result<(), AuthError> {
+    pub async fn save(&self, secret_store: &Database) -> Result<(), AuthError> {
         secret_store
-            .set(Self::SECRET_KEY, &serde_json::to_string(&self)?)
+            .set_secret(Self::SECRET_KEY, &serde_json::to_string(&self)?)
             .await?;
         Ok(())
     }
@@ -314,8 +284,8 @@ impl BuilderIdToken {
     }
 
     /// Load the token from the keychain, refresh the token if it is expired and return it
-    pub async fn load(database: &mut Database) -> Result<Option<Self>, AuthError> {
-        match database.secret_store.get(Self::SECRET_KEY).await {
+    pub async fn load(database: &Database) -> Result<Option<Self>, AuthError> {
+        match database.get_secret(Self::SECRET_KEY).await {
             Ok(Some(secret)) => {
                 let token: Option<Self> = serde_json::from_str(&secret.0)?;
                 match token {
@@ -325,7 +295,7 @@ impl BuilderIdToken {
                         let client = client(region.clone());
                         // if token is expired try to refresh
                         if token.is_expired() {
-                            token.refresh_token(&client, &database.secret_store, &region).await
+                            token.refresh_token(&client, database, &region).await
                         } else {
                             Ok(Some(token))
                         }
@@ -345,19 +315,19 @@ impl BuilderIdToken {
     pub async fn refresh_token(
         &self,
         client: &Client,
-        secret_store: &SecretStore,
+        database: &Database,
         region: &Region,
     ) -> Result<Option<Self>, AuthError> {
         let Some(refresh_token) = &self.refresh_token else {
             // if the token is expired and has no refresh token, delete it
-            if let Err(err) = self.delete(secret_store).await {
+            if let Err(err) = self.delete(database).await {
                 error!(?err, "Failed to delete builder id token");
             }
 
             return Ok(None);
         };
 
-        let registration = match DeviceRegistration::load_from_secret_store(secret_store, region).await? {
+        let registration = match DeviceRegistration::load_from_secret_store(database, region).await? {
             Some(registration) if registration.oauth_flow == self.oauth_flow => registration,
             // If the OIDC client registration is for a different oauth flow or doesn't exist, then
             // we can't refresh the token.
@@ -394,7 +364,7 @@ impl BuilderIdToken {
                 );
                 debug!("Refreshed access token, new token: {:?}", token);
 
-                if let Err(err) = token.save(secret_store).await {
+                if let Err(err) = token.save(database).await {
                     error!(?err, "Failed to store builder id access token");
                 };
 
@@ -407,7 +377,7 @@ impl BuilderIdToken {
                 // 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() {
-                        if let Err(err) = self.delete(secret_store).await {
+                        if let Err(err) = self.delete(database).await {
                             error!(?err, "Failed to delete builder id token");
                         }
                     }
@@ -427,16 +397,16 @@ impl BuilderIdToken {
     }
 
     /// Save the token to the keychain
-    pub async fn save(&self, secret_store: &SecretStore) -> Result<(), AuthError> {
-        secret_store
-            .set(Self::SECRET_KEY, &serde_json::to_string(self)?)
+    pub async fn save(&self, database: &Database) -> Result<(), AuthError> {
+        database
+            .set_secret(Self::SECRET_KEY, &serde_json::to_string(self)?)
             .await?;
         Ok(())
     }
 
     /// Delete the token from the keychain
-    pub async fn delete(&self, secret_store: &SecretStore) -> Result<(), AuthError> {
-        secret_store.delete(Self::SECRET_KEY).await?;
+    pub async fn delete(&self, database: &Database) -> Result<(), AuthError> {
+        database.delete_secret(Self::SECRET_KEY).await?;
         Ok(())
     }
 
@@ -508,7 +478,7 @@ pub async fn poll_create_token(
             let token: BuilderIdToken =
                 BuilderIdToken::from_output(output, region, start_url, OAuthFlow::DeviceCode, scopes);
 
-            if let Err(err) = token.save(&database.secret_store).await {
+            if let Err(err) = token.save(database).await {
                 error!(?err, "Failed to store builder id token");
             };
 
@@ -529,13 +499,13 @@ pub async fn is_logged_in(database: &mut Database) -> bool {
 }
 
 pub async fn logout(database: &mut Database) -> Result<(), AuthError> {
-    let Ok(secret_store) = SecretStore::new().await else {
+    let Ok(secret_store) = Database::new().await else {
         return Ok(());
     };
 
     let (builder_res, device_res) = tokio::join!(
-        secret_store.delete(BuilderIdToken::SECRET_KEY),
-        secret_store.delete(DeviceRegistration::SECRET_KEY),
+        secret_store.delete_secret(BuilderIdToken::SECRET_KEY),
+        secret_store.delete_secret(DeviceRegistration::SECRET_KEY),
     );
 
     let profile_res = database.unset_auth_profile();
@@ -585,20 +555,10 @@ mod tests {
     const US_EAST_1: Region = Region::from_static("us-east-1");
     const US_WEST_2: Region = Region::from_static("us-west-2");
 
-    macro_rules! test_ser_deser {
-        ($ty:ident, $variant:expr, $text:expr) => {
-            let quoted = format!("\"{}\"", $text);
-            assert_eq!(quoted, serde_json::to_string(&$variant).unwrap());
-            assert_eq!($variant, serde_json::from_str(&quoted).unwrap());
-
-            assert_eq!($text, format!("{}", $variant));
-        };
-    }
-
     #[test]
-    fn test_oauth_flow_ser_deser() {
-        test_ser_deser!(OAuthFlow, OAuthFlow::DeviceCode, "DeviceCode");
-        test_ser_deser!(OAuthFlow, OAuthFlow::Pkce, "PKCE");
+    fn test_oauth_flow_deser() {
+        assert_eq!(OAuthFlow::Pkce, serde_json::from_str("\"PKCE\"").unwrap());
+        assert_eq!(OAuthFlow::Pkce, serde_json::from_str("\"Pkce\"").unwrap());
     }
 
     #[tokio::test]

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

@@ -231,7 +231,7 @@ impl PkceRegistration {
     /// then the access and refresh tokens will be saved.
     ///
     /// Only the first connection will be served.
-    pub async fn finish<C: PkceClient>(self, client: &C, database: Option<&Database>) -> Result<(), AuthError> {
+    pub async fn finish<C: PkceClient>(self, client: &C, database: Option<&mut Database>) -> Result<(), AuthError> {
         let code = tokio::select! {
             code = Self::recv_code(self.listener, self.state) => {
                 code?
@@ -270,11 +270,11 @@ impl PkceRegistration {
         );
 
         if let Some(database) = database {
-            if let Err(err) = device_registration.save(&database.secret_store).await {
+            if let Err(err) = device_registration.save(database).await {
                 error!(?err, "Failed to store pkce registration to secret store");
             }
 
-            if let Err(err) = token.save(&database.secret_store).await {
+            if let Err(err) = token.save(database).await {
                 error!(?err, "Failed to store builder id token");
             };
         }

crates/chat-cli/src/cli/debug.rs

@@ -1,7 +1,4 @@
-use clap::{
-    Subcommand,
-    ValueEnum,
-};
+use clap::ValueEnum;
 
 #[derive(Debug, ValueEnum, Clone, PartialEq, Eq)]
 pub enum Build {
@@ -65,7 +62,7 @@ pub enum TISAction {
 use std::path::PathBuf;
 
 #[cfg(target_os = "macos")]
-#[derive(Debug, Subcommand, Clone, PartialEq, Eq)]
+#[derive(Debug, clap::Subcommand, Clone, PartialEq, Eq)]
 pub enum InputMethodDebugAction {
     Install {
         bundle_path: Option<PathBuf>,