refreshes token on expiration

Author: dingfeli

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

@@ -90,6 +90,7 @@ use crate::database::settings::Setting;
 use crate::mcp_client::messenger::Messenger;
 use crate::mcp_client::{
     InitializedMcpClient,
+    InnerService,
     McpClientService,
 };
 use crate::os::Os;
@@ -631,7 +632,9 @@ impl ToolManager {
                         tokio::spawn(async move {
                             match handle.await {
                                 Ok(Ok(client)) => {
-                                    let client = client.inner_service;
+                                    let InnerService::Original(client) = client.inner_service else {
+                                        unreachable!();
+                                    };
                                     match client.cancel().await {
                                         Ok(_) => info!("Server {server_name_clone} evicted due to agent swap"),
                                         Err(e) => error!("Server {server_name_clone} has failed to cancel: {e}"),
@@ -644,7 +647,9 @@ impl ToolManager {
                         });
                     },
                     InitializedMcpClient::Ready(running_service) => {
-                        let client = running_service.inner_service;
+                        let InnerService::Original(client) = running_service.inner_service else {
+                            unreachable!();
+                        };
                         match client.cancel().await {
                             Ok(_) => info!("Server {server_name} evicted due to agent swap"),
                             Err(e) => error!("Server {server_name} has failed to cancel: {e}"),
@@ -886,7 +891,7 @@ impl ToolManager {
                 Tool::Custom(CustomTool {
                     name: tool_name.to_owned(),
                     server_name: server_name.to_owned(),
-                    client: (*running_service).clone(),
+                    client: running_service.clone(),
                     auth_client,
                     params: value.args.as_object().cloned(),
                 })

crates/chat-cli/src/cli/chat/tools/custom_tool.rs

@@ -9,7 +9,6 @@ use crossterm::{
 };
 use eyre::Result;
 use reqwest::Client;
-use rmcp::RoleClient;
 use rmcp::model::CallToolRequestParam;
 use rmcp::transport::auth::AuthClient;
 use schemars::JsonSchema;
@@ -29,6 +28,7 @@ use crate::cli::agent::{
 };
 use crate::cli::chat::CONTINUATION_LINE;
 use crate::cli::chat::token_counter::TokenCounter;
+use crate::mcp_client::RunningService;
 use crate::os::Os;
 use crate::util::MCP_SERVER_TOOL_DELIMITER;
 use crate::util::pattern_matching::matches_any_pattern;
@@ -94,7 +94,7 @@ pub struct CustomTool {
     /// prefixed to the tool name when presented to the model for disambiguation.
     pub server_name: String,
     /// Reference to the client that manages communication with the tool's server process.
-    pub client: rmcp::Peer<RoleClient>,
+    pub client: RunningService,
     /// Optional authentication client for handling authentication with HTTP-based MCP servers.
     /// This is used when the MCP server requires authentication for tool invocation.
     pub auth_client: Option<AuthClient<Client>>,

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

@@ -1,24 +1,11 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
-use std::ops::{
-    Deref,
-    DerefMut,
-};
 use std::process::Stdio;
 
 use regex::Regex;
 use reqwest::Client;
 use rmcp::model::{
-    ErrorCode,
-    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,
@@ -151,30 +138,95 @@ pub enum McpClientError {
     Auth(#[from] crate::auth::AuthError),
 }
 
-pub struct RunningService {
-    pub inner_service: rmcp::service::RunningService<RoleClient, Box<dyn DynService<RoleClient>>>,
-    #[allow(dead_code)]
-    pub auth_dropguard: Option<AuthClientDropGuard>,
+macro_rules! decorate_with_auth_retry {
+    ($param_type:ty, $method_name:ident, $return_type:ty) => {
+        pub async fn $method_name(&self, param: $param_type) -> Result<$return_type, rmcp::ServiceError> {
+            let first_attempt = match &self.inner_service {
+                InnerService::Original(rs) => rs.$method_name(param.clone()).await,
+                InnerService::Peer(peer) => peer.$method_name(param.clone()).await,
+            };
+
+            match first_attempt {
+                Ok(result) => Ok(result),
+                Err(e) => {
+                    // 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;
+                        match refresh_result {
+                            Ok(_) => {
+                                // Retry the operation after token refresh
+                                match &self.inner_service {
+                                    InnerService::Original(rs) => rs.$method_name(param).await,
+                                    InnerService::Peer(peer) => peer.$method_name(param).await,
+                                }
+                            },
+                            Err(_) => {
+                                // If refresh fails, return the original error
+                                Err(e)
+                            }
+                        }
+                    } else {
+                        // No auth client available, return original error
+                        Err(e)
+                    }
+                },
+            }
+        }
+    };
 }
 
-impl RunningService {
-    pub fn get_auth_client(&self) -> Option<AuthClient<Client>> {
-        self.auth_dropguard.as_ref().map(|a| a.auth_client.clone())
+pub enum InnerService {
+    Original(rmcp::service::RunningService<RoleClient, Box<dyn DynService<RoleClient>>>),
+    Peer(rmcp::service::Peer<RoleClient>),
+}
+
+impl std::fmt::Debug for InnerService {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            InnerService::Original(_) => f.debug_tuple("Original").field(&"RunningService<..>").finish(),
+            InnerService::Peer(peer) => f.debug_tuple("Peer").field(peer).finish(),
+        }
     }
 }
 
-impl Deref for RunningService {
-    type Target = rmcp::service::RunningService<RoleClient, Box<dyn DynService<RoleClient>>>;
+impl Clone for InnerService {
+    fn clone(&self) -> Self {
+        match self {
+            InnerService::Original(rs) => InnerService::Peer((*rs).clone()),
+            InnerService::Peer(peer) => InnerService::Peer(peer.clone())
+        }
+    }
+}
 
-    fn deref(&self) -> &Self::Target {
-        &self.inner_service
+#[derive(Debug)]
+pub struct RunningService {
+    pub inner_service: InnerService,
+    auth_dropguard: Option<AuthClientDropGuard>,
+}
+
+impl Clone for RunningService {
+    fn clone(&self) -> Self {
+        let auth_dropguard = self.auth_dropguard.as_ref().map(|dg| {
+            let mut dg = dg.clone();
+            dg.should_write = false;
+            dg
+        });
+
+        RunningService {
+            inner_service: self.inner_service.clone(),
+            auth_dropguard
+        }
     }
 }
 
-impl DerefMut for RunningService {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.inner_service
+impl RunningService {
+    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>);
@@ -397,7 +449,7 @@ impl McpClientService {
             });
 
             Ok(RunningService {
-                inner_service: service,
+                inner_service: InnerService::Original(service),
                 auth_dropguard,
             })
         });

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

@@ -78,6 +78,7 @@ impl Drop for LoopBackDropGuard {
     }
 }
 
+#[derive(Clone, Debug)]
 pub struct AuthClientDropGuard {
     pub should_write: bool,
     pub cred_full_path: PathBuf,