restrict profilie

Author: evanliu048

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

@@ -41,10 +41,7 @@ use aws_smithy_runtime_api::client::identity::{
 };
 use aws_smithy_types::error::display::DisplayErrorContext;
 use aws_types::region::Region;
-use eyre::{
-    Result,
-    eyre,
-};
+use eyre::Result;
 use time::OffsetDateTime;
 use tracing::{
     debug,
@@ -606,7 +603,7 @@ pub async fn is_idc_user(database: &Database) -> Result<bool> {
     if let Ok(Some(token)) = BuilderIdToken::load(database).await {
         Ok(token.token_type() == TokenType::IamIdentityCenter)
     } else {
-        Err(eyre!("No auth token found - is the user signed in?"))
+        Ok(false)
     }
 }
 

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

@@ -33,6 +33,7 @@ use crate::auth::builder_id::{
     BuilderIdToken,
     PollCreateToken,
     TokenType,
+    is_idc_user,
     poll_create_token,
     start_device_authorization,
 };
@@ -338,10 +339,8 @@ pub enum LicenseType {
 }
 
 pub async fn profile(os: &mut Os) -> Result<ExitCode> {
-    if let Ok(Some(token)) = BuilderIdToken::load(&os.database).await {
-        if matches!(token.token_type(), TokenType::BuilderId) {
-            bail!("This command is only available for Pro users");
-        }
+    if !is_idc_user(&os.database).await? {
+        bail!("This command is only available for IAM Identity Center users");
     }
 
     select_profile_interactive(os, false).await?;

crates/q_cli/src/cli/user.rs

@@ -229,13 +229,18 @@ impl RootUserSubcommand {
             Self::Profile => {
                 assert_logged_in().await?;
 
-                if let Ok(Some(token)) = fig_auth::builder_id_token().await {
-                    if matches!(token.token_type(), TokenType::IamIdentityCenter) {
-                        select_profile_interactive(false).await?;
-                    }
+                let is_idc = fig_auth::builder_id_token()
+                    .await
+                    .ok()
+                    .flatten()
+                    .is_some_and(|token| matches!(token.token_type(), TokenType::IamIdentityCenter));
+
+                if !is_idc {
+                    bail!("This command is only available for IAM Identity Center users");
                 }
 
-                bail!("This command is only available for IAM Identity Center users");
+                select_profile_interactive(false).await?;
+                Ok(ExitCode::SUCCESS)
             },
         }
     }

packages/dashboard-app/src/pages/settings/preferences.tsx

@@ -23,6 +23,10 @@ export default function Page() {
   const [profiles, setProfiles] = useState<Profile[] | undefined>(undefined);
 
   useEffect(() => {
+    if (auth.authKind !== "IamIdentityCenter") {
+      return;
+    }
+
     Profile.listAvailableProfiles()
       .then(async (res) => {
         setProfiles(
@@ -35,7 +39,7 @@ export default function Page() {
       .catch((err) => {
         console.error(err);
       });
-  }, []);
+  }, [auth.authKind]);
 
   useEffect(() => {
     State.get("api.codewhisperer.profile").then((profile) => {