1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
// mcp-server/src/prompts.rs
use anyhow::{anyhow, Result};
use mcp_protocol::types::prompt::{Prompt, PromptGetResult, PromptMessage};
use std::collections::HashMap;
use std::sync::{RwLock};
use tokio::sync::broadcast;
/// Handler type for generating prompt messages
pub type PromptHandler = Box<dyn Fn(Option<HashMap<String, String>>) -> Result<Vec<PromptMessage>> + Send + Sync>;
/// Handler type for generating parameter completions
pub type CompletionHandler = Box<dyn Fn(String, Option<String>) -> Result<Vec<String>> + Send + Sync>;
/// Manages prompts for the MCP server
pub struct PromptManager {
/// Map of prompt name to prompt definition
prompts: RwLock<HashMap<String, Prompt>>,
/// Map of prompt name to prompt handler
handlers: RwLock<HashMap<String, PromptHandler>>,
/// Map of prompt name to parameter completion handlers
completion_handlers: RwLock<HashMap<String, HashMap<String, CompletionHandler>>>,
/// Sender for update notifications
update_tx: broadcast::Sender<()>,
}
impl PromptManager {
/// Create a new prompt manager
pub fn new() -> Self {
let (update_tx, _) = broadcast::channel(100);
Self {
prompts: RwLock::new(HashMap::new()),
handlers: RwLock::new(HashMap::new()),
completion_handlers: RwLock::new(HashMap::new()),
update_tx,
}
}
/// Register a prompt with the manager
pub fn register_prompt(
&self,
prompt: Prompt,
handler: impl Fn(Option<HashMap<String, String>>) -> Result<Vec<PromptMessage>> + Send + Sync + 'static,
) {
let name = prompt.name.clone();
// Add prompt to registry
{
let mut prompts = self.prompts.write().unwrap();
prompts.insert(name.clone(), prompt);
}
// Add handler to registry
{
let mut handlers = self.handlers.write().unwrap();
handlers.insert(name, Box::new(handler));
}
// Notify of update
let _ = self.update_tx.send(());
}
/// Register a completion provider for a prompt parameter
pub fn register_completion_provider(
&self,
prompt_name: &str,
param_name: &str,
handler: impl Fn(String, Option<String>) -> Result<Vec<String>> + Send + Sync + 'static,
) {
let mut completion_handlers = self.completion_handlers.write().unwrap();
// Get or create the map for this prompt
let prompt_completions = completion_handlers
.entry(prompt_name.to_string())
.or_insert_with(HashMap::new);
// Register the handler for this parameter
prompt_completions.insert(param_name.to_string(), Box::new(handler));
}
/// Get completions for a prompt parameter
pub async fn get_completions(
&self,
prompt_name: &str,
param_name: &str,
value: Option<String>,
) -> Result<Vec<String>> {
let completion_handlers = self.completion_handlers.read().unwrap();
// Check if we have any completion handlers for this prompt
if let Some(prompt_completions) = completion_handlers.get(prompt_name) {
// Check if we have a handler for this parameter
if let Some(handler) = prompt_completions.get(param_name) {
// Call the handler
return handler(param_name.to_string(), value);
}
}
// If we don't have a handler, return empty results
Ok(Vec::new())
}
/// List all registered prompts with optional pagination
pub async fn list_prompts(&self, cursor: Option<String>) -> (Vec<Prompt>, Option<String>) {
let prompts = self.prompts.read().unwrap();
// Get all prompts in a vector
let mut prompt_list: Vec<Prompt> = prompts.values().cloned().collect();
// Sort by name for consistent ordering
prompt_list.sort_by(|a, b| a.name.cmp(&b.name));
// Simple pagination implementation
if let Some(cursor) = cursor {
if !cursor.is_empty() {
// Skip items before the cursor
prompt_list = prompt_list
.into_iter()
.skip_while(|p| p.name != cursor)
.skip(1) // Skip the cursor item itself
.collect();
}
}
// For simplicity, we'll return at most 50 items per page
let page_size = 50;
let next_cursor = if prompt_list.len() > page_size {
// If we have more than page_size, return the next cursor
prompt_list[page_size - 1].name.clone()
} else {
// No more pages
return (prompt_list, None);
};
// Return the current page and the next cursor
(prompt_list.into_iter().take(page_size).collect(), Some(next_cursor))
}
/// Get a prompt by name and generate its content with the provided arguments
pub async fn get_prompt(&self, name: &str, arguments: Option<HashMap<String, String>>) -> Result<PromptGetResult> {
// Get prompt definition
let prompt = {
let prompts = self.prompts.read().unwrap();
prompts.get(name).cloned().ok_or_else(|| anyhow!("Prompt not found: {}", name))?
};
// Validate arguments against the prompt definition
self.validate_arguments(&prompt, &arguments)?;
// Get handler and execute it
let messages = {
let handlers = self.handlers.read().unwrap();
if let Some(handler) = handlers.get(name) {
// Execute handler
handler(arguments.clone())?
} else {
return Err(anyhow!("Handler not found for prompt: {}", name));
}
};
// Construct result
let result = PromptGetResult {
description: prompt.description,
messages,
};
Ok(result)
}
/// Subscribe to prompt list updates
pub fn subscribe_to_updates(&self) -> broadcast::Receiver<()> {
self.update_tx.subscribe()
}
/// Add an annotation to a prompt
pub async fn add_annotation(&self, name: &str, key: &str, value: serde_json::Value) -> Result<()> {
let mut prompts = self.prompts.write().unwrap();
if let Some(prompt) = prompts.get_mut(name) {
// Initialize annotations if not present
if prompt.annotations.is_none() {
prompt.annotations = Some(HashMap::new());
}
// Add or update annotation
if let Some(annotations) = &mut prompt.annotations {
annotations.insert(key.to_string(), value);
}
// Notify of update
let _ = self.update_tx.send(());
Ok(())
} else {
Err(anyhow!("Prompt not found: {}", name))
}
}
/// Get an annotation from a prompt
pub async fn get_annotation(&self, name: &str, key: &str) -> Result<Option<serde_json::Value>> {
let prompts = self.prompts.read().unwrap();
if let Some(prompt) = prompts.get(name) {
if let Some(annotations) = &prompt.annotations {
Ok(annotations.get(key).cloned())
} else {
Ok(None)
}
} else {
Err(anyhow!("Prompt not found: {}", name))
}
}
/// Validate prompt arguments against the prompt definition
fn validate_arguments(
&self,
prompt: &Prompt,
arguments: &Option<HashMap<String, String>>
) -> Result<()> {
// Check for required arguments
if let Some(prompt_args) = &prompt.arguments {
for arg in prompt_args {
if arg.required.unwrap_or(false) {
match arguments {
Some(args) => {
if !args.contains_key(&arg.name) {
return Err(anyhow!("Missing required argument: {}", arg.name));
}
// Check for empty values
if let Some(value) = args.get(&arg.name) {
if value.trim().is_empty() {
return Err(anyhow!("Required argument cannot be empty: {}", arg.name));
}
}
},
None => return Err(anyhow!("Missing required arguments")),
}
}
}
// Check for unexpected arguments
if let Some(args) = arguments {
for arg_name in args.keys() {
if !prompt_args.iter().any(|a| &a.name == arg_name) {
return Err(anyhow!("Unexpected argument: {}", arg_name));
}
}
}
}
Ok(())
}
}