Skip to main content

mcp_attr/server/
builder.rs

1use std::pin::Pin;
2
3use derive_ex::Ex;
4pub use mcp_attr_macros::{prompt, resource, route, tool};
5use uri_template_ex::{Captures, UriTemplate};
6
7use crate::{
8    Result,
9    schema::{
10        CallToolRequestParams, CallToolResult, GetPromptRequestParams, GetPromptResult,
11        Implementation, ListPromptsRequestParams, ListPromptsResult,
12        ListResourceTemplatesRequestParams, ListResourcesRequestParams, ListResourcesResult,
13        ListToolsRequestParams, ListToolsResult, Prompt, ReadResourceRequestParams,
14        ReadResourceResult, Resource, ResourceTemplate, Tool,
15    },
16    server::errors::{prompt_not_found, resource_not_found, tool_not_found},
17};
18
19use super::{McpServer, RequestContext};
20
21struct CustomServer {
22    route: Route,
23    instructions: Option<String>,
24    server_info: Implementation,
25}
26impl McpServer for CustomServer {
27    fn capabilities(&self) -> crate::schema::ServerCapabilities {
28        let mut c = crate::schema::ServerCapabilities::default();
29        if !self.route.tools.is_empty() {
30            c.tools = Some(crate::schema::ServerCapabilitiesTools {
31                ..Default::default()
32            });
33        }
34        if !self.route.prompts.is_empty() {
35            c.prompts = Some(crate::schema::ServerCapabilitiesPrompts {
36                ..Default::default()
37            });
38        }
39        if !self.route.resources.is_empty() {
40            c.resources = Some(crate::schema::ServerCapabilitiesResources {
41                ..Default::default()
42            });
43        }
44        c
45    }
46    fn server_info(&self) -> Implementation {
47        self.server_info.clone()
48    }
49    fn instructions(&self) -> Option<String> {
50        self.instructions.clone()
51    }
52    async fn prompts_list(
53        &self,
54        _p: ListPromptsRequestParams,
55        _cx: &mut RequestContext,
56    ) -> Result<ListPromptsResult> {
57        let prompts: Vec<Prompt> = self
58            .route
59            .prompts
60            .iter()
61            .map(|p| p.prompt.clone())
62            .collect();
63        Ok(prompts.into())
64    }
65    async fn prompts_get(
66        &self,
67        p: GetPromptRequestParams,
68        cx: &mut RequestContext,
69    ) -> Result<GetPromptResult> {
70        for prompt in &self.route.prompts {
71            if prompt.prompt.name == p.name {
72                return (prompt.f)(&p, cx).await;
73            }
74        }
75        Err(prompt_not_found(&p.name))
76    }
77    async fn resources_list(
78        &self,
79        _p: ListResourcesRequestParams,
80        _cx: &mut RequestContext,
81    ) -> Result<ListResourcesResult> {
82        let resources: Vec<Resource> = self
83            .route
84            .resources
85            .iter()
86            .filter_map(|r| r.to_resource())
87            .collect();
88        Ok(resources.into())
89    }
90    async fn resources_templates_list(
91        &self,
92        _p: ListResourceTemplatesRequestParams,
93        _cx: &mut RequestContext,
94    ) -> Result<crate::schema::ListResourceTemplatesResult> {
95        let templates: Vec<ResourceTemplate> = self
96            .route
97            .resources
98            .iter()
99            .filter_map(|r| r.to_resource_template())
100            .collect();
101        Ok(templates.into())
102    }
103
104    async fn resources_read(
105        &self,
106        p: ReadResourceRequestParams,
107        cx: &mut RequestContext,
108    ) -> Result<ReadResourceResult> {
109        for resource in &self.route.resources {
110            if let Some(c) = resource.captures(&p.uri) {
111                return (resource.f)(&p, &c, cx).await;
112            }
113        }
114        Err(resource_not_found(&p.uri))
115    }
116    async fn tools_list(
117        &self,
118        _p: ListToolsRequestParams,
119        _cx: &mut RequestContext,
120    ) -> Result<ListToolsResult> {
121        let tools: Vec<Tool> = self.route.tools.iter().map(|t| t.tool.clone()).collect();
122        Ok(tools.into())
123    }
124    async fn tools_call(
125        &self,
126        p: CallToolRequestParams,
127        cx: &mut RequestContext,
128    ) -> Result<CallToolResult> {
129        for tool in &self.route.tools {
130            if tool.tool.name == p.name {
131                return (tool.f)(&p, cx).await;
132            }
133        }
134        Err(tool_not_found(&p.name))
135    }
136}
137
138#[derive(Ex)]
139#[derive_ex(Default)]
140#[default(Self::new())]
141pub struct McpServerBuilder {
142    route: Route,
143    instructions: Option<String>,
144    server_info: Implementation,
145}
146impl McpServerBuilder {
147    pub fn new() -> Self {
148        Self {
149            route: Route::default(),
150            instructions: None,
151            server_info: Implementation::from_compile_time_env(),
152        }
153    }
154    pub fn route(mut self, route: impl Into<Route>) -> Self {
155        self.route.extend(route);
156        self
157    }
158    pub fn instructions(mut self, instructions: &str) -> Self {
159        self.instructions = Some(instructions.to_string());
160        self
161    }
162    pub fn server_info(mut self, server_info: Implementation) -> Self {
163        self.server_info = server_info;
164        self
165    }
166    pub fn build(self) -> impl McpServer {
167        CustomServer {
168            route: self.route,
169            instructions: self.instructions,
170            server_info: self.server_info,
171        }
172    }
173}
174
175#[derive(Default)]
176pub struct Route {
177    tools: Vec<ToolDefinition>,
178    prompts: Vec<PromptDefinition>,
179    resources: Vec<ResourceDefinition>,
180}
181impl Route {
182    pub fn new() -> Self {
183        Self::default()
184    }
185    pub fn extend(&mut self, route: impl Into<Route>) {
186        let route = route.into();
187        self.tools.extend(route.tools);
188        self.prompts.extend(route.prompts);
189        self.resources.extend(route.resources);
190    }
191}
192impl<T> FromIterator<T> for Route
193where
194    T: Into<Route>,
195{
196    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
197        let mut route = Route::new();
198        for r in iter {
199            let r = r.into();
200            route.tools.extend(r.tools);
201            route.prompts.extend(r.prompts);
202            route.resources.extend(r.resources);
203        }
204        route
205    }
206}
207
208impl<T> From<Vec<T>> for Route
209where
210    T: Into<Route>,
211{
212    fn from(value: Vec<T>) -> Self {
213        value.into_iter().collect()
214    }
215}
216impl<T, const N: usize> From<[T; N]> for Route
217where
218    T: Into<Route>,
219{
220    fn from(value: [T; N]) -> Self {
221        value.into_iter().collect()
222    }
223}
224
225type PromptResultFuture<'a> =
226    Pin<Box<dyn Future<Output = Result<GetPromptResult>> + Send + Sync + 'a>>;
227
228#[doc(hidden)]
229pub struct PromptDefinition {
230    prompt: Prompt,
231    #[allow(clippy::type_complexity)]
232    f: Box<
233        dyn for<'a> Fn(&'a GetPromptRequestParams, &'a RequestContext) -> PromptResultFuture<'a>
234            + Send
235            + Sync,
236    >,
237}
238impl PromptDefinition {
239    pub fn new(
240        prompt: Prompt,
241        f: impl for<'a> Fn(&'a GetPromptRequestParams, &'a RequestContext) -> PromptResultFuture<'a>
242        + Send
243        + Sync
244        + 'static,
245    ) -> Self {
246        let f = Box::new(f);
247        Self { prompt, f }
248    }
249}
250impl From<PromptDefinition> for Route {
251    fn from(value: PromptDefinition) -> Self {
252        Route {
253            prompts: vec![value],
254            ..Default::default()
255        }
256    }
257}
258
259type ResourceResultFuture<'a> =
260    Pin<Box<dyn Future<Output = Result<ReadResourceResult>> + Send + Sync + 'a>>;
261
262#[doc(hidden)]
263pub struct ResourceDefinition {
264    uri: Option<UriTemplate>,
265    #[allow(clippy::type_complexity)]
266    f: Box<
267        dyn for<'a> Fn(
268                &'a ReadResourceRequestParams,
269                &'a Captures<'a>,
270                &'a RequestContext,
271            ) -> ResourceResultFuture<'a>
272            + Send
273            + Sync
274            + 'static,
275    >,
276    name: String,
277    description: Option<String>,
278    mime_type: Option<String>,
279}
280impl ResourceDefinition {
281    pub fn new(
282        name: &str,
283        uri: Option<&str>,
284        f: impl for<'a> Fn(
285            &'a ReadResourceRequestParams,
286            &'a Captures<'a>,
287            &'a RequestContext,
288        ) -> ResourceResultFuture<'a>
289        + Send
290        + Sync
291        + 'static,
292    ) -> Result<Self> {
293        let f = Box::new(f);
294        Ok(Self {
295            uri: uri.map(UriTemplate::new).transpose()?,
296            f,
297            name: name.to_string(),
298            description: None,
299            mime_type: None,
300        })
301    }
302    pub fn with_description(mut self, description: &str) -> Self {
303        self.description = Some(description.to_string());
304        self
305    }
306    pub fn with_mime_type(mut self, mime_type: &str) -> Self {
307        self.mime_type = Some(mime_type.to_string());
308        self
309    }
310    fn to_resource(&self) -> Option<Resource> {
311        let uri = self.uri.as_ref()?;
312        if uri.var_names().count() != 0 {
313            return None;
314        }
315        Some(Resource {
316            name: self.name.clone(),
317            description: self.description.clone(),
318            mime_type: self.mime_type.clone(),
319            uri: uri.to_string(),
320            size: None,
321            annotations: None,
322        })
323    }
324    fn to_resource_template(&self) -> Option<ResourceTemplate> {
325        let uri = self.uri.as_ref()?;
326        if uri.var_names().count() == 0 {
327            return None;
328        }
329        Some(ResourceTemplate {
330            name: self.name.clone(),
331            uri_template: uri.to_string(),
332            description: self.description.clone(),
333            mime_type: self.mime_type.clone(),
334            annotations: None,
335        })
336    }
337    fn captures<'a>(&'a self, input: &'a str) -> Option<Captures<'a>> {
338        if let Some(uri) = self.uri.as_ref() {
339            uri.captures(input)
340        } else {
341            Some(Captures::empty())
342        }
343    }
344}
345impl From<ResourceDefinition> for Route {
346    fn from(value: ResourceDefinition) -> Self {
347        Route {
348            resources: vec![value],
349            ..Default::default()
350        }
351    }
352}
353
354type ToolResultFuture<'a> =
355    Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send + Sync + 'a>>;
356
357#[doc(hidden)]
358pub struct ToolDefinition {
359    tool: Tool,
360    #[allow(clippy::type_complexity)]
361    f: Box<
362        dyn for<'a> Fn(&'a CallToolRequestParams, &'a RequestContext) -> ToolResultFuture<'a>
363            + Send
364            + Sync,
365    >,
366}
367impl ToolDefinition {
368    pub fn new(
369        tool: Tool,
370        f: impl for<'a> Fn(&'a CallToolRequestParams, &'a RequestContext) -> ToolResultFuture<'a>
371        + Send
372        + Sync
373        + 'static,
374    ) -> Self {
375        let f = Box::new(f);
376        Self { tool, f }
377    }
378}
379impl From<ToolDefinition> for Route {
380    fn from(value: ToolDefinition) -> Self {
381        Route {
382            tools: vec![value],
383            ..Default::default()
384        }
385    }
386}