diff --git a/src/config/mod.rs b/src/config/mod.rs index a1f190f..08992fd 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -55,7 +55,9 @@ pub use self::skill_policy::SkillPolicy; 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}; +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/tool_scope.rs b/src/config/tool_scope.rs index 06402cc..6db7a16 100644 --- a/src/config/tool_scope.rs +++ b/src/config/tool_scope.rs @@ -420,16 +420,68 @@ 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!("{} (required)", arg.name) + format!("{name} (required)") } else { - arg.name.clone() + 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 catalog_key(item: &CatalogItem) -> String { let id = item.uri.as_deref().unwrap_or(&item.name); format!("{}:{id}", item.kind) @@ -483,9 +535,9 @@ fn resource_template_catalog_item(server: &str, template: ResourceTemplate) -> C fn prompt_catalog_item(server: &str, prompt: Prompt) -> CatalogItem { CatalogItem { kind: CatalogItemKind::Prompt, - name: prompt.name, + name: sanitize_display_text(&prompt.name), server: server.to_string(), - description: prompt.description.unwrap_or_default(), + description: sanitize_display_text(&prompt.description.unwrap_or_default()), arguments: prompt.arguments, ..Default::default() } @@ -510,10 +562,10 @@ pub(crate) mod test_fixtures { use base64::Engine; use base64::engine::general_purpose::STANDARD; use rmcp::model::{ - ErrorData, GetPromptResponse, ListPromptsResult, ListResourceTemplatesResult, - ListResourcesResult, ListToolsResult, PaginatedRequestParams, PromptsCapability, - ReadResourceResponse, ResourceContents, ResourcesCapability, ServerCapabilities, - ServerInfo, + CallToolResponse, ErrorData, GetPromptResponse, ListPromptsResult, + ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams, + PromptsCapability, ReadResourceResponse, ResourceContents, ResourcesCapability, + ServerCapabilities, ServerInfo, }; use rmcp::service::{RequestContext, RunningService}; use rmcp::{RoleServer, ServerHandler, ServiceExt}; @@ -537,6 +589,7 @@ pub(crate) mod test_fixtures { 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, @@ -544,6 +597,7 @@ pub(crate) mod test_fixtures { 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 { @@ -552,6 +606,7 @@ pub(crate) mod test_fixtures { tools_capability: true, resources_capability: false, prompts_capability: false, + hostile_prompt: false, fail_resource_listings: false, fail_prompt_listings: false, fail_get_prompt: false, @@ -559,6 +614,7 @@ pub(crate) mod test_fixtures { list_resources_calls: Arc::default(), list_prompts_calls: Arc::default(), get_prompt_calls: Arc::default(), + call_tool_calls: Arc::default(), } } } @@ -594,6 +650,18 @@ pub(crate) mod test_fixtures { )])) } + async fn call_tool( + &self, + _request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + self.call_tool_calls.fetch_add(1, Ordering::SeqCst); + Err(ErrorData::internal_error( + "call_tool should not be reached", + None, + )) + } + async fn list_resources( &self, _request: Option, @@ -673,7 +741,7 @@ pub(crate) mod test_fixtures { if self.fail_prompt_listings { return Err(ErrorData::internal_error("prompt listing exploded", None)); } - Ok(ListPromptsResult::with_all_items(vec![Prompt::new( + let mut prompts = vec![Prompt::new( "summarize", Some("Summarize a document"), Some(vec![ @@ -682,7 +750,19 @@ pub(crate) mod test_fixtures { .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( @@ -864,6 +944,32 @@ mod tests { 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(); @@ -1245,6 +1351,47 @@ mod tests { 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 { @@ -1265,6 +1412,28 @@ mod tests { 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(); diff --git a/src/function/mod.rs b/src/function/mod.rs index 47ff850..093df37 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -7,7 +7,7 @@ pub(crate) mod user_interaction; use crate::{ client::ThinkingBlock, - config::{Agent, RequestContext}, + config::{Agent, RequestContext, flatten_prompt_messages, resolve_prompt_args}, graph, utils::*, }; @@ -753,6 +753,23 @@ impl Functions { }, ); + 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); @@ -760,6 +777,7 @@ impl Functions { format!("{}_{server}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX); let invoke_function_name = format!("{}_{server}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); 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 => { @@ -834,8 +852,26 @@ impl Functions { agent: false, }); } - // The declaration is added alongside its handler. - MCP_PROMPT_META_FUNCTION_NAME_PREFIX => {} + 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}"), } } @@ -1359,6 +1395,14 @@ impl ToolCall { eprintln!("{}", muted_warning_text(&format!("⚠️ {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(&format!("⚠️ {error_msg} ⚠️"))); + json!({"tool_call_error": error_msg}) + }) } else { Self::invoke_mcp_tool(ctx, cmd_name, &json_data) .await @@ -1429,6 +1473,15 @@ impl ToolCall { 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(&format!("⚠️ {error_msg} ⚠️"))); + json!({"tool_call_error": error_msg}) + }) + } _ if cmd_name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) => { Self::invoke_mcp_tool(ctx, &cmd_name, &json_data) .await @@ -1659,6 +1712,66 @@ impl ToolCall { } } + 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( &self, functions: &Functions, @@ -2617,20 +2730,22 @@ mod tests { 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_four() { + 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(), 4); + 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] @@ -2859,6 +2974,25 @@ mod tests { .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"))]); @@ -3035,6 +3169,99 @@ mod tests { 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 { diff --git a/src/repl/completer.rs b/src/repl/completer.rs index e71fbc3..ee1807e 100644 --- a/src/repl/completer.rs +++ b/src/repl/completer.rs @@ -1,6 +1,6 @@ use super::{REPL_COMMANDS, ReplCommand}; -use crate::config::{McpPromptCompletion, RequestContext}; +use crate::config::{McpPromptCompletion, RequestContext, sanitize_display_text}; use crate::mcp::ConnectedServer; use crate::utils::fuzzy_filter; @@ -173,7 +173,14 @@ fn complete_prompt_stage( McpPromptCompletion::PromptNames { server } => list_prompts_blocking(server, rpc_timeout) .unwrap_or_default() .into_iter() - .map(|prompt| (prompt.name, prompt.description)) + .map(|prompt| { + ( + sanitize_display_text(&prompt.name), + prompt + .description + .map(|description| sanitize_display_text(&description)), + ) + }) .collect(), McpPromptCompletion::ArgumentKeys { server, @@ -188,12 +195,18 @@ fn complete_prompt_stage( .into_iter() .filter(|arg| !typed_keys.contains(&arg.name)) .map(|arg| { - let description = match (arg.required == Some(true), arg.description) { + 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!("{}=", arg.name), description) + ( + format!("{}=", sanitize_display_text(&arg.name)), + description, + ) }) .collect(), }; @@ -358,6 +371,45 @@ mod prompt_completion_tests { 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 { diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 6a4a7db..95a04e3 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, flatten_prompt_messages, 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}; @@ -39,7 +39,6 @@ use reedline::{ default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings, }; use reedline::{MenuBuilder, Signal}; -use rmcp::model::PromptArgument; use std::collections::HashMap; use std::sync::LazyLock; use std::{env, process, sync::Arc}; @@ -798,8 +797,9 @@ pub async fn run_repl_command( .unwrap_or_default(); let (mut arguments, missing) = resolve_prompt_args(&declared, provided); for key in missing { - let value = - Text::new(&format!("{key}:")).prompt().with_context(|| { + 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); @@ -1804,16 +1804,13 @@ fn unquote_prompt_value(value: &str) -> &str { } } -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) +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) { @@ -1999,20 +1996,15 @@ mod tests { } #[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()); + 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]