feat(mcp): add mcp_prompt meta-tool and harden prompt display rendering
Emit an mcp_prompt_<server> declaration for servers advertising the prompts capability, execute prompts via McpRuntime::prompt on both tool dispatch chains, and return the flattened prompt text as the tool result. Sanitize server-controlled prompt names, descriptions, and argument names before terminal rendering, and attribute the .prompt argument inquire label to its server and prompt. Per plans/mcp-resources-prompts-design.md §5.2 (T7).
This commit is contained in:
+56
-4
@@ -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 {
|
||||
|
||||
+20
-28
@@ -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<String, String>,
|
||||
) -> (HashMap<String, String>, Vec<String>) {
|
||||
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<String>, &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]
|
||||
|
||||
Reference in New Issue
Block a user