adds remote mcp support in agent crate

Author: dingfeli

crates/agent/Cargo.toml

@@ -51,7 +51,7 @@ r2d2_sqlite.workspace = true
 rand.workspace = true
 regex.workspace = true
 reqwest.workspace = true
-rmcp = { version = "0.8.0", features = ["client", "transport-async-rw", "transport-child-process", "transport-io"] }
+rmcp.workspace = true
 rusqlite.workspace = true
 rustls.workspace = true
 rustls-native-certs.workspace = true

crates/agent/src/agent/agent_config/definitions.rs

@@ -12,6 +12,7 @@ use serde::{
 use super::types::ResourcePath;
 use crate::agent::consts::DEFAULT_AGENT_NAME;
 use crate::agent::tools::BuiltInToolName;
+use crate::mcp::oauth_util::OAuthConfig;
 
 #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
 #[serde(untagged)]
@@ -218,7 +219,7 @@ pub struct McpServers {
 #[serde(untagged)]
 pub enum McpServerConfig {
     Local(LocalMcpServerConfig),
-    StreamableHTTP(StreamableHTTPMcpServerConfig),
+    Remote(RemoteMcpServerConfig),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
@@ -241,7 +242,7 @@ pub struct LocalMcpServerConfig {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
-pub struct StreamableHTTPMcpServerConfig {
+pub struct RemoteMcpServerConfig {
     /// The URL endpoint for HTTP-based MCP servers
     pub url: String,
     /// HTTP headers to include when communicating with HTTP-based MCP servers
@@ -251,6 +252,11 @@ pub struct StreamableHTTPMcpServerConfig {
     #[serde(alias = "timeout")]
     #[serde(default = "default_timeout")]
     pub timeout_ms: u64,
+    /// OAuth scopes required for authentication with the remote MCP server
+    pub oauth_scope: Vec<String>,
+    /// OAuth configuration for this server
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub oauth: Option<OAuthConfig>,
 }
 
 pub fn default_timeout() -> u64 {

crates/agent/src/agent/consts.rs

@@ -13,6 +13,8 @@ pub const MAX_TOOL_NAME_LEN: usize = 64;
 
 pub const MAX_TOOL_SPEC_DESCRIPTION_LEN: usize = 10_004;
 
+pub const DEFAULT_MCP_CREDENTIAL_PATH: &str = "~/.aws/sso/cache";
+
 /// 10 MB
 pub const MAX_IMAGE_SIZE_BYTES: u64 = 10 * 1024 * 1024;
 

crates/agent/src/agent/mcp/actor.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -44,20 +45,16 @@ pub enum McpMessage {
     Tools(Result<Vec<RmcpTool>, ServiceError>),
     Prompts(Result<Vec<RmcpPrompt>, ServiceError>),
     ExecuteTool { request_id: u32, result: ExecuteToolResult },
+    OauthRequest { oauth_url: String },
 }
 
 #[derive(Debug)]
 pub struct McpServerActorHandle {
     _server_name: String,
     sender: RequestSender<McpServerActorRequest, McpServerActorResponse, McpServerActorError>,
-    event_rx: mpsc::Receiver<McpServerActorEvent>,
 }
 
 impl McpServerActorHandle {
-    pub async fn recv(&mut self) -> Option<McpServerActorEvent> {
-        self.event_rx.recv().await
-    }
-
     pub async fn get_tool_specs(&self) -> Result<Vec<ToolSpec>, McpServerActorError> {
         match self
             .sender
@@ -153,6 +150,7 @@ impl From<ServiceError> for McpServerActorError {
 pub enum McpServerActorEvent {
     /// The MCP server has launched successfully
     Initialized {
+        server_name: String,
         /// Time taken to launch the server
         serve_duration: Duration,
         /// Time taken to list all tools.
@@ -165,7 +163,9 @@ pub enum McpServerActorEvent {
         list_prompts_duration: Option<Duration>,
     },
     /// The MCP server failed to initialize successfully
-    InitializeError(String),
+    InitializeError { server_name: String, error: String },
+    /// An OAuth authentication request from the MCP server
+    OauthRequest { server_name: String, oauth_url: String },
 }
 
 #[derive(Debug)]
@@ -195,34 +195,38 @@ pub struct McpServerActor {
 
 impl McpServerActor {
     /// Spawns an actor to manage the MCP server, returning a [McpServerActorHandle].
-    pub fn spawn(server_name: String, config: McpServerConfig) -> McpServerActorHandle {
-        let (event_tx, event_rx) = mpsc::channel(32);
+    pub fn spawn(
+        server_name: String,
+        config: McpServerConfig,
+        cred_path: PathBuf,
+        event_tx: mpsc::Sender<McpServerActorEvent>,
+    ) -> McpServerActorHandle {
         let (req_tx, req_rx) = new_request_channel();
 
         let server_name_clone = server_name.clone();
-        tokio::spawn(async move { Self::launch(server_name_clone, config, req_rx, event_tx).await });
+        tokio::spawn(async move { Self::launch(server_name_clone, config, cred_path, req_rx, event_tx).await });
 
         McpServerActorHandle {
             _server_name: server_name,
             sender: req_tx,
-            event_rx,
         }
     }
 
     async fn launch(
         server_name: String,
         config: McpServerConfig,
+        cred_path: PathBuf,
         req_rx: RequestReceiver<McpServerActorRequest, McpServerActorResponse, McpServerActorError>,
         event_tx: mpsc::Sender<McpServerActorEvent>,
     ) {
         let (message_tx, message_rx) = mpsc::channel(32);
-        match McpService::new(server_name.clone(), config.clone(), message_tx.clone())
+        match McpService::new(server_name.clone(), config.clone(), cred_path, message_tx.clone())
             .launch()
             .await
         {
             Ok((service_handle, launch_md)) => {
                 let s = Self {
-                    server_name,
+                    server_name: server_name.clone(),
                     _config: config,
                     tools: launch_md.tools.unwrap_or_default(),
                     prompts: launch_md.prompts.unwrap_or_default(),
@@ -237,6 +241,7 @@ impl McpServerActor {
                 let _ = s
                     .event_tx
                     .send(McpServerActorEvent::Initialized {
+                        server_name,
                         serve_duration: launch_md.serve_time_taken,
                         list_tools_duration: launch_md.list_tools_duration,
                         list_prompts_duration: launch_md.list_prompts_duration,
@@ -246,7 +251,10 @@ impl McpServerActor {
             },
             Err(err) => {
                 let _ = event_tx
-                    .send(McpServerActorEvent::InitializeError(err.to_string()))
+                    .send(McpServerActorEvent::InitializeError {
+                        server_name,
+                        error: err.to_string(),
+                    })
                     .await;
             },
         }
@@ -331,6 +339,18 @@ impl McpServerActor {
                     );
                 },
             },
+            McpMessage::OauthRequest { oauth_url } => {
+                if let Err(err) = self
+                    .event_tx
+                    .send(McpServerActorEvent::OauthRequest {
+                        server_name: self.server_name.clone(),
+                        oauth_url,
+                    })
+                    .await
+                {
+                    error!(?self.server_name, ?err, "failed to send oauth request");
+                }
+            },
         }
     }
 
@@ -340,7 +360,7 @@ impl McpServerActor {
         let service_handle = self.service_handle.clone();
         let tx = self.message_tx.clone();
         tokio::spawn(async move {
-            let res = service_handle.list_tools().await;
+            let res = service_handle.list_all_tools().await;
             let _ = tx.send(McpMessage::Tools(res)).await;
         });
     }
@@ -351,7 +371,7 @@ impl McpServerActor {
         let service_handle = self.service_handle.clone();
         let tx = self.message_tx.clone();
         tokio::spawn(async move {
-            let res = service_handle.list_prompts().await;
+            let res = service_handle.list_all_prompts().await;
             let _ = tx.send(McpMessage::Prompts(res)).await;
         });
     }