use async_trait::async_trait;
use pmcp::types::{Content, GetPromptResult, PromptArgument, PromptInfo, PromptMessage};
use pmcp::PromptHandler;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::error::ToolkitError;
use crate::resources::StaticResourceHandler;
pub const CODE_MODE_PROMPT_NAME: &str = "start_code_mode";
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PromptConfig {
pub name: String,
pub description: String,
#[serde(default)]
pub include_resources: Vec<String>,
}
type PromptInfoOut = pmcp::types::PromptInfo;
impl PromptConfig {
pub fn to_prompt_info(&self) -> PromptInfoOut {
let info = PromptInfo::new(&self.name);
info.with_description(&self.description)
}
}
pub struct StaticPromptHandler {
name: String,
description: Option<String>,
arguments: Vec<PromptArgument>,
body: String,
}
impl StaticPromptHandler {
pub fn new(
name: impl Into<String>,
description: Option<impl Into<String>>,
arguments: Vec<PromptArgument>,
body: impl Into<String>,
) -> Self {
Self {
name: name.into(),
description: description.map(Into::into),
arguments,
body: body.into(),
}
}
pub fn from_configs(
prompts: &[PromptConfig],
resources: &StaticResourceHandler,
) -> Vec<(String, Self)> {
prompts
.iter()
.map(|p| {
let body = Self::resolve_body(p, resources);
let handler = Self::new(
&p.name,
Some(p.description.clone()),
Vec::new(), body,
);
(p.name.clone(), handler)
})
.collect()
}
fn resolve_body(prompt: &PromptConfig, resources: &StaticResourceHandler) -> String {
let mut content_parts: Vec<String> = Vec::new();
for resource_uri in &prompt.include_resources {
if let Some(resource) = resources.get(resource_uri) {
content_parts.push(resource.content.clone());
} else {
tracing::warn!(
uri = %resource_uri,
prompt = %prompt.name,
"Resource not found for prompt",
);
}
}
if content_parts.is_empty() {
format!("(No resources found for prompt '{}')", prompt.name)
} else {
content_parts.join("\n\n---\n\n")
}
}
}
#[async_trait]
impl PromptHandler for StaticPromptHandler {
async fn handle(
&self,
args: HashMap<String, String>,
_extra: pmcp::RequestHandlerExtra,
) -> pmcp::Result<GetPromptResult> {
for arg in &self.arguments {
if arg.required && !args.contains_key(&arg.name) {
return Err(pmcp::Error::validation(format!(
"Required argument '{}' is missing",
arg.name
)));
}
}
Ok(GetPromptResult::new(
vec![PromptMessage::user(Content::text(self.body.clone()))],
self.description.clone(),
))
}
fn metadata(&self) -> Option<PromptInfoOut> {
let mut info = PromptInfo::new(&self.name);
if let Some(desc) = &self.description {
info = info.with_description(desc);
}
if !self.arguments.is_empty() {
info = info.with_arguments(self.arguments.clone());
}
Some(info)
}
}
pub fn prompt_handlers_from_config(
cfg: &crate::config::ServerConfig,
) -> Vec<(String, StaticPromptHandler)> {
let resource_handler = crate::resources::StaticResourceHandler::from(cfg);
let configs: Vec<PromptConfig> = cfg
.prompts
.iter()
.map(|p| PromptConfig {
name: p.name.clone(),
description: p.description.clone().unwrap_or_default(),
include_resources: p.include_resources.clone(),
})
.collect();
StaticPromptHandler::from_configs(&configs, &resource_handler)
}
impl From<&crate::config::ServerConfig> for StaticPromptHandler {
fn from(cfg: &crate::config::ServerConfig) -> Self {
let mut pairs = prompt_handlers_from_config(cfg);
if pairs.is_empty() {
StaticPromptHandler::new(
"<no-prompts>",
Some("config declared no [[prompts]] entries"),
Vec::new(),
String::new(),
)
} else {
pairs.remove(0).1
}
}
}
pub fn resolve_extra_prompt_content(
prompts: &[PromptConfig],
resources: &[crate::resources::ResourceConfig],
) -> Vec<String> {
const AUTO_GENERATED: &[&str] = &["code-mode://instructions", "code-mode://policies"];
let prompt = prompts.iter().find(|p| p.name == CODE_MODE_PROMPT_NAME);
let Some(prompt) = prompt else {
return vec![];
};
prompt
.include_resources
.iter()
.filter(|uri| !AUTO_GENERATED.contains(&uri.as_str()))
.filter_map(|uri| {
resources
.iter()
.find(|r| r.uri == *uri)
.and_then(|r| r.content.clone())
})
.filter(|c| !c.is_empty())
.collect()
}
#[allow(dead_code)]
fn _ensure_error_path_kept() -> Option<ToolkitError> {
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::resources::ResourceConfig;
use pmcp::types::Content;
use pmcp::RequestHandlerExtra;
fn mk_extra() -> RequestHandlerExtra {
RequestHandlerExtra::default()
}
#[test]
fn prompt_config_to_info() {
let config = PromptConfig {
name: "test-prompt".to_string(),
description: "A test prompt".to_string(),
include_resources: vec!["docs://test".to_string()],
};
let info = config.to_prompt_info();
assert_eq!(info.name, "test-prompt");
assert_eq!(info.description, Some("A test prompt".to_string()));
assert!(info.arguments.is_none());
}
#[tokio::test]
async fn handle_with_all_required_args_succeeds() {
let handler = StaticPromptHandler::new(
"needs-foo",
Some("requires foo"),
vec![PromptArgument::new("foo").required()],
"Hello {{foo}}",
);
let args = HashMap::from([("foo".to_string(), "world".to_string())]);
let result = handler.handle(args, mk_extra()).await.unwrap();
assert_eq!(result.messages.len(), 1);
assert_eq!(result.description.as_deref(), Some("requires foo"));
match &result.messages[0].content {
Content::Text { text } => assert_eq!(text, "Hello {{foo}}"),
other => panic!("expected text content, got {:?}", other),
}
}
#[tokio::test]
async fn handle_missing_required_arg_returns_validation_err() {
let handler = StaticPromptHandler::new(
"needs-foo",
Some("requires foo"),
vec![PromptArgument::new("foo").required()],
"Hello {{foo}}",
);
let result = handler.handle(HashMap::new(), mk_extra()).await;
let err = result.expect_err("expected validation error");
let msg = err.to_string();
assert!(
msg.contains("foo"),
"error message should mention the missing argument 'foo': {msg}",
);
assert!(
msg.to_lowercase().contains("missing") || msg.to_lowercase().contains("required"),
"error message should indicate the missing-required-arg path: {msg}",
);
}
#[tokio::test]
async fn metadata_returns_some_promptinfo_with_description_and_args() {
let handler = StaticPromptHandler::new(
"with-meta",
Some("a described prompt"),
vec![
PromptArgument::new("a").required(),
PromptArgument::new("b"),
],
"body",
);
let info = handler.metadata().expect("metadata should return Some");
assert_eq!(info.name, "with-meta");
assert_eq!(info.description.as_deref(), Some("a described prompt"));
let args = info.arguments.expect("arguments should be populated");
assert_eq!(args.len(), 2);
assert_eq!(args[0].name, "a");
assert!(args[0].required);
assert_eq!(args[1].name, "b");
assert!(!args[1].required);
}
#[test]
fn metadata_with_no_arguments_omits_arguments_field() {
let handler = StaticPromptHandler::new("plain", Some("d"), vec![], "body");
let info = handler.metadata().unwrap();
assert!(info.arguments.is_none());
}
#[tokio::test]
async fn from_configs_resolves_resource_bodies_deterministically() {
let resource_configs = vec![ResourceConfig {
uri: "docs://test".to_string(),
name: "Test Resource".to_string(),
description: None,
mime_type: "text/plain".to_string(),
content: Some("Hello from resource".to_string()),
content_file: None,
meta: None,
}];
let resources =
crate::resources::StaticResourceHandler::from_configs(&resource_configs).unwrap();
let prompts = vec![
PromptConfig {
name: "p1".to_string(),
description: "first".to_string(),
include_resources: vec!["docs://test".to_string()],
},
PromptConfig {
name: "p2".to_string(),
description: "second".to_string(),
include_resources: vec![],
},
];
let mut materialized = StaticPromptHandler::from_configs(&prompts, &resources);
assert_eq!(materialized.len(), 2);
assert_eq!(materialized[0].0, "p1");
assert_eq!(materialized[1].0, "p2");
let (_, p1_handler) = materialized.remove(0);
let result = p1_handler.handle(HashMap::new(), mk_extra()).await.unwrap();
match &result.messages[0].content {
Content::Text { text } => assert_eq!(text, "Hello from resource"),
other => panic!("expected text, got {:?}", other),
}
let (_, p2_handler) = materialized.remove(0);
let result = p2_handler.handle(HashMap::new(), mk_extra()).await.unwrap();
match &result.messages[0].content {
Content::Text { text } => assert!(text.contains("p2")),
other => panic!("expected text, got {:?}", other),
}
}
#[test]
fn resolve_extra_prompt_content_filters_auto_generated() {
let prompts = vec![PromptConfig {
name: CODE_MODE_PROMPT_NAME.to_string(),
description: "code mode".to_string(),
include_resources: vec![
"code-mode://instructions".to_string(), "docs://learnings".to_string(),
],
}];
let resources = vec![
ResourceConfig {
uri: "code-mode://instructions".to_string(),
name: "auto".to_string(),
description: None,
mime_type: "text/markdown".to_string(),
content: Some("AUTO".to_string()),
content_file: None,
meta: None,
},
ResourceConfig {
uri: "docs://learnings".to_string(),
name: "learnings".to_string(),
description: None,
mime_type: "text/markdown".to_string(),
content: Some("LEARNED".to_string()),
content_file: None,
meta: None,
},
];
let extras = resolve_extra_prompt_content(&prompts, &resources);
assert_eq!(extras, vec!["LEARNED".to_string()]);
}
}