enfore login when chat

Author: evanliu048

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

@@ -123,17 +123,23 @@ impl RootSubcommand {
         matches!(self, Self::Chat(_) | Self::Login(_) | Self::Profile | Self::Issue(_))
     }
 
-    pub fn requires_auth(&self) -> bool {
-        matches!(self, Self::Chat(_) | Self::Profile)
-    }
-
     pub async fn execute(self, os: &mut Os) -> Result<ExitCode> {
         // Check for auth on subcommands that require it.
-        if self.requires_auth() && !is_logged_in(&mut os.database).await {
-            bail!(
-                "You are not logged in, please log in with {}",
-                format!("{CLI_BINARY_NAME} login").bold()
-            );
+        if !is_logged_in(&mut os.database).await {
+            if matches!(self, Self::Chat(_)) {
+                let options = ["Yes", "No"];
+                match crate::util::choose(" You are not logged in. Login now?", &options)? {
+                    Some(0) => {},
+                    _ => bail!("Login is required to use chat"),
+                }
+
+                LoginArgs::default().execute(os).await?;
+            } else if matches!(self, Self::Profile) {
+                bail!(
+                    "You are not logged in, please log in with {}",
+                    format!("{CLI_BINARY_NAME} login").bold()
+                );
+            }
         }
 
         // Send executed telemetry.

crates/q_cli/src/cli/mod.rs

@@ -20,7 +20,6 @@ mod translate;
 mod uninstall;
 mod update;
 mod user;
-
 use std::io::{
     Write as _,
     stdout,
@@ -87,7 +86,7 @@ use crate::util::desktop::{
 };
 use crate::util::{
     CliContext,
-    assert_logged_in,
+    is_logged_in_check,
 };
 
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)]
@@ -375,30 +374,23 @@ impl Cli {
 
                     Self::execute_chat("chat", Some(args), true).await
                 },
-                CliRootCommands::Mcp { args } => {
-                    if args.iter().any(|arg| ["--help", "-h"].contains(&arg.as_str())) {
-                        return Self::execute_chat("mcp", Some(args), false).await;
-                    }
-
-                    Self::execute_chat("mcp", Some(args), true).await
-                },
+                CliRootCommands::Mcp { args } => Self::execute_chat("mcp", Some(args), false).await,
                 CliRootCommands::Inline(subcommand) => subcommand.execute(&cli_context).await,
-                CliRootCommands::Agent { args } => {
-                    if args.iter().any(|arg| ["--help", "-h"].contains(&arg.as_str())) {
-                        return Self::execute_chat("agent", Some(args), false).await;
-                    }
-
-                    Self::execute_chat("agent", Some(args), true).await
-                },
+                CliRootCommands::Agent { args } => Self::execute_chat("agent", Some(args), false).await,
             },
             // Root command
             None => Self::execute_chat("chat", None, true).await,
         }
     }
 
     pub async fn execute_chat(subcmd: &str, args: Option<Vec<String>>, enforce_login: bool) -> Result<ExitCode> {
-        if enforce_login {
-            assert_logged_in().await?;
+        if enforce_login && !is_logged_in_check().await {
+            let options = ["Yes", "No"];
+            match crate::util::choose(" You are not logged in. Login now?", &options)? {
+                Some(0) => {},
+                _ => bail!("Login is required to use chat"),
+            }
+            crate::cli::user::login_interactive(Default::default()).await?;
         }
 
         // Save credentials from the macOS keychain to sqlite.

crates/q_cli/src/util/mod.rs

@@ -290,8 +290,12 @@ pub fn dialoguer_theme() -> ColorfulTheme {
     }
 }
 
+pub async fn is_logged_in_check() -> bool {
+    std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty()) || fig_auth::is_logged_in().await
+}
+
 pub async fn assert_logged_in() -> Result<(), Error> {
-    if !(std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty()) || fig_auth::is_logged_in().await) {
+    if !is_logged_in_check().await {
         bail!(
             "You are not logged in, please log in with {}",
             format!("{CLI_BINARY_NAME} login",).bold()