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:
+3
-1
@@ -55,7 +55,9 @@ pub use self::skill_policy::SkillPolicy;
|
|||||||
pub use self::skill_registry::SkillRegistry;
|
pub use self::skill_registry::SkillRegistry;
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(crate) use self::tool_scope::test_fixtures;
|
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;
|
pub use self::update::run_self_update;
|
||||||
use crate::client::{
|
use crate::client::{
|
||||||
self, ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
|
self, ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
|
||||||
|
|||||||
+179
-10
@@ -420,16 +420,68 @@ pub fn format_prompt_arguments(arguments: &[PromptArgument]) -> String {
|
|||||||
arguments
|
arguments
|
||||||
.iter()
|
.iter()
|
||||||
.map(|arg| {
|
.map(|arg| {
|
||||||
|
let name = sanitize_display_text(&arg.name);
|
||||||
if arg.required == Some(true) {
|
if arg.required == Some(true) {
|
||||||
format!("{} (required)", arg.name)
|
format!("{name} (required)")
|
||||||
} else {
|
} else {
|
||||||
arg.name.clone()
|
name
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect::<Vec<_>>()
|
.collect::<Vec<_>>()
|
||||||
.join(", ")
|
.join(", ")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
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 {
|
fn catalog_key(item: &CatalogItem) -> String {
|
||||||
let id = item.uri.as_deref().unwrap_or(&item.name);
|
let id = item.uri.as_deref().unwrap_or(&item.name);
|
||||||
format!("{}:{id}", item.kind)
|
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 {
|
fn prompt_catalog_item(server: &str, prompt: Prompt) -> CatalogItem {
|
||||||
CatalogItem {
|
CatalogItem {
|
||||||
kind: CatalogItemKind::Prompt,
|
kind: CatalogItemKind::Prompt,
|
||||||
name: prompt.name,
|
name: sanitize_display_text(&prompt.name),
|
||||||
server: server.to_string(),
|
server: server.to_string(),
|
||||||
description: prompt.description.unwrap_or_default(),
|
description: sanitize_display_text(&prompt.description.unwrap_or_default()),
|
||||||
arguments: prompt.arguments,
|
arguments: prompt.arguments,
|
||||||
..Default::default()
|
..Default::default()
|
||||||
}
|
}
|
||||||
@@ -510,10 +562,10 @@ pub(crate) mod test_fixtures {
|
|||||||
use base64::Engine;
|
use base64::Engine;
|
||||||
use base64::engine::general_purpose::STANDARD;
|
use base64::engine::general_purpose::STANDARD;
|
||||||
use rmcp::model::{
|
use rmcp::model::{
|
||||||
ErrorData, GetPromptResponse, ListPromptsResult, ListResourceTemplatesResult,
|
CallToolResponse, ErrorData, GetPromptResponse, ListPromptsResult,
|
||||||
ListResourcesResult, ListToolsResult, PaginatedRequestParams, PromptsCapability,
|
ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams,
|
||||||
ReadResourceResponse, ResourceContents, ResourcesCapability, ServerCapabilities,
|
PromptsCapability, ReadResourceResponse, ResourceContents, ResourcesCapability,
|
||||||
ServerInfo,
|
ServerCapabilities, ServerInfo,
|
||||||
};
|
};
|
||||||
use rmcp::service::{RequestContext, RunningService};
|
use rmcp::service::{RequestContext, RunningService};
|
||||||
use rmcp::{RoleServer, ServerHandler, ServiceExt};
|
use rmcp::{RoleServer, ServerHandler, ServiceExt};
|
||||||
@@ -537,6 +589,7 @@ pub(crate) mod test_fixtures {
|
|||||||
pub(crate) tools_capability: bool,
|
pub(crate) tools_capability: bool,
|
||||||
pub(crate) resources_capability: bool,
|
pub(crate) resources_capability: bool,
|
||||||
pub(crate) prompts_capability: bool,
|
pub(crate) prompts_capability: bool,
|
||||||
|
pub(crate) hostile_prompt: bool,
|
||||||
pub(crate) fail_resource_listings: bool,
|
pub(crate) fail_resource_listings: bool,
|
||||||
pub(crate) fail_prompt_listings: bool,
|
pub(crate) fail_prompt_listings: bool,
|
||||||
pub(crate) fail_get_prompt: bool,
|
pub(crate) fail_get_prompt: bool,
|
||||||
@@ -544,6 +597,7 @@ pub(crate) mod test_fixtures {
|
|||||||
pub(crate) list_resources_calls: Arc<AtomicUsize>,
|
pub(crate) list_resources_calls: Arc<AtomicUsize>,
|
||||||
pub(crate) list_prompts_calls: Arc<AtomicUsize>,
|
pub(crate) list_prompts_calls: Arc<AtomicUsize>,
|
||||||
pub(crate) get_prompt_calls: Arc<AtomicUsize>,
|
pub(crate) get_prompt_calls: Arc<AtomicUsize>,
|
||||||
|
pub(crate) call_tool_calls: Arc<AtomicUsize>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for FixtureServer {
|
impl Default for FixtureServer {
|
||||||
@@ -552,6 +606,7 @@ pub(crate) mod test_fixtures {
|
|||||||
tools_capability: true,
|
tools_capability: true,
|
||||||
resources_capability: false,
|
resources_capability: false,
|
||||||
prompts_capability: false,
|
prompts_capability: false,
|
||||||
|
hostile_prompt: false,
|
||||||
fail_resource_listings: false,
|
fail_resource_listings: false,
|
||||||
fail_prompt_listings: false,
|
fail_prompt_listings: false,
|
||||||
fail_get_prompt: false,
|
fail_get_prompt: false,
|
||||||
@@ -559,6 +614,7 @@ pub(crate) mod test_fixtures {
|
|||||||
list_resources_calls: Arc::default(),
|
list_resources_calls: Arc::default(),
|
||||||
list_prompts_calls: Arc::default(),
|
list_prompts_calls: Arc::default(),
|
||||||
get_prompt_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<RoleServer>,
|
||||||
|
) -> Result<CallToolResponse, ErrorData> {
|
||||||
|
self.call_tool_calls.fetch_add(1, Ordering::SeqCst);
|
||||||
|
Err(ErrorData::internal_error(
|
||||||
|
"call_tool should not be reached",
|
||||||
|
None,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
async fn list_resources(
|
async fn list_resources(
|
||||||
&self,
|
&self,
|
||||||
_request: Option<PaginatedRequestParams>,
|
_request: Option<PaginatedRequestParams>,
|
||||||
@@ -673,7 +741,7 @@ pub(crate) mod test_fixtures {
|
|||||||
if self.fail_prompt_listings {
|
if self.fail_prompt_listings {
|
||||||
return Err(ErrorData::internal_error("prompt listing exploded", None));
|
return Err(ErrorData::internal_error("prompt listing exploded", None));
|
||||||
}
|
}
|
||||||
Ok(ListPromptsResult::with_all_items(vec![Prompt::new(
|
let mut prompts = vec![Prompt::new(
|
||||||
"summarize",
|
"summarize",
|
||||||
Some("Summarize a document"),
|
Some("Summarize a document"),
|
||||||
Some(vec![
|
Some(vec![
|
||||||
@@ -682,7 +750,19 @@ pub(crate) mod test_fixtures {
|
|||||||
.with_required(true),
|
.with_required(true),
|
||||||
PromptArgument::new("style"),
|
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(
|
async fn get_prompt(
|
||||||
@@ -864,6 +944,32 @@ mod tests {
|
|||||||
assert!(functions.contains("mcp_read_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]
|
#[test]
|
||||||
fn tool_scope_default_has_empty_mcp_runtime() {
|
fn tool_scope_default_has_empty_mcp_runtime() {
|
||||||
let scope = ToolScope::default();
|
let scope = ToolScope::default();
|
||||||
@@ -1245,6 +1351,47 @@ mod tests {
|
|||||||
assert_eq!(format_prompt_arguments(&[]), "");
|
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]
|
#[tokio::test]
|
||||||
async fn prompt_catalog_carries_arguments() {
|
async fn prompt_catalog_carries_arguments() {
|
||||||
let fixture = FixtureServer {
|
let fixture = FixtureServer {
|
||||||
@@ -1265,6 +1412,28 @@ mod tests {
|
|||||||
assert_eq!(format_prompt_arguments(arguments), "path (required), style");
|
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]
|
#[tokio::test]
|
||||||
async fn prompt_catalog_degrades_when_one_server_fails() {
|
async fn prompt_catalog_degrades_when_one_server_fails() {
|
||||||
install_warn_collector();
|
install_warn_collector();
|
||||||
|
|||||||
+232
-5
@@ -7,7 +7,7 @@ pub(crate) mod user_interaction;
|
|||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
client::ThinkingBlock,
|
client::ThinkingBlock,
|
||||||
config::{Agent, RequestContext},
|
config::{Agent, RequestContext, flatten_prompt_messages, resolve_prompt_args},
|
||||||
graph,
|
graph,
|
||||||
utils::*,
|
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 {
|
for features in mcp_servers {
|
||||||
let server = &features.name;
|
let server = &features.name;
|
||||||
let search_function_name = format!("{}_{server}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
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);
|
format!("{}_{server}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX);
|
||||||
let invoke_function_name = format!("{}_{server}", MCP_INVOKE_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 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) {
|
for prefix in gated_meta_function_prefixes(&features) {
|
||||||
match prefix {
|
match prefix {
|
||||||
MCP_INVOKE_META_FUNCTION_NAME_PREFIX => {
|
MCP_INVOKE_META_FUNCTION_NAME_PREFIX => {
|
||||||
@@ -834,8 +852,26 @@ impl Functions {
|
|||||||
agent: false,
|
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}"),
|
_ => debug_assert!(false, "unhandled MCP meta-function prefix: {prefix}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1359,6 +1395,14 @@ impl ToolCall {
|
|||||||
eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️")));
|
eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️")));
|
||||||
json!({"tool_call_error": 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 {
|
} else {
|
||||||
Self::invoke_mcp_tool(ctx, cmd_name, &json_data)
|
Self::invoke_mcp_tool(ctx, cmd_name, &json_data)
|
||||||
.await
|
.await
|
||||||
@@ -1429,6 +1473,15 @@ impl ToolCall {
|
|||||||
json!({"tool_call_error": 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(&format!("⚠️ {error_msg} ⚠️")));
|
||||||
|
json!({"tool_call_error": error_msg})
|
||||||
|
})
|
||||||
|
}
|
||||||
_ if cmd_name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) => {
|
_ if cmd_name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) => {
|
||||||
Self::invoke_mcp_tool(ctx, &cmd_name, &json_data)
|
Self::invoke_mcp_tool(ctx, &cmd_name, &json_data)
|
||||||
.await
|
.await
|
||||||
@@ -1659,6 +1712,66 @@ impl ToolCall {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn get_mcp_prompt(
|
||||||
|
ctx: &RequestContext,
|
||||||
|
cmd_name: &str,
|
||||||
|
json_data: &Value,
|
||||||
|
) -> Result<Value> {
|
||||||
|
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(
|
fn extract_call_config_from_agent(
|
||||||
&self,
|
&self,
|
||||||
functions: &Functions,
|
functions: &Functions,
|
||||||
@@ -2617,20 +2730,22 @@ mod tests {
|
|||||||
f.append_mcp_meta_functions(vec![mcp_features("res", false, true, false)]);
|
f.append_mcp_meta_functions(vec![mcp_features("res", false, true, false)]);
|
||||||
assert_eq!(f.declarations().len(), 3);
|
assert_eq!(f.declarations().len(), 3);
|
||||||
assert!(!f.contains("mcp_invoke_res"));
|
assert!(!f.contains("mcp_invoke_res"));
|
||||||
|
assert!(!f.contains("mcp_prompt_res"));
|
||||||
assert!(f.contains("mcp_search_res"));
|
assert!(f.contains("mcp_search_res"));
|
||||||
assert!(f.contains("mcp_describe_res"));
|
assert!(f.contains("mcp_describe_res"));
|
||||||
assert!(f.contains("mcp_read_res"));
|
assert!(f.contains("mcp_read_res"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn functions_append_mcp_meta_all_capabilities_emits_four() {
|
fn functions_append_mcp_meta_all_capabilities_emits_five() {
|
||||||
let mut f = Functions::default();
|
let mut f = Functions::default();
|
||||||
f.append_mcp_meta_functions(vec![mcp_features("srv", true, true, true)]);
|
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_invoke_srv"));
|
||||||
assert!(f.contains("mcp_search_srv"));
|
assert!(f.contains("mcp_search_srv"));
|
||||||
assert!(f.contains("mcp_describe_srv"));
|
assert!(f.contains("mcp_describe_srv"));
|
||||||
assert!(f.contains("mcp_read_srv"));
|
assert!(f.contains("mcp_read_srv"));
|
||||||
|
assert!(f.contains("mcp_prompt_srv"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -2859,6 +2974,25 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn prompts_fixture() -> FixtureServer {
|
||||||
|
FixtureServer {
|
||||||
|
prompts_capability: true,
|
||||||
|
..Default::default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn eval_mcp_prompt(args: Value) -> Result<Value> {
|
||||||
|
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]
|
#[test]
|
||||||
fn expand_uri_template_substitutes_simple_vars() {
|
fn expand_uri_template_substitutes_simple_vars() {
|
||||||
let args = template_args(&[("path", json!("docs")), ("name", json!("readme"))]);
|
let args = template_args(&[("path", json!("docs")), ("name", json!("readme"))]);
|
||||||
@@ -3035,6 +3169,99 @@ mod tests {
|
|||||||
assert_eq!(output["text"], FIXTURE_LOG_TEXT);
|
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]
|
#[test]
|
||||||
fn eval_mcp_read_pages_text_with_offset() {
|
fn eval_mcp_read_pages_text_with_offset() {
|
||||||
let (page1, page2) = run_async(async {
|
let (page1, page2) = run_async(async {
|
||||||
|
|||||||
+56
-4
@@ -1,6 +1,6 @@
|
|||||||
use super::{REPL_COMMANDS, ReplCommand};
|
use super::{REPL_COMMANDS, ReplCommand};
|
||||||
|
|
||||||
use crate::config::{McpPromptCompletion, RequestContext};
|
use crate::config::{McpPromptCompletion, RequestContext, sanitize_display_text};
|
||||||
use crate::mcp::ConnectedServer;
|
use crate::mcp::ConnectedServer;
|
||||||
use crate::utils::fuzzy_filter;
|
use crate::utils::fuzzy_filter;
|
||||||
|
|
||||||
@@ -173,7 +173,14 @@ fn complete_prompt_stage(
|
|||||||
McpPromptCompletion::PromptNames { server } => list_prompts_blocking(server, rpc_timeout)
|
McpPromptCompletion::PromptNames { server } => list_prompts_blocking(server, rpc_timeout)
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|prompt| (prompt.name, prompt.description))
|
.map(|prompt| {
|
||||||
|
(
|
||||||
|
sanitize_display_text(&prompt.name),
|
||||||
|
prompt
|
||||||
|
.description
|
||||||
|
.map(|description| sanitize_display_text(&description)),
|
||||||
|
)
|
||||||
|
})
|
||||||
.collect(),
|
.collect(),
|
||||||
McpPromptCompletion::ArgumentKeys {
|
McpPromptCompletion::ArgumentKeys {
|
||||||
server,
|
server,
|
||||||
@@ -188,12 +195,18 @@ fn complete_prompt_stage(
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|arg| !typed_keys.contains(&arg.name))
|
.filter(|arg| !typed_keys.contains(&arg.name))
|
||||||
.map(|arg| {
|
.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, Some(description)) => Some(format!("{description} (required)")),
|
||||||
(true, None) => Some("(required)".to_string()),
|
(true, None) => Some("(required)".to_string()),
|
||||||
(false, description) => description,
|
(false, description) => description,
|
||||||
};
|
};
|
||||||
(format!("{}=", arg.name), description)
|
(
|
||||||
|
format!("{}=", sanitize_display_text(&arg.name)),
|
||||||
|
description,
|
||||||
|
)
|
||||||
})
|
})
|
||||||
.collect(),
|
.collect(),
|
||||||
};
|
};
|
||||||
@@ -358,6 +371,45 @@ mod prompt_completion_tests {
|
|||||||
assert!(values.is_empty());
|
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)]
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||||
async fn slow_listing_times_out_to_empty() {
|
async fn slow_listing_times_out_to_empty() {
|
||||||
let fixture = FixtureServer {
|
let fixture = FixtureServer {
|
||||||
|
|||||||
+20
-28
@@ -13,7 +13,7 @@ use crate::client::{
|
|||||||
};
|
};
|
||||||
use crate::config::{
|
use crate::config::{
|
||||||
AgentVariables, AppConfig, AssertState, Input, LastMessage, MacroState, RequestContext,
|
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::config::{AssetCategory, paths};
|
||||||
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
|
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,
|
default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
|
||||||
};
|
};
|
||||||
use reedline::{MenuBuilder, Signal};
|
use reedline::{MenuBuilder, Signal};
|
||||||
use rmcp::model::PromptArgument;
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
use std::{env, process, sync::Arc};
|
use std::{env, process, sync::Arc};
|
||||||
@@ -798,8 +797,9 @@ pub async fn run_repl_command(
|
|||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
let (mut arguments, missing) = resolve_prompt_args(&declared, provided);
|
let (mut arguments, missing) = resolve_prompt_args(&declared, provided);
|
||||||
for key in missing {
|
for key in missing {
|
||||||
let value =
|
let value = Text::new(&prompt_arg_inquire_label(server, name, &key))
|
||||||
Text::new(&format!("{key}:")).prompt().with_context(|| {
|
.prompt()
|
||||||
|
.with_context(|| {
|
||||||
format!("Failed to read prompt argument '{key}'")
|
format!("Failed to read prompt argument '{key}'")
|
||||||
})?;
|
})?;
|
||||||
arguments.insert(key, value);
|
arguments.insert(key, value);
|
||||||
@@ -1804,16 +1804,13 @@ fn unquote_prompt_value(value: &str) -> &str {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn resolve_prompt_args(
|
fn prompt_arg_inquire_label(server: &str, prompt: &str, arg: &str) -> String {
|
||||||
declared: &[PromptArgument],
|
format!(
|
||||||
provided: HashMap<String, String>,
|
"Prompt '{}' on '{}' requires '{}':",
|
||||||
) -> (HashMap<String, String>, Vec<String>) {
|
sanitize_display_text(prompt),
|
||||||
let missing = declared
|
sanitize_display_text(server),
|
||||||
.iter()
|
sanitize_display_text(arg)
|
||||||
.filter(|arg| arg.required == Some(true) && !provided.contains_key(&arg.name))
|
)
|
||||||
.map(|arg| arg.name.clone())
|
|
||||||
.collect();
|
|
||||||
(provided, missing)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
|
pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
|
||||||
@@ -1999,20 +1996,15 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn resolve_prompt_args_reports_missing_required_only() {
|
fn prompt_arg_inquire_label_sanitizes_all_components() {
|
||||||
let declared = vec![
|
assert_eq!(
|
||||||
PromptArgument::new("path").with_required(true),
|
prompt_arg_inquire_label("srv", "summarize", "path"),
|
||||||
PromptArgument::new("style"),
|
"Prompt 'summarize' on 'srv' requires 'path':"
|
||||||
];
|
);
|
||||||
|
assert_eq!(
|
||||||
let (resolved, missing) = resolve_prompt_args(&declared, HashMap::new());
|
prompt_arg_inquire_label("s\u{1b}[31mrv", "sum\u{1b}]0;x\u{7}marize", "pa\u{7}th"),
|
||||||
assert!(resolved.is_empty());
|
"Prompt 'summarize' on 'srv' requires 'pa th':"
|
||||||
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]
|
#[test]
|
||||||
|
|||||||
Reference in New Issue
Block a user