diff --git a/Cargo.toml b/Cargo.toml index 8b39d6f..d8473a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -141,6 +141,7 @@ arboard = { version = "3.3.0", default-features = false } [dev-dependencies] pretty_assertions = "1.4.0" +rmcp = { version = "3.1.2", features = ["server"] } serial_test = "3" [[bin]] diff --git a/README.md b/README.md index d484d9c..0d32e56 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,9 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g * [Create Custom TypeScript Tools](https://github.com/Dark-Alex-17/coyote/wiki/Custom-Tools#custom-typescript-based-tools) * [Create Custom Bash Tools](https://github.com/Dark-Alex-17/coyote/wiki/Custom-Bash-Tools) * [Bash Prompt Utilities](https://github.com/Dark-Alex-17/coyote/wiki/Bash-Prompt-Helpers) -* [First-Class MCP Server Support](https://github.com/Dark-Alex-17/coyote/wiki/MCP-Servers): Easily connect and interact with MCP servers for advanced functionality. +* [First-Class MCP Server Support](https://github.com/Dark-Alex-17/coyote/wiki/MCP-Servers): Easily connect and interact with MCP servers for advanced functionality. Coyote supports all three MCP capabilities: tools, resources, and prompts. + * Models interact with each server through a compact set of capability-gated meta-tools: `mcp_search`/`mcp_describe` for discovery across tools, resources, and prompts, `mcp_invoke` for tool calls, `mcp_read` for paged and regex-filterable resource reads, and `mcp_prompt` for server-defined prompts. Binary content is spilled to disk instead of inlined, and oversized tool results are bounded before they reach the model. + * Invoke server prompts yourself with `.prompt [key=value ...]` in the REPL, with live staged tab-completion (servers, then prompt names, then `key=` arguments), and discover them with `.list prompts`. * [Macros](https://github.com/Dark-Alex-17/coyote/wiki/Macros): Automate repetitive tasks and workflows with Coyote "scripts" (macros). Macros are Coyote's custom commands: invoke any macro directly by name (e.g. `.review main`), with tab-completion, right alongside the built-in REPL commands. * Give a macro a `description` (shown in `.list macros` and completions) and set `isolated: false` to run its steps on the live session, exactly as if you typed them. Note that non-isolated steps are recorded in the session, and mutating steps (`.role`, `.model`, ...) persist after the macro ends — by design. Steps are fail-fast: an error aborts the remaining steps, but completed steps' effects remain. A non-isolated macro step cannot invoke another macro, and a `.exit` step never exits the REPL. * Commit project-specific macros to `.coyote/macros/` in your repo — they shadow same-named global macros (opt out with `--no-workspace-macros`). diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 1114217..6c4e456 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -47,7 +47,7 @@ pub enum McpScopeArg { .args(["sandbox", "fresh"]) .multiple(true) .conflicts_with_all([ - "model", "prompt", "role", "session", "agent", "rag", "rebuild_rag", + "model", "temp_role", "role", "session", "agent", "rag", "rebuild_rag", "macro_name", "execute", "code", "file", "no_stream", "no_memory", "init_memory", "dry_run", "info", "build_tools", "install", "install_builtins", "sync_models", "list_models", "list_roles", @@ -70,9 +70,9 @@ pub struct Cli { /// Select a LLM model #[arg(short, long, add = ArgValueCompleter::new(model_completer))] pub model: Option, - /// Use the system prompt + /// Set a temporary role (an ad-hoc system prompt) for this invocation #[arg(long)] - pub prompt: Option, + pub temp_role: Option, /// Select a role #[arg(short, long, add = ArgValueCompleter::new(role_completer))] pub role: Option, @@ -705,9 +705,9 @@ mod tests { } #[test] - fn parse_prompt_flag() { - let cli = parse(&["--prompt", "be a pirate"]); - assert_eq!(cli.prompt, Some("be a pirate".to_string())); + fn parse_temp_role_flag() { + let cli = parse(&["--temp-role", "be a pirate"]); + assert_eq!(cli.temp_role, Some("be a pirate".to_string())); } #[test] diff --git a/src/config/agent.rs b/src/config/agent.rs index 9a6b83f..40771fa 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -15,6 +15,7 @@ use crate::config::prompts::{ }; use crate::graph::types::RagNode; use crate::graph::{Graph, GraphParser, NodeType}; +use crate::mcp::McpServerFeatures; use crate::rag::RagInitConfig; use crate::vault::SECRET_RE; use anyhow::{Context, Result}; @@ -380,7 +381,7 @@ impl Agent { self.graph_rags.get(node_id).cloned() } - pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { + pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { self.functions.append_mcp_meta_functions(mcp_servers); } diff --git a/src/config/app_state.rs b/src/config/app_state.rs index c1f4769..7a6cb0e 100644 --- a/src/config/app_state.rs +++ b/src/config/app_state.rs @@ -70,7 +70,7 @@ impl AppState { let mut functions = Functions::init(config.visible_tools.as_ref().unwrap_or(&Vec::new()))?; if !mcp_registry.is_empty() && config.mcp_server_support { - functions.append_mcp_meta_functions(mcp_registry.list_started_servers()); + functions.append_mcp_meta_functions(mcp_registry.server_features()); } let mcp_registry = if mcp_registry.is_empty() { diff --git a/src/config/bundles.rs b/src/config/bundles.rs index a378f7e..69b181d 100644 --- a/src/config/bundles.rs +++ b/src/config/bundles.rs @@ -442,7 +442,9 @@ impl BundleStore { self.ensure_bundle_exists(bundle)?; for (name, record) in self.bundles.iter_mut() { if name != bundle { - record.files.retain(|owned| owned.path != file.path); + record + .files + .retain(|owned| !same_installed_path(&owned.path, &file.path)); } } let record = self @@ -450,7 +452,9 @@ impl BundleStore { .get_mut(bundle) .expect("bundle existence checked above"); - record.files.retain(|owned| owned.path != file.path); + record + .files + .retain(|owned| !same_installed_path(&owned.path, &file.path)); record.files.push(file); self.save() @@ -577,6 +581,17 @@ pub(crate) struct BundleListRow { pub(crate) drift: DriftSummary, } +/// NTFS and default APFS resolve file names case-insensitively, so records +/// differing only in case denote the same physical file there. Linux keeps +/// exact matching because case variants are genuinely distinct files. +fn same_installed_path(a: &str, b: &str) -> bool { + if cfg!(any(windows, target_os = "macos")) { + a.eq_ignore_ascii_case(b) + } else { + a == b + } +} + /// An unreadable file counts as locally modified: it exists but its integrity /// cannot be verified. pub(crate) fn bundle_list_rows(store: &BundleStore, config_dir: &Path) -> Vec { @@ -1489,4 +1504,53 @@ mod tests { assert!(rows.is_empty()); } + + #[test] + fn same_installed_path_matches_filesystem_case_semantics() { + assert!(same_installed_path("macros/a.yaml", "macros/a.yaml")); + assert!(!same_installed_path("macros/a.yaml", "macros/b.yaml")); + assert_eq!( + same_installed_path("macros/Foo.yaml", "macros/foo.yaml"), + cfg!(any(windows, target_os = "macos")) + ); + } + + #[test] + fn record_file_transfers_case_variant_ownership_on_case_insensitive_hosts() { + let dir = TempStoreDir::new("bundles-case-variant"); + let mut store = dir.store(); + store + .upsert_bundle("alpha", metadata("https://x/a", "aaa")) + .unwrap(); + store + .upsert_bundle("beta", metadata("https://x/b", "bbb")) + .unwrap(); + store + .record_file("alpha", file_record("macros/Shared.yaml", "one")) + .unwrap(); + + store + .record_file("beta", file_record("macros/shared.yaml", "two")) + .unwrap(); + + let alpha_still_owns = store + .get("alpha") + .unwrap() + .files + .iter() + .any(|f| f.path == "macros/Shared.yaml"); + assert_eq!( + alpha_still_owns, + !cfg!(any(windows, target_os = "macos")), + "case-variant paths are one physical file on case-insensitive hosts" + ); + assert!( + store + .get("beta") + .unwrap() + .files + .iter() + .any(|f| f.path == "macros/shared.yaml") + ); + } } diff --git a/src/config/install_remote.rs b/src/config/install_remote.rs index 6a43966..caaadcb 100644 --- a/src/config/install_remote.rs +++ b/src/config/install_remote.rs @@ -752,9 +752,29 @@ fn select_uninstall_candidate(store: &BundleStore, spec: &str) -> Result bool { let recorded = Path::new(path); !recorded.is_absolute() - && recorded - .components() - .all(|c| matches!(c, Component::Normal(_))) + && recorded.components().all(|c| match c { + Component::Normal(name) => is_safe_component(&name.to_string_lossy()), + _ => false, + }) +} + +/// Rejects names Windows refuses or silently rewrites (alternate data stream +/// colons, reserved device names, trailing dots or spaces) so a recorded path +/// denotes the same regular file on every platform. +fn is_safe_component(name: &str) -> bool { + !name.contains(':') + && !name.ends_with('.') + && !name.ends_with(' ') + && !is_windows_reserved_name(name) +} + +fn is_windows_reserved_name(name: &str) -> bool { + let stem = name.split('.').next().unwrap_or(""); + let lower = stem.to_ascii_lowercase(); + matches!(lower.as_str(), "con" | "prn" | "aux" | "nul") + || (lower.len() == 4 + && (lower.starts_with("com") || lower.starts_with("lpt")) + && matches!(lower.as_bytes()[3], b'1'..=b'9')) } fn uninstall_owned_files( @@ -1067,7 +1087,12 @@ impl TempRepoDir { impl Drop for TempRepoDir { fn drop(&mut self) { - let _ = fs::remove_dir_all(&self.path); + if let Err(error) = fs::remove_dir_all(&self.path) { + log::warn!( + "failed to remove temp clone {}: {error}", + self.path.display() + ); + } } } @@ -1079,13 +1104,37 @@ fn is_commit_sha(reference: &str) -> bool { fn clone_to_temp(url: &str, reference: Option<&str>) -> Result { let dest = utils::temp_file("coyote-remote-install-", ""); + match clone_into(&dest, url, reference) { + Ok(head_sha) => Ok(TempRepoDir { + path: dest, + head_sha, + }), + Err(error) => { + let _ = fs::remove_dir_all(&dest); + Err(error) + } + } +} + +/// Checked-out bytes must not depend on the machine's git configuration: +/// recorded sha256 provenance would otherwise drift with autocrlf settings. +/// Long paths are opted into for deep bundle trees on Windows. +fn git_content_config() -> Vec { + ["core.autocrlf=false", "core.eol=lf", "core.longpaths=true"] + .iter() + .flat_map(|setting| ["-c".into(), (*setting).into()]) + .collect() +} + +fn clone_into(dest: &Path, url: &str, reference: Option<&str>) -> Result { let dest_arg: OsString = dest.as_os_str().into(); let is_sha = reference.is_some_and(is_commit_sha); match reference { Some(r) if !is_sha => { - run_git(vec![ + let mut args = git_content_config(); + args.extend([ "clone".into(), "--depth".into(), "1".into(), @@ -1094,26 +1143,28 @@ fn clone_to_temp(url: &str, reference: Option<&str>) -> Result { "--".into(), url.into(), dest_arg, - ])?; + ]); + run_git(args)?; } Some(r) => { - run_git(vec![ - "clone".into(), - "--".into(), - url.into(), - dest_arg.clone(), - ])?; - run_git(vec!["-C".into(), dest_arg, "checkout".into(), r.into()])?; + let mut args = git_content_config(); + args.extend(["clone".into(), "--".into(), url.into(), dest_arg.clone()]); + run_git(args)?; + let mut args = git_content_config(); + args.extend(["-C".into(), dest_arg, "checkout".into(), r.into()]); + run_git(args)?; } None => { - run_git(vec![ + let mut args = git_content_config(); + args.extend([ "clone".into(), "--depth".into(), "1".into(), "--".into(), url.into(), dest_arg, - ])?; + ]); + run_git(args)?; } } @@ -1123,11 +1174,7 @@ fn clone_to_temp(url: &str, reference: Option<&str>) -> Result { "rev-parse".into(), "HEAD".into(), ])?; - - Ok(TempRepoDir { - path: dest, - head_sha, - }) + Ok(head_sha) } fn run_git(args: Vec) -> Result<()> { @@ -1816,7 +1863,17 @@ fn record_written_file( } fn provenance_path(dst: &Path) -> String { - let rel = dst.strip_prefix(paths::config_dir()).unwrap_or(dst); + let rel = match dst.strip_prefix(paths::config_dir()) { + Ok(rel) => rel, + Err(_) => { + log::warn!( + "bundle file {} lies outside the config dir (an asset dir override?); \ + it will not be uninstallable and drift checks may misreport it", + dst.display() + ); + dst + } + }; rel.to_string_lossy().replace('\\', "/") } @@ -2242,6 +2299,53 @@ mod tests { use std::env; use std::time::{SystemTime, UNIX_EPOCH}; + #[test] + fn safe_relative_path_accepts_plain_portable_components() { + assert!(is_safe_relative_path("macros/a.yaml")); + assert!(is_safe_relative_path("skills/deep/nested/file.md")); + assert!(is_safe_relative_path("functions/tools/console.sh")); + assert!(is_safe_relative_path("roles/common.md")); + } + + #[test] + fn safe_relative_path_rejects_escapes_and_windows_hazards() { + assert!(!is_safe_relative_path("../outside.yaml")); + assert!(!is_safe_relative_path("/abs/path.yaml")); + assert!(!is_safe_relative_path("macros/../../evil.yaml")); + assert!(!is_safe_relative_path("macros/a.yaml:stream")); + assert!(!is_safe_relative_path("macros/trailing.")); + assert!(!is_safe_relative_path("macros/trailing ")); + assert!(!is_safe_relative_path("macros/nul")); + assert!(!is_safe_relative_path("macros/NUL.yaml")); + assert!(!is_safe_relative_path("con/a.yaml")); + assert!(!is_safe_relative_path("macros/COM1.txt")); + assert!(!is_safe_relative_path("macros/lpt9")); + } + + #[test] + fn windows_reserved_name_check_is_stem_based() { + assert!(is_windows_reserved_name("nul")); + assert!(is_windows_reserved_name("NUL.txt")); + assert!(is_windows_reserved_name("com1")); + assert!(!is_windows_reserved_name("com0")); + assert!(!is_windows_reserved_name("com10")); + assert!(!is_windows_reserved_name("console")); + assert!(!is_windows_reserved_name("nullable.yaml")); + } + + #[test] + fn git_content_config_pins_line_endings_and_long_paths() { + let args = git_content_config(); + let rendered: Vec = args + .iter() + .map(|a| a.to_string_lossy().into_owned()) + .collect(); + assert_eq!(rendered.len(), 6); + assert!(rendered.contains(&"core.autocrlf=false".to_string())); + assert!(rendered.contains(&"core.eol=lf".to_string())); + assert!(rendered.contains(&"core.longpaths=true".to_string())); + } + struct TestVaultConfigGuard { dir_key: String, file_key: String, @@ -4286,6 +4390,12 @@ mod tests { #[test] #[serial] fn uninstall_shorthand_with_multiple_matches_bails_non_interactively() { + if *IS_STDOUT_TERMINAL { + eprintln!( + "Skipping uninstall_shorthand_with_multiple_matches_bails_non_interactively: requires non-TTY stdout" + ); + return; + } let _guard = TestVaultConfigGuard::new("uninst-short-multi"); let mut store = BundleStore::load().unwrap(); store diff --git a/src/config/macro_policy.rs b/src/config/macro_policy.rs index 4c08ed7..f9b97cc 100644 --- a/src/config/macro_policy.rs +++ b/src/config/macro_policy.rs @@ -287,6 +287,7 @@ fn discover_macros_in(dirs: &[(MacroSource, PathBuf)]) -> Vec { #[cfg(test)] mod tests { use super::*; + use crate::repl; use crate::utils::get_env_name; use serial_test::serial; use std::path::Path; @@ -565,6 +566,34 @@ mod tests { assert_eq!(state_of(&policy, "a"), &MacroState::Enabled); } + #[test] + fn prompt_macro_is_shadowed_by_builtin() { + let policy = MacroPolicy::effective_with( + globals(&["prompt"]), + None, + None, + None, + None, + &repl::builtin_command_names(), + ); + + assert_eq!(state_of(&policy, "prompt"), &MacroState::ShadowedBuiltin); + } + + #[test] + fn temp_role_macro_is_shadowed_by_builtin() { + let policy = MacroPolicy::effective_with( + globals(&["temp-role"]), + None, + None, + None, + None, + &repl::builtin_command_names(), + ); + + assert_eq!(state_of(&policy, "temp-role"), &MacroState::ShadowedBuiltin); + } + #[test] fn locked_wins_over_shadowed_builtin() { let l = list(&["a"]); diff --git a/src/config/mod.rs b/src/config/mod.rs index f1749ce..08992fd 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -22,6 +22,8 @@ pub(crate) mod todo; mod tool_scope; mod update; +#[cfg(test)] +pub(crate) use self::agent::AgentConfig; pub use self::agent::{ Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents, list_agents_with_descriptions, @@ -51,6 +53,11 @@ pub use self::skill::Skill; pub use self::skill_policy::SkillPolicy; #[allow(unused_imports)] pub use self::skill_registry::SkillRegistry; +#[cfg(test)] +pub(crate) use self::tool_scope::test_fixtures; +pub use self::tool_scope::{ + McpPromptCompletion, flatten_prompt_messages, resolve_prompt_args, sanitize_display_text, +}; pub use self::update::run_self_update; use crate::client::{ self, ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS, diff --git a/src/config/request_context.rs b/src/config/request_context.rs index bd04c47..2be3d37 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -5,7 +5,7 @@ use super::skill::{SKILL_SCAFFOLD, Skill}; use super::skill_policy::SkillPolicy; use super::skill_registry::SkillRegistry; use super::todo::TodoList; -use super::tool_scope::{McpRuntime, ToolScope}; +use super::tool_scope::{McpPromptCompletion, McpRuntime, ToolScope, format_prompt_arguments}; use super::{ AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE, Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, MacroAllowlistLevel, @@ -23,8 +23,8 @@ use crate::function::{ user_interaction::USER_FUNCTION_PREFIX, }; use crate::mcp::{ - MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, - MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error, + CatalogItem, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, + McpServersConfig, is_auth_required_error, is_mcp_meta_function, mcp_meta_function_names, }; use crate::rag::Rag; use crate::supervisor::Supervisor; @@ -77,6 +77,31 @@ fn installed_bundle_names() -> Vec { } } +pub(crate) fn expand_enabled_mcp_server_ids( + app: &AppConfig, + mcp_config: &McpServersConfig, + enabled_mcp_servers: &[String], +) -> Vec { + if enabled_mcp_servers.iter().any(|s| s.trim() == "all") { + return mcp_config.mcp_servers.keys().cloned().collect(); + } + + let mut ids = Vec::new(); + for item in enabled_mcp_servers.iter().map(|s| s.trim()) { + if mcp_config.mcp_servers.contains_key(item) { + ids.push(item.to_string()); + } else if let Some(mapped) = app.mapping_mcp_servers.get(item) { + for mapped_id in mapped.split(',').map(|s| s.trim()) { + if mcp_config.mcp_servers.contains_key(mapped_id) { + ids.push(mapped_id.to_string()); + } + } + } + } + + ids +} + pub struct AutoContinueConfig { pub enabled: bool, pub max_continues: usize, @@ -139,6 +164,20 @@ pub(crate) fn asset_table(header: &[&str]) -> Table { table } +fn mcp_prompt_rows(items: &[CatalogItem]) -> Vec<[String; 4]> { + items + .iter() + .map(|item| { + [ + item.server.clone(), + item.name.clone(), + item.description.clone(), + format_prompt_arguments(item.arguments.as_deref().unwrap_or_default()), + ] + }) + .collect() +} + fn complete_skills_with_descriptions(names: Vec) -> Vec<(String, Option)> { names .into_iter() @@ -2011,11 +2050,7 @@ impl RequestContext { .functions .declarations() .iter() - .filter(|v| { - !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| !is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()) .collect(); @@ -2025,11 +2060,7 @@ impl RequestContext { .functions() .declarations() .iter() - .filter(|v| { - !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| !is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()), ); } @@ -2102,11 +2133,7 @@ impl RequestContext { .declarations() .to_vec() .into_iter() - .filter(|v| { - !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| !is_mcp_meta_function(&v.name)) .collect(); if let Some(ref tool_names) = role_filter { @@ -2155,11 +2182,7 @@ impl RequestContext { .functions .declarations() .iter() - .filter(|v| { - v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()) .collect(); if let Some(agent) = &self.agent { @@ -2168,12 +2191,7 @@ impl RequestContext { .functions() .declarations() .iter() - .filter(|v| { - v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || v.name - .starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()), ); } @@ -2188,41 +2206,17 @@ impl RequestContext { continue; } - let item_invoke_name = - format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); let item_search_name = format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX); - let item_describe_name = - format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX); if let Some(values) = app.mapping_mcp_servers.get(item) { server_names.extend( values .split(',') - .flat_map(|v| { - vec![ - format!( - "{}_{}", - MCP_INVOKE_META_FUNCTION_NAME_PREFIX, - v.to_string() - ), - format!( - "{}_{}", - MCP_SEARCH_META_FUNCTION_NAME_PREFIX, - v.to_string() - ), - format!( - "{}_{}", - MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, - v.to_string() - ), - ] - }) + .flat_map(mcp_meta_function_names) .filter(|v| mcp_declaration_names.contains(v)), ) - } else if mcp_declaration_names.contains(&item_invoke_name) { - server_names.insert(item_invoke_name); - server_names.insert(item_search_name); - server_names.insert(item_describe_name); + } else if mcp_declaration_names.contains(&item_search_name) { + server_names.extend(mcp_meta_function_names(item)); } } } @@ -2251,11 +2245,7 @@ impl RequestContext { .declarations() .to_vec() .into_iter() - .filter(|v| { - v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| is_mcp_meta_function(&v.name)) .collect(); if let Some(ref server_names) = role_filter { @@ -2351,7 +2341,7 @@ impl RequestContext { Ok(()) } - pub fn use_prompt(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> { + pub fn use_temp_role(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> { let mut role = Role::new(TEMP_ROLE_NAME, prompt); role.set_model(self.current_model().clone()); self.use_role_obj(role) @@ -2828,7 +2818,7 @@ impl RequestContext { } "bundles" => bundles::list_installed_bundles(), _ => bail!( - "Unknown kind '{kind}'. Valid kinds: roles, sessions, agents, rags, macros, skills, tools, mcp-servers, bundles" + "Unknown kind '{kind}'. Valid kinds: roles, sessions, agents, rags, macros, skills, prompts, tools, mcp-servers, bundles" ), } } @@ -3396,6 +3386,7 @@ impl RequestContext { "rags", "macros", "skills", + "prompts", "tools", "mcp-servers", "bundles", @@ -3720,6 +3711,40 @@ impl RequestContext { fuzzy_filter(values, |v| v.0.as_str(), filter) } + pub fn mcp_prompt_completion(&self, args: &[&str]) -> McpPromptCompletion { + let app = self.app.config.as_ref(); + let enabled_ids = match &self.app.mcp_config { + Some(mcp_config) => { + let mut servers = self + .enabled_mcp_servers_for_current_scope(app, true) + .unwrap_or_default(); + servers.extend(self.skill_registry.loaded_mcp_servers()); + expand_enabled_mcp_server_ids(app, mcp_config, &servers) + } + None => vec![], + }; + self.tool_scope + .mcp_runtime + .prompt_completion(&enabled_ids, args) + } + + pub async fn list_mcp_prompts(&self) -> Result<()> { + let items = self.tool_scope.mcp_runtime.prompt_catalog().await; + if items.is_empty() { + println!("No prompts found."); + return Ok(()); + } + + let mut table = asset_table(&["server", "name", "description", "args"]); + for row in mcp_prompt_rows(&items) { + table.add_row(row.to_vec()); + } + + println!("Prompts:"); + println!("{table}"); + Ok(()) + } + async fn rebuild_tool_scope( &mut self, app: &AppConfig, @@ -3763,24 +3788,7 @@ impl RequestContext { && let Some(mcp_config) = &self.app.mcp_config { let server_ids: Vec = match &enabled_mcp_servers { - Some(servers) if servers.iter().any(|s| s.trim() == "all") => { - mcp_config.mcp_servers.keys().cloned().collect() - } - Some(servers) => { - let mut ids = Vec::new(); - for item in servers.iter().map(|s| s.trim()) { - if mcp_config.mcp_servers.contains_key(item) { - ids.push(item.to_string()); - } else if let Some(mapped) = app.mapping_mcp_servers.get(item) { - for mapped_id in mapped.split(',').map(|s| s.trim()) { - if mcp_config.mcp_servers.contains_key(mapped_id) { - ids.push(mapped_id.to_string()); - } - } - } - } - ids - } + Some(servers) => expand_enabled_mcp_server_ids(app, mcp_config, servers), None => vec![], }; @@ -3836,7 +3844,7 @@ impl RequestContext { functions.append_todo_functions(); } if !mcp_runtime.is_empty() { - functions.append_mcp_meta_functions(mcp_runtime.server_names()); + functions.append_mcp_meta_functions(mcp_runtime.server_features()); } if app.function_calling_support && policy.skills_enabled { functions.append_skill_functions(); @@ -4685,10 +4693,11 @@ mod tests { use crate::config::AppState; use crate::config::agent::AgentConfig; use crate::function::{ToolCall, skill}; - use crate::mcp::{McpServer, McpServersConfig, McpTransportType}; + use crate::mcp::{McpServer, McpServerFeatures, McpServersConfig, McpTransportType}; use crate::utils; use crate::utils::get_env_name; use crate::vault::Vault; + use rmcp::model::PromptArgument; use serde_json::json; use serial_test::serial; use std::env; @@ -4746,6 +4755,15 @@ mod tests { RequestContext::new(default_app_state(), WorkingMode::Cmd) } + fn tools_only_features(name: &str) -> McpServerFeatures { + McpServerFeatures { + name: name.to_string(), + tools: true, + resources: false, + prompts: false, + } + } + #[test] fn new_creates_clean_state() { let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); @@ -4834,10 +4852,10 @@ mod tests { } #[test] - fn use_prompt_creates_temp_role() { + fn use_temp_role_creates_temp_role() { let mut ctx = create_test_ctx(); let app = ctx.app.config.clone(); - ctx.use_prompt(&app, "you are a pirate").unwrap(); + ctx.use_temp_role(&app, "you are a pirate").unwrap(); assert!(ctx.role.is_some()); assert_eq!(ctx.role.as_ref().unwrap().name(), "temp"); assert!( @@ -4849,6 +4867,41 @@ mod tests { ); } + #[test] + fn mcp_prompt_rows_assembles_columns() { + let items = vec![ + CatalogItem { + name: "summarize".to_string(), + server: "docs".to_string(), + description: "Summarize a document".to_string(), + arguments: Some(vec![ + PromptArgument::new("path").with_required(true), + PromptArgument::new("style"), + ]), + ..Default::default() + }, + CatalogItem { + name: "greet".to_string(), + server: "misc".to_string(), + ..Default::default() + }, + ]; + + let rows = mcp_prompt_rows(&items); + + assert_eq!(rows.len(), 2); + assert_eq!( + rows[0], + [ + "docs", + "summarize", + "Summarize a document", + "path (required), style" + ] + ); + assert_eq!(rows[1], ["misc", "greet", "", ""]); + } + #[test] fn extract_role_returns_standalone_role() { let mut ctx = create_test_ctx(); @@ -5753,9 +5806,10 @@ mod tests { #[test] fn select_enabled_mcp_servers_all_returns_all_mcp_functions() { let mut ctx = create_test_ctx(); - ctx.tool_scope - .functions - .append_mcp_meta_functions(vec!["github".into(), "slack".into()]); + ctx.tool_scope.functions.append_mcp_meta_functions(vec![ + tools_only_features("github"), + tools_only_features("slack"), + ]); let mut role = Role::new("r", "p"); role.set_enabled_mcp_servers(Some(vec!["all".to_string()])); @@ -5770,9 +5824,10 @@ mod tests { #[test] fn select_enabled_mcp_servers_comma_filters() { let mut ctx = create_test_ctx(); - ctx.tool_scope - .functions - .append_mcp_meta_functions(vec!["github".into(), "slack".into()]); + ctx.tool_scope.functions.append_mcp_meta_functions(vec![ + tools_only_features("github"), + tools_only_features("slack"), + ]); let mut role = Role::new("r", "p"); role.set_enabled_mcp_servers(Some(vec!["github".to_string()])); @@ -5783,6 +5838,28 @@ mod tests { assert!(!names.contains(&"mcp_invoke_slack")); } + #[test] + fn select_enabled_mcp_servers_keeps_resources_only_server() { + let mut ctx = create_test_ctx(); + ctx.tool_scope + .functions + .append_mcp_meta_functions(vec![McpServerFeatures { + name: "res".to_string(), + tools: false, + resources: true, + prompts: false, + }]); + + let mut role = Role::new("r", "p"); + role.set_enabled_mcp_servers(Some(vec!["res".to_string()])); + + let fns = ctx.select_enabled_mcp_servers(&role); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!(names.contains(&"mcp_search_res")); + assert!(names.contains(&"mcp_describe_res")); + assert!(!names.contains(&"mcp_invoke_res")); + } + #[test] fn state_empty_context_has_no_context_flags() { let ctx = create_test_ctx(); diff --git a/src/config/tool_scope.rs b/src/config/tool_scope.rs index 9836ca5..9939bd7 100644 --- a/src/config/tool_scope.rs +++ b/src/config/tool_scope.rs @@ -1,9 +1,13 @@ use crate::function::{Functions, ToolCallTracker}; -use crate::mcp::{CatalogItem, ConnectedServer, McpRegistry}; +use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry, McpServerFeatures}; use anyhow::{Context, Result, anyhow}; use bm25::{Document, Language, SearchEngineBuilder}; -use rmcp::model::{CallToolRequestParams, CallToolResult}; +use rmcp::model::{ + Annotations, CallToolRequestParams, CallToolResult, ContentBlock, GetPromptRequestParams, + GetPromptResult, Prompt, PromptArgument, PromptMessage, ReadResourceRequestParams, + ReadResourceResult, Resource, ResourceTemplate, Role, Tool, +}; use serde_json::{Value, json}; use std::collections::HashMap; use std::sync::Arc; @@ -25,6 +29,18 @@ impl Default for ToolScope { } } +pub enum McpPromptCompletion { + Ready(Vec<(String, Option)>), + PromptNames { + server: Arc, + }, + ArgumentKeys { + server: Arc, + prompt: String, + typed_keys: Vec, + }, +} + #[derive(Default, Clone)] pub struct McpRuntime { pub servers: HashMap>, @@ -47,8 +63,20 @@ impl McpRuntime { self.servers.get(name) } - pub fn server_names(&self) -> Vec { - self.servers.keys().cloned().collect() + pub fn server_features(&self) -> Vec { + let mut features: Vec = self + .servers + .iter() + .map(|(name, handle)| { + let info = handle.peer_info(); + McpServerFeatures::from_capabilities( + name.as_str(), + info.as_ref().map(|info| &info.capabilities), + ) + }) + .collect(); + features.sort_by(|a, b| a.name.cmp(&b.name)); + features } pub fn sync_from_registry(&mut self, registry: &McpRegistry) { @@ -63,16 +91,58 @@ impl McpRuntime { .get(server) .cloned() .with_context(|| format!("{server} MCP server not found in runtime"))?; - let tools = server_handle.list_tools(None).await?; + let info = server_handle.peer_info(); + let features = McpServerFeatures::from_capabilities( + server, + info.as_ref().map(|info| &info.capabilities), + ); let mut items = HashMap::new(); - for tool in tools.tools { - let item = CatalogItem { - name: tool.name.to_string(), - server: server.to_string(), - description: tool.description.unwrap_or_default().to_string(), - }; - items.insert(item.name.clone(), item); + if features.tools { + match server_handle.list_all_tools().await { + Ok(tools) => merge_catalog_items( + &mut items, + tools + .into_iter() + .map(|tool| tool_catalog_item(server, tool)), + ), + Err(e) => warn!("Failed to list tools on MCP server {server}: {e}"), + } + } + + if features.resources { + match server_handle.list_all_resources().await { + Ok(resources) => merge_catalog_items( + &mut items, + resources + .into_iter() + .map(|resource| resource_catalog_item(server, resource)), + ), + Err(e) => warn!("Failed to list resources on MCP server {server}: {e}"), + } + match server_handle.list_all_resource_templates().await { + Ok(templates) => merge_catalog_items( + &mut items, + templates + .into_iter() + .map(|template| resource_template_catalog_item(server, template)), + ), + Err(e) => { + warn!("Failed to list resource templates on MCP server {server}: {e}") + } + } + } + + if features.prompts { + match server_handle.list_all_prompts().await { + Ok(prompts) => merge_catalog_items( + &mut items, + prompts + .into_iter() + .map(|prompt| prompt_catalog_item(server, prompt)), + ), + Err(e) => warn!("Failed to list prompts on MCP server {server}: {e}"), + } } Ok(items) @@ -85,11 +155,15 @@ impl McpRuntime { top_k: usize, ) -> Result> { let items = self.catalog_items(server).await?; - let docs = items.values().map(|item| Document { - id: item.name.clone(), + let docs = items.iter().map(|(key, item)| Document { + id: key.clone(), contents: format!( - "{}\n{}\nserver:{}", - item.name, item.description, item.server + "{}\n{}\n{}\n{}\nserver:{}", + item.name, + item.description, + item.kind, + item.uri.as_deref().unwrap_or_default(), + item.server ), }); let engine = SearchEngineBuilder::::with_documents(Language::English, docs).build(); @@ -103,30 +177,113 @@ impl McpRuntime { .collect()) } - pub async fn describe(&self, server: &str, tool: &str) -> Result { + /// Best-effort audience lookup for a catalog resource: a listing failure + /// or an unknown uri (e.g. template-expanded) yields `None`. + pub async fn resource_audience(&self, server: &str, uri: &str) -> Option> { + let items = self.catalog_items(server).await.ok()?; + items + .into_values() + .find(|item| item.uri.as_deref() == Some(uri)) + .and_then(|item| item.audience) + } + + pub async fn describe(&self, server: &str, kind: &str, tool: &str) -> Result { let server_handle = self .get(server) .cloned() .with_context(|| format!("{server} MCP server not found in runtime"))?; - let tool_schema = server_handle - .list_tools(None) - .await? - .tools - .into_iter() - .find(|item| item.name == tool) - .ok_or_else(|| anyhow!("{tool} not found in {server} MCP server catalog"))? - .input_schema; + match kind { + "tool" => { + let tool_schema = server_handle + .list_all_tools() + .await? + .into_iter() + .find(|item| item.name == tool) + .ok_or_else(|| anyhow!("{tool} not found in {server} MCP server catalog"))? + .input_schema; - Ok(json!({ - "type": "object", - "properties": { - "tool": { - "type": "string", - }, - "arguments": tool_schema + Ok(json!({ + "type": "object", + "properties": { + "tool": { + "type": "string", + }, + "arguments": tool_schema + } + })) } - })) + "resource" => { + let resource = server_handle + .list_all_resources() + .await? + .into_iter() + .find(|item| item.uri == tool) + .ok_or_else(|| { + anyhow!("{tool} not found in {server} MCP server resource catalog") + })?; + + Ok(json!({ + "uri": resource.uri, + "name": resource.name, + "title": resource.title, + "description": resource.description, + "mime_type": resource.mime_type, + "size": resource.size, + })) + } + "resource_template" => { + let template = server_handle + .list_all_resource_templates() + .await? + .into_iter() + .find(|item| item.uri_template == tool) + .ok_or_else(|| { + anyhow!("{tool} not found in {server} MCP server resource template catalog") + })?; + + Ok(json!({ + "uri_template": template.uri_template, + "name": template.name, + "title": template.title, + "description": template.description, + "mime_type": template.mime_type, + "variables": uri_template_variables(&template.uri_template), + })) + } + "prompt" => { + let prompt = server_handle + .list_all_prompts() + .await? + .into_iter() + .find(|item| item.name == tool) + .ok_or_else(|| { + anyhow!("{tool} not found in {server} MCP server prompt catalog") + })?; + + let arguments: Vec = prompt + .arguments + .unwrap_or_default() + .into_iter() + .map(|arg| { + json!({ + "name": arg.name, + "description": arg.description, + "required": arg.required, + }) + }) + .collect(); + + Ok(json!({ + "name": prompt.name, + "description": prompt.description, + "arguments": arguments, + })) + } + other => Err(anyhow!( + "Unknown kind '{other}'. Valid kinds: tool, resource, resource_template, prompt" + )), + } } pub async fn invoke( @@ -145,18 +302,616 @@ impl McpRuntime { server_handle.call_tool(request).await.map_err(Into::into) } + + pub async fn read(&self, server: &str, uri: &str) -> Result { + let server_handle = self + .get(server) + .cloned() + .with_context(|| format!("Read MCP server does not exist: {server}"))?; + + server_handle + .read_resource(ReadResourceRequestParams::new(uri)) + .await + .map_err(Into::into) + } + + pub async fn list_prompts(&self, server: &str) -> Result> { + let server_handle = self + .get(server) + .cloned() + .with_context(|| format!("Prompt MCP server does not exist: {server}"))?; + + server_handle.list_all_prompts().await.map_err(Into::into) + } + + pub async fn prompt( + &self, + server: &str, + name: &str, + arguments: HashMap, + ) -> Result { + let server_handle = self + .get(server) + .cloned() + .with_context(|| format!("Prompt MCP server does not exist: {server}"))?; + + let mut request = GetPromptRequestParams::new(name.to_owned()); + if !arguments.is_empty() { + request.arguments = Some( + arguments + .into_iter() + .map(|(key, value)| (key, Value::String(value))) + .collect(), + ); + } + + server_handle.get_prompt(request).await.map_err(Into::into) + } + + pub async fn prompt_catalog(&self) -> Vec { + let mut items = Vec::new(); + for features in self.server_features() { + if !features.prompts { + continue; + } + let Some(server_handle) = self.get(&features.name) else { + continue; + }; + match server_handle.list_all_prompts().await { + Ok(prompts) => items.extend( + prompts + .into_iter() + .map(|prompt| prompt_catalog_item(&features.name, prompt)), + ), + Err(e) => warn!( + "Failed to list prompts on MCP server {}: {e}", + features.name + ), + } + } + items + } + + pub fn prompt_completion( + &self, + enabled_servers: &[String], + args: &[&str], + ) -> McpPromptCompletion { + match args { + [] | [_] => McpPromptCompletion::Ready( + self.server_features() + .into_iter() + .filter(|features| features.prompts && enabled_servers.contains(&features.name)) + .map(|features| (features.name, None)) + .collect(), + ), + [server, _] => match self.get(server) { + Some(handle) => McpPromptCompletion::PromptNames { + server: Arc::clone(handle), + }, + None => McpPromptCompletion::Ready(vec![]), + }, + [server, prompt, rest @ ..] => match self.get(server) { + Some(handle) => McpPromptCompletion::ArgumentKeys { + server: Arc::clone(handle), + prompt: (*prompt).to_string(), + typed_keys: rest + .iter() + .filter_map(|arg| arg.split_once('=').map(|(key, _)| key.to_string())) + .collect(), + }, + None => McpPromptCompletion::Ready(vec![]), + }, + } + } +} + +pub fn flatten_prompt_messages(messages: &[PromptMessage]) -> String { + messages + .iter() + .map(|message| { + let label = match message.role { + Role::User => "[user]", + Role::Assistant => "[assistant]", + }; + let content = match &message.content { + ContentBlock::Text(text) => text.text.clone(), + ContentBlock::Image(_) => "[image content omitted]".to_string(), + ContentBlock::Audio(_) => "[audio content omitted]".to_string(), + _ => "[resource content omitted]".to_string(), + }; + format!("{label}\n{content}") + }) + .collect::>() + .join("\n\n") +} + +pub fn format_prompt_arguments(arguments: &[PromptArgument]) -> String { + arguments + .iter() + .map(|arg| { + let name = sanitize_display_text(&arg.name); + if arg.required == Some(true) { + format!("{name} (required)") + } else { + name + } + }) + .collect::>() + .join(", ") +} + +pub fn resolve_prompt_args( + declared: &[PromptArgument], + provided: HashMap, +) -> (HashMap, Vec) { + let missing = declared + .iter() + .filter(|arg| arg.required == Some(true) && !provided.contains_key(&arg.name)) + .map(|arg| arg.name.clone()) + .collect(); + (provided, missing) +} + +pub fn sanitize_display_text(text: &str) -> String { + let mut sanitized = String::with_capacity(text.len()); + let mut chars = text.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == '\u{1b}' { + match chars.next() { + // CSI: skip everything up to and including the final byte. + Some('[') => { + for next in chars.by_ref() { + if matches!(next, '\u{40}'..='\u{7e}') { + break; + } + } + } + // OSC: skip until BEL or the ESC \ string terminator. + Some(']') => { + while let Some(next) = chars.next() { + if next == '\u{07}' { + break; + } + if next == '\u{1b}' { + if chars.peek() == Some(&'\\') { + chars.next(); + } + break; + } + } + } + _ => {} + } + } else if ch.is_control() { + sanitized.push(' '); + } else { + sanitized.push(ch); + } + } + sanitized +} + +fn audience_strings(annotations: Option) -> Option> { + let audience = annotations?.audience?; + Some( + audience + .into_iter() + .map(|role| match role { + Role::User => "user".to_string(), + Role::Assistant => "assistant".to_string(), + }) + .collect(), + ) +} + +fn catalog_key(item: &CatalogItem) -> String { + let id = item.uri.as_deref().unwrap_or(&item.name); + format!("{}:{id}", item.kind) +} + +fn merge_catalog_items( + items: &mut HashMap, + new_items: impl IntoIterator, +) { + for item in new_items { + items.insert(catalog_key(&item), item); + } +} + +fn tool_catalog_item(server: &str, tool: Tool) -> CatalogItem { + CatalogItem { + kind: CatalogItemKind::Tool, + name: tool.name.to_string(), + server: server.to_string(), + description: tool.description.unwrap_or_default().to_string(), + ..Default::default() + } +} + +fn resource_catalog_item(server: &str, resource: Resource) -> CatalogItem { + CatalogItem { + kind: CatalogItemKind::Resource, + name: resource.name, + server: server.to_string(), + description: resource.description.unwrap_or_default(), + uri: Some(resource.uri), + mime_type: resource.mime_type, + size: resource.size, + arguments: None, + audience: audience_strings(resource.annotations), + } +} + +fn resource_template_catalog_item(server: &str, template: ResourceTemplate) -> CatalogItem { + CatalogItem { + kind: CatalogItemKind::ResourceTemplate, + name: template.name, + server: server.to_string(), + description: template.description.unwrap_or_default(), + uri: Some(template.uri_template), + mime_type: template.mime_type, + size: None, + arguments: None, + audience: audience_strings(template.annotations), + } +} + +fn prompt_catalog_item(server: &str, prompt: Prompt) -> CatalogItem { + CatalogItem { + kind: CatalogItemKind::Prompt, + name: sanitize_display_text(&prompt.name), + server: server.to_string(), + description: sanitize_display_text(&prompt.description.unwrap_or_default()), + arguments: prompt.arguments, + ..Default::default() + } +} + +fn uri_template_variables(template: &str) -> Vec { + let mut variables = Vec::new(); + let mut rest = template; + while let Some(start) = rest.find('{') { + let Some(len) = rest[start + 1..].find('}') else { + break; + }; + variables.push(rest[start + 1..start + 1 + len].to_string()); + rest = &rest[start + len + 2..]; + } + variables +} + +#[cfg(test)] +pub(crate) mod test_fixtures { + use super::*; + use base64::Engine; + use base64::engine::general_purpose::STANDARD; + use rmcp::model::{ + CallToolResponse, ErrorData, GetPromptResponse, ListPromptsResult, + ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams, + PromptsCapability, ReadResourceResponse, ResourceContents, ResourcesCapability, + ServerCapabilities, ServerInfo, + }; + use rmcp::service::{RequestContext, RunningService}; + use rmcp::{RoleServer, ServerHandler, ServiceExt}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + pub(crate) const FIXTURE_LOG_URI: &str = "file:///app.log"; + pub(crate) const FIXTURE_LOG_TEXT: &str = "début of the log\n\ + second line\n\ + ERROR: disk full\n\ + fourth line\n\ + fifth line\n\ + ERROR: café overheated\n\ + seventh line\n\ + eighth line"; + pub(crate) const FIXTURE_BLOB_URI: &str = "file:///report.pdf"; + pub(crate) const FIXTURE_BLOB_BYTES: &[u8] = &[0xff, 0xfe, 0x00, 0x88, 0x01]; + pub(crate) const FIXTURE_ANNOTATED_URI: &str = "file:///annotated"; + pub(crate) const FIXTURE_ANNOTATED_TEXT: &str = "annotated body"; + + #[derive(Clone)] + pub(crate) struct FixtureServer { + pub(crate) tools_capability: bool, + pub(crate) resources_capability: bool, + pub(crate) prompts_capability: bool, + pub(crate) hostile_prompt: bool, + pub(crate) fail_resource_listings: bool, + pub(crate) fail_prompt_listings: bool, + pub(crate) fail_get_prompt: bool, + pub(crate) prompt_delay: Option, + pub(crate) tool_result: Option, + pub(crate) list_resources_calls: Arc, + pub(crate) list_prompts_calls: Arc, + pub(crate) get_prompt_calls: Arc, + pub(crate) call_tool_calls: Arc, + } + + impl Default for FixtureServer { + fn default() -> Self { + Self { + tools_capability: true, + resources_capability: false, + prompts_capability: false, + hostile_prompt: false, + fail_resource_listings: false, + fail_prompt_listings: false, + fail_get_prompt: false, + prompt_delay: None, + tool_result: None, + list_resources_calls: Arc::default(), + list_prompts_calls: Arc::default(), + get_prompt_calls: Arc::default(), + call_tool_calls: Arc::default(), + } + } + } + + impl ServerHandler for FixtureServer { + fn get_info(&self) -> ServerInfo { + let mut capabilities = if self.tools_capability { + ServerCapabilities::builder().enable_tools().build() + } else { + ServerCapabilities::builder().build() + }; + capabilities.resources = self.resources_capability.then(ResourcesCapability::default); + capabilities.prompts = self.prompts_capability.then(PromptsCapability::default); + ServerInfo::new(capabilities) + } + + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + let schema = json!({ + "type": "object", + "properties": { "q": { "type": "string" } } + }) + .as_object() + .cloned() + .unwrap(); + Ok(ListToolsResult::with_all_items(vec![Tool::new( + "dup", + "Duplicate-named tool", + schema, + )])) + } + + async fn call_tool( + &self, + _request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + self.call_tool_calls.fetch_add(1, Ordering::SeqCst); + match &self.tool_result { + Some(result) => Ok(CallToolResponse::Complete(result.clone())), + None => Err(ErrorData::internal_error( + "call_tool should not be reached", + None, + )), + } + } + + async fn list_resources( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + self.list_resources_calls.fetch_add(1, Ordering::SeqCst); + if self.fail_resource_listings { + return Err(ErrorData::internal_error("resource listing exploded", None)); + } + Ok(ListResourcesResult::with_all_items(vec![ + Resource::new("dup", "dup-resource") + .with_description("Duplicate-named resource") + .with_mime_type("text/plain") + .with_size(42), + Resource::new(FIXTURE_ANNOTATED_URI, "annotated-notes") + .with_description("Notes with an audience annotation") + .with_mime_type("text/plain") + .with_annotations(Annotations::default().with_audience(vec![Role::User])), + ])) + } + + async fn list_resource_templates( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + if self.fail_resource_listings { + return Err(ErrorData::internal_error("template listing exploded", None)); + } + Ok(ListResourceTemplatesResult::with_all_items(vec![ + ResourceTemplate::new("file:///{path}/{name}", "file-template") + .with_description("Read a file") + .with_mime_type("text/plain"), + ])) + } + + async fn read_resource( + &self, + request: ReadResourceRequestParams, + _context: RequestContext, + ) -> Result { + let uri = request.uri.as_str(); + let contents = match uri { + FIXTURE_LOG_URI => vec![ResourceContents::text(FIXTURE_LOG_TEXT, uri)], + FIXTURE_BLOB_URI => vec![ + ResourceContents::blob(STANDARD.encode(FIXTURE_BLOB_BYTES), uri) + .with_mime_type("application/pdf"), + ], + FIXTURE_ANNOTATED_URI => vec![ResourceContents::text(FIXTURE_ANNOTATED_TEXT, uri)], + "file:///multi" => vec![ + ResourceContents::text("first", "file:///multi/0"), + ResourceContents::text("second", "file:///multi/1"), + ResourceContents::text("third", "file:///multi/2"), + ], + "file:///huge" => (0..3) + .map(|i| { + ResourceContents::text("x".repeat(150 * 1024), format!("file:///huge/{i}")) + }) + .collect(), + "file:///docs/readme" => vec![ResourceContents::text("readme body", uri)], + _ => { + return Err(ErrorData::resource_not_found( + format!("Unknown resource: {uri}"), + None, + )); + } + }; + Ok(ReadResourceResponse::Complete(ReadResourceResult::new( + contents, + ))) + } + + async fn list_prompts( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + self.list_prompts_calls.fetch_add(1, Ordering::SeqCst); + if let Some(delay) = self.prompt_delay { + tokio::time::sleep(delay).await; + } + if self.fail_prompt_listings { + return Err(ErrorData::internal_error("prompt listing exploded", None)); + } + let mut prompts = vec![Prompt::new( + "summarize", + Some("Summarize a document"), + Some(vec![ + PromptArgument::new("path") + .with_description("Document path") + .with_required(true), + PromptArgument::new("style"), + ]), + )]; + if self.hostile_prompt { + prompts.push(Prompt::new( + "sum\u{1b}[31mmarize-evil", + Some("Runs\u{1b}]0;pwn\u{7} hostile\ttext"), + Some(vec![ + PromptArgument::new("pa\u{1b}[1mth") + .with_description("Doc\u{1b}[4m path") + .with_required(true), + ]), + )); + } + Ok(ListPromptsResult::with_all_items(prompts)) + } + + async fn get_prompt( + &self, + request: GetPromptRequestParams, + _context: RequestContext, + ) -> Result { + self.get_prompt_calls.fetch_add(1, Ordering::SeqCst); + if let Some(delay) = self.prompt_delay { + tokio::time::sleep(delay).await; + } + if self.fail_get_prompt { + return Err(ErrorData::internal_error("get_prompt exploded", None)); + } + let messages = match request.name.as_str() { + "summarize" => { + let path = request + .arguments + .as_ref() + .and_then(|args| args.get("path")) + .and_then(|value| value.as_str()) + .unwrap_or_default(); + vec![ + PromptMessage::new_text(Role::User, format!("Summarize {path}")), + PromptMessage::new_text(Role::Assistant, "In which style?"), + PromptMessage::new_text(Role::User, "Concise."), + ] + } + "hostile" => vec![ + PromptMessage::new_text(Role::User, "!rm -rf /"), + PromptMessage::new_text(Role::User, ".session hijack"), + ], + other => { + return Err(ErrorData::invalid_params( + format!("Unknown prompt: {other}"), + None, + )); + } + }; + Ok(GetPromptResponse::Complete(GetPromptResult::new(messages))) + } + } + + pub(crate) async fn add_fixture_server( + runtime: &mut McpRuntime, + name: &str, + fixture: FixtureServer, + ) -> RunningService { + let (client_io, server_io) = tokio::io::duplex(4096); + let (server, client) = tokio::join!(fixture.serve(server_io), ().serve(client_io)); + runtime.insert(name.to_string(), Arc::new(client.unwrap())); + server.unwrap() + } + + pub(crate) async fn fixture_runtime( + fixture: FixtureServer, + ) -> (McpRuntime, RunningService) { + let mut runtime = McpRuntime::new(); + let server = add_fixture_server(&mut runtime, "fixture", fixture).await; + (runtime, server) + } } #[cfg(test)] mod tests { + use super::test_fixtures::{ + FIXTURE_ANNOTATED_URI, FixtureServer, add_fixture_server, fixture_runtime, + }; use super::*; use crate::function::ToolCall; + use log::{Level, LevelFilter, Log, Metadata, Record}; + use std::sync::atomic::Ordering; + use std::sync::{Mutex, Once, OnceLock}; + + struct WarnCollector; + + static WARN_MESSAGES: OnceLock>> = OnceLock::new(); + + fn warn_messages() -> &'static Mutex> { + WARN_MESSAGES.get_or_init(Mutex::default) + } + + impl Log for WarnCollector { + fn enabled(&self, metadata: &Metadata) -> bool { + metadata.level() <= Level::Warn + } + + fn log(&self, record: &Record) { + if self.enabled(record.metadata()) { + warn_messages() + .lock() + .unwrap() + .push(record.args().to_string()); + } + } + + fn flush(&self) {} + } + + fn install_warn_collector() { + static INSTALL: Once = Once::new(); + INSTALL.call_once(|| { + log::set_logger(&WarnCollector).expect("no other logger should be installed"); + log::set_max_level(LevelFilter::Warn); + }); + } #[test] fn mcp_runtime_new_is_empty() { let runtime = McpRuntime::new(); assert!(runtime.is_empty()); - assert!(runtime.server_names().is_empty()); + assert!(runtime.server_features().is_empty()); } #[test] @@ -171,6 +926,89 @@ mod tests { assert!(runtime.get("nonexistent").is_none()); } + #[tokio::test] + async fn server_features_reports_fixture_capabilities() { + let (runtime, _server) = fixture_runtime(FixtureServer { + resources_capability: true, + ..Default::default() + }) + .await; + + let features = runtime.server_features(); + assert_eq!( + features, + vec![McpServerFeatures { + name: "fixture".to_string(), + tools: true, + resources: true, + prompts: false, + }] + ); + + let mut functions = Functions::default(); + functions.append_mcp_meta_functions(features); + assert_eq!(functions.declarations().len(), 4); + assert!(functions.contains("mcp_invoke_fixture")); + assert!(functions.contains("mcp_search_fixture")); + assert!(functions.contains("mcp_describe_fixture")); + assert!(functions.contains("mcp_read_fixture")); + } + + #[tokio::test] + async fn server_features_without_tools_capability_gates_invoke() { + let (runtime, _server) = fixture_runtime(FixtureServer { + tools_capability: false, + resources_capability: true, + ..Default::default() + }) + .await; + + let features = runtime.server_features(); + assert_eq!( + features, + vec![McpServerFeatures { + name: "fixture".to_string(), + tools: false, + resources: true, + prompts: false, + }] + ); + + let mut functions = Functions::default(); + functions.append_mcp_meta_functions(features); + assert_eq!(functions.declarations().len(), 3); + assert!(!functions.contains("mcp_invoke_fixture")); + assert!(functions.contains("mcp_search_fixture")); + assert!(functions.contains("mcp_describe_fixture")); + assert!(functions.contains("mcp_read_fixture")); + } + + #[tokio::test] + async fn server_features_with_prompts_capability_emits_prompt_meta_function() { + let (runtime, _server) = fixture_runtime(FixtureServer { + resources_capability: true, + prompts_capability: true, + ..Default::default() + }) + .await; + + let features = runtime.server_features(); + assert_eq!( + features, + vec![McpServerFeatures { + name: "fixture".to_string(), + tools: true, + resources: true, + prompts: true, + }] + ); + + let mut functions = Functions::default(); + functions.append_mcp_meta_functions(features); + assert_eq!(functions.declarations().len(), 5); + assert!(functions.contains("mcp_prompt_fixture")); + } + #[test] fn tool_scope_default_has_empty_mcp_runtime() { let scope = ToolScope::default(); @@ -189,4 +1027,620 @@ mod tests { let dummy_call = ToolCall::default(); assert!(scope.tool_tracker.check_loop(&dummy_call).is_none()); } + + #[test] + fn uri_template_variables_extracts_placeholders() { + assert_eq!( + uri_template_variables("file:///{path}/{name}"), + vec!["path", "name"] + ); + assert!(uri_template_variables("file:///static").is_empty()); + } + + #[tokio::test] + async fn catalog_items_keeps_tool_and_resource_with_same_id() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.catalog_items("fixture").await.unwrap(); + + assert!(items.contains_key("tool:dup")); + assert!(items.contains_key("resource:dup")); + assert!(items.contains_key("resource_template:file:///{path}/{name}")); + } + + #[tokio::test] + async fn catalog_items_degrades_when_resource_listing_fails() { + let fixture = FixtureServer { + resources_capability: true, + fail_resource_listings: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.catalog_items("fixture").await.unwrap(); + + assert!(items.contains_key("tool:dup")); + assert!(!items.keys().any(|key| key.starts_with("resource"))); + } + + #[tokio::test] + async fn catalog_items_warns_when_resource_listing_fails() { + install_warn_collector(); + let fixture = FixtureServer { + resources_capability: true, + fail_resource_listings: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + runtime.catalog_items("fixture").await.unwrap(); + + let messages = warn_messages().lock().unwrap(); + assert!( + messages + .iter() + .any(|msg| msg.contains("Failed to list resources on MCP server fixture")), + "missing resource-listing warning in: {messages:?}" + ); + } + + #[tokio::test] + async fn catalog_items_skips_unadvertised_capabilities() { + let fixture = FixtureServer::default(); + let resources_calls = Arc::clone(&fixture.list_resources_calls); + let prompts_calls = Arc::clone(&fixture.list_prompts_calls); + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.catalog_items("fixture").await.unwrap(); + + assert_eq!(items.keys().collect::>(), vec!["tool:dup"]); + assert_eq!(resources_calls.load(Ordering::SeqCst), 0); + assert_eq!(prompts_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn catalog_items_includes_prompts_when_advertised() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.catalog_items("fixture").await.unwrap(); + + let prompt = items.get("prompt:summarize").unwrap(); + assert_eq!(prompt.kind, CatalogItemKind::Prompt); + assert_eq!(prompt.description, "Summarize a document"); + } + + #[tokio::test] + async fn search_results_carry_kind_and_uri() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let results = runtime.search("fixture", "dup", 10).await.unwrap(); + + let values: Vec = results + .iter() + .map(|item| serde_json::to_value(item).unwrap()) + .collect(); + let tool = values.iter().find(|v| v["kind"] == "tool").unwrap(); + assert!(tool.get("uri").is_none()); + let resource = values.iter().find(|v| v["kind"] == "resource").unwrap(); + assert_eq!(resource["uri"], "dup"); + assert_eq!(resource["mime_type"], "text/plain"); + assert_eq!(resource["size"], 42); + assert!(resource.get("audience").is_none()); + } + + #[tokio::test] + async fn search_results_carry_resource_audience() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let results = runtime + .search("fixture", "annotated notes", 10) + .await + .unwrap(); + + let values: Vec = results + .iter() + .map(|item| serde_json::to_value(item).unwrap()) + .collect(); + let resource = values + .iter() + .find(|v| v["uri"] == FIXTURE_ANNOTATED_URI) + .unwrap(); + assert_eq!(resource["audience"], json!(["user"])); + } + + #[tokio::test] + async fn resource_audience_returns_annotated_roles() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let audience = runtime + .resource_audience("fixture", FIXTURE_ANNOTATED_URI) + .await; + + assert_eq!(audience, Some(vec!["user".to_string()])); + } + + #[tokio::test] + async fn resource_audience_is_none_for_unknown_uri_or_server() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + assert!(runtime.resource_audience("fixture", "dup").await.is_none()); + assert!( + runtime + .resource_audience("fixture", "file:///unknown") + .await + .is_none() + ); + assert!(runtime.resource_audience("ghost", "dup").await.is_none()); + } + + #[tokio::test] + async fn describe_tool_keeps_existing_schema_shape() { + let fixture = FixtureServer::default(); + let (runtime, _server) = fixture_runtime(fixture).await; + + let result = runtime.describe("fixture", "tool", "dup").await.unwrap(); + + assert_eq!( + result, + json!({ + "type": "object", + "properties": { + "tool": { "type": "string" }, + "arguments": { + "type": "object", + "properties": { "q": { "type": "string" } } + } + } + }) + ); + } + + #[tokio::test] + async fn describe_resource_returns_metadata() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let result = runtime + .describe("fixture", "resource", "dup") + .await + .unwrap(); + + assert_eq!(result["uri"], "dup"); + assert_eq!(result["name"], "dup-resource"); + assert_eq!(result["description"], "Duplicate-named resource"); + assert_eq!(result["mime_type"], "text/plain"); + assert_eq!(result["size"], 42); + } + + #[tokio::test] + async fn describe_resource_template_returns_variables() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let result = runtime + .describe("fixture", "resource_template", "file:///{path}/{name}") + .await + .unwrap(); + + assert_eq!(result["uri_template"], "file:///{path}/{name}"); + assert_eq!(result["name"], "file-template"); + assert_eq!(result["variables"], json!(["path", "name"])); + } + + #[tokio::test] + async fn describe_prompt_returns_arguments() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let result = runtime + .describe("fixture", "prompt", "summarize") + .await + .unwrap(); + + assert_eq!(result["name"], "summarize"); + assert_eq!(result["description"], "Summarize a document"); + assert_eq!(result["arguments"][0]["name"], "path"); + assert_eq!(result["arguments"][0]["description"], "Document path"); + assert_eq!(result["arguments"][0]["required"], true); + assert_eq!(result["arguments"][1]["name"], "style"); + assert_eq!(result["arguments"][1]["required"], Value::Null); + } + + #[tokio::test] + async fn describe_unknown_kind_lists_valid_kinds() { + let fixture = FixtureServer::default(); + let (runtime, _server) = fixture_runtime(fixture).await; + + let err = runtime + .describe("fixture", "widget", "dup") + .await + .unwrap_err() + .to_string(); + + assert!(err.contains("widget")); + for kind in ["tool", "resource", "resource_template", "prompt"] { + assert!(err.contains(kind), "missing {kind} in: {err}"); + } + } + + #[tokio::test] + async fn describe_missing_resource_names_kind_in_error() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let err = runtime + .describe("fixture", "resource", "file:///missing") + .await + .unwrap_err() + .to_string(); + + assert_eq!( + err, + "file:///missing not found in fixture MCP server resource catalog" + ); + } + + #[tokio::test] + async fn prompt_returns_result_and_counts_calls() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let get_prompt_calls = Arc::clone(&fixture.get_prompt_calls); + let (runtime, _server) = fixture_runtime(fixture).await; + + let result = runtime + .prompt( + "fixture", + "summarize", + HashMap::from([("path".to_string(), "notes.txt".to_string())]), + ) + .await + .unwrap(); + + assert_eq!(result.messages.len(), 3); + assert_eq!( + result.messages[0] + .content + .as_text() + .map(|t| t.text.as_str()), + Some("Summarize notes.txt") + ); + assert_eq!(get_prompt_calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn prompt_missing_server_errors() { + let runtime = McpRuntime::new(); + + let err = runtime + .prompt("ghost", "summarize", HashMap::new()) + .await + .unwrap_err() + .to_string(); + + assert_eq!(err, "Prompt MCP server does not exist: ghost"); + } + + #[tokio::test] + async fn prompt_surfaces_server_failure() { + let fixture = FixtureServer { + prompts_capability: true, + fail_get_prompt: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let err = runtime + .prompt("fixture", "summarize", HashMap::new()) + .await + .unwrap_err() + .to_string(); + + assert!( + err.contains("get_prompt exploded"), + "unexpected error: {err}" + ); + } + + #[test] + fn flatten_prompt_messages_labels_every_message() { + let messages = vec![ + PromptMessage::new_text(Role::User, "Summarize notes.txt"), + PromptMessage::new_text(Role::Assistant, "In which style?"), + PromptMessage::new_text(Role::User, "Concise."), + ]; + + assert_eq!( + flatten_prompt_messages(&messages), + "[user]\nSummarize notes.txt\n\n[assistant]\nIn which style?\n\n[user]\nConcise." + ); + } + + #[test] + fn flatten_prompt_messages_labels_single_message() { + let messages = vec![PromptMessage::new_text(Role::User, "hello")]; + + assert_eq!(flatten_prompt_messages(&messages), "[user]\nhello"); + } + + #[test] + fn flatten_prompt_messages_replaces_non_text_content() { + let messages = vec![PromptMessage::new( + Role::User, + ContentBlock::image("aGk=", "image/png"), + )]; + + assert_eq!( + flatten_prompt_messages(&messages), + "[user]\n[image content omitted]" + ); + } + + #[tokio::test] + async fn flattened_hostile_prompt_cannot_start_a_command_line() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let result = runtime + .prompt("fixture", "hostile", HashMap::new()) + .await + .unwrap(); + let flattened = flatten_prompt_messages(&result.messages); + + assert!(flattened.starts_with("[user]")); + assert!(!flattened.starts_with('!')); + assert!(!flattened.starts_with('.')); + assert!(flattened.contains("!rm -rf /")); + assert!(flattened.contains(".session hijack")); + } + + #[test] + fn format_prompt_arguments_marks_required() { + let arguments = vec![ + PromptArgument::new("path").with_required(true), + PromptArgument::new("style"), + ]; + + assert_eq!( + format_prompt_arguments(&arguments), + "path (required), style" + ); + assert_eq!(format_prompt_arguments(&[]), ""); + } + + #[test] + fn resolve_prompt_args_reports_missing_required_only() { + let declared = vec![ + PromptArgument::new("path").with_required(true), + PromptArgument::new("style"), + ]; + + let (resolved, missing) = resolve_prompt_args(&declared, HashMap::new()); + assert!(resolved.is_empty()); + assert_eq!(missing, vec!["path".to_string()]); + + let provided = HashMap::from([("path".to_string(), "notes.txt".to_string())]); + let (resolved, missing) = resolve_prompt_args(&declared, provided); + assert_eq!(resolved["path"], "notes.txt"); + assert!(missing.is_empty()); + } + + #[test] + fn sanitize_display_text_strips_csi_sequences_entirely() { + assert_eq!(sanitize_display_text("a\u{1b}[31mred\u{1b}[0mb"), "aredb"); + } + + #[test] + fn sanitize_display_text_strips_osc_sequences_entirely() { + assert_eq!(sanitize_display_text("a\u{1b}]0;title\u{7}b"), "ab"); + assert_eq!(sanitize_display_text("a\u{1b}]0;title\u{1b}\\b"), "ab"); + } + + #[test] + fn sanitize_display_text_keeps_plain_text() { + assert_eq!( + sanitize_display_text("path (required), café"), + "path (required), café" + ); + } + + #[test] + fn sanitize_display_text_maps_control_chars_to_spaces() { + assert_eq!(sanitize_display_text("a\nb\tc\rd\u{7}e"), "a b c d e"); + } + + #[tokio::test] + async fn prompt_catalog_carries_arguments() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.prompt_catalog().await; + + assert_eq!(items.len(), 1); + let item = &items[0]; + assert_eq!(item.kind, CatalogItemKind::Prompt); + assert_eq!(item.server, "fixture"); + assert_eq!(item.name, "summarize"); + assert_eq!(item.description, "Summarize a document"); + let arguments = item.arguments.as_deref().unwrap(); + assert_eq!(format_prompt_arguments(arguments), "path (required), style"); + } + + #[tokio::test] + async fn prompt_catalog_sanitizes_hostile_display_strings() { + let fixture = FixtureServer { + prompts_capability: true, + hostile_prompt: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.prompt_catalog().await; + + let item = items + .iter() + .find(|item| item.name == "summarize-evil") + .unwrap(); + assert_eq!(item.description, "Runs hostile text"); + assert_eq!( + format_prompt_arguments(item.arguments.as_deref().unwrap()), + "path (required)" + ); + } + + #[tokio::test] + async fn prompt_catalog_degrades_when_one_server_fails() { + install_warn_collector(); + let healthy = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let failing = FixtureServer { + prompts_capability: true, + fail_prompt_listings: true, + ..Default::default() + }; + let (mut runtime, _server) = fixture_runtime(healthy).await; + let _failing_server = add_fixture_server(&mut runtime, "broken", failing).await; + + let items = runtime.prompt_catalog().await; + + assert_eq!(items.len(), 1); + assert_eq!(items[0].server, "fixture"); + let messages = warn_messages().lock().unwrap(); + assert!( + messages + .iter() + .any(|msg| msg.contains("Failed to list prompts on MCP server broken")), + "missing prompt-listing warning in: {messages:?}" + ); + } + + #[tokio::test] + async fn prompt_catalog_skips_servers_without_prompts_capability() { + let fixture = FixtureServer::default(); + let prompts_calls = Arc::clone(&fixture.list_prompts_calls); + let (runtime, _server) = fixture_runtime(fixture).await; + + let items = runtime.prompt_catalog().await; + + assert!(items.is_empty()); + assert_eq!(prompts_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn prompt_completion_stage_one_uses_local_state_only() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let list_prompts_calls = Arc::clone(&fixture.list_prompts_calls); + let get_prompt_calls = Arc::clone(&fixture.get_prompt_calls); + let (mut runtime, _server) = fixture_runtime(fixture).await; + let _tools_only = + add_fixture_server(&mut runtime, "tools-only", FixtureServer::default()).await; + + let stage = + runtime.prompt_completion(&["fixture".to_string(), "tools-only".to_string()], &[""]); + + let McpPromptCompletion::Ready(values) = stage else { + panic!("stage one must not require an RPC"); + }; + assert_eq!(values, vec![("fixture".to_string(), None)]); + assert_eq!(list_prompts_calls.load(Ordering::SeqCst), 0); + assert_eq!(get_prompt_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn prompt_completion_stage_one_excludes_disabled_servers() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let McpPromptCompletion::Ready(values) = runtime.prompt_completion(&[], &[""]) else { + panic!("stage one must not require an RPC"); + }; + + assert!(values.is_empty()); + } + + #[tokio::test] + async fn prompt_completion_unknown_server_is_empty() { + let (runtime, _server) = fixture_runtime(FixtureServer::default()).await; + + for args in [ + ["ghost", ""].as_slice(), + ["ghost", "summarize", ""].as_slice(), + ] { + let McpPromptCompletion::Ready(values) = runtime.prompt_completion(&[], args) else { + panic!("unknown server must degrade to empty suggestions"); + }; + assert!(values.is_empty()); + } + } + + #[tokio::test] + async fn prompt_completion_later_stages_carry_typed_keys() { + let fixture = FixtureServer { + prompts_capability: true, + ..Default::default() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + + let stage = runtime.prompt_completion(&[], &["fixture", "summarize", "path=x", "sty"]); + + let McpPromptCompletion::ArgumentKeys { + prompt, typed_keys, .. + } = stage + else { + panic!("expected the argument-key stage"); + }; + assert_eq!(prompt, "summarize"); + assert_eq!(typed_keys, vec!["path".to_string()]); + } } diff --git a/src/function/mod.rs b/src/function/mod.rs index d17f042..d32a24c 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -7,7 +7,9 @@ pub(crate) mod user_interaction; use crate::{ client::ThinkingBlock, - config::{Agent, RequestContext}, + config::{ + Agent, RequestContext, flatten_prompt_messages, resolve_prompt_args, sanitize_display_text, + }, graph, utils::*, }; @@ -16,7 +18,9 @@ use crate::config::ensure_parent_exists; use crate::config::paths; use crate::mcp::{ MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, - MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpServersConfig, + MCP_META_FUNCTION_PREFIXES, MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpServerFeatures, + McpServersConfig, is_mcp_meta_function, render, }; use crate::parsers::{bash, python, typescript}; use anyhow::{Context, Result, anyhow, bail}; @@ -284,14 +288,9 @@ pub async fn eval_tool_calls( } } - let (mcp_calls, sequential_calls): (Vec<_>, Vec<_>) = - to_execute.into_iter().partition(|(_, call)| { - call.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || call.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || call - .name - .starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }); + let (mcp_calls, sequential_calls): (Vec<_>, Vec<_>) = to_execute + .into_iter() + .partition(|(_, call)| is_mcp_meta_function(&call.name)); if !mcp_calls.is_empty() { let ctx_ref: &RequestContext = ctx; @@ -414,6 +413,18 @@ impl ToolResult { } } +fn gated_meta_function_prefixes(features: &McpServerFeatures) -> Vec<&'static str> { + MCP_META_FUNCTION_PREFIXES + .into_iter() + .filter(|&prefix| match prefix { + MCP_INVOKE_META_FUNCTION_NAME_PREFIX => features.tools, + MCP_READ_META_FUNCTION_NAME_PREFIX => features.resources, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX => features.prompts, + _ => true, + }) + .collect() +} + #[derive(Debug, Clone, Default)] pub struct Functions { declarations: Vec, @@ -629,7 +640,7 @@ impl Functions { .retain(|f| !f.name.starts_with(RAG_FUNCTION_PREFIX)); } - pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { + pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { let mut invoke_function_properties = IndexMap::new(); invoke_function_properties.insert( "tool".to_string(), @@ -674,58 +685,198 @@ impl Functions { ..Default::default() }, ); + describe_function_properties.insert( + "kind".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "Catalog item kind: tool (default), resource, resource_template, or prompt" + .into(), + ), + default: Some(Value::from("tool")), + ..Default::default() + }, + ); - for server in mcp_servers { + let mut read_function_properties = IndexMap::new(); + read_function_properties.insert( + "uri".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "Resource URI, or a resource template with {var} placeholders".into(), + ), + ..Default::default() + }, + ); + read_function_properties.insert( + "arguments".to_string(), + JsonSchema { + type_value: Some("object".to_string()), + description: Some("Template variable values (RFC 6570 Level 1 only)".into()), + ..Default::default() + }, + ); + read_function_properties.insert( + "pattern".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "Optional regex; returns only matching lines (with context) from text content" + .into(), + ), + ..Default::default() + }, + ); + read_function_properties.insert( + "offset".to_string(), + JsonSchema { + type_value: Some("integer".to_string()), + description: Some( + "Byte offset for paging text. When pattern is set, offsets (and \ + next_offset/total_bytes in the result) refer to the filtered stream, not \ + the raw resource" + .into(), + ), + default: Some(Value::from(0usize)), + ..Default::default() + }, + ); + read_function_properties.insert( + "max_bytes".to_string(), + JsonSchema { + type_value: Some("integer".to_string()), + description: Some(format!( + "Max text bytes to return (clamped to {})", + render::TEXT_MAX_BYTES_CLAMP + )), + default: Some(Value::from(render::DEFAULT_TEXT_MAX_BYTES)), + ..Default::default() + }, + ); + + let mut prompt_function_properties = IndexMap::new(); + prompt_function_properties.insert( + "prompt".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + ..Default::default() + }, + ); + prompt_function_properties.insert( + "arguments".to_string(), + JsonSchema { + type_value: Some("object".to_string()), + description: Some("String values only; prompt arguments have no schemas".into()), + ..Default::default() + }, + ); + + for features in mcp_servers { + let server = &features.name; let search_function_name = format!("{}_{server}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX); let describe_function_name = format!("{}_{server}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX); let invoke_function_name = format!("{}_{server}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); - let invoke_function_declaration = FunctionDeclaration { - name: invoke_function_name.clone(), - description: formatdoc!( - r#" - Invoke the specified tool on the {server} MCP server. Always call {describe_function_name} first to - find the correct invocation schema for the given tool. - "# - ), - parameters: JsonSchema { - type_value: Some("object".to_string()), - properties: Some(invoke_function_properties.clone()), - required: Some(vec!["tool".to_string()]), - ..Default::default() - }, - agent: false, - }; - let search_functions_declaration = FunctionDeclaration { - name: search_function_name.clone(), - description: formatdoc!( - r#" - Find candidate tools by keywords for the {server} MCP server. Returns small suggestions; fetch - schemas with {describe_function_name}. - "# - ), - parameters: JsonSchema { - type_value: Some("object".to_string()), - properties: Some(search_function_properties.clone()), - required: Some(vec!["query".to_string()]), - ..Default::default() - }, - agent: false, - }; - let describe_functions_declaration = FunctionDeclaration { - name: describe_function_name.clone(), - description: "Get the full JSON schema for exactly one MCP tool.".to_string(), - parameters: JsonSchema { - type_value: Some("object".to_string()), - properties: Some(describe_function_properties.clone()), - required: Some(vec!["tool".to_string()]), - ..Default::default() - }, - agent: false, - }; - self.declarations.push(invoke_function_declaration); - self.declarations.push(search_functions_declaration); - self.declarations.push(describe_functions_declaration); + let read_function_name = format!("{}_{server}", MCP_READ_META_FUNCTION_NAME_PREFIX); + let prompt_function_name = format!("{}_{server}", MCP_PROMPT_META_FUNCTION_NAME_PREFIX); + for prefix in gated_meta_function_prefixes(&features) { + match prefix { + MCP_INVOKE_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: invoke_function_name.clone(), + description: formatdoc!( + r#" + Invoke the specified tool on the {server} MCP server. Always call {describe_function_name} first to + find the correct invocation schema for the given tool. + "# + ), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(invoke_function_properties.clone()), + required: Some(vec!["tool".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + MCP_SEARCH_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: search_function_name.clone(), + description: formatdoc!( + r#" + Find candidate tools by keywords for the {server} MCP server. Returns small suggestions; fetch + schemas with {describe_function_name}. + "# + ), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(search_function_properties.clone()), + required: Some(vec!["query".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: describe_function_name.clone(), + description: "Get the full schema or metadata for exactly one MCP \ + catalog item: a tool, resource, resource template, or \ + prompt." + .to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(describe_function_properties.clone()), + required: Some(vec!["tool".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + MCP_READ_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: read_function_name.clone(), + description: formatdoc!( + r#" + Read a resource, or expand a resource template, from the {server} MCP server. Call + {describe_function_name} with kind "resource" or "resource_template" to find URIs and + template variables. Text content is paged via offset/max_bytes and can be filtered + with pattern; binary content is spilled to disk and its metadata returned. + "# + ), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(read_function_properties.clone()), + required: Some(vec!["uri".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + MCP_PROMPT_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: prompt_function_name.clone(), + description: formatdoc!( + r#" + Fetch a prompt from the {server} MCP server, rendered with the given arguments. Call + {describe_function_name} with kind "prompt" to discover prompt names and their + arguments. The result is the prompt text, labeled per message; fold it into your + reasoning. + "# + ), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(prompt_function_properties.clone()), + required: Some(vec!["prompt".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + _ => debug_assert!(false, "unhandled MCP meta-function prefix: {prefix}"), + } + } } } @@ -1227,7 +1378,7 @@ impl ToolCall { .await .unwrap_or_else(|e| { let error_msg = format!("MCP search failed: {e}"); - eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); json!({"tool_call_error": error_msg}) }) } else if cmd_name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) { @@ -1235,7 +1386,23 @@ impl ToolCall { .await .unwrap_or_else(|e| { let error_msg = format!("MCP describe failed: {e}"); - eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); + json!({"tool_call_error": error_msg}) + }) + } else if cmd_name.starts_with(MCP_READ_META_FUNCTION_NAME_PREFIX) { + Self::read_mcp_resource(ctx, cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("MCP read failed: {e}"); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); + json!({"tool_call_error": error_msg}) + }) + } else if cmd_name.starts_with(MCP_PROMPT_META_FUNCTION_NAME_PREFIX) { + Self::get_mcp_prompt(ctx, cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("MCP prompt failed: {e}"); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); json!({"tool_call_error": error_msg}) }) } else { @@ -1243,7 +1410,7 @@ impl ToolCall { .await .unwrap_or_else(|e| { let error_msg = format!("MCP tool invocation failed: {e}"); - eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); json!({"tool_call_error": error_msg}) }) }; @@ -1286,7 +1453,7 @@ impl ToolCall { .await .unwrap_or_else(|e| { let error_msg = format!("MCP search failed: {e}"); - eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); json!({"tool_call_error": error_msg}) }) } @@ -1295,7 +1462,25 @@ impl ToolCall { .await .unwrap_or_else(|e| { let error_msg = format!("MCP describe failed: {e}"); - eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); + json!({"tool_call_error": error_msg}) + }) + } + _ if cmd_name.starts_with(MCP_READ_META_FUNCTION_NAME_PREFIX) => { + Self::read_mcp_resource(ctx, &cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("MCP read failed: {e}"); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); + json!({"tool_call_error": error_msg}) + }) + } + _ if cmd_name.starts_with(MCP_PROMPT_META_FUNCTION_NAME_PREFIX) => { + Self::get_mcp_prompt(ctx, &cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("MCP prompt failed: {e}"); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); json!({"tool_call_error": error_msg}) }) } @@ -1304,7 +1489,7 @@ impl ToolCall { .await .unwrap_or_else(|e| { let error_msg = format!("MCP tool invocation failed: {e}"); - eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); json!({"tool_call_error": error_msg}) }) } @@ -1383,10 +1568,16 @@ impl ToolCall { .ok_or_else(|| anyhow!("Missing 'tool' in arguments"))? .as_str() .ok_or_else(|| anyhow!("Invalid 'tool' in arguments"))?; + let kind = match json_data.get("kind") { + Some(value) => value + .as_str() + .ok_or_else(|| anyhow!("Invalid 'kind' in arguments"))?, + None => "tool", + }; let result = ctx .tool_scope .mcp_runtime - .describe(&server_id, tool) + .describe(&server_id, kind, tool) .await?; Ok(serde_json::to_value(result)?) } @@ -1441,7 +1632,154 @@ impl ToolCall { .mcp_runtime .invoke(&server, tool, arguments) .await?; - Ok(serde_json::to_value(result)?) + render_tool_result(serde_json::to_value(result)?, &server) + } + + async fn read_mcp_resource( + ctx: &RequestContext, + cmd_name: &str, + json_data: &Value, + ) -> Result { + let server = cmd_name + .strip_prefix(&format!("{MCP_READ_META_FUNCTION_NAME_PREFIX}_")) + .ok_or_else(|| anyhow!("Malformed MCP read function name: {cmd_name}"))?; + let uri = json_data + .get("uri") + .ok_or_else(|| anyhow!("Missing 'uri' in arguments"))? + .as_str() + .ok_or_else(|| anyhow!("Invalid 'uri' in arguments"))?; + let pattern = match json_data.get("pattern") { + Some(value) => Some( + value + .as_str() + .ok_or_else(|| anyhow!("Invalid 'pattern' in arguments"))?, + ), + None => None, + }; + let offset = match json_data.get("offset") { + Some(value) => value + .as_u64() + .ok_or_else(|| anyhow!("Invalid 'offset' in arguments"))? + as usize, + None => 0, + }; + let max_bytes = match json_data.get("max_bytes") { + Some(value) => Some( + value + .as_u64() + .ok_or_else(|| anyhow!("Invalid 'max_bytes' in arguments"))? + as usize, + ), + None => None, + }; + let uri = match json_data.get("arguments").and_then(Value::as_object) { + Some(args) if !args.is_empty() => expand_uri_template(uri, args)?, + _ => uri.to_string(), + }; + + let result = ctx.tool_scope.mcp_runtime.read(server, &uri).await?; + let audience = ctx + .tool_scope + .mcp_runtime + .resource_audience(server, &uri) + .await; + let items: Vec = result + .contents + .iter() + .map(serde_json::to_value) + .collect::>()?; + + let mut rendered_items = Vec::with_capacity(items.len()); + let mut total_size = 0usize; + for (index, item) in items.iter().enumerate() { + let mut rendered = render_resource_content(item, pattern, offset, max_bytes, server)?; + if let (Some(audience), Some(map)) = (&audience, rendered.as_object_mut()) { + map.insert("audience".to_string(), json!(audience)); + } + let size = rendered.to_string().len(); + // Bound the overall response; the first item is always included. + if index > 0 && total_size + size > render::TEXT_MAX_BYTES_CLAMP { + let omitted = items.len() - index; + rendered_items.push(json!({ + "truncated": true, + "omitted_items": omitted, + "note": format!( + "{omitted} content item(s) omitted: the combined response would exceed \ + {} bytes", + render::TEXT_MAX_BYTES_CLAMP + ), + })); + break; + } + total_size += size; + rendered_items.push(rendered); + } + + if rendered_items.len() == 1 { + Ok(rendered_items.remove(0)) + } else { + Ok(Value::Array(rendered_items)) + } + } + + async fn get_mcp_prompt( + ctx: &RequestContext, + cmd_name: &str, + json_data: &Value, + ) -> Result { + let server = cmd_name + .strip_prefix(&format!("{MCP_PROMPT_META_FUNCTION_NAME_PREFIX}_")) + .ok_or_else(|| anyhow!("Malformed MCP prompt function name: {cmd_name}"))?; + let prompt = json_data + .get("prompt") + .ok_or_else(|| anyhow!("Missing 'prompt' in arguments"))? + .as_str() + .ok_or_else(|| anyhow!("Invalid 'prompt' in arguments"))?; + let mut provided = HashMap::new(); + if let Some(value) = json_data.get("arguments") { + let entries = value + .as_object() + .ok_or_else(|| anyhow!("Invalid 'arguments' in arguments"))?; + for (key, value) in entries { + let value = value.as_str().ok_or_else(|| { + anyhow!( + "Invalid value for prompt argument '{key}': prompt arguments are strings" + ) + })?; + provided.insert(key.clone(), value.to_string()); + } + } + + let declared = ctx + .tool_scope + .mcp_runtime + .list_prompts(server) + .await? + .into_iter() + .find(|candidate| candidate.name == prompt) + .ok_or_else(|| { + anyhow!( + "Prompt '{prompt}' not found on MCP server '{server}'; call the describe \ + meta-tool with kind \"prompt\" to list available prompts" + ) + })? + .arguments + .unwrap_or_default(); + let (arguments, missing) = resolve_prompt_args(&declared, provided); + if !missing.is_empty() { + bail!( + "Missing required prompt argument(s): {}. Provide them as string values in \ + 'arguments'.", + missing.join(", ") + ); + } + + let result = ctx + .tool_scope + .mcp_runtime + .prompt(server, prompt, arguments) + .await?; + Ok(Value::String(flatten_prompt_messages(&result.messages))) } fn extract_call_config_from_agent( @@ -1487,6 +1825,308 @@ impl ToolCall { } } +fn expand_uri_template(template: &str, args: &serde_json::Map) -> Result { + let mut expanded = String::with_capacity(template.len()); + let mut rest = template; + while let Some(start) = rest.find('{') { + expanded.push_str(&rest[..start]); + let after = &rest[start + 1..]; + let Some(len) = after.find('}') else { + bail!("Unclosed '{{' in URI template: {template}"); + }; + expanded.push_str(&expand_uri_template_variable(&after[..len], args)?); + rest = &after[len + 1..]; + } + expanded.push_str(rest); + Ok(expanded) +} + +fn expand_uri_template_variable( + expr: &str, + args: &serde_json::Map, +) -> Result { + const LEVEL_1_ONLY: &str = "only RFC 6570 Level 1 simple substitution {var} is supported"; + if let Some(operator) = expr.chars().next().filter(|c| "+#./;?&".contains(*c)) { + let name = match operator { + '+' => "reserved-expansion", + '#' => "fragment-expansion", + '.' => "label-expansion", + '/' => "path-segment-expansion", + ';' => "path-style-parameter-expansion", + '?' => "form-style-query-expansion", + _ => "form-style-query-continuation", + }; + bail!( + "The '{operator}' {name} operator in '{{{expr}}}' requires RFC 6570 Level 2 or \ + higher; {LEVEL_1_ONLY}" + ); + } + if expr.contains(',') { + bail!( + "The ',' multi-variable expression '{{{expr}}}' requires RFC 6570 Level 3; {LEVEL_1_ONLY}" + ); + } + if expr.contains(':') { + bail!("The ':' prefix modifier in '{{{expr}}}' requires RFC 6570 Level 4; {LEVEL_1_ONLY}"); + } + if expr.ends_with('*') { + bail!("The '*' explode modifier in '{{{expr}}}' requires RFC 6570 Level 4; {LEVEL_1_ONLY}"); + } + if expr.is_empty() + || !expr + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '.') + { + bail!("Invalid variable name '{{{expr}}}' in URI template; expected [A-Za-z0-9_.]+"); + } + let value = args + .get(expr) + .ok_or_else(|| anyhow!("URI template variable '{expr}' is missing from 'arguments'"))?; + let text = match value { + Value::String(text) => text.clone(), + Value::Number(number) => number.to_string(), + Value::Bool(boolean) => boolean.to_string(), + other => bail!( + "URI template variable '{expr}' must be a string, number, or boolean; got {other}" + ), + }; + Ok(urlencoding::encode(&text).into_owned()) +} + +#[derive(Debug)] +enum ResourceContentBody { + Text(String), + Blob(String), +} + +// rmcp's untagged ResourceContents enum cannot represent malformed items +// (both or neither of text/blob), so classification happens on the raw Value. +fn parse_resource_content(item: &Value) -> Result { + let text = item.get("text"); + let blob = item.get("blob"); + match (text, blob) { + (Some(text), None) => Ok(ResourceContentBody::Text( + text.as_str() + .ok_or_else(|| anyhow!("Resource content 'text' is not a string"))? + .to_string(), + )), + (None, Some(blob)) => Ok(ResourceContentBody::Blob( + blob.as_str() + .ok_or_else(|| anyhow!("Resource content 'blob' is not a string"))? + .to_string(), + )), + (Some(_), Some(_)) => { + bail!("Resource content item has both 'text' and 'blob'; expected exactly one") + } + (None, None) => { + bail!("Resource content item has neither 'text' nor 'blob'; expected exactly one") + } + } +} + +fn render_resource_content( + item: &Value, + pattern: Option<&str>, + offset: usize, + max_bytes: Option, + server: &str, +) -> Result { + let uri = item + .get("uri") + .and_then(Value::as_str) + .map(render::clamp_metadata); + let mime_type = item + .get("mimeType") + .and_then(Value::as_str) + .map(render::clamp_metadata); + let text = match parse_resource_content(item)? { + ResourceContentBody::Text(text) => text, + ResourceContentBody::Blob(blob) => { + match render::render_blob(&blob, mime_type.as_deref(), server)? { + render::RenderedBlob::Text(text) => text, + render::RenderedBlob::Spilled(meta) => { + let mut value = serde_json::to_value(meta)?; + if let Some(map) = value.as_object_mut() { + map.insert("uri".to_string(), json!(uri)); + } + return Ok(value); + } + } + } + }; + let rendered = render::render_text(&text, pattern, offset, max_bytes)?; + let mut value = json!({ + "uri": uri, + "mime_type": mime_type, + "text": rendered.text, + "truncated": rendered.truncated, + "total_bytes": rendered.total_bytes, + "next_offset": rendered.next_offset, + }); + + if let Some(next_offset) = rendered.next_offset { + value["note"] = json!(format!( + "Content truncated; re-call with offset={next_offset} to continue (max_bytes is \ + clamped to {})", + render::TEXT_MAX_BYTES_CLAMP + )); + } + + Ok(value) +} + +// Terminal-only rendering of an MCP dispatch error: escape sequences are +// stripped so a hostile server cannot drive the terminal, while the JSON +// payload keeps the raw message. +fn mcp_error_display(error_msg: &str) -> String { + sanitize_display_text(&format!("⚠️ {error_msg} ⚠️")) +} + +/// Bounds a raw `CallToolResult` JSON value: oversized text is sliced, +/// base64 blob content is routed through the blob renderer instead of +/// reaching model context, and oversized structured content is replaced with +/// a truncation marker. In-bounds results pass through unchanged. +fn render_tool_result(mut result: Value, server: &str) -> Result { + let Some(map) = result.as_object_mut() else { + return Ok(result); + }; + if let Some(items) = map.get_mut("content").and_then(Value::as_array_mut) { + for item in items { + render_tool_content_item(item, server)?; + } + } + let oversized_structured = map + .get("structuredContent") + .is_some_and(|structured| structured.to_string().len() > render::TEXT_MAX_BYTES_CLAMP); + if oversized_structured { + map.insert( + "structuredContent".to_string(), + json!({ + "truncated": true, + "note": format!( + "structuredContent omitted: its serialized form exceeds \ + TEXT_MAX_BYTES_CLAMP ({} bytes); re-call the tool with narrower \ + arguments", + render::TEXT_MAX_BYTES_CLAMP + ), + }), + ); + } + + Ok(result) +} + +fn render_tool_content_item(item: &mut Value, server: &str) -> Result<()> { + match item.get("type").and_then(Value::as_str) { + Some("text") => clamp_tool_text(item), + Some("image") | Some("audio") => { + let mime_type = item + .get("mimeType") + .and_then(Value::as_str) + .map(render::clamp_metadata); + if let Some(data) = item.get("data").and_then(Value::as_str) { + let replacement = render_tool_blob(data, mime_type, None, server)?; + *item = replacement; + } + } + Some("resource") => { + let Some(resource) = item.get("resource") else { + return Ok(()); + }; + let uri = resource + .get("uri") + .and_then(Value::as_str) + .map(render::clamp_metadata); + let mime_type = resource + .get("mimeType") + .and_then(Value::as_str) + .map(render::clamp_metadata); + if let Some(blob) = resource.get("blob").and_then(Value::as_str) { + let replacement = render_tool_blob(blob, mime_type, uri, server)?; + *item = replacement; + } else if let Some(resource) = item.get_mut("resource") { + clamp_tool_text(resource); + clamp_metadata_field(resource, "uri"); + clamp_metadata_field(resource, "mimeType"); + } + } + _ => {} + } + + Ok(()) +} + +fn render_tool_blob( + b64: &str, + mime_type: Option, + uri: Option, + server: &str, +) -> Result { + let mut value = match render::render_blob(b64, mime_type.as_deref(), server) { + Ok(render::RenderedBlob::Text(text)) => { + let mut item = json!({ "type": "text", "text": text }); + clamp_tool_text(&mut item); + item + } + Ok(render::RenderedBlob::Spilled(meta)) => serde_json::to_value(meta)?, + // One undecodable item must not sink the rest of the result. + Err(error) => json!({ "error": format!("Failed to render blob content: {error}") }), + }; + + if let Some(map) = value.as_object_mut() { + if let Some(mime_type) = mime_type + && !map.contains_key("mime_type") + { + map.insert("mime_type".to_string(), json!(mime_type)); + } + if let Some(uri) = uri { + map.insert("uri".to_string(), json!(uri)); + } + } + + Ok(value) +} + +fn clamp_tool_text(container: &mut Value) { + let Some(text) = container.get("text").and_then(Value::as_str) else { + return; + }; + + if text.len() <= render::TEXT_MAX_BYTES_CLAMP { + return; + } + + let total_bytes = text.len(); + let clamped = render::truncate_utf8(text, render::TEXT_MAX_BYTES_CLAMP).to_string(); + let Some(map) = container.as_object_mut() else { + return; + }; + map.insert("text".to_string(), json!(clamped)); + map.insert("truncated".to_string(), json!(true)); + map.insert("total_bytes".to_string(), json!(total_bytes)); + map.insert( + "note".to_string(), + json!(format!( + "Text truncated; re-call the tool with narrower arguments (text is clamped to \ + TEXT_MAX_BYTES_CLAMP = {} bytes)", + render::TEXT_MAX_BYTES_CLAMP + )), + ); +} + +fn clamp_metadata_field(object: &mut Value, key: &str) { + let Some(text) = object.get(key).and_then(Value::as_str) else { + return; + }; + + if text.len() <= render::METADATA_MAX_BYTES { + return; + } + + let clamped = render::clamp_metadata(text); + object[key] = json!(clamped); +} + pub fn run_llm_function( cmd_name: String, cmd_args: Vec, @@ -1830,9 +2470,18 @@ fn format_json_colored_keys(value: &serde_json::Value) -> String { #[cfg(test)] mod tests { use super::*; + use crate::config::test_fixtures::{ + FIXTURE_ANNOTATED_TEXT, FIXTURE_ANNOTATED_URI, FIXTURE_BLOB_BYTES, FIXTURE_BLOB_URI, + FIXTURE_LOG_TEXT, FIXTURE_LOG_URI, FixtureServer, fixture_runtime, + }; use crate::config::{AppState, WorkingMode}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; + use base64::Engine; + use base64::engine::general_purpose::STANDARD; + use rmcp::model::{CallToolResult, ContentBlock}; use serde_json::json; + use serial_test::serial; + use std::process; use std::sync::Arc; fn call(name: &str, id: Option<&str>) -> ToolCall { @@ -1843,6 +2492,19 @@ mod tests { ToolCall::new(name.to_string(), args, Some("id1".to_string())) } + fn mcp_features(name: &str, tools: bool, resources: bool, prompts: bool) -> McpServerFeatures { + McpServerFeatures { + name: name.to_string(), + tools, + resources, + prompts, + } + } + + fn tools_only(name: &str) -> McpServerFeatures { + mcp_features(name, true, false, false) + } + fn run_async(f: F) -> F::Output { tokio::runtime::Builder::new_current_thread() .enable_all() @@ -2128,6 +2790,8 @@ mod tests { assert_eq!(MCP_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke"); assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search"); assert_eq!(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, "mcp_describe"); + assert_eq!(MCP_READ_META_FUNCTION_NAME_PREFIX, "mcp_read"); + assert_eq!(MCP_PROMPT_META_FUNCTION_NAME_PREFIX, "mcp_prompt"); } #[test] @@ -2210,7 +2874,9 @@ mod tests { #[test] fn functions_append_mcp_meta_creates_three_per_server() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["github".to_string()]); + + f.append_mcp_meta_functions(vec![tools_only("github")]); + assert_eq!(f.declarations().len(), 3); assert!(f.contains("mcp_invoke_github")); assert!(f.contains("mcp_search_github")); @@ -2220,7 +2886,9 @@ mod tests { #[test] fn functions_append_mcp_meta_multiple_servers() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["github".into(), "slack".into()]); + + f.append_mcp_meta_functions(vec![tools_only("github"), tools_only("slack")]); + assert_eq!(f.declarations().len(), 6); assert!(f.contains("mcp_invoke_github")); assert!(f.contains("mcp_invoke_slack")); @@ -2233,6 +2901,140 @@ mod tests { assert!(f.is_empty()); } + #[test] + fn functions_append_mcp_meta_resources_only_omits_invoke() { + let mut f = Functions::default(); + + f.append_mcp_meta_functions(vec![mcp_features("res", false, true, false)]); + + assert_eq!(f.declarations().len(), 3); + assert!(!f.contains("mcp_invoke_res")); + assert!(!f.contains("mcp_prompt_res")); + assert!(f.contains("mcp_search_res")); + assert!(f.contains("mcp_describe_res")); + assert!(f.contains("mcp_read_res")); + } + + #[test] + fn functions_append_mcp_meta_all_capabilities_emits_five() { + let mut f = Functions::default(); + + f.append_mcp_meta_functions(vec![mcp_features("srv", true, true, true)]); + + assert_eq!(f.declarations().len(), 5); + assert!(f.contains("mcp_invoke_srv")); + assert!(f.contains("mcp_search_srv")); + assert!(f.contains("mcp_describe_srv")); + assert!(f.contains("mcp_read_srv")); + assert!(f.contains("mcp_prompt_srv")); + } + + #[test] + fn features_from_missing_capabilities_fail_open_for_tools() { + let features = McpServerFeatures::from_capabilities("srv", None); + assert!(features.tools); + assert!(!features.resources); + assert!(!features.prompts); + + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![features]); + assert!(f.contains("mcp_invoke_srv")); + } + + #[test] + fn gated_prefixes_tools_only() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, false, false)), + vec![ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_tools_and_resources_include_read() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, true, false)), + vec![ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_tools_and_prompts_include_prompt() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, false, true)), + vec![ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_all_capabilities_include_all() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, true, true)), + MCP_META_FUNCTION_PREFIXES.to_vec() + ); + } + + #[test] + fn gated_prefixes_resources_only_omit_invoke() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, true, false)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_prompts_only_omit_invoke_and_read() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, false, true)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_resources_and_prompts_omit_invoke() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, true, true)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_no_capabilities_keep_search_and_describe() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, false, false)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + ] + ); + } + #[test] fn functions_find_returns_declaration() { let mut f = Functions::default(); @@ -2264,7 +3066,7 @@ mod tests { #[test] fn functions_mcp_invoke_declaration_has_tool_and_arguments_params() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_invoke_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); assert!(props.contains_key("tool")); @@ -2276,7 +3078,7 @@ mod tests { #[test] fn functions_mcp_search_declaration_has_query_and_top_k_params() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_search_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); assert!(props.contains_key("query")); @@ -2286,12 +3088,707 @@ mod tests { #[test] fn functions_mcp_describe_declaration_has_tool_param() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_describe_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); assert!(props.contains_key("tool")); } + #[test] + fn functions_mcp_describe_declaration_has_optional_kind_param() { + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![tools_only("srv")]); + let decl = f.find("mcp_describe_srv").unwrap(); + let props = decl.parameters.properties.as_ref().unwrap(); + let kind = props.get("kind").unwrap(); + assert_eq!(kind.default, Some(Value::from("tool"))); + let required = decl.parameters.required.as_ref().unwrap(); + assert_eq!(required, &vec!["tool".to_string()]); + } + + #[test] + fn eval_mcp_describe_without_kind_defaults_to_tool() { + let output = run_async(async { + let (runtime, _server) = fixture_runtime(FixtureServer::default()).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.mcp_runtime = runtime; + let call = call_with_args("mcp_describe_fixture", json!({"tool": "dup"})); + call.eval_mcp(&ctx).await + }) + .unwrap(); + + assert_eq!( + output, + json!({ + "type": "object", + "properties": { + "tool": { "type": "string" }, + "arguments": { + "type": "object", + "properties": { "q": { "type": "string" } } + } + } + }) + ); + } + + fn template_args(pairs: &[(&str, Value)]) -> serde_json::Map { + pairs + .iter() + .map(|(key, value)| (key.to_string(), value.clone())) + .collect() + } + + fn resources_fixture() -> FixtureServer { + FixtureServer { + resources_capability: true, + ..Default::default() + } + } + + async fn eval_mcp_read(args: Value) -> Result { + let (runtime, _server) = fixture_runtime(resources_fixture()).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.mcp_runtime = runtime; + call_with_args("mcp_read_fixture", args) + .eval_mcp(&ctx) + .await + } + + fn prompts_fixture() -> FixtureServer { + FixtureServer { + prompts_capability: true, + ..Default::default() + } + } + + async fn eval_mcp_prompt(args: Value) -> Result { + let (runtime, _server) = fixture_runtime(prompts_fixture()).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.mcp_runtime = runtime; + call_with_args("mcp_prompt_fixture", args) + .eval_mcp(&ctx) + .await + } + + const FLATTENED_SUMMARIZE_PROMPT: &str = + "[user]\nSummarize notes.txt\n\n[assistant]\nIn which style?\n\n[user]\nConcise."; + + #[test] + fn expand_uri_template_substitutes_simple_vars() { + let args = template_args(&[("path", json!("docs")), ("name", json!("readme"))]); + assert_eq!( + expand_uri_template("file:///{path}/{name}", &args).unwrap(), + "file:///docs/readme" + ); + } + + #[test] + fn expand_uri_template_stringifies_numbers_and_bools() { + let args = template_args(&[("id", json!(42)), ("flag", json!(true))]); + assert_eq!( + expand_uri_template("item://{id}/{flag}", &args).unwrap(), + "item://42/true" + ); + } + + #[test] + fn expand_uri_template_percent_encodes_values() { + let args = template_args(&[("q", json!("a b/c✓"))]); + assert_eq!( + expand_uri_template("search://{q}", &args).unwrap(), + "search://a%20b%2Fc%E2%9C%93" + ); + } + + #[test] + fn expand_uri_template_without_placeholders_is_noop() { + let args = template_args(&[("unused", json!("x"))]); + assert_eq!( + expand_uri_template("file:///static", &args).unwrap(), + "file:///static" + ); + } + + #[test] + fn expand_uri_template_missing_variable_names_it() { + let err = expand_uri_template("file:///{path}", &template_args(&[])) + .unwrap_err() + .to_string(); + assert!(err.contains("'path'"), "{err}"); + assert!(err.contains("missing"), "{err}"); + } + + #[test] + fn expand_uri_template_rejects_higher_level_operators() { + let args = template_args(&[("var", json!("v"))]); + for operator in ["+", "#", ".", "/", ";", "?", "&"] { + let err = expand_uri_template(&format!("x://{{{operator}var}}"), &args) + .unwrap_err() + .to_string(); + assert!(err.contains(&format!("'{operator}'")), "{err}"); + assert!(err.contains("Level 1 simple substitution"), "{err}"); + } + } + + #[test] + fn expand_uri_template_rejects_modifiers_and_multi_vars() { + let args = template_args(&[("a", json!("v")), ("b", json!("w")), ("var", json!("v"))]); + for (template, construct) in [ + ("x://{var*}", "'*' explode modifier"), + ("x://{var:3}", "':' prefix modifier"), + ("x://{a,b}", "',' multi-variable"), + ] { + let err = expand_uri_template(template, &args) + .unwrap_err() + .to_string(); + assert!(err.contains(construct), "{err}"); + assert!(err.contains("Level 1 simple substitution"), "{err}"); + } + } + + #[test] + fn expand_uri_template_unclosed_brace_errors() { + let err = expand_uri_template("file:///{path", &template_args(&[])) + .unwrap_err() + .to_string(); + assert!(err.contains("Unclosed"), "{err}"); + } + + #[test] + fn expand_uri_template_rejects_invalid_variable_names() { + let err = expand_uri_template("x://{va r}", &template_args(&[])) + .unwrap_err() + .to_string(); + assert!(err.contains("Invalid variable name"), "{err}"); + } + + #[test] + fn expand_uri_template_rejects_non_scalar_values() { + for value in [json!(null), json!(["a"]), json!({"k": "v"})] { + let args = template_args(&[("v", value)]); + let err = expand_uri_template("x://{v}", &args) + .unwrap_err() + .to_string(); + assert!(err.contains("string, number, or boolean"), "{err}"); + } + } + + #[test] + fn parse_resource_content_classifies_by_field_presence() { + assert!(matches!( + parse_resource_content(&json!({"uri": "u", "text": "hi"})).unwrap(), + ResourceContentBody::Text(text) if text == "hi" + )); + assert!(matches!( + parse_resource_content(&json!({"uri": "u", "blob": "aGk="})).unwrap(), + ResourceContentBody::Blob(blob) if blob == "aGk=" + )); + } + + #[test] + fn parse_resource_content_rejects_both_and_neither() { + let err = parse_resource_content(&json!({"text": "t", "blob": "b"})) + .unwrap_err() + .to_string(); + assert!(err.contains("both"), "{err}"); + + let err = parse_resource_content(&json!({"uri": "u"})) + .unwrap_err() + .to_string(); + assert!(err.contains("neither"), "{err}"); + } + + #[test] + fn functions_mcp_read_declaration_has_paging_params() { + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![mcp_features("srv", false, true, false)]); + let decl = f.find("mcp_read_srv").unwrap(); + let props = decl.parameters.properties.as_ref().unwrap(); + for param in ["uri", "arguments", "pattern", "offset", "max_bytes"] { + assert!(props.contains_key(param), "missing {param}"); + } + assert_eq!(props.len(), 5); + assert_eq!(props["offset"].default, Some(Value::from(0usize))); + assert_eq!( + props["max_bytes"].default, + Some(Value::from(render::DEFAULT_TEXT_MAX_BYTES)) + ); + assert_eq!(decl.parameters.required, Some(vec!["uri".to_string()])); + } + + #[test] + fn mcp_read_routes_through_the_concurrent_mcp_path() { + assert!(is_mcp_meta_function("mcp_read_x")); + } + + #[test] + fn eval_mcp_read_returns_rendered_text() { + let output = run_async(eval_mcp_read(json!({"uri": FIXTURE_LOG_URI}))).unwrap(); + + assert_eq!(output["uri"], FIXTURE_LOG_URI); + assert_eq!(output["mime_type"], "text/plain"); + assert_eq!(output["text"], FIXTURE_LOG_TEXT); + assert_eq!(output["truncated"], false); + assert_eq!(output["next_offset"], Value::Null); + } + + #[test] + fn eval_routes_mcp_read_to_resource_handler() { + let output = run_async(async { + let (runtime, _server) = fixture_runtime(resources_fixture()).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope + .functions + .append_mcp_meta_functions(vec![mcp_features("fixture", true, true, false)]); + ctx.tool_scope.mcp_runtime = runtime; + let call = call_with_args("mcp_read_fixture", json!({"uri": FIXTURE_LOG_URI})); + call.eval(&mut ctx).await + }) + .unwrap(); + + assert_eq!(output["text"], FIXTURE_LOG_TEXT); + } + + #[test] + fn functions_mcp_prompt_declaration_has_prompt_and_arguments_params() { + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![mcp_features("srv", false, false, true)]); + let decl = f.find("mcp_prompt_srv").unwrap(); + let props = decl.parameters.properties.as_ref().unwrap(); + assert!(props.contains_key("prompt")); + assert!(props.contains_key("arguments")); + assert_eq!(props.len(), 2); + assert_eq!(decl.parameters.required, Some(vec!["prompt".to_string()])); + } + + #[test] + fn eval_mcp_routes_mcp_prompt_to_prompt_handler() { + let fixture = prompts_fixture(); + let get_prompt_calls = Arc::clone(&fixture.get_prompt_calls); + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + let output = run_async(async { + let (runtime, _server) = fixture_runtime(fixture).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.mcp_runtime = runtime; + let call = call_with_args( + "mcp_prompt_fixture", + json!({"prompt": "summarize", "arguments": {"path": "notes.txt"}}), + ); + call.eval_mcp(&ctx).await + }) + .unwrap(); + + assert_eq!(output, Value::String(FLATTENED_SUMMARIZE_PROMPT.into())); + assert_eq!(get_prompt_calls.load(Ordering::SeqCst), 1); + assert_eq!(call_tool_calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn eval_routes_mcp_prompt_to_prompt_handler() { + let fixture = prompts_fixture(); + let get_prompt_calls = Arc::clone(&fixture.get_prompt_calls); + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + let output = run_async(async { + let (runtime, _server) = fixture_runtime(fixture).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope + .functions + .append_mcp_meta_functions(vec![mcp_features("fixture", true, false, true)]); + ctx.tool_scope.mcp_runtime = runtime; + let call = call_with_args( + "mcp_prompt_fixture", + json!({"prompt": "summarize", "arguments": {"path": "notes.txt"}}), + ); + call.eval(&mut ctx).await + }) + .unwrap(); + + assert_eq!(output, Value::String(FLATTENED_SUMMARIZE_PROMPT.into())); + assert_eq!(get_prompt_calls.load(Ordering::SeqCst), 1); + assert_eq!(call_tool_calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn eval_mcp_prompt_missing_required_arg_returns_teaching_error() { + let output = run_async(eval_mcp_prompt(json!({"prompt": "summarize"}))).unwrap(); + + let err = output["tool_call_error"].as_str().unwrap(); + assert!( + err.contains("Missing required prompt argument(s): path"), + "{err}" + ); + } + + #[test] + fn eval_mcp_prompt_unknown_prompt_returns_teaching_error() { + let output = run_async(eval_mcp_prompt(json!({"prompt": "ghost"}))).unwrap(); + + let err = output["tool_call_error"].as_str().unwrap(); + assert!( + err.contains("Prompt 'ghost' not found on MCP server 'fixture'"), + "{err}" + ); + assert!(err.contains("kind \"prompt\""), "{err}"); + } + + #[test] + fn eval_mcp_prompt_rejects_non_string_argument_values() { + let output = run_async(eval_mcp_prompt( + json!({"prompt": "summarize", "arguments": {"path": 5}}), + )) + .unwrap(); + + let err = output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("prompt arguments are strings"), "{err}"); + } + + #[test] + fn eval_mcp_read_pages_text_with_offset() { + let (page1, page2) = run_async(async { + let (runtime, _server) = fixture_runtime(resources_fixture()).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.mcp_runtime = runtime; + let call = call_with_args( + "mcp_read_fixture", + json!({"uri": FIXTURE_LOG_URI, "max_bytes": 20}), + ); + let page1 = call.eval_mcp(&ctx).await.unwrap(); + let call = call_with_args( + "mcp_read_fixture", + json!({"uri": FIXTURE_LOG_URI, "offset": page1["next_offset"], "max_bytes": 20}), + ); + let page2 = call.eval_mcp(&ctx).await.unwrap(); + (page1, page2) + }); + + assert_eq!(page1["truncated"], true); + assert!(page1["note"].as_str().unwrap().contains("204800")); + let text1 = page1["text"].as_str().unwrap(); + let text2 = page2["text"].as_str().unwrap(); + assert!(!text2.is_empty()); + assert!(FIXTURE_LOG_TEXT.starts_with(&format!("{text1}{text2}"))); + } + + #[test] + fn eval_mcp_read_pattern_filters_lines_with_context() { + let output = run_async(eval_mcp_read( + json!({"uri": FIXTURE_LOG_URI, "pattern": "café"}), + )) + .unwrap(); + + let text = output["text"].as_str().unwrap(); + assert!(text.contains("6:ERROR: café overheated"), "{text}"); + assert!(text.contains("4-fourth line"), "{text}"); + assert!(!text.contains("disk full"), "{text}"); + assert_eq!(output["total_bytes"], text.len()); + } + + #[test] + fn eval_mcp_read_invalid_pattern_returns_teaching_error() { + let output = run_async(eval_mcp_read( + json!({"uri": FIXTURE_LOG_URI, "pattern": "("}), + )) + .unwrap(); + + let err = output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Invalid filter pattern"), "{err}"); + } + + #[test] + #[serial] + fn eval_mcp_read_blob_spills_with_metadata() { + static COUNTER: AtomicU64 = AtomicU64::new(0); + let cache_dir = env::temp_dir().join(format!( + "coyote-read-blob-{}-{}", + process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + )); + fs::create_dir_all(&cache_dir).unwrap(); + let env_name = get_env_name("cache_dir"); + let previous = env::var_os(&env_name); + unsafe { env::set_var(&env_name, &cache_dir) }; + + let output = run_async(eval_mcp_read(json!({"uri": FIXTURE_BLOB_URI}))).unwrap(); + + unsafe { + match previous { + Some(value) => env::set_var(&env_name, value), + None => env::remove_var(&env_name), + } + } + + assert_eq!(output["spilled"], true); + assert_eq!(output["uri"], FIXTURE_BLOB_URI); + assert_eq!(output["mime_type"], "application/pdf"); + assert_eq!(output["sha256"].as_str().unwrap().len(), 64); + let path = PathBuf::from(output["path"].as_str().unwrap()); + assert!(path.starts_with(&cache_dir)); + assert_eq!(fs::read(&path).unwrap(), FIXTURE_BLOB_BYTES); + + fs::remove_dir_all(&cache_dir).unwrap(); + } + + #[test] + fn eval_mcp_read_multi_content_returns_array() { + let output = run_async(eval_mcp_read(json!({"uri": "file:///multi"}))).unwrap(); + + let items = output.as_array().unwrap(); + assert_eq!(items.len(), 3); + assert_eq!(items[0]["text"], "first"); + assert_eq!(items[1]["text"], "second"); + assert_eq!(items[2]["text"], "third"); + assert_eq!(items[1]["uri"], "file:///multi/1"); + } + + #[test] + fn eval_mcp_read_multi_content_enforces_overall_ceiling() { + let output = run_async(eval_mcp_read( + json!({"uri": "file:///huge", "max_bytes": render::TEXT_MAX_BYTES_CLAMP}), + )) + .unwrap(); + + let items = output.as_array().unwrap(); + assert_eq!(items.len(), 2); + assert_eq!(items[0]["text"].as_str().unwrap().len(), 150 * 1024); + let marker = &items[1]; + assert_eq!(marker["truncated"], true); + assert_eq!(marker["omitted_items"], 2); + let note = marker["note"].as_str().unwrap(); + assert!(note.contains("2 content item(s) omitted"), "{note}"); + assert!(note.contains("204800"), "{note}"); + } + + #[test] + fn eval_mcp_read_expands_template_end_to_end() { + let output = run_async(eval_mcp_read(json!({ + "uri": "file:///{path}/{name}", + "arguments": {"path": "docs", "name": "readme"}, + }))) + .unwrap(); + + assert_eq!(output["uri"], "file:///docs/readme"); + assert_eq!(output["text"], "readme body"); + } + + #[test] + fn eval_mcp_read_attaches_catalog_audience() { + let output = run_async(eval_mcp_read(json!({"uri": FIXTURE_ANNOTATED_URI}))).unwrap(); + + assert_eq!(output["audience"], json!(["user"])); + assert_eq!(output["text"], FIXTURE_ANNOTATED_TEXT); + } + + #[test] + fn eval_mcp_read_omits_audience_for_unannotated_uri() { + let output = run_async(eval_mcp_read(json!({"uri": FIXTURE_LOG_URI}))).unwrap(); + + assert!(output.get("audience").is_none()); + } + + #[test] + fn mcp_error_display_strips_terminal_escapes() { + let hostile = "fail\u{1b}[31mred\u{1b}]0;pwn\u{7}end"; + + let display = mcp_error_display(hostile); + + assert!(!display.contains('\u{1b}')); + assert_eq!(display, "⚠️ failredend ⚠️"); + // The payload keeps the raw message; only the terminal string differs. + assert_ne!(display, format!("⚠️ {hostile} ⚠️")); + } + + #[test] + fn render_tool_result_passes_in_bounds_result_through_unchanged() { + let mut result = CallToolResult::success(vec![ContentBlock::text("small text")]); + result.structured_content = Some(json!({"rows": [1, 2, 3]})); + + let bounded = render_tool_result(serde_json::to_value(&result).unwrap(), "srv").unwrap(); + + assert_eq!(bounded, serde_json::to_value(&result).unwrap()); + assert_eq!(bounded["isError"], false); + } + + #[test] + fn render_tool_result_clamps_oversized_text() { + let result = CallToolResult::success(vec![ContentBlock::text( + "x".repeat(render::TEXT_MAX_BYTES_CLAMP + 10), + )]); + + let bounded = render_tool_result(serde_json::to_value(&result).unwrap(), "srv").unwrap(); + + let item = &bounded["content"][0]; + assert_eq!( + item["text"].as_str().unwrap().len(), + render::TEXT_MAX_BYTES_CLAMP + ); + assert_eq!(item["truncated"], true); + assert_eq!(item["total_bytes"], render::TEXT_MAX_BYTES_CLAMP + 10); + assert!( + item["note"] + .as_str() + .unwrap() + .contains("TEXT_MAX_BYTES_CLAMP") + ); + assert_eq!(bounded["isError"], false); + } + + #[test] + #[serial] + fn render_tool_result_spills_blob_content_without_base64() { + static COUNTER: AtomicU64 = AtomicU64::new(0); + let cache_dir = env::temp_dir().join(format!( + "coyote-tool-blob-{}-{}", + process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + )); + fs::create_dir_all(&cache_dir).unwrap(); + let env_name = get_env_name("cache_dir"); + let previous = env::var_os(&env_name); + unsafe { env::set_var(&env_name, &cache_dir) }; + + let b64 = STANDARD.encode(FIXTURE_BLOB_BYTES); + let result = CallToolResult::success(vec![ContentBlock::image(b64.clone(), "image/png")]); + let bounded = render_tool_result(serde_json::to_value(&result).unwrap(), "srv"); + + unsafe { + match previous { + Some(value) => env::set_var(&env_name, value), + None => env::remove_var(&env_name), + } + } + + let bounded = bounded.unwrap(); + let item = &bounded["content"][0]; + assert_eq!(item["spilled"], true); + assert_eq!(item["mime_type"], "image/png"); + assert_eq!(item["sha256"].as_str().unwrap().len(), 64); + assert!( + !bounded.to_string().contains(&b64), + "base64 payload must not reach model context" + ); + + fs::remove_dir_all(&cache_dir).unwrap(); + } + + #[test] + fn render_tool_result_inlines_utf8_blob_as_text() { + let b64 = STANDARD.encode("hello ✓ world"); + let result = CallToolResult::success(vec![ContentBlock::image(b64.clone(), "image/png")]); + + let bounded = render_tool_result(serde_json::to_value(&result).unwrap(), "srv").unwrap(); + + let item = &bounded["content"][0]; + assert_eq!(item["type"], "text"); + assert_eq!(item["text"], "hello ✓ world"); + assert!(item.get("spilled").is_none()); + assert!(!bounded.to_string().contains(&b64)); + } + + #[test] + fn render_tool_result_degrades_undecodable_blob_item() { + let value = json!({ + "content": [ + {"type": "image", "data": "!!!not base64!!!", "mimeType": "image/png"}, + {"type": "text", "text": "still here"}, + ], + }); + + let bounded = render_tool_result(value, "srv").unwrap(); + + let error = bounded["content"][0]["error"].as_str().unwrap(); + assert!(error.contains("base64"), "{error}"); + assert_eq!(bounded["content"][0]["mime_type"], "image/png"); + assert_eq!(bounded["content"][1]["text"], "still here"); + } + + #[test] + fn render_tool_result_replaces_oversized_structured_content() { + let mut result = CallToolResult::success(vec![]); + result.structured_content = + Some(json!({"blob": "x".repeat(render::TEXT_MAX_BYTES_CLAMP + 1)})); + + let bounded = render_tool_result(serde_json::to_value(&result).unwrap(), "srv").unwrap(); + + let structured = &bounded["structuredContent"]; + assert_eq!(structured["truncated"], true); + assert!( + structured["note"] + .as_str() + .unwrap() + .contains("TEXT_MAX_BYTES_CLAMP") + ); + assert!(bounded.to_string().len() < render::TEXT_MAX_BYTES_CLAMP); + } + + #[test] + fn render_tool_result_clamps_embedded_resource_metadata() { + let uri = "u".repeat(render::METADATA_MAX_BYTES + 1); + let value = json!({ + "content": [{"type": "resource", "resource": {"uri": uri, "text": "hi"}}], + }); + + let bounded = render_tool_result(value, "srv").unwrap(); + + let resource = &bounded["content"][0]["resource"]; + assert!( + resource["uri"] + .as_str() + .unwrap() + .contains("METADATA_MAX_BYTES") + ); + assert_eq!(resource["text"], "hi"); + } + + #[test] + fn render_resource_content_clamps_metadata_strings() { + let uri = format!("file:///{}", "u".repeat(render::METADATA_MAX_BYTES)); + let mime = format!("text/{}", "m".repeat(render::METADATA_MAX_BYTES)); + let item = json!({"uri": uri, "mimeType": mime, "text": "hi"}); + + let value = render_resource_content(&item, None, 0, None, "srv").unwrap(); + + assert!( + value["uri"] + .as_str() + .unwrap() + .contains("METADATA_MAX_BYTES") + ); + assert!( + value["mime_type"] + .as_str() + .unwrap() + .contains("METADATA_MAX_BYTES") + ); + assert_eq!(value["text"], "hi"); + } + + #[test] + fn eval_mcp_invoke_bounds_tool_results_end_to_end() { + let mut result = CallToolResult::success(vec![ContentBlock::text("hi")]); + result.structured_content = Some(json!({"ok": true})); + let fixture = FixtureServer { + tool_result: Some(result), + ..Default::default() + }; + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + + let output = run_async(async { + let (runtime, _server) = fixture_runtime(fixture).await; + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.mcp_runtime = runtime; + call_with_args("mcp_invoke_fixture", json!({"tool": "dup"})) + .eval_mcp(&ctx) + .await + }) + .unwrap(); + + assert_eq!(output["content"][0]["text"], "hi"); + assert_eq!(output["structuredContent"], json!({"ok": true})); + assert_eq!(output["isError"], false); + assert_eq!(call_tool_calls.load(Ordering::SeqCst), 1); + } + #[test] fn functions_supervisor_includes_task_queue_tools() { let mut f = Functions::default(); diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index 05cdf65..32c5d48 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -630,14 +630,14 @@ async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[Stri } fn sync_agent_functions_to_ctx(ctx: &mut RequestContext) -> Result<()> { - let server_names = ctx.tool_scope.mcp_runtime.server_names(); + let server_features = ctx.tool_scope.mcp_runtime.server_features(); let functions = { let agent = ctx .agent .as_mut() .with_context(|| "Agent should be initialized")?; - if !server_names.is_empty() { - agent.append_mcp_meta_functions(server_names); + if !server_features.is_empty() { + agent.append_mcp_meta_functions(server_features); } agent.functions().clone() }; @@ -1453,7 +1453,8 @@ async fn summarize_output(ctx: &RequestContext, agent_name: &str, output: &str) #[cfg(test)] mod tests { use super::*; - use crate::config::{AppState, WorkingMode}; + use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; + use crate::config::{AgentConfig, AppState, WorkingMode}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use serde_json::json; use serial_test::serial; @@ -1510,6 +1511,28 @@ mod tests { .block_on(f) } + #[tokio::test] + async fn sync_agent_functions_gates_meta_functions_on_live_capabilities() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + let (runtime, _server) = fixture_runtime(FixtureServer { + tools_capability: false, + resources_capability: true, + ..FixtureServer::default() + }) + .await; + ctx.tool_scope.mcp_runtime = runtime; + + sync_agent_functions_to_ctx(&mut ctx).unwrap(); + + let functions = &ctx.tool_scope.functions; + assert_eq!(functions.declarations().len(), 3); + assert!(functions.contains("mcp_search_fixture")); + assert!(functions.contains("mcp_describe_fixture")); + assert!(functions.contains("mcp_read_fixture")); + assert!(!functions.contains("mcp_invoke_fixture")); + } + #[test] fn handle_list_running_empty_supervisor() { let mut ctx = ctx_with_supervisor(4, 3); diff --git a/src/main.rs b/src/main.rs index 3f98954..3e5d9b0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -380,8 +380,8 @@ async fn run( .await?; } else { let app: Arc = Arc::clone(&ctx.app.config); - if let Some(prompt) = &cli.prompt { - ctx.use_prompt(app.as_ref(), prompt)?; + if let Some(prompt) = &cli.temp_role { + ctx.use_temp_role(app.as_ref(), prompt)?; } else if let Some(name) = &cli.role { ctx.use_role(app.as_ref(), name, abort_signal.clone()) .await?; diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index e6b6222..99d8272 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -1,6 +1,7 @@ mod auth_client; pub(crate) mod manage; pub(crate) mod oauth; +pub(crate) mod render; mod sse_transport; use crate::config::AppConfig; @@ -15,6 +16,7 @@ use futures_util::{StreamExt, TryStreamExt, stream}; use http::{HeaderName, HeaderValue}; use indexmap::IndexMap; use indoc::formatdoc; +use rmcp::model::{PromptArgument, ServerCapabilities}; use rmcp::service::RunningService; use rmcp::transport::StreamableHttpClientTransport; use rmcp::transport::TokioChildProcess; @@ -34,27 +36,97 @@ use tokio::process::Command; pub const MCP_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke"; pub const MCP_SEARCH_META_FUNCTION_NAME_PREFIX: &str = "mcp_search"; pub const MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX: &str = "mcp_describe"; +pub const MCP_READ_META_FUNCTION_NAME_PREFIX: &str = "mcp_read"; +pub const MCP_PROMPT_META_FUNCTION_NAME_PREFIX: &str = "mcp_prompt"; + +pub const MCP_META_FUNCTION_PREFIXES: [&str; 5] = [ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, +]; + +pub fn is_mcp_meta_function(name: &str) -> bool { + MCP_META_FUNCTION_PREFIXES + .iter() + .any(|prefix| name.starts_with(prefix)) +} + +pub fn mcp_meta_function_names(server: &str) -> Vec { + MCP_META_FUNCTION_PREFIXES + .iter() + .map(|prefix| format!("{prefix}_{server}")) + .collect() +} pub type ConnectedServer = RunningService; +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct McpServerFeatures { + pub name: String, + pub tools: bool, + pub resources: bool, + pub prompts: bool, +} + +impl McpServerFeatures { + pub fn from_capabilities( + name: impl Into, + capabilities: Option<&ServerCapabilities>, + ) -> Self { + Self { + name: name.into(), + tools: capabilities.is_none_or(|c| c.tools.is_some()), + resources: capabilities.is_some_and(|c| c.resources.is_some()), + prompts: capabilities.is_some_and(|c| c.prompts.is_some()), + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum CatalogItemKind { + #[default] + Tool, + Resource, + ResourceTemplate, + Prompt, +} + +impl CatalogItemKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Tool => "tool", + Self::Resource => "resource", + Self::ResourceTemplate => "resource_template", + Self::Prompt => "prompt", + } + } +} + +impl Display for CatalogItemKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + #[derive(Clone, Debug, Default, Serialize)] pub struct CatalogItem { + pub kind: CatalogItemKind, pub name: String, pub server: String, pub description: String, -} - -#[derive(Debug)] -struct ServerCatalog { - items: HashMap, -} - -impl Clone for ServerCatalog { - fn clone(&self) -> Self { - Self { - items: self.items.clone(), - } - } + #[serde(skip_serializing_if = "Option::is_none")] + pub uri: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub mime_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub size: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub arguments: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub audience: Option>, } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -160,7 +232,6 @@ pub struct McpRegistry { log_path: Option, config: Option, servers: HashMap>, - catalogs: HashMap, } impl McpRegistry { @@ -303,7 +374,7 @@ impl McpRegistry { debug!("Starting selected MCP servers: {:?}", ids_to_start); - let results: Vec, ServerCatalog)>> = stream::iter( + let results: Vec)>> = stream::iter( ids_to_start .into_iter() .map(|id| async { self.start_server(id).await }), @@ -312,18 +383,14 @@ impl McpRegistry { .try_collect() .await?; - for (id, server, catalog) in results.into_iter().flatten() { - self.servers.insert(id.clone(), server); - self.catalogs.insert(id, catalog); + for (id, server) in results.into_iter().flatten() { + self.servers.insert(id, server); } Ok(()) } - async fn start_server( - &self, - id: String, - ) -> Result, ServerCatalog)>> { + async fn start_server(&self, id: String) -> Result)>> { let spec = self .config .as_ref() @@ -347,30 +414,9 @@ impl McpRegistry { Err(e) => return Err(e), }; - let tools = service.list_tools(None).await?; - debug!("Available tools for MCP server {id}: {tools:?}"); - - let mut items_vec = Vec::new(); - for t in tools.tools { - let name = t.name.to_string(); - let description = t.description.unwrap_or_default().to_string(); - items_vec.push(CatalogItem { - name, - server: id.clone(), - description, - }); - } - - let mut items_map = HashMap::new(); - items_vec.into_iter().for_each(|it| { - items_map.insert(it.name.clone(), it); - }); - - let catalog = ServerCatalog { items: items_map }; - info!("Started MCP server: {id}"); - Ok(Some((id.to_string(), service, catalog))) + Ok(Some((id, service))) } fn resolve_server_ids(&self, enabled_mcp_servers: Option>) -> Vec { @@ -398,8 +444,21 @@ impl McpRegistry { &self.servers } - pub fn list_started_servers(&self) -> Vec { - self.servers.keys().cloned().collect() + pub fn server_features(&self) -> Vec { + let mut features: Vec = self + .servers + .iter() + .map(|(name, handle)| { + let info = handle.peer_info(); + McpServerFeatures::from_capabilities( + name.as_str(), + info.as_ref().map(|info| &info.capabilities), + ) + }) + .collect(); + + features.sort_by(|a, b| a.name.cmp(&b.name)); + features } pub fn is_empty(&self) -> bool { @@ -1161,7 +1220,7 @@ mod tests { let registry = McpRegistry::default(); assert!(registry.is_empty()); - assert!(registry.list_started_servers().is_empty()); + assert!(registry.server_features().is_empty()); assert!(registry.mcp_config().is_none()); assert!(registry.log_path().is_none()); } @@ -1185,6 +1244,51 @@ mod tests { assert_eq!(MCP_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke"); assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search"); assert_eq!(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, "mcp_describe"); + assert_eq!(MCP_READ_META_FUNCTION_NAME_PREFIX, "mcp_read"); + assert_eq!(MCP_PROMPT_META_FUNCTION_NAME_PREFIX, "mcp_prompt"); + } + + #[test] + fn is_mcp_meta_function_classifies_names() { + assert!(is_mcp_meta_function("mcp_invoke_github")); + assert!(is_mcp_meta_function("mcp_search_github")); + assert!(is_mcp_meta_function("mcp_describe_github")); + assert!(is_mcp_meta_function("mcp_read_github")); + assert!(is_mcp_meta_function("mcp_prompt_github")); + assert!(!is_mcp_meta_function("mcp_gateway_tool")); + assert!(!is_mcp_meta_function("fs_read")); + assert!(!is_mcp_meta_function("")); + assert!(!is_mcp_meta_function("mcp_")); + } + + #[test] + fn meta_function_prefixes_are_not_prefixes_of_each_other() { + for (i, a) in MCP_META_FUNCTION_PREFIXES.iter().enumerate() { + for (j, b) in MCP_META_FUNCTION_PREFIXES.iter().enumerate() { + if i != j { + assert!(!b.starts_with(a), "{a} is a prefix of {b}"); + } + } + } + } + + #[test] + fn is_mcp_meta_function_preserves_lax_prefix_matching() { + assert!(is_mcp_meta_function("mcp_invoker_x")); + } + + #[test] + fn mcp_meta_function_names_returns_all_prefixes_in_order() { + assert_eq!( + mcp_meta_function_names("github"), + vec![ + "mcp_invoke_github", + "mcp_search_github", + "mcp_describe_github", + "mcp_read_github", + "mcp_prompt_github", + ] + ); } #[test] diff --git a/src/mcp/render.rs b/src/mcp/render.rs new file mode 100644 index 0000000..6c1da0c --- /dev/null +++ b/src/mcp/render.rs @@ -0,0 +1,896 @@ +//! Content policy for MCP resource and tool content: UTF-8-boundary-safe text +//! paging, grep-style pattern filtering, and spill-to-disk for binary blobs. + +use crate::config::paths; +use base64::engine::general_purpose::STANDARD; +use base64::read::DecoderReader; +use fancy_regex::Regex; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use std::error::Error; +use std::fs::{self, OpenOptions}; +use std::io::{ErrorKind, Read, Write}; +#[cfg(unix)] +use std::os::unix::fs::OpenOptionsExt; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::SystemTime; +use std::{fmt, io}; + +/// Default page size when the caller does not specify `max_bytes`. +pub const DEFAULT_TEXT_MAX_BYTES: usize = 51_200; +/// Hard upper bound on a single text slice regardless of requested `max_bytes`. +pub const TEXT_MAX_BYTES_CLAMP: usize = 204_800; +/// Maximum decoded size of a base64 blob before rendering is refused. +pub const BLOB_DECODE_CEILING_BYTES: usize = 50 * 1024 * 1024; +/// Total size bound for the spill tree; oldest files are evicted beyond it. +pub const SPILL_DIR_MAX_BYTES: u64 = 512 * 1024 * 1024; +/// Byte bound on server-supplied metadata strings (uri, mime type) copied into output. +pub const METADATA_MAX_BYTES: usize = 4096; + +const PATTERN_CONTEXT_LINES: usize = 2; +const HUNK_SEPARATOR: &str = "--"; + +const MIME_EXTENSIONS: &[(&str, &str)] = &[ + ("application/gzip", "gz"), + ("application/json", "json"), + ("application/pdf", "pdf"), + ("application/zip", "zip"), + ("audio/mpeg", "mp3"), + ("image/gif", "gif"), + ("image/jpeg", "jpg"), + ("image/png", "png"), + ("image/webp", "webp"), + ("text/csv", "csv"), + ("video/mp4", "mp4"), +]; + +#[derive(Debug)] +pub enum RenderError { + InvalidPattern { pattern: String, error: String }, + DecodedSizeExceeded, + InvalidBase64(String), + Io(io::Error), +} + +impl fmt::Display for RenderError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidPattern { pattern, error } => write!( + f, + "Invalid filter pattern '{pattern}': {error}. Provide a valid regex; \ + lines matching it are returned with {PATTERN_CONTEXT_LINES} lines of context." + ), + Self::DecodedSizeExceeded => write!( + f, + "Decoded blob exceeds BLOB_DECODE_CEILING_BYTES ({} MiB); refusing to render it", + BLOB_DECODE_CEILING_BYTES / (1024 * 1024) + ), + Self::InvalidBase64(error) => write!(f, "Invalid base64 in blob content: {error}"), + Self::Io(error) => write!(f, "Failed to spill blob to disk: {error}"), + } + } +} + +impl Error for RenderError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + Self::Io(error) => Some(error), + _ => None, + } + } +} + +impl From for RenderError { + fn from(error: io::Error) -> Self { + Self::Io(error) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RenderedText { + pub text: String, + pub truncated: bool, + pub total_bytes: usize, + pub next_offset: Option, +} + +#[derive(Debug)] +pub enum RenderedBlob { + Text(String), + Spilled(SpillMetadata), +} + +#[derive(Debug, Serialize)] +pub struct SpillMetadata { + pub spilled: bool, + pub path: PathBuf, + pub mime_type: Option, + pub sniffed: bool, + pub size_bytes: u64, + pub sha256: String, +} + +/// Pages `text` with UTF-8-boundary-safe slicing. When `pattern` is set, the +/// text is first reduced to matching lines plus context (grep-style, with +/// 1-based line-number prefixes), and all offset/size math operates on that +/// filtered stream. +pub fn render_text( + text: &str, + pattern: Option<&str>, + offset: usize, + max_bytes: Option, +) -> Result { + let filtered = match pattern { + Some(pattern) => Some(filter_lines(text, pattern)?), + None => None, + }; + let stream = filtered.as_deref().unwrap_or(text); + let max_bytes = max_bytes + .unwrap_or(DEFAULT_TEXT_MAX_BYTES) + .min(TEXT_MAX_BYTES_CLAMP); + let total_bytes = stream.len(); + let mut start = offset.min(total_bytes); + while !stream.is_char_boundary(start) { + start += 1; + } + let mut end = start.saturating_add(max_bytes).min(total_bytes); + while !stream.is_char_boundary(end) { + end -= 1; + } + // A max_bytes smaller than one codepoint would produce an empty page with + // next_offset == offset, stalling paging; always advance by at least one. + if end == start && start < total_bytes { + end += 1; + while !stream.is_char_boundary(end) { + end += 1; + } + } + + let truncated = end < total_bytes; + Ok(RenderedText { + text: stream[start..end].to_string(), + truncated, + total_bytes, + next_offset: truncated.then_some(end), + }) +} + +/// Decodes a base64 blob, returning it as text when it is valid UTF-8 and +/// spilling it under `cache_dir()/mcp-resources//` otherwise. +pub fn render_blob( + b64: &str, + claimed_mime: Option<&str>, + server: &str, +) -> Result { + let spill_base = paths::cache_dir().join("mcp-resources"); + render_blob_at(b64, claimed_mime, server, &spill_base) +} + +pub fn render_blob_at( + b64: &str, + claimed_mime: Option<&str>, + server: &str, + spill_base: &Path, +) -> Result { + let decoded = decode_base64_bounded(b64)?; + let decoded = match String::from_utf8(decoded) { + Ok(text) => return Ok(RenderedBlob::Text(text)), + Err(error) => error.into_bytes(), + }; + let sha256 = format!("{:x}", Sha256::digest(&decoded)); + let dir = spill_base.join(sanitize_server(server)); + fs::create_dir_all(&dir)?; + let path = dir.join(format!("{sha256}.{}", extension_for_mime(claimed_mime))); + + // Writes land in a temp file and are renamed into place, so a visible + // file at the final path is always complete and the dedup check below is + // race-safe across processes (same sha means same content). + if !path.exists() { + static TEMP_COUNTER: AtomicUsize = AtomicUsize::new(0); + let temp = dir.join(format!( + "{sha256}.tmp-{}-{}", + std::process::id(), + TEMP_COUNTER.fetch_add(1, Ordering::Relaxed) + )); + let mut options = OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + options.mode(0o600); + let written = options + .open(&temp) + .and_then(|mut file| file.write_all(&decoded)) + .and_then(|()| fs::rename(&temp, &path)); + if let Err(error) = written { + let _ = fs::remove_file(&temp); + return Err(RenderError::Io(error)); + } + } + enforce_spill_bound(spill_base, SPILL_DIR_MAX_BYTES, &path); + + Ok(RenderedBlob::Spilled(SpillMetadata { + spilled: true, + path, + mime_type: claimed_mime.map(str::to_string), + sniffed: false, + size_bytes: decoded.len() as u64, + sha256, + })) +} + +/// Truncates `text` to at most `max_bytes`, rounding the cut point back to a +/// UTF-8 character boundary. +pub fn truncate_utf8(text: &str, max_bytes: usize) -> &str { + if text.len() <= max_bytes { + return text; + } + let mut end = max_bytes; + while !text.is_char_boundary(end) { + end -= 1; + } + &text[..end] +} + +/// Bounds a server-supplied metadata string to [`METADATA_MAX_BYTES`], +/// appending a marker citing the constant when the input is truncated. +pub fn clamp_metadata(text: &str) -> String { + if text.len() <= METADATA_MAX_BYTES { + return text.to_string(); + } + + let clamped = truncate_utf8(text, METADATA_MAX_BYTES); + format!("{clamped} [truncated: exceeds METADATA_MAX_BYTES ({METADATA_MAX_BYTES} bytes)]") +} + +fn filter_lines(text: &str, pattern: &str) -> Result { + let regex = Regex::new(pattern).map_err(|error| RenderError::InvalidPattern { + pattern: pattern.to_string(), + error: error.to_string(), + })?; + let lines: Vec<&str> = text.lines().collect(); + // fancy_regex can also fail at match time (backtracking limits); treat + // that as a non-match rather than failing the whole render. + let is_match: Vec = lines + .iter() + .map(|line| regex.is_match(line).unwrap_or(false)) + .collect(); + + let mut keep = vec![false; lines.len()]; + for (i, _) in is_match.iter().enumerate().filter(|&(_, matched)| *matched) { + let start = i.saturating_sub(PATTERN_CONTEXT_LINES); + let end = (i + PATTERN_CONTEXT_LINES).min(lines.len() - 1); + keep[start..=end].fill(true); + } + + let mut out: Vec = Vec::new(); + let mut prev_kept: Option = None; + for (i, line) in lines.iter().enumerate() { + if !keep[i] { + continue; + } + if prev_kept.is_some_and(|prev| i > prev + 1) { + out.push(HUNK_SEPARATOR.to_string()); + } + let marker = if is_match[i] { ':' } else { '-' }; + out.push(format!("{}{marker}{line}", i + 1)); + prev_kept = Some(i); + } + + Ok(out.join("\n")) +} + +fn decode_base64_bounded(b64: &str) -> Result, RenderError> { + // The encoded length puts a lower bound on the decoded size; reject + // inputs that bound already proves oversized before decoding anything. + let min_decoded = (b64.len() / 4).saturating_mul(3).saturating_sub(2); + if min_decoded > BLOB_DECODE_CEILING_BYTES { + return Err(RenderError::DecodedSizeExceeded); + } + + let mut reader = DecoderReader::new(b64.as_bytes(), &STANDARD); + let mut decoded = Vec::new(); + let mut chunk = [0u8; 8192]; + loop { + match reader.read(&mut chunk) { + Ok(0) => return Ok(decoded), + Ok(n) => { + if decoded.len() + n > BLOB_DECODE_CEILING_BYTES { + return Err(RenderError::DecodedSizeExceeded); + } + decoded.extend_from_slice(&chunk[..n]); + } + Err(error) => return Err(RenderError::InvalidBase64(error.to_string())), + } + } +} + +/// Maps a server-controlled mime type to a spill-file extension via an exact +/// allowlist lookup; anything unrecognized falls back to `bin`. +fn extension_for_mime(mime: Option<&str>) -> &'static str { + let Some(mime) = mime else { + return "bin"; + }; + let bare = mime + .split(';') + .next() + .unwrap_or("") + .trim() + .to_ascii_lowercase(); + let ext = MIME_EXTENSIONS + .iter() + .find(|(known, _)| *known == bare) + .map(|(_, ext)| *ext) + .unwrap_or("bin"); + let safe = !ext.is_empty() + && ext.len() <= 8 + && ext + .bytes() + .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit()); + + if safe { ext } else { "bin" } +} + +fn sanitize_server(server: &str) -> String { + let mut sanitized: String = server + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') { + c + } else { + '_' + } + }) + .take(64) + .collect(); + // Windows strips trailing dots at create time, which would make the + // constructed path disagree with the on-disk name. + while sanitized.ends_with('.') { + sanitized.pop(); + } + if sanitized.is_empty() { + return "_".to_string(); + } + // Windows reserves device names (bare or with any extension). + let stem = sanitized.split('.').next().unwrap_or(""); + if is_windows_reserved(stem) { + sanitized.insert(0, '_'); + } + sanitized +} + +fn is_windows_reserved(stem: &str) -> bool { + let lower = stem.to_ascii_lowercase(); + matches!(lower.as_str(), "con" | "prn" | "aux" | "nul") + || (lower.len() == 4 + && (lower.starts_with("com") || lower.starts_with("lpt")) + && matches!(lower.as_bytes()[3], b'1'..=b'9')) +} + +struct SpillEntry { + path: PathBuf, + size: u64, + modified: SystemTime, +} + +fn enforce_spill_bound(base: &Path, max_total: u64, protect: &Path) { + let mut entries = Vec::new(); + collect_spill_files(base, &mut entries); + evict_oldest(entries, max_total, protect); +} + +/// Best-effort eviction: the spill dir is shared across processes, so a file +/// vanishing underneath us (`NotFound`) is expected and never fails the spill. +fn evict_oldest(mut entries: Vec, max_total: u64, protect: &Path) { + let mut total: u64 = entries.iter().map(|entry| entry.size).sum(); + if total <= max_total { + return; + } + entries.sort_by_key(|entry| entry.modified); + for entry in &entries { + if total <= max_total { + break; + } + + // Filenames are content-hashed, so name equality is sufficient and + // survives filesystems that normalize directory names (case folding, + // trailing-dot stripping) where a full-path comparison would miss. + if entry.path.file_name() == protect.file_name() { + continue; + } + + match fs::remove_file(&entry.path) { + Ok(()) => total -= entry.size, + Err(error) if error.kind() == ErrorKind::NotFound => total -= entry.size, + Err(_) => {} + } + } +} + +fn collect_spill_files(dir: &Path, out: &mut Vec) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + for entry in entries.flatten() { + let path = entry.path(); + let Ok(metadata) = entry.metadata() else { + continue; + }; + if metadata.is_dir() { + collect_spill_files(&path, out); + } else if metadata.is_file() { + out.push(SpillEntry { + path, + size: metadata.len(), + modified: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH), + }); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use base64::Engine; + use std::env; + #[cfg(unix)] + use std::os::unix::fs::PermissionsExt; + use std::process; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + fn with_spill_base(f: F) { + static COUNTER: AtomicUsize = AtomicUsize::new(0); + let unique = format!( + "{}-{}", + process::id(), + COUNTER.fetch_add(1, Ordering::Relaxed) + ); + let base = env::temp_dir().join(format!("coyote-render-test-{unique}")); + fs::create_dir_all(&base).unwrap(); + f(&base); + let _ = fs::remove_dir_all(&base); + } + + fn set_mtime(path: &Path, secs_after_epoch: u64) { + let file = OpenOptions::new().write(true).open(path).unwrap(); + file.set_modified(SystemTime::UNIX_EPOCH + Duration::from_secs(secs_after_epoch)) + .unwrap(); + } + + fn write_spill_file(dir: &Path, name: &str, len: usize, mtime_secs: u64) -> PathBuf { + let path = dir.join(name); + fs::write(&path, vec![0u8; len]).unwrap(); + set_mtime(&path, mtime_secs); + path + } + + const TEN_LINES: &str = "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten"; + + #[test] + fn slices_basic_ascii_page() { + let rendered = render_text("hello world", None, 0, Some(5)).unwrap(); + + assert_eq!(rendered.text, "hello"); + assert!(rendered.truncated); + assert_eq!(rendered.total_bytes, 11); + assert_eq!(rendered.next_offset, Some(5)); + } + + #[test] + fn offset_mid_codepoint_rounds_forward() { + // 'é' occupies bytes 1..3; offset 2 lands inside it. + let rendered = render_text("héllo", None, 2, None).unwrap(); + + assert_eq!(rendered.text, "llo"); + assert!(!rendered.truncated); + assert_eq!(rendered.next_offset, None); + } + + #[test] + fn end_mid_codepoint_rounds_backward() { + // 'é' occupies bytes 1..3; offset 0 + max_bytes 2 lands inside it. + let rendered = render_text("aé", None, 0, Some(2)).unwrap(); + + assert_eq!(rendered.text, "a"); + assert!(rendered.truncated); + assert_eq!(rendered.total_bytes, 3); + assert_eq!(rendered.next_offset, Some(1)); + + let rest = render_text("aé", None, 1, Some(2)).unwrap(); + + assert_eq!(rest.text, "é"); + assert!(!rest.truncated); + } + + #[test] + fn max_bytes_below_one_codepoint_still_advances() { + // 'é' is 2 bytes; max_bytes 1 must not stall at next_offset == offset. + let rendered = render_text("éa", None, 0, Some(1)).unwrap(); + + assert_eq!(rendered.text, "é"); + assert!(rendered.truncated); + assert_eq!(rendered.total_bytes, 3); + assert_eq!(rendered.next_offset, Some(2)); + } + + #[test] + fn offset_past_eof_returns_empty() { + let rendered = render_text("short", None, 100, None).unwrap(); + + assert_eq!(rendered.text, ""); + assert!(!rendered.truncated); + assert_eq!(rendered.total_bytes, 5); + assert_eq!(rendered.next_offset, None); + } + + #[test] + fn exact_fit_is_not_truncated() { + let rendered = render_text("exact", None, 0, Some(5)).unwrap(); + + assert_eq!(rendered.text, "exact"); + assert!(!rendered.truncated); + assert_eq!(rendered.next_offset, None); + } + + #[test] + fn default_max_bytes_is_default_text_max_bytes() { + let text = "a".repeat(DEFAULT_TEXT_MAX_BYTES + 1); + + let rendered = render_text(&text, None, 0, None).unwrap(); + + assert_eq!(rendered.text.len(), DEFAULT_TEXT_MAX_BYTES); + assert!(rendered.truncated); + assert_eq!(rendered.next_offset, Some(DEFAULT_TEXT_MAX_BYTES)); + } + + #[test] + fn max_bytes_above_clamp_is_clamped() { + let text = "a".repeat(TEXT_MAX_BYTES_CLAMP + 1); + + let rendered = render_text(&text, None, 0, Some(usize::MAX)).unwrap(); + + assert_eq!(rendered.text.len(), TEXT_MAX_BYTES_CLAMP); + assert!(rendered.truncated); + assert_eq!(rendered.next_offset, Some(TEXT_MAX_BYTES_CLAMP)); + } + + #[test] + fn truncate_utf8_rounds_back_to_char_boundary() { + // 'é' occupies bytes 1..3; a cut at byte 2 lands inside it. + assert_eq!(truncate_utf8("aé", 2), "a"); + assert_eq!(truncate_utf8("aé", 3), "aé"); + assert_eq!(truncate_utf8("abc", 10), "abc"); + assert_eq!(truncate_utf8("abc", 0), ""); + } + + #[test] + fn clamp_metadata_appends_marker_only_when_oversized() { + assert_eq!(clamp_metadata("text/plain"), "text/plain"); + + let long = "u".repeat(METADATA_MAX_BYTES + 1); + + let clamped = clamp_metadata(&long); + + assert!(clamped.starts_with(&"u".repeat(METADATA_MAX_BYTES))); + assert!(clamped.contains("METADATA_MAX_BYTES")); + assert!(clamped.contains(&METADATA_MAX_BYTES.to_string())); + } + + #[test] + fn pattern_emits_matches_with_context_and_line_numbers() { + let rendered = render_text(TEN_LINES, Some("^five$"), 0, None).unwrap(); + + assert_eq!(rendered.text, "3-three\n4-four\n5:five\n6-six\n7-seven"); + assert!(!rendered.truncated); + assert_eq!(rendered.total_bytes, rendered.text.len()); + } + + #[test] + fn pattern_separates_disjoint_hunks() { + let rendered = render_text(TEN_LINES, Some("^(two|nine)$"), 0, None).unwrap(); + + assert_eq!( + rendered.text, + "1-one\n2:two\n3-three\n4-four\n--\n7-seven\n8-eight\n9:nine\n10-ten" + ); + } + + #[test] + fn pattern_merges_adjacent_hunks_without_duplicates() { + let rendered = render_text(TEN_LINES, Some("^(two|six)$"), 0, None).unwrap(); + + assert_eq!( + rendered.text, + "1-one\n2:two\n3-three\n4-four\n5-five\n6:six\n7-seven\n8-eight" + ); + assert!(!rendered.text.contains(HUNK_SEPARATOR)); + } + + #[test] + fn pattern_paging_walks_the_filtered_stream() { + let full = render_text(TEN_LINES, Some("^t"), 0, None).unwrap(); + assert!(!full.truncated); + + let mut assembled = String::new(); + let mut offset = 0; + loop { + let page = render_text(TEN_LINES, Some("^t"), offset, Some(7)).unwrap(); + assert_eq!(page.total_bytes, full.text.len()); + assembled.push_str(&page.text); + match page.next_offset { + Some(next) => offset = next, + None => break, + } + } + + assert_eq!(assembled, full.text); + } + + #[test] + fn pattern_with_no_matches_returns_empty() { + let rendered = render_text(TEN_LINES, Some("^zebra$"), 0, None).unwrap(); + + assert_eq!(rendered.text, ""); + assert_eq!(rendered.total_bytes, 0); + assert!(!rendered.truncated); + assert_eq!(rendered.next_offset, None); + } + + #[test] + fn invalid_pattern_is_a_teaching_error() { + let parse_error = Regex::new("(").unwrap_err().to_string(); + + let err = render_text("text", Some("("), 0, None).unwrap_err(); + + assert!(matches!(err, RenderError::InvalidPattern { .. })); + let message = err.to_string(); + assert!(message.contains("'('")); + assert!(message.contains(&parse_error)); + } + + #[test] + fn utf8_blob_decodes_to_text_without_spilling() { + with_spill_base(|base| { + let b64 = STANDARD.encode("hello ✓ world"); + + let rendered = render_blob_at(&b64, Some("text/plain"), "srv", base).unwrap(); + + let RenderedBlob::Text(text) = rendered else { + panic!("expected text variant"); + }; + assert_eq!(text, "hello ✓ world"); + assert_eq!(fs::read_dir(base).unwrap().count(), 0); + }); + } + + #[test] + fn binary_blob_spills_with_metadata_and_0600_perms() { + with_spill_base(|base| { + let data: &[u8] = &[0xff, 0xfe, 0x00, 0x88, 0x01]; + let b64 = STANDARD.encode(data); + + let rendered = render_blob_at(&b64, Some("application/pdf"), "docs", base).unwrap(); + + let RenderedBlob::Spilled(meta) = rendered else { + panic!("expected spilled variant"); + }; + let expected_sha = format!("{:x}", Sha256::digest(data)); + assert_eq!(meta.sha256, expected_sha); + assert_eq!( + meta.path, + base.join("docs").join(format!("{expected_sha}.pdf")) + ); + assert_eq!(meta.size_bytes, data.len() as u64); + assert_eq!(meta.mime_type.as_deref(), Some("application/pdf")); + assert!(!meta.sniffed); + assert!(meta.spilled); + assert_eq!(fs::read(&meta.path).unwrap(), data); + #[cfg(unix)] + { + let mode = fs::metadata(&meta.path).unwrap().permissions().mode(); + assert_eq!(mode & 0o777, 0o600); + } + }); + } + + #[test] + fn decode_ceiling_rejects_oversized_blob() { + with_spill_base(|base| { + // base64 of 51 MiB of zero bytes is just a repeated-'A' string. + let encoded = "A".repeat(51 * 1024 * 1024 / 3 * 4); + + let err = render_blob_at(&encoded, None, "srv", base).unwrap_err(); + + assert!(matches!(err, RenderError::DecodedSizeExceeded)); + assert!(err.to_string().contains("BLOB_DECODE_CEILING_BYTES")); + }); + } + + #[test] + fn malformed_base64_is_rejected() { + with_spill_base(|base| { + let err = render_blob_at("!!!not base64!!!", None, "srv", base).unwrap_err(); + + assert!(matches!(err, RenderError::InvalidBase64(_))); + }); + } + + #[test] + fn spill_dedup_returns_same_path_without_rewriting() { + with_spill_base(|base| { + let data: &[u8] = &[0xff, 0x01, 0x02]; + let b64 = STANDARD.encode(data); + + let RenderedBlob::Spilled(first) = render_blob_at(&b64, None, "srv", base).unwrap() + else { + panic!("expected spilled variant"); + }; + fs::write(&first.path, b"sentinel").unwrap(); + + let RenderedBlob::Spilled(second) = render_blob_at(&b64, None, "srv", base).unwrap() + else { + panic!("expected spilled variant"); + }; + + assert_eq!(second.path, first.path); + assert_eq!(second.sha256, first.sha256); + assert_eq!(fs::read(&second.path).unwrap(), b"sentinel"); + }); + } + + #[test] + fn spill_metadata_serializes_spilled_true() { + with_spill_base(|base| { + let b64 = STANDARD.encode([0xffu8, 0x00]); + + let RenderedBlob::Spilled(meta) = + render_blob_at(&b64, Some("image/png"), "srv", base).unwrap() + else { + panic!("expected spilled variant"); + }; + + let value = serde_json::to_value(&meta).unwrap(); + assert_eq!(value["spilled"], serde_json::Value::Bool(true)); + assert_eq!(value["sniffed"], serde_json::Value::Bool(false)); + assert_eq!(value["sha256"].as_str(), Some(meta.sha256.as_str())); + assert_eq!(value["mime_type"].as_str(), Some("image/png")); + }); + } + + #[test] + fn extension_allowlist_normalizes_and_defaults_to_bin() { + assert_eq!(extension_for_mime(Some("application/pdf")), "pdf"); + assert_eq!(extension_for_mime(Some("image/png")), "png"); + assert_eq!(extension_for_mime(Some(" TEXT/CSV ; charset=utf-8")), "csv"); + assert_eq!(extension_for_mime(Some("../../evil")), "bin"); + assert_eq!(extension_for_mime(Some("image/png/../../x")), "bin"); + assert_eq!(extension_for_mime(Some("application/x-∞")), "bin"); + assert_eq!(extension_for_mime(Some("text/plain")), "bin"); + assert_eq!(extension_for_mime(None), "bin"); + } + + #[test] + fn sanitize_server_strips_path_separators() { + assert_eq!(sanitize_server("../evil/srv"), ".._evil_srv"); + assert_eq!(sanitize_server("srv name!"), "srv_name_"); + assert_eq!(sanitize_server(""), "_"); + assert_eq!(sanitize_server("."), "_"); + assert_eq!(sanitize_server(".."), "_"); + assert_eq!(sanitize_server("good-server_1.0"), "good-server_1.0"); + } + + #[test] + fn sanitize_server_escapes_windows_reserved_names() { + assert_eq!(sanitize_server("con"), "_con"); + assert_eq!(sanitize_server("CON"), "_CON"); + assert_eq!(sanitize_server("nul.txt"), "_nul.txt"); + assert_eq!(sanitize_server("COM1"), "_COM1"); + assert_eq!(sanitize_server("lpt9"), "_lpt9"); + assert_eq!(sanitize_server("com0"), "com0"); + assert_eq!(sanitize_server("com10"), "com10"); + assert_eq!(sanitize_server("consul"), "consul"); + } + + #[test] + fn sanitize_server_strips_trailing_dots_and_caps_length() { + assert_eq!(sanitize_server("srv."), "srv"); + assert_eq!(sanitize_server("srv..."), "srv"); + assert_eq!(sanitize_server("..."), "_"); + let long = "a".repeat(100); + assert_eq!(sanitize_server(&long).len(), 64); + } + + #[test] + fn spill_path_confines_crafted_server_and_mime() { + with_spill_base(|base| { + let b64 = STANDARD.encode([0xffu8, 0x00, 0x11]); + + let RenderedBlob::Spilled(meta) = + render_blob_at(&b64, Some("../../evil"), "../evil/srv", base).unwrap() + else { + panic!("expected spilled variant"); + }; + + assert!(meta.path.starts_with(base)); + let dir_name = meta.path.parent().unwrap().file_name().unwrap(); + assert_eq!(dir_name, ".._evil_srv"); + assert_eq!(meta.path.extension().unwrap(), "bin"); + }); + } + + #[test] + fn eviction_removes_oldest_files_first_across_server_dirs() { + with_spill_base(|base| { + let srv_a = base.join("srv-a"); + let srv_b = base.join("srv-b"); + fs::create_dir_all(&srv_a).unwrap(); + fs::create_dir_all(&srv_b).unwrap(); + let oldest = write_spill_file(&srv_a, "a.bin", 100, 100); + let middle = write_spill_file(&srv_b, "b.bin", 100, 200); + let newest = write_spill_file(&srv_b, "c.bin", 100, 300); + + enforce_spill_bound(base, 150, &newest); + + assert!(!oldest.exists()); + assert!(!middle.exists()); + assert!(newest.exists()); + }); + } + + #[test] + fn eviction_skips_protected_file() { + with_spill_base(|base| { + let srv = base.join("srv"); + fs::create_dir_all(&srv).unwrap(); + let oldest = write_spill_file(&srv, "a.bin", 100, 100); + let middle = write_spill_file(&srv, "b.bin", 100, 200); + let newest = write_spill_file(&srv, "c.bin", 100, 300); + + enforce_spill_bound(base, 250, &oldest); + + assert!(oldest.exists()); + assert!(!middle.exists()); + assert!(newest.exists()); + }); + } + + #[test] + fn eviction_under_bound_is_noop() { + with_spill_base(|base| { + let srv = base.join("srv"); + fs::create_dir_all(&srv).unwrap(); + let first = write_spill_file(&srv, "a.bin", 100, 100); + let second = write_spill_file(&srv, "b.bin", 100, 200); + + enforce_spill_bound(base, 1000, &second); + + assert!(first.exists()); + assert!(second.exists()); + }); + } + + #[test] + fn eviction_tolerates_already_removed_entries() { + with_spill_base(|base| { + let srv = base.join("srv"); + fs::create_dir_all(&srv).unwrap(); + let real = write_spill_file(&srv, "real.bin", 100, 200); + let entries = vec![ + SpillEntry { + path: srv.join("ghost.bin"), + size: 100, + modified: SystemTime::UNIX_EPOCH + Duration::from_secs(100), + }, + SpillEntry { + path: real.clone(), + size: 100, + modified: SystemTime::UNIX_EPOCH + Duration::from_secs(200), + }, + ]; + + evict_oldest(entries, 50, &base.join("untouched")); + + assert!(!real.exists()); + }); + } +} diff --git a/src/repl/completer.rs b/src/repl/completer.rs index 440edbb..be439fd 100644 --- a/src/repl/completer.rs +++ b/src/repl/completer.rs @@ -1,11 +1,17 @@ use super::{REPL_COMMANDS, ReplCommand}; -use crate::{config::RequestContext, utils::fuzzy_filter}; +use crate::config::{McpPromptCompletion, RequestContext, sanitize_display_text}; +use crate::mcp::ConnectedServer; +use crate::utils::fuzzy_filter; use parking_lot::RwLock; use reedline::{Completer, Span, Suggestion}; +use rmcp::model::Prompt; use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; + +const PROMPT_COMPLETION_RPC_TIMEOUT: Duration = Duration::from_secs(2); impl Completer for ReplCompleter { fn complete(&mut self, line: &str, pos: usize) -> Vec { @@ -29,6 +35,22 @@ impl Completer for ReplCompleter { return suggestions; } + if cmd == ".prompt" && parts_len > 1 { + let span = Span::new(parts[parts_len - 1].1, pos); + let args: Vec<&str> = parts.iter().skip(1).map(|(v, _)| *v).collect(); + let filter = args.last().copied().unwrap_or_default().to_string(); + let stage = { + let ctx = self.ctx.read(); + ctx.mcp_prompt_completion(&args) + }; + return complete_prompt_stage(stage, &filter, PROMPT_COMPLETION_RPC_TIMEOUT) + .iter() + .map(|(value, description)| { + create_suggestion(value, description.as_deref().unwrap_or_default(), span) + }) + .collect(); + } + let ctx = self.ctx.read(); let state = ctx.state(); let model_has_reasoning = !ctx.current_model().reasoning_levels().is_empty(); @@ -141,6 +163,71 @@ fn create_suggestion(value: &str, description: &str, span: Span) -> Suggestion { } } +fn complete_prompt_stage( + stage: McpPromptCompletion, + filter: &str, + rpc_timeout: Duration, +) -> Vec<(String, Option)> { + let values = match stage { + McpPromptCompletion::Ready(values) => values, + McpPromptCompletion::PromptNames { server } => list_prompts_blocking(server, rpc_timeout) + .unwrap_or_default() + .into_iter() + .map(|prompt| { + ( + sanitize_display_text(&prompt.name), + prompt + .description + .map(|description| sanitize_display_text(&description)), + ) + }) + .collect(), + McpPromptCompletion::ArgumentKeys { + server, + prompt, + typed_keys, + } => list_prompts_blocking(server, rpc_timeout) + .unwrap_or_default() + .into_iter() + .find(|candidate| candidate.name == prompt) + .and_then(|candidate| candidate.arguments) + .unwrap_or_default() + .into_iter() + .filter(|arg| !typed_keys.contains(&arg.name)) + .map(|arg| { + let description = arg + .description + .map(|description| sanitize_display_text(&description)); + let description = match (arg.required == Some(true), description) { + (true, Some(description)) => Some(format!("{description} (required)")), + (true, None) => Some("(required)".to_string()), + (false, description) => description, + }; + ( + format!("{}=", sanitize_display_text(&arg.name)), + description, + ) + }) + .collect(), + }; + fuzzy_filter(values, |(value, _)| value.as_str(), filter) +} + +fn list_prompts_blocking( + server: Arc, + rpc_timeout: Duration, +) -> Option> { + let fut = async move { tokio::time::timeout(rpc_timeout, server.list_all_prompts()).await }; + // block_in_place is only sound because the REPL's read_line runs inside the + // main-thread block_on of the multi-thread runtime. + let result = match tokio::runtime::Handle::try_current().ok() { + Some(handle) => tokio::task::block_in_place(|| handle.block_on(fut)), + None => tokio::runtime::Runtime::new().ok()?.block_on(fut), + }; + + result.ok()?.ok() +} + fn split_line(line: &str) -> Vec<(&str, usize)> { let mut parts = vec![]; let mut part_start = None; @@ -178,3 +265,208 @@ fn test_split_line() { vec![(".set", 0), ("highlight", 5), ("t", 15)], ); } + +#[cfg(test)] +mod prompt_completion_tests { + use super::*; + use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; + use std::sync::atomic::Ordering; + + fn prompts_fixture() -> FixtureServer { + FixtureServer { + prompts_capability: true, + ..Default::default() + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn stage_two_lists_prompt_names_with_descriptions() { + let (runtime, _server) = fixture_runtime(prompts_fixture()).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::PromptNames { server }, + "", + Duration::from_secs(2), + ); + + assert_eq!( + values, + vec![( + "summarize".to_string(), + Some("Summarize a document".to_string()) + )] + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn stage_three_suggests_argument_keys_with_required_marker() { + let (runtime, _server) = fixture_runtime(prompts_fixture()).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::ArgumentKeys { + server, + prompt: "summarize".to_string(), + typed_keys: vec![], + }, + "", + Duration::from_secs(2), + ); + + assert_eq!( + values, + vec![ + ( + "path=".to_string(), + Some("Document path (required)".to_string()) + ), + ("style=".to_string(), None), + ] + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn stage_three_excludes_typed_keys_and_fuzzy_filters() { + let (runtime, _server) = fixture_runtime(prompts_fixture()).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::ArgumentKeys { + server: Arc::clone(&server), + prompt: "summarize".to_string(), + typed_keys: vec!["path".to_string()], + }, + "", + Duration::from_secs(2), + ); + assert_eq!(values, vec![("style=".to_string(), None)]); + + let values = complete_prompt_stage( + McpPromptCompletion::ArgumentKeys { + server, + prompt: "summarize".to_string(), + typed_keys: vec![], + }, + "sty", + Duration::from_secs(2), + ); + assert_eq!(values, vec![("style=".to_string(), None)]); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn stage_three_unknown_prompt_is_empty() { + let (runtime, _server) = fixture_runtime(prompts_fixture()).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::ArgumentKeys { + server, + prompt: "ghost".to_string(), + typed_keys: vec![], + }, + "", + Duration::from_secs(2), + ); + + assert!(values.is_empty()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn hostile_prompt_strings_are_sanitized_in_suggestions() { + let fixture = FixtureServer { + hostile_prompt: true, + ..prompts_fixture() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::PromptNames { + server: Arc::clone(&server), + }, + "evil", + Duration::from_secs(2), + ); + assert_eq!( + values, + vec![( + "summarize-evil".to_string(), + Some("Runs hostile text".to_string()) + )] + ); + + let values = complete_prompt_stage( + McpPromptCompletion::ArgumentKeys { + server, + prompt: "sum\u{1b}[31mmarize-evil".to_string(), + typed_keys: vec![], + }, + "", + Duration::from_secs(2), + ); + assert_eq!( + values, + vec![("path=".to_string(), Some("Doc path (required)".to_string()))] + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn slow_listing_times_out_to_empty() { + let fixture = FixtureServer { + prompt_delay: Some(Duration::from_millis(200)), + ..prompts_fixture() + }; + let (runtime, _server) = fixture_runtime(fixture).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::PromptNames { server }, + "", + Duration::from_millis(20), + ); + + assert!(values.is_empty()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn failed_listing_is_swallowed_without_retry() { + let fixture = FixtureServer { + fail_prompt_listings: true, + ..prompts_fixture() + }; + let list_prompts_calls = Arc::clone(&fixture.list_prompts_calls); + let (runtime, _server) = fixture_runtime(fixture).await; + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::PromptNames { server }, + "", + Duration::from_secs(2), + ); + + assert!(values.is_empty()); + assert_eq!(list_prompts_calls.load(Ordering::SeqCst), 1); + } + + #[test] + fn bridge_without_ambient_runtime_uses_fallback_runtime() { + let rt = tokio::runtime::Runtime::new().unwrap(); + let (runtime, _server) = rt.block_on(fixture_runtime(prompts_fixture())); + let server = runtime.get("fixture").cloned().unwrap(); + + let values = complete_prompt_stage( + McpPromptCompletion::PromptNames { server }, + "", + Duration::from_secs(2), + ); + + assert_eq!( + values, + vec![( + "summarize".to_string(), + Some("Summarize a document".to_string()) + )] + ); + } +} diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 3c159f2..f3b0d0f 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -13,7 +13,7 @@ use crate::client::{ }; use crate::config::{ AgentVariables, AppConfig, AssertState, Input, LastMessage, MacroState, RequestContext, - StateFlags, macro_execute, + StateFlags, flatten_prompt_messages, macro_execute, resolve_prompt_args, sanitize_display_text, }; use crate::config::{AssetCategory, paths}; use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail}; @@ -29,6 +29,7 @@ use anyhow::{Context, Result, bail}; use crossterm::cursor::SetCursorStyle; use fancy_regex::Regex; use indoc::indoc; +use inquire::Text; use log::warn; use parking_lot::RwLock; use reedline::CursorConfig; @@ -38,6 +39,7 @@ use reedline::{ default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings, }; use reedline::{MenuBuilder, Signal}; +use std::collections::HashMap; use std::sync::LazyLock; use std::{env, process, sync::Arc}; use tokio::task; @@ -53,7 +55,7 @@ pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {" 4. Continue with the next pending item now. Call tools immediately." }; -static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| { +static REPL_COMMANDS: LazyLock<[ReplCommand; 62]> = LazyLock::new(|| { [ ReplCommand::new(".help", "Show this help guide", AssertState::pass()), ReplCommand::new(".info", "Show system info", AssertState::pass()), @@ -105,6 +107,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| { ReplCommand::new(".model", "Switch LLM model", AssertState::pass()), ReplCommand::new( ".prompt", + "Invoke an MCP prompt and submit the result as chat input", + AssertState::pass(), + ), + ReplCommand::new( + ".temp-role", "Set a temporary role using a prompt", AssertState::False(StateFlags::SESSION | StateFlags::AGENT), ), @@ -307,7 +314,7 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| { ), ReplCommand::new( ".list", - "List roles, sessions, agents, RAGs, macros, skills, tools, MCP servers, or bundles", + "List roles, sessions, agents, RAGs, macros, skills, prompts, tools, MCP servers, or bundles", AssertState::pass(), ), ReplCommand::new( @@ -774,12 +781,47 @@ pub async fn run_repl_command( .tool disable # Disable a single tool in the current context"# ), }, - ".prompt" => match args { + ".prompt" => { + let (words, _) = split_args_text(args.unwrap_or_default(), cfg!(windows)); + match words.as_slice() { + [server, name, rest @ ..] => { + let provided = parse_prompt_call_args(rest)?; + let prompts = ctx.tool_scope.mcp_runtime.list_prompts(server).await?; + let declared = prompts + .into_iter() + .find(|prompt| prompt.name == *name) + .with_context(|| { + format!("Prompt '{name}' not found on MCP server '{server}'") + })? + .arguments + .unwrap_or_default(); + let (mut arguments, missing) = resolve_prompt_args(&declared, provided); + for key in missing { + let value = Text::new(&prompt_arg_inquire_label(server, name, &key)) + .prompt() + .with_context(|| { + format!("Failed to read prompt argument '{key}'") + })?; + arguments.insert(key, value); + } + let result = ctx + .tool_scope + .mcp_runtime + .prompt(server, name, arguments) + .await?; + let flattened = flatten_prompt_messages(&result.messages); + let input = Input::from_str(ctx, &flattened, None)?; + ask(ctx, abort_signal.clone(), input, true).await?; + } + _ => println!("Usage: .prompt [key=value ...]"), + } + } + ".temp-role" => match args { Some(text) => { let app = Arc::clone(&ctx.app.config); - ctx.use_prompt(app.as_ref(), text)?; + ctx.use_temp_role(app.as_ref(), text)?; } - None => println!("Usage: .prompt ..."), + None => println!("Usage: .temp-role ..."), }, ".role" => match args { Some(args) => match args.split_once(['\n', ' ']) { @@ -1207,13 +1249,16 @@ pub async fn run_repl_command( println!("Usage: .uninstall [--yes] (see `.uninstall --help`)") } }, - ".list" => match args { + ".list" => match args.map(str::trim) { + Some("prompts") => { + ctx.list_mcp_prompts().await?; + } Some(args) => { - ctx.list_assets(args.trim())?; + ctx.list_assets(args)?; } _ => { println!( - "Usage: .list " + "Usage: .list " ) } }, @@ -1737,6 +1782,37 @@ fn split_first_arg(args: Option<&str>) -> Option<(&str, Option<&str>)> { }) } +fn parse_prompt_call_args(words: &[String]) -> Result> { + let mut args = HashMap::new(); + for word in words { + let Some((key, value)) = word.split_once('=') else { + bail!("Invalid prompt argument '{word}': arguments must be key=value pairs"); + }; + args.insert(key.to_string(), unquote_prompt_value(value).to_string()); + } + Ok(args) +} + +fn unquote_prompt_value(value: &str) -> &str { + let quoted = value.len() >= 2 + && ((value.starts_with('"') && value.ends_with('"')) + || (value.starts_with('\'') && value.ends_with('\''))); + if quoted { + &value[1..value.len() - 1] + } else { + value + } +} + +fn prompt_arg_inquire_label(server: &str, prompt: &str, arg: &str) -> String { + format!( + "Prompt '{}' on '{}' requires '{}':", + sanitize_display_text(prompt), + sanitize_display_text(server), + sanitize_display_text(arg) + ) +} + pub fn split_args_text(line: &str, is_win: bool) -> (Vec, &str) { let mut words = Vec::new(); let mut word = String::new(); @@ -1888,8 +1964,47 @@ mod tests { } #[test] - fn repl_commands_has_61_entries() { - assert_eq!(REPL_COMMANDS.len(), 61); + fn repl_commands_has_62_entries() { + assert_eq!(REPL_COMMANDS.len(), 62); + } + + #[test] + fn parse_prompt_call_args_splits_on_first_equals_and_unquotes() { + let words = vec![ + "path=notes.txt".to_string(), + r#"style="a b""#.to_string(), + "expr=a=b".to_string(), + ]; + + let args = parse_prompt_call_args(&words).unwrap(); + + assert_eq!(args["path"], "notes.txt"); + assert_eq!(args["style"], "a b"); + assert_eq!(args["expr"], "a=b"); + } + + #[test] + fn parse_prompt_call_args_rejects_words_without_equals() { + let err = parse_prompt_call_args(&["positional".to_string()]) + .unwrap_err() + .to_string(); + + assert_eq!( + err, + "Invalid prompt argument 'positional': arguments must be key=value pairs" + ); + } + + #[test] + fn prompt_arg_inquire_label_sanitizes_all_components() { + assert_eq!( + prompt_arg_inquire_label("srv", "summarize", "path"), + "Prompt 'summarize' on 'srv' requires 'path':" + ); + assert_eq!( + prompt_arg_inquire_label("s\u{1b}[31mrv", "sum\u{1b}]0;x\u{7}marize", "pa\u{7}th"), + "Prompt 'summarize' on 'srv' requires 'pa th':" + ); } #[test] @@ -2105,10 +2220,22 @@ mod tests { } #[test] - fn repl_commands_prompt_blocked_in_session_or_agent() { + fn repl_commands_prompt_always_available() { let cmd = REPL_COMMANDS.iter().find(|c| c.name == ".prompt").unwrap(); assert!(cmd.is_valid(StateFlags::empty())); assert!(cmd.is_valid(StateFlags::ROLE)); + assert!(cmd.is_valid(StateFlags::SESSION)); + assert!(cmd.is_valid(StateFlags::AGENT)); + } + + #[test] + fn repl_commands_temp_role_blocked_in_session_or_agent() { + let cmd = REPL_COMMANDS + .iter() + .find(|c| c.name == ".temp-role") + .unwrap(); + assert!(cmd.is_valid(StateFlags::empty())); + assert!(cmd.is_valid(StateFlags::ROLE)); assert!(!cmd.is_valid(StateFlags::SESSION)); assert!(!cmd.is_valid(StateFlags::AGENT)); }