feat: Add prompt support (#351)
* feat: add prompt support with typed argument handling - Implement #[prompt], #[prompt_router], and #[prompt_handler] macros - Add automatic JSON schema generation from Rust types for arguments - Support flexible async handler signatures with automatic adaptation - Create PromptRouter for efficient prompt dispatch - Include comprehensive tests and example implementation This enables MCP servers to provide reusable prompt templates that LLMs can discover and invoke with strongly-typed parameters, similar to the existing tool system but optimized for prompt use cases. 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com> * refactor: unify parameter handling between tools and prompts - Replace Arguments<T> with Parameters<T> for consistent API - Create shared common module for tool/prompt utilities - Modernize async handling with futures::future::BoxFuture - Move cached_schema_for_type to common module for reuse - Update error types from rmcp::Error to rmcp::ErrorData - Add comprehensive trait implementations for parameter extraction 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com> * refactor: extract Parameters wrapper to shared module and unify trait usage - Move Parameters type from common.rs to dedicated wrapper/parameters.rs module - Unify extractor traits by using FromContextPart for both tools and prompts - Remove duplicate FromToolCallContextPart and FromPromptContextPart traits - Add structured output support to Json wrapper with IntoCallToolResult impl - Improve error messages with more descriptive panic for schema serialization - Update all imports across codebase to use new module path - Clean up trailing whitespace and formatting inconsistencies This consolidates parameter extraction logic and reduces code duplication between tool and prompt handlers while maintaining backward compatibility. * chore: remove committed .egg-info directory and update gitignore * docs: fix documentation formatting
This commit is contained in:
Michael Assaf
GitHub
parent
c0452f93a4
commit
9349f5cb26
3 changed files with 3191 additions and 464 deletions
|
|
@ -2,7 +2,10 @@
|
|||
|
||||
use rmcp::{
|
||||
ServerHandler,
|
||||
handler::server::{router::tool::ToolRouter, tool::Parameters, wrapper::Json},
|
||||
handler::server::{
|
||||
router::tool::ToolRouter,
|
||||
wrapper::{Json, Parameters},
|
||||
},
|
||||
model::{ServerCapabilities, ServerInfo},
|
||||
schemars, tool, tool_handler, tool_router,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -3,9 +3,12 @@ use std::sync::Arc;
|
|||
|
||||
use rmcp::{
|
||||
ErrorData as McpError, RoleServer, ServerHandler,
|
||||
handler::server::{router::tool::ToolRouter, tool::Parameters},
|
||||
handler::server::{
|
||||
router::{prompt::PromptRouter, tool::ToolRouter},
|
||||
wrapper::Parameters,
|
||||
},
|
||||
model::*,
|
||||
schemars,
|
||||
prompt, prompt_handler, prompt_router, schemars,
|
||||
service::RequestContext,
|
||||
tool, tool_handler, tool_router,
|
||||
};
|
||||
|
|
@ -18,10 +21,26 @@ pub struct StructRequest {
|
|||
pub b: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct ExamplePromptArgs {
|
||||
/// A message to put in the prompt
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct CounterAnalysisArgs {
|
||||
/// The target value you're trying to reach
|
||||
pub goal: i32,
|
||||
/// Preferred strategy: 'fast' or 'careful'
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub strategy: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Counter {
|
||||
counter: Arc<Mutex<i32>>,
|
||||
tool_router: ToolRouter<Counter>,
|
||||
prompt_router: PromptRouter<Counter>,
|
||||
}
|
||||
|
||||
#[tool_router]
|
||||
|
|
@ -31,6 +50,7 @@ impl Counter {
|
|||
Self {
|
||||
counter: Arc::new(Mutex::new(0)),
|
||||
tool_router: Self::tool_router(),
|
||||
prompt_router: Self::prompt_router(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -86,7 +106,63 @@ impl Counter {
|
|||
)]))
|
||||
}
|
||||
}
|
||||
|
||||
#[prompt_router]
|
||||
impl Counter {
|
||||
/// This is an example prompt that takes one required argument, message
|
||||
#[prompt(name = "example_prompt")]
|
||||
async fn example_prompt(
|
||||
&self,
|
||||
Parameters(args): Parameters<ExamplePromptArgs>,
|
||||
_ctx: RequestContext<RoleServer>,
|
||||
) -> Result<Vec<PromptMessage>, McpError> {
|
||||
let prompt = format!(
|
||||
"This is an example prompt with your message here: '{}'",
|
||||
args.message
|
||||
);
|
||||
Ok(vec![PromptMessage {
|
||||
role: PromptMessageRole::User,
|
||||
content: PromptMessageContent::text(prompt),
|
||||
}])
|
||||
}
|
||||
|
||||
/// Analyze the current counter value and suggest next steps
|
||||
#[prompt(name = "counter_analysis")]
|
||||
async fn counter_analysis(
|
||||
&self,
|
||||
Parameters(args): Parameters<CounterAnalysisArgs>,
|
||||
_ctx: RequestContext<RoleServer>,
|
||||
) -> Result<GetPromptResult, McpError> {
|
||||
let strategy = args.strategy.unwrap_or_else(|| "careful".to_string());
|
||||
let current_value = *self.counter.lock().await;
|
||||
let difference = args.goal - current_value;
|
||||
|
||||
let messages = vec![
|
||||
PromptMessage::new_text(
|
||||
PromptMessageRole::Assistant,
|
||||
"I'll analyze the counter situation and suggest the best approach.",
|
||||
),
|
||||
PromptMessage::new_text(
|
||||
PromptMessageRole::User,
|
||||
format!(
|
||||
"Current counter value: {}\nGoal value: {}\nDifference: {}\nStrategy preference: {}\n\nPlease analyze the situation and suggest the best approach to reach the goal.",
|
||||
current_value, args.goal, difference, strategy
|
||||
),
|
||||
),
|
||||
];
|
||||
|
||||
Ok(GetPromptResult {
|
||||
description: Some(format!(
|
||||
"Counter analysis for reaching {} from {}",
|
||||
args.goal, current_value
|
||||
)),
|
||||
messages,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tool_handler]
|
||||
#[prompt_handler]
|
||||
impl ServerHandler for Counter {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo {
|
||||
|
|
@ -97,7 +173,7 @@ impl ServerHandler for Counter {
|
|||
.enable_tools()
|
||||
.build(),
|
||||
server_info: Implementation::from_build_env(),
|
||||
instructions: Some("This server provides a counter tool that can increment and decrement values. The counter starts at 0 and can be modified using the 'increment' and 'decrement' tools. Use 'get_value' to check the current count.".to_string()),
|
||||
instructions: Some("This server provides counter tools and prompts. Tools: increment, decrement, get_value, say_hello, echo, sum. Prompts: example_prompt (takes a message), counter_analysis (analyzes counter state with a goal).".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -142,52 +218,6 @@ impl ServerHandler for Counter {
|
|||
}
|
||||
}
|
||||
|
||||
async fn list_prompts(
|
||||
&self,
|
||||
_request: Option<PaginatedRequestParam>,
|
||||
_: RequestContext<RoleServer>,
|
||||
) -> Result<ListPromptsResult, McpError> {
|
||||
Ok(ListPromptsResult {
|
||||
next_cursor: None,
|
||||
prompts: vec![Prompt::new(
|
||||
"example_prompt",
|
||||
Some("This is an example prompt that takes one required argument, message"),
|
||||
Some(vec![PromptArgument {
|
||||
name: "message".to_string(),
|
||||
description: Some("A message to put in the prompt".to_string()),
|
||||
required: Some(true),
|
||||
}]),
|
||||
)],
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_prompt(
|
||||
&self,
|
||||
GetPromptRequestParam { name, arguments }: GetPromptRequestParam,
|
||||
_: RequestContext<RoleServer>,
|
||||
) -> Result<GetPromptResult, McpError> {
|
||||
match name.as_str() {
|
||||
"example_prompt" => {
|
||||
let message = arguments
|
||||
.and_then(|json| json.get("message")?.as_str().map(|s| s.to_string()))
|
||||
.ok_or_else(|| {
|
||||
McpError::invalid_params("No message provided to example_prompt", None)
|
||||
})?;
|
||||
|
||||
let prompt =
|
||||
format!("This is an example prompt with your message here: '{message}'");
|
||||
Ok(GetPromptResult {
|
||||
description: None,
|
||||
messages: vec![PromptMessage {
|
||||
role: PromptMessageRole::User,
|
||||
content: PromptMessageContent::text(prompt),
|
||||
}],
|
||||
})
|
||||
}
|
||||
_ => Err(McpError::invalid_params("prompt not found", None)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_resource_templates(
|
||||
&self,
|
||||
_request: Option<PaginatedRequestParam>,
|
||||
|
|
@ -212,3 +242,76 @@ impl ServerHandler for Counter {
|
|||
Ok(self.get_info())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_prompt_attributes_generated() {
|
||||
// Verify that the prompt macros generate the expected attributes
|
||||
let example_attr = Counter::example_prompt_prompt_attr();
|
||||
assert_eq!(example_attr.name, "example_prompt");
|
||||
assert!(example_attr.description.is_some());
|
||||
assert!(example_attr.arguments.is_some());
|
||||
|
||||
let args = example_attr.arguments.unwrap();
|
||||
assert_eq!(args.len(), 1);
|
||||
assert_eq!(args[0].name, "message");
|
||||
assert_eq!(args[0].required, Some(true));
|
||||
|
||||
let analysis_attr = Counter::counter_analysis_prompt_attr();
|
||||
assert_eq!(analysis_attr.name, "counter_analysis");
|
||||
assert!(analysis_attr.description.is_some());
|
||||
assert!(analysis_attr.arguments.is_some());
|
||||
|
||||
let args = analysis_attr.arguments.unwrap();
|
||||
assert_eq!(args.len(), 2);
|
||||
assert_eq!(args[0].name, "goal");
|
||||
assert_eq!(args[0].required, Some(true));
|
||||
assert_eq!(args[1].name, "strategy");
|
||||
assert_eq!(args[1].required, Some(false));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_prompt_router_has_routes() {
|
||||
let router = Counter::prompt_router();
|
||||
assert!(router.has_route("example_prompt"));
|
||||
assert!(router.has_route("counter_analysis"));
|
||||
|
||||
let prompts = router.list_all();
|
||||
assert_eq!(prompts.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_example_prompt_execution() {
|
||||
let counter = Counter::new();
|
||||
let context = rmcp::handler::server::prompt::PromptContext::new(
|
||||
&counter,
|
||||
"example_prompt".to_string(),
|
||||
Some({
|
||||
let mut map = serde_json::Map::new();
|
||||
map.insert(
|
||||
"message".to_string(),
|
||||
serde_json::Value::String("Test message".to_string()),
|
||||
);
|
||||
map
|
||||
}),
|
||||
RequestContext {
|
||||
meta: Default::default(),
|
||||
ct: tokio_util::sync::CancellationToken::new(),
|
||||
id: rmcp::model::NumberOrString::String("test-1".to_string()),
|
||||
peer: Default::default(),
|
||||
extensions: Default::default(),
|
||||
},
|
||||
);
|
||||
|
||||
let router = Counter::prompt_router();
|
||||
let result = router.get_prompt(context).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
let prompt_result = result.unwrap();
|
||||
assert_eq!(prompt_result.messages.len(), 1);
|
||||
assert_eq!(prompt_result.messages[0].role, PromptMessageRole::User);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use std::sync::Arc;
|
|||
|
||||
use rmcp::{
|
||||
ServerHandler,
|
||||
handler::server::{router::tool::ToolRouter, tool::Parameters},
|
||||
handler::server::{router::tool::ToolRouter, wrapper::Parameters},
|
||||
model::{ServerCapabilities, ServerInfo},
|
||||
schemars, tool, tool_handler, tool_router,
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in a new issue