| 1 | //! An [`McpToolset`] as [Ed](ed_agent) tools. Behind the `agent` feature. |
| 2 | |
| 3 | use std::collections::{HashMap, HashSet}; |
| 4 | use std::sync::Arc; |
| 5 | |
| 6 | use async_trait::async_trait; |
| 7 | use ed_agent::{AgentToolDefinition, Ed, EventSink, ToolExecutor, ToolRisk}; |
| 8 | use serde_json::Value; |
| 9 | use tokio::sync::{Mutex, RwLock}; |
| 10 | use tokio_util::sync::CancellationToken; |
| 11 | |
| 12 | use crate::{McpTool, McpToolset}; |
| 13 | |
| 14 | /// How MCP tools are presented to Ed. |
| 15 | #[derive(Debug, Clone, Default)] |
| 16 | pub struct ToolPolicy { |
| 17 | read_only: HashSet<String>, |
| 18 | unprefixed: bool, |
| 19 | } |
| 20 | |
| 21 | impl ToolPolicy { |
| 22 | /// Default policy: tool names are `mcp__<server>__<tool>`, and a tool is read-only only when |
| 23 | /// its server sets `readOnlyHint: true`. |
| 24 | pub fn new() -> Self { |
| 25 | Self::default() |
| 26 | } |
| 27 | |
| 28 | /// Treat these tools as read-only (no approval) whatever their server says. Names are the |
| 29 | /// server's own tool names, or the names the model sees. For servers that do not set |
| 30 | /// `readOnlyHint`. |
| 31 | pub fn read_only(mut self, names: &[&str]) -> Self { |
| 32 | self.read_only.extend(names.iter().map(|n| (*n).to_owned())); |
| 33 | self |
| 34 | } |
| 35 | |
| 36 | /// Show the model the server's own tool names, without the `mcp__<server>__` prefix. For apps |
| 37 | /// with one server whose tool names are already in the system prompt. If two servers use the |
| 38 | /// same tool name the first one wins. |
| 39 | pub fn unprefixed(mut self) -> Self { |
| 40 | self.unprefixed = true; |
| 41 | self |
| 42 | } |
| 43 | |
| 44 | fn risk(&self, tool: &McpTool, exposed: &str) -> ToolRisk { |
| 45 | if tool.read_only || self.read_only.contains(exposed) || self.read_only.contains(&tool.tool) |
| 46 | { |
| 47 | ToolRisk::ReadOnly |
| 48 | } else { |
| 49 | ToolRisk::Mutating |
| 50 | } |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | struct Loaded { |
| 55 | toolset: Arc<McpToolset>, |
| 56 | /// Name the model sees -> qualified name in the toolset. |
| 57 | names: HashMap<String, String>, |
| 58 | definitions: Vec<AgentToolDefinition>, |
| 59 | } |
| 60 | |
| 61 | /// Runs the tools of an [`McpToolset`] for Ed. |
| 62 | /// |
| 63 | /// Cheap to clone. Hand one clone to [`Ed::with_agent`] as the executor, and keep another to |
| 64 | /// connect and disconnect servers as the user signs in and out: |
| 65 | /// |
| 66 | /// ```ignore |
| 67 | /// let mcp = McpExecutor::new(ToolPolicy::new().read_only(&["list_messages"]).unprefixed()); |
| 68 | /// let ed = Ed::with_agent(sink, Arc::new(mcp.clone()), approvals, config); |
| 69 | /// mcp.set_toolset(Some(toolset)); |
| 70 | /// mcp.sync(&ed).await; |
| 71 | /// ``` |
| 72 | #[derive(Clone)] |
| 73 | pub struct McpExecutor { |
| 74 | policy: Arc<ToolPolicy>, |
| 75 | loaded: Arc<RwLock<Option<Loaded>>>, |
| 76 | registered: Arc<Mutex<Vec<String>>>, |
| 77 | } |
| 78 | |
| 79 | impl McpExecutor { |
| 80 | pub fn new(policy: ToolPolicy) -> Self { |
| 81 | Self { |
| 82 | policy: Arc::new(policy), |
| 83 | loaded: Arc::default(), |
| 84 | registered: Arc::default(), |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | /// Use `toolset`, or none, and shut the previous toolset down. Call |
| 89 | /// [`sync`](McpExecutor::sync) afterwards. A tool call still running on the old toolset |
| 90 | /// finishes first; its connections close when it does. |
| 91 | pub async fn set_toolset(&self, toolset: Option<McpToolset>) { |
| 92 | let new = toolset.map(|toolset| { |
| 93 | let mut names = HashMap::new(); |
| 94 | let mut definitions = Vec::new(); |
| 95 | for tool in toolset.tools() { |
| 96 | let exposed = if self.policy.unprefixed { |
| 97 | tool.tool.clone() |
| 98 | } else { |
| 99 | tool.name.clone() |
| 100 | }; |
| 101 | if names.contains_key(&exposed) { |
| 102 | tracing::warn!( |
| 103 | "duplicate tool name `{exposed}` after prefix removal; keeping the first" |
| 104 | ); |
| 105 | continue; |
| 106 | } |
| 107 | let description = if self.policy.unprefixed { |
| 108 | tool.description |
| 109 | .strip_prefix(&format!("[{} MCP server] ", tool.server)) |
| 110 | .unwrap_or(&tool.description) |
| 111 | .to_owned() |
| 112 | } else { |
| 113 | tool.description.clone() |
| 114 | }; |
| 115 | let description = if description.trim().is_empty() { |
| 116 | format!("The `{}` tool.", tool.tool) |
| 117 | } else { |
| 118 | description |
| 119 | }; |
| 120 | let schema = tool.parameters.to_string(); |
| 121 | definitions.push(match self.policy.risk(tool, &exposed) { |
| 122 | ToolRisk::ReadOnly => { |
| 123 | AgentToolDefinition::read_only(exposed.clone(), description, schema) |
| 124 | } |
| 125 | ToolRisk::Mutating => { |
| 126 | AgentToolDefinition::mutating(exposed.clone(), description, schema) |
| 127 | } |
| 128 | }); |
| 129 | names.insert(exposed, tool.name.clone()); |
| 130 | } |
| 131 | Loaded { |
| 132 | toolset: Arc::new(toolset), |
| 133 | names, |
| 134 | definitions, |
| 135 | } |
| 136 | }); |
| 137 | let old = std::mem::replace(&mut *self.loaded.write().await, new); |
| 138 | if let Some(old) = old |
| 139 | && let Some(toolset) = Arc::into_inner(old.toolset) |
| 140 | { |
| 141 | toolset.shutdown().await; |
| 142 | } |
| 143 | } |
| 144 | |
| 145 | /// The tools Ed should offer, for the connected toolset (empty when none). |
| 146 | pub async fn definitions(&self) -> Vec<AgentToolDefinition> { |
| 147 | self.loaded |
| 148 | .read() |
| 149 | .await |
| 150 | .as_ref() |
| 151 | .map(|l| l.definitions.clone()) |
| 152 | .unwrap_or_default() |
| 153 | } |
| 154 | |
| 155 | /// Make `ed`'s registered tools match the connected toolset: the ones this executor registered |
| 156 | /// before are removed, the current ones added. |
| 157 | pub async fn sync<S: EventSink>(&self, ed: &Ed<S>) { |
| 158 | let mut registered = self.registered.lock().await; |
| 159 | for name in registered.drain(..) { |
| 160 | ed.unregister_tool(&name).await; |
| 161 | } |
| 162 | for definition in self.definitions().await { |
| 163 | let name = definition.name.clone(); |
| 164 | match ed.register_tool(definition).await { |
| 165 | Ok(()) => registered.push(name), |
| 166 | Err(error) => tracing::warn!("could not register MCP tool `{name}`: {error}"), |
| 167 | } |
| 168 | } |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | #[async_trait] |
| 173 | impl ToolExecutor for McpExecutor { |
| 174 | async fn execute(&self, tool_name: &str, arguments: &str) -> Result<String, String> { |
| 175 | let args: Value = |
| 176 | serde_json::from_str(arguments).map_err(|e| format!("bad arguments: {e}"))?; |
| 177 | // Clone out of the lock: a call can run for a long time and must not hold up a |
| 178 | // disconnect. |
| 179 | let (toolset, qualified) = { |
| 180 | let guard = self.loaded.read().await; |
| 181 | let loaded = guard.as_ref().ok_or("no MCP server is connected")?; |
| 182 | let qualified = loaded |
| 183 | .names |
| 184 | .get(tool_name) |
| 185 | .ok_or_else(|| format!("unknown tool `{tool_name}`"))? |
| 186 | .clone(); |
| 187 | (loaded.toolset.clone(), qualified) |
| 188 | }; |
| 189 | let outcome = toolset |
| 190 | .call(&qualified, args, &CancellationToken::new(), &|_| {}) |
| 191 | .await |
| 192 | .map_err(|e| e.to_string())?; |
| 193 | if outcome.failed { |
| 194 | Err(outcome.text) |
| 195 | } else { |
| 196 | Ok(outcome.text) |
| 197 | } |
| 198 | } |
| 199 | } |