| 1 | //! The tool layer: what a tool gets ([`ToolCtx`]), what it returns ([`ToolOutcome`]), the |
| 2 | //! approval flow, ACP path rules, and dispatch to product [`Toolset`]s and MCP servers. |
| 3 | |
| 4 | use std::collections::HashSet; |
| 5 | use std::path::{Component, Path, PathBuf}; |
| 6 | use std::sync::atomic::{AtomicBool, Ordering}; |
| 7 | use std::sync::{Arc, Mutex}; |
| 8 | |
| 9 | use agent_client_protocol::schema::v1::{ |
| 10 | ClientCapabilities, PermissionOption, PermissionOptionKind, RequestPermissionOutcome, |
| 11 | RequestPermissionRequest, SessionId, SessionNotification, SessionUpdate, ToolCallContent, |
| 12 | ToolCallLocation, ToolCallUpdate, ToolCallUpdateFields, ToolKind, |
| 13 | }; |
| 14 | use agent_client_protocol::{Client, ConnectionTo}; |
| 15 | use anyhow::{Result, anyhow, bail}; |
| 16 | use async_trait::async_trait; |
| 17 | use ed_mcp::McpToolset; |
| 18 | use serde_json::{Value, json}; |
| 19 | use tokio_util::sync::CancellationToken; |
| 20 | |
| 21 | use crate::{McpCall, Profile}; |
| 22 | |
| 23 | /// Cap on tool output sent to the model and on client terminal output. |
| 24 | pub const MAX_OUTPUT_BYTES: usize = 32 * 1024; |
| 25 | |
| 26 | /// Everything a tool needs to act on behalf of one session. |
| 27 | pub struct ToolCtx { |
| 28 | pub connection: ConnectionTo<Client>, |
| 29 | pub session_id: SessionId, |
| 30 | /// Primary working directory: the base for relative paths and `run_command`. |
| 31 | pub cwd: PathBuf, |
| 32 | /// Additional workspace roots the model may address with the `root` parameter. |
| 33 | pub roots: Vec<PathBuf>, |
| 34 | pub caps: ClientCapabilities, |
| 35 | pub cancel: CancellationToken, |
| 36 | /// Approve tool calls without asking: `--yolo`, or the session's `auto_approve` option. |
| 37 | pub auto_approve: Arc<AtomicBool>, |
| 38 | /// Tool names the user chose "always allow" for in this session. |
| 39 | pub always_allowed: Arc<Mutex<HashSet<String>>>, |
| 40 | /// Tool names the user chose "always reject" for in this session. |
| 41 | pub always_rejected: Arc<Mutex<HashSet<String>>>, |
| 42 | /// Tools from MCP servers available to this session. |
| 43 | pub mcp: Arc<McpToolset>, |
| 44 | /// How long `run_command` may run: `<PREFIX>_COMMAND_TIMEOUT_SECS`, default 120 seconds. |
| 45 | pub command_timeout: std::time::Duration, |
| 46 | } |
| 47 | |
| 48 | /// What a tool call produced. |
| 49 | pub struct ToolOutcome { |
| 50 | /// Text returned to the model. |
| 51 | pub text: String, |
| 52 | /// Rich content shown to the user in the client. |
| 53 | pub content: Vec<ToolCallContent>, |
| 54 | pub failed: bool, |
| 55 | /// Files the tool produced or touched, surfaced as tool-call locations. |
| 56 | pub locations: Vec<PathBuf>, |
| 57 | } |
| 58 | |
| 59 | impl ToolOutcome { |
| 60 | pub fn ok(text: impl Into<String>) -> Self { |
| 61 | let text = text.into(); |
| 62 | Self { |
| 63 | content: vec![ToolCallContent::from(text.clone())], |
| 64 | text, |
| 65 | failed: false, |
| 66 | locations: Vec::new(), |
| 67 | } |
| 68 | } |
| 69 | pub fn err(text: impl Into<String>) -> Self { |
| 70 | let text = text.into(); |
| 71 | Self { |
| 72 | content: vec![ToolCallContent::from(text.clone())], |
| 73 | text, |
| 74 | failed: true, |
| 75 | locations: Vec::new(), |
| 76 | } |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | /// An OpenAI `tools` entry. |
| 81 | pub fn function_def(name: &str, desc: &str, params: Value) -> Value { |
| 82 | json!({ "type": "function", "function": { "name": name, "description": desc, "parameters": params } }) |
| 83 | } |
| 84 | |
| 85 | /// The session's directories, for describing a tool call without running it (also used when |
| 86 | /// `session/load` replays history). |
| 87 | pub struct DescribeCtx<'a> { |
| 88 | pub cwd: &'a Path, |
| 89 | pub roots: &'a [PathBuf], |
| 90 | } |
| 91 | |
| 92 | /// A group of tools a product offers to the model. |
| 93 | #[async_trait] |
| 94 | pub trait Toolset: Send + Sync { |
| 95 | /// OpenAI `tools` entries, built with [`function_def`]. |
| 96 | fn definitions(&self) -> Vec<Value>; |
| 97 | /// Whether `name` is one of this toolset's tools. |
| 98 | fn handles(&self, name: &str) -> bool; |
| 99 | /// Title, kind and affected locations shown in the client when the call starts. Locations |
| 100 | /// must be absolute ([`absolutize`]). |
| 101 | fn describe( |
| 102 | &self, |
| 103 | name: &str, |
| 104 | args: &Value, |
| 105 | ctx: &DescribeCtx<'_>, |
| 106 | ) -> (String, ToolKind, Vec<ToolCallLocation>); |
| 107 | /// Run the tool. A mutating tool asks first with [`ToolCtx::permit`]. An `Err` is shown to |
| 108 | /// the model as `Error: …` and the call is marked failed. |
| 109 | async fn execute( |
| 110 | &self, |
| 111 | ctx: &ToolCtx, |
| 112 | tool_call_id: &str, |
| 113 | name: &str, |
| 114 | args: Value, |
| 115 | ) -> Result<ToolOutcome>; |
| 116 | } |
| 117 | |
| 118 | /// All tool definitions for one session, in OpenAI `tools` format: product toolsets first, |
| 119 | /// then MCP tools sorted by name. |
| 120 | pub fn definitions(toolsets: &[Arc<dyn Toolset>], mcp: &McpToolset) -> Value { |
| 121 | let mut defs: Vec<Value> = toolsets.iter().flat_map(|t| t.definitions()).collect(); |
| 122 | defs.extend( |
| 123 | mcp.tools() |
| 124 | .into_iter() |
| 125 | .map(|t| function_def(&t.name, &t.description, t.parameters.clone())), |
| 126 | ); |
| 127 | Value::Array(defs) |
| 128 | } |
| 129 | |
| 130 | /// Title, kind and affected locations shown in the client when the call starts. |
| 131 | pub fn describe( |
| 132 | toolsets: &[Arc<dyn Toolset>], |
| 133 | cwd: &Path, |
| 134 | roots: &[PathBuf], |
| 135 | name: &str, |
| 136 | args: &Value, |
| 137 | ) -> (String, ToolKind, Vec<ToolCallLocation>) { |
| 138 | if let Some((server, tool)) = ed_mcp::split_qualified(name) { |
| 139 | let loc = ["input", "path", "file"] |
| 140 | .iter() |
| 141 | .filter_map(|k| args.get(*k).and_then(Value::as_str)) |
| 142 | .map(PathBuf::from) |
| 143 | .filter(|p| p.is_absolute()) |
| 144 | .map(|p| ToolCallLocation::new(absolutize(&p))) |
| 145 | .collect(); |
| 146 | return (format!("{server}: {tool}"), ToolKind::Other, loc); |
| 147 | } |
| 148 | let ctx = DescribeCtx { cwd, roots }; |
| 149 | match toolsets.iter().find(|t| t.handles(name)) { |
| 150 | Some(t) => t.describe(name, args, &ctx), |
| 151 | None => (name.to_string(), ToolKind::Other, vec![]), |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | /// Run one tool call: a product toolset's, or an MCP server's. Never fails; errors become a |
| 156 | /// failed outcome the model sees. |
| 157 | pub async fn execute( |
| 158 | ctx: &ToolCtx, |
| 159 | toolsets: &[Arc<dyn Toolset>], |
| 160 | profile: &dyn Profile, |
| 161 | tool_call_id: &str, |
| 162 | name: &str, |
| 163 | args: Value, |
| 164 | ) -> ToolOutcome { |
| 165 | let result = if ed_mcp::split_qualified(name).is_some() { |
| 166 | execute_mcp(ctx, profile, tool_call_id, name, args).await |
| 167 | } else if let Some(t) = toolsets.iter().find(|t| t.handles(name)) { |
| 168 | t.execute(ctx, tool_call_id, name, args).await |
| 169 | } else { |
| 170 | Err(anyhow!("unknown tool `{name}`")) |
| 171 | }; |
| 172 | result.unwrap_or_else(|e| ToolOutcome::err(format!("Error: {e:#}"))) |
| 173 | } |
| 174 | |
| 175 | async fn execute_mcp( |
| 176 | ctx: &ToolCtx, |
| 177 | profile: &dyn Profile, |
| 178 | tool_call_id: &str, |
| 179 | name: &str, |
| 180 | args: Value, |
| 181 | ) -> Result<ToolOutcome> { |
| 182 | let Some(tool) = ctx.mcp.tool(name) else { |
| 183 | bail!("unknown tool `{name}`"); |
| 184 | }; |
| 185 | if !tool.read_only { |
| 186 | let call = McpCall { |
| 187 | tool, |
| 188 | args: &args, |
| 189 | connection: ctx.mcp.connection(name), |
| 190 | }; |
| 191 | let preview = match profile.mcp_permission_preview(call).await { |
| 192 | Some(text) => text, |
| 193 | None => default_mcp_preview(tool, &args), |
| 194 | }; |
| 195 | let preview = vec![ToolCallContent::from(preview)]; |
| 196 | if !ctx.permit(tool_call_id, name, preview).await? { |
| 197 | return Ok(ToolOutcome::err("User rejected running this tool.")); |
| 198 | } |
| 199 | } |
| 200 | let on_progress = |line: String| { |
| 201 | let _ = ctx.update( |
| 202 | tool_call_id, |
| 203 | ToolCallUpdateFields::new().content(vec![ToolCallContent::from(line)]), |
| 204 | ); |
| 205 | }; |
| 206 | let out = ctx.mcp.call(name, args, &ctx.cancel, &on_progress).await?; |
| 207 | let mut outcome = if out.failed { |
| 208 | ToolOutcome::err(truncate(out.text)) |
| 209 | } else { |
| 210 | ToolOutcome::ok(truncate(out.text)) |
| 211 | }; |
| 212 | if let Some(structured) = &out.structured { |
| 213 | outcome.locations = profile |
| 214 | .mcp_result_locations(tool, structured) |
| 215 | .into_iter() |
| 216 | .filter(|p| p.is_absolute()) |
| 217 | .map(|p| absolutize(&p)) |
| 218 | .collect(); |
| 219 | } |
| 220 | Ok(outcome) |
| 221 | } |
| 222 | |
| 223 | /// The permission prompt for an MCP tool when the product has nothing better to say. |
| 224 | fn default_mcp_preview(tool: &ed_mcp::McpTool, args: &Value) -> String { |
| 225 | let mut text = format!("Run `{}` on the {} MCP server.", tool.tool, tool.server); |
| 226 | if args.as_object().is_some_and(|o| !o.is_empty()) { |
| 227 | let shown = serde_json::to_string_pretty(args).unwrap_or_default(); |
| 228 | text.push_str(&format!("\nArguments:\n```json\n{}\n```", truncate(shown))); |
| 229 | } |
| 230 | text |
| 231 | } |
| 232 | |
| 233 | /// The base directory for a tool call's `root` parameter: a workspace root or any directory |
| 234 | /// beneath one (so the model can pass e.g. `root/crate/src`). Paths that escape a root via |
| 235 | /// `..` are rejected. |
| 236 | pub fn base_for(cwd: &Path, roots: &[PathBuf], root: Option<&str>) -> Result<PathBuf> { |
| 237 | let Some(root) = root else { |
| 238 | return Ok(cwd.to_path_buf()); |
| 239 | }; |
| 240 | // Lexical normalization resolves `..`, so `root/../elsewhere` can't pass the prefix check. |
| 241 | let candidate = normalize(Path::new(root)); |
| 242 | let all_roots = all_roots(cwd, roots); |
| 243 | for ws_root in &all_roots { |
| 244 | if candidate == *ws_root { |
| 245 | return Ok(ws_root.clone()); |
| 246 | } |
| 247 | // Allow any directory below a root (e.g. a crate dir like `root/crate/src`), |
| 248 | // as long as it doesn't escape the root via `..`. |
| 249 | if let Ok(rest) = candidate.strip_prefix(ws_root) |
| 250 | && rest.components().all(|c| matches!(c, Component::Normal(_))) |
| 251 | { |
| 252 | return Ok(candidate); |
| 253 | } |
| 254 | } |
| 255 | let known = all_roots |
| 256 | .iter() |
| 257 | .map(|r| r.display().to_string()) |
| 258 | .collect::<Vec<_>>() |
| 259 | .join(", "); |
| 260 | bail!( |
| 261 | "unknown root `{}` (workspace roots: {known})", |
| 262 | candidate.display() |
| 263 | ) |
| 264 | } |
| 265 | |
| 266 | fn all_roots(cwd: &Path, roots: &[PathBuf]) -> Vec<PathBuf> { |
| 267 | std::iter::once(cwd.to_path_buf()) |
| 268 | .chain(roots.iter().filter(|r| *r != cwd).cloned()) |
| 269 | .collect() |
| 270 | } |
| 271 | |
| 272 | impl ToolCtx { |
| 273 | pub fn base_for(&self, root: Option<&str>) -> Result<PathBuf> { |
| 274 | base_for(&self.cwd, &self.roots, root) |
| 275 | } |
| 276 | |
| 277 | pub fn resolve_with(&self, root: Option<&str>, path: &str) -> Result<PathBuf> { |
| 278 | Ok(resolve_in(&self.base_for(root)?, path)) |
| 279 | } |
| 280 | |
| 281 | /// Ask the user for permission. Returns `Ok(false)` if rejected or cancelled. Honours |
| 282 | /// "always allow" and "always reject" choices for `tool` and the session's auto-approve. |
| 283 | pub async fn permit( |
| 284 | &self, |
| 285 | id: &str, |
| 286 | tool: &str, |
| 287 | content: Vec<ToolCallContent>, |
| 288 | ) -> Result<bool> { |
| 289 | if self.always_rejected.lock().unwrap().contains(tool) { |
| 290 | return Ok(false); |
| 291 | } |
| 292 | if self.auto_approve.load(Ordering::Relaxed) |
| 293 | || self.always_allowed.lock().unwrap().contains(tool) |
| 294 | { |
| 295 | return Ok(true); |
| 296 | } |
| 297 | let mut fields = ToolCallUpdateFields::new(); |
| 298 | if !content.is_empty() { |
| 299 | fields = fields.content(content); |
| 300 | } |
| 301 | let req = RequestPermissionRequest::new( |
| 302 | self.session_id.clone(), |
| 303 | ToolCallUpdate::new(id.to_string(), fields), |
| 304 | vec![ |
| 305 | PermissionOption::new("allow_once", "Allow", PermissionOptionKind::AllowOnce), |
| 306 | PermissionOption::new( |
| 307 | "allow_always", |
| 308 | "Always allow", |
| 309 | PermissionOptionKind::AllowAlways, |
| 310 | ), |
| 311 | PermissionOption::new("reject_once", "Reject", PermissionOptionKind::RejectOnce), |
| 312 | PermissionOption::new( |
| 313 | "reject_always", |
| 314 | "Always reject", |
| 315 | PermissionOptionKind::RejectAlways, |
| 316 | ), |
| 317 | ], |
| 318 | ); |
| 319 | let resp = tokio::select! { |
| 320 | r = self.connection.send_request(req).block_task() => r?, |
| 321 | () = self.cancel.cancelled() => return Ok(false), |
| 322 | }; |
| 323 | Ok(match resp.outcome { |
| 324 | RequestPermissionOutcome::Selected(sel) => match &*sel.option_id.0 { |
| 325 | "allow_always" => { |
| 326 | self.always_allowed.lock().unwrap().insert(tool.to_string()); |
| 327 | true |
| 328 | } |
| 329 | "allow_once" => true, |
| 330 | "reject_always" => { |
| 331 | self.always_rejected |
| 332 | .lock() |
| 333 | .unwrap() |
| 334 | .insert(tool.to_string()); |
| 335 | false |
| 336 | } |
| 337 | _ => false, |
| 338 | }, |
| 339 | _ => false, |
| 340 | }) |
| 341 | } |
| 342 | |
| 343 | pub fn update(&self, id: &str, fields: ToolCallUpdateFields) -> Result<()> { |
| 344 | let update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(id.to_string(), fields)); |
| 345 | self.connection |
| 346 | .send_notification(SessionNotification::new(self.session_id.clone(), update))?; |
| 347 | Ok(()) |
| 348 | } |
| 349 | } |
| 350 | |
| 351 | /// Join `path` onto `base`, keeping absolute paths as-is. |
| 352 | pub fn resolve_in(base: &Path, path: &str) -> PathBuf { |
| 353 | let p = Path::new(path); |
| 354 | if p.is_absolute() { |
| 355 | p.to_path_buf() |
| 356 | } else { |
| 357 | base.join(p) |
| 358 | } |
| 359 | } |
| 360 | |
| 361 | /// Cut `s` to [`MAX_OUTPUT_BYTES`] on a character boundary, marking the cut. |
| 362 | pub fn truncate(mut s: String) -> String { |
| 363 | if s.len() > MAX_OUTPUT_BYTES { |
| 364 | let mut cut = MAX_OUTPUT_BYTES; |
| 365 | while !s.is_char_boundary(cut) { |
| 366 | cut -= 1; |
| 367 | } |
| 368 | s.truncate(cut); |
| 369 | s.push_str("\n[truncated]"); |
| 370 | } |
| 371 | s |
| 372 | } |
| 373 | |
| 374 | /// Make `path` absolute and free of `.`/`..` components, as the ACP v1 spec requires |
| 375 | /// absolute paths everywhere. Existing paths are canonicalized (resolving symlinks |
| 376 | /// and giving the client the true on-disk location); paths that don't exist yet |
| 377 | /// are normalized lexically with the existing parent canonicalized when possible. |
| 378 | pub fn absolutize(path: &Path) -> PathBuf { |
| 379 | if let Ok(c) = std::fs::canonicalize(path) { |
| 380 | return c; |
| 381 | } |
| 382 | // Not on disk (yet): normalize lexically, anchoring at the canonicalized parent. |
| 383 | let parent = path.parent().filter(|p| !p.as_os_str().is_empty()); |
| 384 | let base = parent.and_then(|p| std::fs::canonicalize(p).ok()); |
| 385 | let file_name = path.file_name().map(|f| f.to_os_string()); |
| 386 | match (base, file_name) { |
| 387 | (Some(mut b), Some(f)) => { |
| 388 | b.push(f); |
| 389 | b |
| 390 | } |
| 391 | _ => normalize(path), |
| 392 | } |
| 393 | } |
| 394 | |
| 395 | /// Lexically normalize a path: resolve `.` and `..` without touching the filesystem. |
| 396 | /// Absolute paths stay absolute; relative paths are made absolute against the cwd. |
| 397 | pub fn normalize(path: &Path) -> PathBuf { |
| 398 | let absolute = if path.is_absolute() { |
| 399 | path.to_path_buf() |
| 400 | } else { |
| 401 | std::env::current_dir() |
| 402 | .unwrap_or_else(|_| PathBuf::from(".")) |
| 403 | .join(path) |
| 404 | }; |
| 405 | let mut out: Vec<std::ffi::OsString> = Vec::new(); |
| 406 | for c in absolute.components() { |
| 407 | match c { |
| 408 | Component::CurDir => {} |
| 409 | Component::ParentDir => { |
| 410 | out.pop(); |
| 411 | } |
| 412 | other => out.push(other.as_os_str().to_os_string()), |
| 413 | } |
| 414 | } |
| 415 | out.into_iter().collect() |
| 416 | } |
| 417 | |
| 418 | #[cfg(test)] |
| 419 | mod fs_path_tests { |
| 420 | use super::*; |
| 421 | |
| 422 | #[test] |
| 423 | fn normalize_resolves_dot_dot_lexically() { |
| 424 | assert_eq!(normalize(Path::new("/a/b/../c")), PathBuf::from("/a/c")); |
| 425 | assert_eq!(normalize(Path::new("/a/./b")), PathBuf::from("/a/b")); |
| 426 | assert_eq!(normalize(Path::new("/a/b/..")), PathBuf::from("/a")); |
| 427 | } |
| 428 | |
| 429 | #[test] |
| 430 | fn normalize_makes_relative_paths_absolute() { |
| 431 | let n = normalize(Path::new("src/../README.md")); |
| 432 | assert!(n.is_absolute()); |
| 433 | assert!(n.ends_with("README.md")); |
| 434 | assert!(!n.components().any(|c| c == Component::ParentDir)); |
| 435 | } |
| 436 | |
| 437 | #[test] |
| 438 | fn absolutize_canonicalizes_existing_paths() { |
| 439 | let n = absolutize(&std::env::temp_dir()); |
| 440 | assert!(n.is_absolute()); |
| 441 | assert!(!n.components().any(|c| c == Component::ParentDir)); |
| 442 | } |
| 443 | |
| 444 | #[test] |
| 445 | fn absolutize_normalizes_missing_paths() { |
| 446 | let n = absolutize(Path::new("/definitely/does/not/../exist.txt")); |
| 447 | assert_eq!(n, PathBuf::from("/definitely/does/exist.txt")); |
| 448 | } |
| 449 | |
| 450 | #[test] |
| 451 | fn base_for_rejects_escapes_and_accepts_subdirs() { |
| 452 | let cwd = PathBuf::from("/ws/a"); |
| 453 | let roots = vec![PathBuf::from("/ws/b")]; |
| 454 | assert_eq!(base_for(&cwd, &roots, None).unwrap(), cwd); |
| 455 | assert_eq!( |
| 456 | base_for(&cwd, &roots, Some("/ws/b/sub/../x")).unwrap(), |
| 457 | PathBuf::from("/ws/b/x") |
| 458 | ); |
| 459 | assert!(base_for(&cwd, &roots, Some("/ws/b/../c")).is_err()); |
| 460 | assert!(base_for(&cwd, &roots, Some("/etc")).is_err()); |
| 461 | } |
| 462 | } |