reformats sign in messages

Author: dingfeli

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

@@ -181,10 +181,15 @@ impl Default for Agent {
                 set.extend(default_approve);
                 set
             },
-            resources: vec!["file://AmazonQ.md", "file://AGENTS.md", "file://README.md", "file://.amazonq/rules/**/*.md"]
-                .into_iter()
-                .map(Into::into)
-                .collect::<Vec<_>>(),
+            resources: vec![
+                "file://AmazonQ.md",
+                "file://AGENTS.md",
+                "file://README.md",
+                "file://.amazonq/rules/**/*.md",
+            ]
+            .into_iter()
+            .map(Into::into)
+            .collect::<Vec<_>>(),
             hooks: Default::default(),
             tools_settings: Default::default(),
             use_legacy_mcp_json: true,

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

@@ -138,6 +138,11 @@ enum LoadingMsg {
     /// This is sent when all tool initialization is complete or when the application is shutting
     /// down.
     Terminate { still_loading: Vec<String> },
+    /// Indicates that a server requires user authentication and provides a sign-in link.
+    /// This message is used to notify the user about authentication requirements for MCP servers
+    /// that need OAuth or other authentication methods. Contains the server name and the
+    /// authentication message (typically a URL or instructions).
+    SignInNotice { name: String },
 }
 
 /// Used to denote the loading outcome associated with a server.
@@ -1183,6 +1188,15 @@ fn spawn_display_task(
                                 execute!(output, style::Print("\n"),)?;
                                 break;
                             },
+                            LoadingMsg::SignInNotice { name } => {
+                                execute!(
+                                    output,
+                                    cursor::MoveToColumn(0),
+                                    cursor::MoveUp(1),
+                                    terminal::Clear(terminal::ClearType::CurrentLine),
+                                )?;
+                                queue_oauth_message(&name, &mut output)?;
+                            },
                         },
                         Err(_e) => {
                             spinner_logo_idx = (spinner_logo_idx + 1) % SPINNER_CHARS.len();
@@ -1611,7 +1625,7 @@ fn spawn_orchestrator_task(
                 UpdateEventMessage::OauthLink { server_name, link } => {
                     let mut buf_writer = BufWriter::new(&mut *record_temp_buf);
                     let msg = eyre::eyre!(link);
-                    let _ = queue_oauth_message(server_name.as_str(), &msg, &mut buf_writer);
+                    let _ = queue_oauth_message_with_link(server_name.as_str(), &msg, &mut buf_writer);
                     let _ = buf_writer.flush();
                     drop(buf_writer);
                     let record_str = String::from_utf8_lossy(record_temp_buf).to_string();
@@ -1625,10 +1639,8 @@ fn spawn_orchestrator_task(
                         })
                         .or_insert(vec![record]);
                     if let Some(sender) = &loading_status_sender {
-                        let msg = LoadingMsg::Warn {
+                        let msg = LoadingMsg::SignInNotice {
                             name: server_name.clone(),
-                            msg: eyre::eyre!("{}", record_str),
-                            time: "0.0".to_string(),
                         };
                         if let Err(e) = sender.send(msg).await {
                             warn!(
@@ -1920,18 +1932,31 @@ fn queue_failure_message(
     )?)
 }
 
-fn queue_oauth_message(name: &str, msg: &eyre::Report, output: &mut impl Write) -> eyre::Result<()> {
+fn queue_oauth_message(name: &str, output: &mut impl Write) -> eyre::Result<()> {
+    Ok(queue!(
+        output,
+        style::SetForegroundColor(style::Color::Yellow),
+        style::Print("⚠ "),
+        style::SetForegroundColor(style::Color::Blue),
+        style::Print(name),
+        style::ResetColor,
+        style::Print(" requires OAuth authentication. Use /mcp to see the auth link\n"),
+    )?)
+}
+
+fn queue_oauth_message_with_link(name: &str, msg: &eyre::Report, output: &mut impl Write) -> eyre::Result<()> {
     Ok(queue!(
         output,
         style::SetForegroundColor(style::Color::Yellow),
         style::Print("⚠ "),
         style::SetForegroundColor(style::Color::Blue),
         style::Print(name),
         style::ResetColor,
-        style::Print(" requires OAuth authentication. Please visit:\n"),
-        style::SetForegroundColor(style::Color::Cyan),
+        style::Print(" requires OAuth authentication. Follow this link to proceed: \n"),
+        style::SetForegroundColor(style::Color::Yellow),
         style::Print(msg),
         style::ResetColor,
+        style::Print("\n")
     )?)
 }
 

crates/chat-cli/src/mcp_client/client.rs

@@ -5,7 +5,20 @@ use std::process::Stdio;
 use regex::Regex;
 use reqwest::Client;
 use rmcp::model::{
-    CallToolRequestParam, CallToolResult, ErrorCode, GetPromptRequestParam, GetPromptResult, Implementation, InitializeRequestParam, ListPromptsResult, ListToolsResult, LoggingLevel, LoggingMessageNotificationParam, PaginatedRequestParam, ServerNotification, ServerRequest
+    CallToolRequestParam,
+    CallToolResult,
+    ErrorCode,
+    GetPromptRequestParam,
+    GetPromptResult,
+    Implementation,
+    InitializeRequestParam,
+    ListPromptsResult,
+    ListToolsResult,
+    LoggingLevel,
+    LoggingMessageNotificationParam,
+    PaginatedRequestParam,
+    ServerNotification,
+    ServerRequest,
 };
 use rmcp::service::{
     ClientInitializeError,
@@ -31,6 +44,7 @@ use tokio::process::{
 };
 use tokio::task::JoinHandle;
 use tracing::{
+    debug,
     error,
     info,
 };
@@ -152,7 +166,7 @@ macro_rules! decorate_with_auth_retry {
                     // TODO: discern error type prior to retrying
                     // Not entirely sure what is thrown when auth is required
                     if let Some(auth_client) = self.get_auth_client() {
-                        let refresh_result = auth_client.auth_manager.lock().await.refresh_token().await;
+                        let refresh_result = auth_client.get_access_token().await;
                         match refresh_result {
                             Ok(_) => {
                                 // Retry the operation after token refresh
@@ -163,8 +177,11 @@ macro_rules! decorate_with_auth_retry {
                             },
                             Err(_) => {
                                 // If refresh fails, return the original error
+                                // Currently our event loop just does not allow us easy ways to
+                                // reauth entirely once a session starts since this would mean
+                                // swapping of transport (which also means swapping of client)
                                 Err(e)
-                            }
+                            },
                         }
                     } else {
                         // No auth client available, return original error
@@ -176,6 +193,11 @@ macro_rules! decorate_with_auth_retry {
     };
 }
 
+/// Wrapper around rmcp service types to enable cloning.
+///
+/// This exists because `rmcp::service::RunningService` is not directly cloneable as it is a
+/// pointer type to `Peer<C>`. This enum allows us to hold either the original service or its
+/// peer representation, enabling cloning by converting the original service to a peer when needed.
 pub enum InnerService {
     Original(rmcp::service::RunningService<RoleClient, Box<dyn DynService<RoleClient>>>),
     Peer(rmcp::service::Peer<RoleClient>),
@@ -194,11 +216,22 @@ impl Clone for InnerService {
     fn clone(&self) -> Self {
         match self {
             InnerService::Original(rs) => InnerService::Peer((*rs).clone()),
-            InnerService::Peer(peer) => InnerService::Peer(peer.clone())
+            InnerService::Peer(peer) => InnerService::Peer(peer.clone()),
         }
     }
 }
 
+/// A wrapper around MCP (Model Context Protocol) service instances that manages
+/// authentication and enables cloning functionality.
+///
+/// This struct holds either an original `RunningService` or its peer representation,
+/// along with an optional authentication drop guard for managing OAuth tokens.
+/// The authentication drop guard handles token lifecycle and cleanup when the
+/// service is dropped.
+///
+/// # Fields
+/// * `inner_service` - The underlying MCP service instance (original or peer)
+/// * `auth_dropguard` - Optional authentication manager for OAuth token handling
 #[derive(Debug)]
 pub struct RunningService {
     pub inner_service: InnerService,
@@ -215,18 +248,19 @@ impl Clone for RunningService {
 
         RunningService {
             inner_service: self.inner_service.clone(),
-            auth_dropguard
+            auth_dropguard,
         }
     }
 }
 
 impl RunningService {
+    decorate_with_auth_retry!(CallToolRequestParam, call_tool, CallToolResult);
+
+    decorate_with_auth_retry!(GetPromptRequestParam, get_prompt, GetPromptResult);
+
     pub fn get_auth_client(&self) -> Option<AuthClient<Client>> {
         self.auth_dropguard.as_ref().map(|a| a.auth_client.clone())
     }
-
-    decorate_with_auth_retry!(CallToolRequestParam, call_tool, CallToolResult);
-    decorate_with_auth_retry!(GetPromptRequestParam, get_prompt, GetPromptResult);
 }
 
 pub type StdioTransport = (TokioChildProcess, Option<ChildStderr>);
@@ -304,16 +338,17 @@ impl McpClientService {
                                 let service = match self.into_dyn().serve(transport).await.map_err(Box::new) {
                                     Ok(service) => service,
                                     Err(e) if matches!(*e, ClientInitializeError::ConnectionClosed(_)) => {
+                                        debug!("## mcp: first hand shake attempt failed: {:?}", e);
                                         let refresh_res =
-                                            auth_dg.auth_client.auth_manager.lock().await.refresh_token().await;
+                                            auth_dg.auth_client.get_access_token().await;
                                         let new_self = McpClientService::new(
                                             server_name.clone(),
                                             backup_config,
                                             messenger_clone.clone(),
                                         );
 
                                         let new_transport =
-                                            get_http_transport(&os_clone, true, &url, &*messenger_dup).await?;
+                                            get_http_transport(&os_clone, true, &url, Some(auth_dg.auth_client.clone()), &*messenger_dup).await?;
 
                                         match new_transport {
                                             HttpTransport::WithAuth((new_transport, new_auth_dg)) => {
@@ -325,13 +360,14 @@ impl McpClientService {
                                                         new_self.into_dyn().serve(new_transport).await.map_err(Box::new)?
                                                     },
                                                     Err(e) => {
+                                                        error!("## mcp: token refresh attempt failed: {:?}", e);
                                                         info!("Retry for http transport failed {e}. Possible reauth needed");
                                                         // This could be because the refresh token is expired, in which
                                                         // case we would need to have user go through the auth flow
                                                         // again
                                                         let new_transport  =
-                                                            get_http_transport(&os_clone, true, &url, &*messenger_dup).await?;
-                                                        
+                                                            get_http_transport(&os_clone, true, &url, None, &*messenger_dup).await?;
+
                                                         match new_transport {
                                                             HttpTransport::WithAuth((new_transport, new_auth_dg)) => {
                                                                 auth_dg = new_auth_dg;
@@ -345,7 +381,7 @@ impl McpClientService {
                                                     },
                                                 }
                                             },
-                                            HttpTransport::WithoutAuth(new_transport) => 
+                                            HttpTransport::WithoutAuth(new_transport) =>
                                                 new_self.into_dyn().serve(new_transport).await.map_err(Box::new)?,
                                         }
                                     },
@@ -487,7 +523,7 @@ impl McpClientService {
                 Ok(Transport::Stdio((tokio_child_process, child_stderr)))
             },
             TransportType::Http => {
-                let http_transport = get_http_transport(os, false, url, messenger).await?;
+                let http_transport = get_http_transport(os, false, url, None, messenger).await?;
 
                 Ok(Transport::Http(http_transport))
             },

crates/chat-cli/src/mcp_client/oauth_util.rs

@@ -33,6 +33,7 @@ use sha2::{
 use tokio::sync::oneshot::Sender;
 use tokio_util::sync::CancellationToken;
 use tracing::{
+    debug,
     error,
     info,
 };
@@ -78,6 +79,13 @@ impl Drop for LoopBackDropGuard {
     }
 }
 
+/// A guard that manages the lifecycle of an authenticated MCP client and automatically
+/// persists OAuth credentials when dropped.
+///
+/// This struct wraps an `AuthClient` and ensures that OAuth tokens are written to disk
+/// when the guard goes out of scope, unless explicitly disabled via `should_write`.
+/// This provides automatic credential caching for MCP server connections that require
+/// OAuth authentication.
 #[derive(Clone, Debug)]
 pub struct AuthClientDropGuard {
     pub should_write: bool,
@@ -136,6 +144,16 @@ impl Drop for AuthClientDropGuard {
     }
 }
 
+/// HTTP transport wrapper that handles both authenticated and non-authenticated MCP connections.
+///
+/// This enum provides two variants for different authentication scenarios:
+/// - `WithAuth`: Used when the MCP server requires OAuth authentication, containing both the
+///   transport worker and an auth client guard that manages credential persistence
+/// - `WithoutAuth`: Used for servers that don't require authentication, containing only the basic
+///   transport worker
+///
+/// The appropriate variant is automatically selected based on the server's response to
+/// an initial probe request during transport creation.
 pub enum HttpTransport {
     WithAuth(
         (
@@ -150,6 +168,7 @@ pub async fn get_http_transport(
     os: &Os,
     delete_cache: bool,
     url: &str,
+    auth_client: Option<AuthClient<Client>>,
     messenger: &dyn Messenger,
 ) -> Result<HttpTransport, OauthUtilError> {
     let cred_dir = get_mcp_auth_dir(os)?;
@@ -165,8 +184,14 @@ pub async fn get_http_transport(
     let probe_resp = reqwest_client.get(url.clone()).send().await?;
     match probe_resp.status() {
         StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
-            let am = get_auth_manager(url.clone(), cred_full_path.clone(), messenger).await?;
-            let auth_client = AuthClient::new(reqwest_client, am);
+            debug!("## mcp: requires auth, auth client passed in is {:?}", auth_client);
+            let auth_client = match auth_client {
+                Some(auth_client) => auth_client,
+                None => {
+                    let am = get_auth_manager(url.clone(), cred_full_path.clone(), messenger).await?;
+                    AuthClient::new(reqwest_client, am)
+                },
+            };
             let transport =
                 StreamableHttpClientTransport::with_client(auth_client.clone(), StreamableHttpClientTransportConfig {
                     uri: url.as_str().into(),
@@ -175,6 +200,7 @@ pub async fn get_http_transport(
                 });
 
             let auth_dg = AuthClientDropGuard::new(cred_full_path, auth_client);
+            debug!("## mcp: transport obtained");
 
             Ok(HttpTransport::WithAuth((transport, auth_dg)))
         },
@@ -191,7 +217,7 @@ async fn get_auth_manager(
     cred_full_path: PathBuf,
     messenger: &dyn Messenger,
 ) -> Result<AuthorizationManager, OauthUtilError> {
-    let content_as_bytes = tokio::fs::read(cred_full_path).await;
+    let content_as_bytes = tokio::fs::read(&cred_full_path).await;
     let mut oauth_state = OAuthState::new(url, None).await?;
 
     match content_as_bytes {
@@ -200,12 +226,15 @@ async fn get_auth_manager(
 
             oauth_state.set_credentials("id", token).await?;
 
+            debug!("## mcp: credentials set with cache");
+
             Ok(oauth_state
                 .into_authorization_manager()
                 .ok_or(OauthUtilError::MissingAuthorizationManager)?)
         },
         Err(e) => {
             info!("Error reading cached credentials: {e}");
+            debug!("## mcp: cache read failed. constructing auth manager from scratch");
             get_auth_manager_impl(oauth_state, messenger).await
         },
     }