mcpkit_server/capability/
prompts.rs

1//! Prompt capability implementation.
2//!
3//! This module provides utilities for managing and rendering prompts
4//! in an MCP server.
5
6use crate::context::Context;
7use crate::handler::PromptHandler;
8use mcpkit_core::error::McpError;
9use mcpkit_core::types::prompt::{GetPromptResult, Prompt, PromptArgument, PromptMessage};
10use serde_json::Value;
11use std::collections::HashMap;
12use std::future::Future;
13use std::pin::Pin;
14
15/// A boxed async function for prompt rendering.
16pub type BoxedPromptFn = Box<
17    dyn for<'a> Fn(
18            Option<Value>,
19            &'a Context<'a>,
20        )
21            -> Pin<Box<dyn Future<Output = Result<GetPromptResult, McpError>> + Send + 'a>>
22        + Send
23        + Sync,
24>;
25
26/// A registered prompt with metadata and handler.
27pub struct RegisteredPrompt {
28    /// Prompt metadata.
29    pub prompt: Prompt,
30    /// Handler function for rendering.
31    pub handler: BoxedPromptFn,
32}
33
34/// Service for managing prompts.
35///
36/// This provides a registry for prompts and handles rendering
37/// them with arguments.
38pub struct PromptService {
39    prompts: HashMap<String, RegisteredPrompt>,
40}
41
42impl Default for PromptService {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl PromptService {
49    /// Create a new empty prompt service.
50    #[must_use]
51    pub fn new() -> Self {
52        Self {
53            prompts: HashMap::new(),
54        }
55    }
56
57    /// Register a prompt with a handler function.
58    pub fn register<F, Fut>(&mut self, prompt: Prompt, handler: F)
59    where
60        F: Fn(Option<Value>, &Context<'_>) -> Fut + Send + Sync + 'static,
61        Fut: Future<Output = Result<GetPromptResult, McpError>> + Send + 'static,
62    {
63        let name = prompt.name.clone();
64        let boxed: BoxedPromptFn = Box::new(move |args, ctx| Box::pin(handler(args, ctx)));
65        self.prompts.insert(
66            name,
67            RegisteredPrompt {
68                prompt,
69                handler: boxed,
70            },
71        );
72    }
73
74    /// Get a prompt by name.
75    #[must_use]
76    pub fn get(&self, name: &str) -> Option<&RegisteredPrompt> {
77        self.prompts.get(name)
78    }
79
80    /// Check if a prompt exists.
81    #[must_use]
82    pub fn contains(&self, name: &str) -> bool {
83        self.prompts.contains_key(name)
84    }
85
86    /// List all registered prompts.
87    #[must_use]
88    pub fn list(&self) -> Vec<&Prompt> {
89        self.prompts.values().map(|r| &r.prompt).collect()
90    }
91
92    /// Get the number of registered prompts.
93    #[must_use]
94    pub fn len(&self) -> usize {
95        self.prompts.len()
96    }
97
98    /// Check if the service has no prompts.
99    #[must_use]
100    pub fn is_empty(&self) -> bool {
101        self.prompts.is_empty()
102    }
103
104    /// Render a prompt by name with arguments.
105    pub async fn render(
106        &self,
107        name: &str,
108        arguments: Option<Value>,
109        ctx: &Context<'_>,
110    ) -> Result<GetPromptResult, McpError> {
111        let registered = self.prompts.get(name).ok_or_else(|| {
112            McpError::invalid_params("prompts/get", format!("Unknown prompt: {name}"))
113        })?;
114
115        (registered.handler)(arguments, ctx).await
116    }
117}
118
119impl PromptHandler for PromptService {
120    async fn list_prompts(&self, _ctx: &Context<'_>) -> Result<Vec<Prompt>, McpError> {
121        Ok(self.list().into_iter().cloned().collect())
122    }
123
124    async fn get_prompt(
125        &self,
126        name: &str,
127        arguments: Option<serde_json::Map<String, Value>>,
128        ctx: &Context<'_>,
129    ) -> Result<GetPromptResult, McpError> {
130        let args = arguments.map(Value::Object);
131        self.render(name, args, ctx).await
132    }
133}
134
135/// Builder for creating prompts with a fluent API.
136pub struct PromptBuilder {
137    name: String,
138    description: Option<String>,
139    arguments: Vec<PromptArgument>,
140}
141
142impl PromptBuilder {
143    /// Create a new prompt builder.
144    pub fn new(name: impl Into<String>) -> Self {
145        Self {
146            name: name.into(),
147            description: None,
148            arguments: Vec::new(),
149        }
150    }
151
152    /// Set the prompt description.
153    pub fn description(mut self, desc: impl Into<String>) -> Self {
154        self.description = Some(desc.into());
155        self
156    }
157
158    /// Add a required argument.
159    pub fn required_arg(mut self, name: impl Into<String>, description: impl Into<String>) -> Self {
160        self.arguments.push(PromptArgument {
161            name: name.into(),
162            title: None,
163            description: Some(description.into()),
164            required: Some(true),
165        });
166        self
167    }
168
169    /// Add an optional argument.
170    pub fn optional_arg(mut self, name: impl Into<String>, description: impl Into<String>) -> Self {
171        self.arguments.push(PromptArgument {
172            name: name.into(),
173            title: None,
174            description: Some(description.into()),
175            required: Some(false),
176        });
177        self
178    }
179
180    /// Add a custom argument.
181    #[must_use]
182    pub fn argument(mut self, arg: PromptArgument) -> Self {
183        self.arguments.push(arg);
184        self
185    }
186
187    /// Build the prompt.
188    #[must_use]
189    pub fn build(self) -> Prompt {
190        Prompt {
191            name: self.name,
192            title: None,
193            description: self.description,
194            icons: None,
195            arguments: if self.arguments.is_empty() {
196                None
197            } else {
198                Some(self.arguments)
199            },
200            meta: None,
201        }
202    }
203}
204
205/// Builder for creating prompt results.
206pub struct PromptResultBuilder {
207    description: Option<String>,
208    messages: Vec<PromptMessage>,
209}
210
211impl Default for PromptResultBuilder {
212    fn default() -> Self {
213        Self::new()
214    }
215}
216
217impl PromptResultBuilder {
218    /// Create a new result builder.
219    #[must_use]
220    pub const fn new() -> Self {
221        Self {
222            description: None,
223            messages: Vec::new(),
224        }
225    }
226
227    /// Set the result description.
228    pub fn description(mut self, desc: impl Into<String>) -> Self {
229        self.description = Some(desc.into());
230        self
231    }
232
233    /// Add a user message with text content.
234    pub fn user_text(mut self, text: impl Into<String>) -> Self {
235        self.messages.push(PromptMessage::user(text.into()));
236        self
237    }
238
239    /// Add an assistant message with text content.
240    pub fn assistant_text(mut self, text: impl Into<String>) -> Self {
241        self.messages.push(PromptMessage::assistant(text.into()));
242        self
243    }
244
245    /// Add a custom message.
246    #[must_use]
247    pub fn message(mut self, msg: PromptMessage) -> Self {
248        self.messages.push(msg);
249        self
250    }
251
252    /// Build the result.
253    #[must_use]
254    pub fn build(self) -> GetPromptResult {
255        GetPromptResult {
256            description: self.description,
257            messages: self.messages,
258            meta: None,
259        }
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266    use crate::context::{Context, NoOpPeer};
267    use mcpkit_core::capability::{ClientCapabilities, ServerCapabilities};
268    use mcpkit_core::protocol::RequestId;
269    use mcpkit_core::protocol_version::ProtocolVersion;
270
271    fn make_context() -> (
272        RequestId,
273        ClientCapabilities,
274        ServerCapabilities,
275        ProtocolVersion,
276        NoOpPeer,
277    ) {
278        (
279            RequestId::Number(1),
280            ClientCapabilities::default(),
281            ServerCapabilities::default(),
282            ProtocolVersion::LATEST,
283            NoOpPeer,
284        )
285    }
286
287    #[test]
288    fn test_prompt_builder() {
289        let prompt = PromptBuilder::new("code-review")
290            .description("Review code for issues")
291            .required_arg("code", "The code to review")
292            .optional_arg("language", "Programming language")
293            .build();
294
295        assert_eq!(prompt.name, "code-review");
296        assert_eq!(
297            prompt.description.as_deref(),
298            Some("Review code for issues")
299        );
300        assert_eq!(prompt.arguments.as_ref().map(std::vec::Vec::len), Some(2));
301    }
302
303    #[test]
304    fn test_prompt_result_builder() {
305        let result = PromptResultBuilder::new()
306            .description("Generated review")
307            .user_text("Please review this code")
308            .assistant_text("I'll analyze the code...")
309            .build();
310
311        assert_eq!(result.description.as_deref(), Some("Generated review"));
312        assert_eq!(result.messages.len(), 2);
313    }
314
315    #[tokio::test]
316    async fn test_prompt_service() -> Result<(), Box<dyn std::error::Error>> {
317        let mut service = PromptService::new();
318
319        let prompt = PromptBuilder::new("greeting")
320            .description("Generate a greeting")
321            .required_arg("name", "Name to greet")
322            .build();
323
324        service.register(prompt, |args, _ctx| async move {
325            let name = args
326                .and_then(|v| v.get("name").and_then(|n| n.as_str()).map(String::from))
327                .unwrap_or_else(|| "World".to_string());
328
329            Ok(PromptResultBuilder::new()
330                .user_text(format!("Generate a greeting for {name}"))
331                .build())
332        });
333
334        assert!(service.contains("greeting"));
335        assert_eq!(service.len(), 1);
336
337        let (req_id, client_caps, server_caps, protocol_version, peer) = make_context();
338        let ctx = Context::new(
339            &req_id,
340            None,
341            &client_caps,
342            &server_caps,
343            protocol_version,
344            &peer,
345        );
346
347        let result = service
348            .render("greeting", Some(serde_json::json!({"name": "Alice"})), &ctx)
349            .await?;
350
351        assert!(!result.messages.is_empty());
352
353        Ok(())
354    }
355}
mcpkit_server/capability/prompts.rs

mcpkit_server/capability/
prompts.rs
