mcp_attr/
schema_ext.rs

1#![allow(missing_docs)]
2
3use jsoncall::{ErrorCode, bail_public};
4use schemars::{JsonSchema, schema::Metadata, schema_for};
5use serde::Serialize;
6use serde_json::{Value, to_value};
7use url::Url;
8
9use crate::{
10    Result,
11    schema::{
12        BlobResourceContents, CallToolRequestParams, CallToolResult, CallToolResultContentItem,
13        CompleteRequestParams, CompleteRequestParamsArgument, CompleteRequestParamsRef,
14        CompleteResult, CompleteResultCompletion, EmbeddedResource, EmbeddedResourceResource,
15        GetPromptRequestParams, GetPromptResult, ImageContent, Implementation, ListPromptsResult,
16        ListResourceTemplatesResult, ListResourcesResult, ListRootsResult, ListToolsResult, Prompt,
17        PromptArgument, PromptMessage, PromptMessageContent, PromptReference,
18        ReadResourceRequestParams, ReadResourceResult, ReadResourceResultContentsItem, Resource,
19        ResourceAnnotations, ResourceReference, ResourceTemplate, ResourceTemplateAnnotations,
20        Role, Root, TextContent, TextResourceContents, Tool, ToolInputSchema,
21    },
22    utils::Base64Bytes,
23};
24use std::{
25    collections::BTreeMap,
26    path::{Path, PathBuf},
27};
28use std::{fmt::Display, str::FromStr};
29
30impl From<Vec<Prompt>> for ListPromptsResult {
31    fn from(prompts: Vec<Prompt>) -> Self {
32        ListPromptsResult {
33            prompts,
34            next_cursor: None,
35            meta: Default::default(),
36        }
37    }
38}
39impl<T: Into<PromptMessage>> From<Vec<T>> for GetPromptResult {
40    fn from(messages: Vec<T>) -> Self {
41        GetPromptResult {
42            description: None,
43            messages: messages.into_iter().map(|m| m.into()).collect(),
44            meta: Default::default(),
45        }
46    }
47}
48impl From<PromptMessage> for GetPromptResult {
49    fn from(message: PromptMessage) -> Self {
50        vec![message].into()
51    }
52}
53impl From<PromptMessageContent> for PromptMessage {
54    fn from(content: PromptMessageContent) -> Self {
55        PromptMessage {
56            content,
57            role: Role::User,
58        }
59    }
60}
61impl From<Vec<Resource>> for ListResourcesResult {
62    fn from(resources: Vec<Resource>) -> Self {
63        ListResourcesResult {
64            resources,
65            next_cursor: None,
66            meta: Default::default(),
67        }
68    }
69}
70impl From<Vec<ResourceTemplate>> for ListResourceTemplatesResult {
71    fn from(resource_templates: Vec<ResourceTemplate>) -> Self {
72        ListResourceTemplatesResult {
73            resource_templates,
74            next_cursor: None,
75            meta: Default::default(),
76        }
77    }
78}
79impl From<Vec<ReadResourceResultContentsItem>> for ReadResourceResult {
80    fn from(contents: Vec<ReadResourceResultContentsItem>) -> Self {
81        ReadResourceResult {
82            contents,
83            meta: Default::default(),
84        }
85    }
86}
87impl From<ReadResourceResultContentsItem> for ReadResourceResult {
88    fn from(content: ReadResourceResultContentsItem) -> Self {
89        ReadResourceResult {
90            contents: vec![content],
91            meta: Default::default(),
92        }
93    }
94}
95
96impl From<Vec<Tool>> for ListToolsResult {
97    fn from(tools: Vec<Tool>) -> Self {
98        ListToolsResult {
99            tools,
100            next_cursor: None,
101            meta: Default::default(),
102        }
103    }
104}
105impl<T: Into<CallToolResultContentItem>> From<Vec<T>> for CallToolResult {
106    fn from(content: Vec<T>) -> Self {
107        CallToolResult {
108            content: content.into_iter().map(|c| c.into()).collect(),
109            is_error: None,
110            meta: Default::default(),
111        }
112    }
113}
114impl From<()> for CallToolResult {
115    fn from(_: ()) -> Self {
116        Vec::<CallToolResultContentItem>::new().into()
117    }
118}
119
120impl From<CallToolResultContentItem> for CallToolResult {
121    fn from(content: CallToolResultContentItem) -> Self {
122        vec![content].into()
123    }
124}
125impl GetPromptRequestParams {
126    pub fn new(name: &str) -> Self {
127        GetPromptRequestParams {
128            name: name.to_string(),
129            arguments: BTreeMap::new(),
130        }
131    }
132    pub fn with_arguments<K, V>(mut self, arguments: impl IntoIterator<Item = (K, V)>) -> Self
133    where
134        K: Display,
135        V: Display,
136    {
137        self.arguments = arguments
138            .into_iter()
139            .map(|(k, v)| (k.to_string(), v.to_string()))
140            .collect();
141        self
142    }
143}
144
145impl Prompt {
146    pub fn new(name: &str) -> Self {
147        Prompt {
148            name: name.to_string(),
149            arguments: vec![],
150            description: None,
151        }
152    }
153    pub fn with_description(mut self, description: &str) -> Self {
154        self.description = Some(description.to_string());
155        self
156    }
157    pub fn with_arguments(mut self, arguments: Vec<PromptArgument>) -> Self {
158        self.arguments = arguments;
159        self
160    }
161}
162impl PromptArgument {
163    pub fn new(name: &str, required: bool) -> Self {
164        PromptArgument {
165            name: name.to_string(),
166            description: None,
167            required: Some(required),
168        }
169    }
170    pub fn with_description(mut self, description: &str) -> Self {
171        self.description = Some(description.to_string());
172        self
173    }
174    pub fn with_required(mut self, required: bool) -> Self {
175        self.required = Some(required);
176        self
177    }
178}
179
180impl Resource {
181    pub fn new(uri: &str, name: &str) -> Self {
182        Resource {
183            uri: uri.to_string(),
184            name: name.to_string(),
185            description: None,
186            mime_type: None,
187            annotations: None,
188            size: None,
189        }
190    }
191    pub fn with_description(mut self, description: &str) -> Self {
192        self.description = Some(description.to_string());
193        self
194    }
195    pub fn with_mime_type(mut self, mime_type: &str) -> Self {
196        self.mime_type = Some(mime_type.to_string());
197        self
198    }
199    pub fn with_annotations(mut self, annotations: impl Into<ResourceAnnotations>) -> Self {
200        self.annotations = Some(annotations.into());
201        self
202    }
203    pub fn with_size(mut self, size: i64) -> Self {
204        self.size = Some(size);
205        self
206    }
207}
208impl ResourceTemplate {
209    pub fn new(uri_template: &str, name: &str) -> Self {
210        ResourceTemplate {
211            uri_template: uri_template.to_string(),
212            name: name.to_string(),
213            annotations: None,
214            description: None,
215            mime_type: None,
216        }
217    }
218    pub fn with_description(mut self, description: &str) -> Self {
219        self.description = Some(description.to_string());
220        self
221    }
222    pub fn with_mime_type(mut self, mime_type: &str) -> Self {
223        self.mime_type = Some(mime_type.to_string());
224        self
225    }
226    pub fn with_annotations(mut self, annotations: impl Into<ResourceTemplateAnnotations>) -> Self {
227        self.annotations = Some(annotations.into());
228        self
229    }
230}
231impl ReadResourceRequestParams {
232    pub fn new(uri: &str) -> Self {
233        ReadResourceRequestParams {
234            uri: uri.to_string(),
235        }
236    }
237}
238
239impl Tool {
240    pub fn new(name: &str, input_schema: ToolInputSchema) -> Self {
241        Tool {
242            name: name.to_string(),
243            description: None,
244            input_schema,
245        }
246    }
247    pub fn with_description(mut self, description: &str) -> Self {
248        self.description = Some(description.to_string());
249        self
250    }
251}
252
253impl ToolInputSchema {
254    pub fn new() -> Self {
255        Self {
256            properties: BTreeMap::new(),
257            required: vec![],
258            type_: "object".to_string(),
259        }
260    }
261    pub fn insert_property<T: JsonSchema>(
262        &mut self,
263        name: &str,
264        description: &str,
265        required: bool,
266    ) -> Result<()> {
267        let mut root = schema_for!(T);
268        if !description.is_empty() {
269            let metadata = root
270                .schema
271                .metadata
272                .get_or_insert(Box::new(Metadata::default()));
273            metadata.description = Some(description.to_string());
274        }
275        let value = to_value(root.schema)?;
276        let Value::Object(obj) = value else {
277            bail_public!(
278                ErrorCode::INVALID_PARAMS,
279                "schema for `{name}` is not an object"
280            );
281        };
282        self.properties.insert(name.to_string(), obj);
283        if required {
284            self.required.push(name.to_string());
285        }
286        Ok(())
287    }
288    pub fn with_property<T: JsonSchema>(
289        mut self,
290        name: &str,
291        description: &str,
292        required: bool,
293    ) -> Result<Self> {
294        self.insert_property::<T>(name, description, required)?;
295        Ok(self)
296    }
297}
298impl Default for ToolInputSchema {
299    fn default() -> Self {
300        Self::new()
301    }
302}
303impl CallToolRequestParams {
304    pub fn new(name: &str) -> Self {
305        CallToolRequestParams {
306            name: name.to_string(),
307            arguments: None,
308        }
309    }
310    pub fn with_argument(mut self, name: &str, value: impl Serialize) -> Result<Self> {
311        let mut arguments = self.arguments.unwrap_or_default();
312        arguments.insert(name.to_string(), to_value(value)?);
313        self.arguments = Some(arguments);
314        Ok(self)
315    }
316}
317
318impl TextContent {
319    pub fn new(text: impl std::fmt::Display) -> Self {
320        Self {
321            text: text.to_string(),
322            annotations: None,
323            type_: "text".to_string(),
324        }
325    }
326}
327impl From<String> for TextContent {
328    fn from(text: String) -> Self {
329        Self::new(text)
330    }
331}
332impl From<&str> for TextContent {
333    fn from(text: &str) -> Self {
334        Self::new(text)
335    }
336}
337
338impl ImageContent {
339    pub fn new(data: Base64Bytes, mime_type: &str) -> Self {
340        Self {
341            data,
342            mime_type: mime_type.to_string(),
343            annotations: None,
344            type_: "image".to_string(),
345        }
346    }
347}
348
349impl EmbeddedResource {
350    pub fn new(resource: impl Into<EmbeddedResourceResource>) -> Self {
351        Self {
352            annotations: None,
353            resource: resource.into(),
354            type_: "resource".to_string(),
355        }
356    }
357}
358
359impl From<String> for TextResourceContents {
360    fn from(text: String) -> Self {
361        TextResourceContents {
362            text,
363            ..Default::default()
364        }
365    }
366}
367impl From<&str> for TextResourceContents {
368    fn from(text: &str) -> Self {
369        text.to_string().into()
370    }
371}
372
373impl From<Base64Bytes> for BlobResourceContents {
374    fn from(blob: Base64Bytes) -> Self {
375        BlobResourceContents {
376            blob,
377            ..Default::default()
378        }
379    }
380}
381
382impl Implementation {
383    pub fn new(name: &str, version: &str) -> Self {
384        Self {
385            name: name.to_string(),
386            version: version.to_string(),
387        }
388    }
389    pub fn from_compile_time_env() -> Self {
390        Self::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"))
391    }
392}
393
394impl Root {
395    pub fn new(uri: &str) -> Self {
396        Self {
397            uri: uri.to_string(),
398            name: None,
399        }
400    }
401    pub fn with_name(mut self, name: impl Display) -> Self {
402        self.name = Some(name.to_string());
403        self
404    }
405
406    pub fn from_file_path(path: impl AsRef<Path>) -> Option<Self> {
407        Some(Self::new(Url::from_file_path(path).ok()?.as_str()))
408    }
409    pub fn to_file_path(&self) -> Option<PathBuf> {
410        Url::from_str(&self.uri).ok()?.to_file_path().ok()
411    }
412}
413
414impl From<Vec<Root>> for ListRootsResult {
415    fn from(roots: Vec<Root>) -> Self {
416        ListRootsResult {
417            roots,
418            meta: Default::default(),
419        }
420    }
421}
422
423impl From<CompleteResultCompletion> for CompleteResult {
424    fn from(completion: CompleteResultCompletion) -> Self {
425        Self {
426            completion,
427            meta: Default::default(),
428        }
429    }
430}
431impl CompleteResultCompletion {
432    pub const MAX_VALUES: usize = 100;
433}
434
435impl From<Vec<String>> for CompleteResultCompletion {
436    fn from(mut values: Vec<String>) -> Self {
437        let total = Some(values.len() as i64);
438        let has_more = if values.len() > Self::MAX_VALUES {
439            values.truncate(Self::MAX_VALUES);
440            Some(true)
441        } else {
442            None
443        };
444        Self {
445            has_more,
446            total,
447            values,
448        }
449    }
450}
451impl From<&[&str]> for CompleteResultCompletion {
452    fn from(values: &[&str]) -> Self {
453        values
454            .iter()
455            .map(|s| s.to_string())
456            .collect::<Vec<String>>()
457            .into()
458    }
459}
460
461impl CompleteRequestParams {
462    pub fn new(r: CompleteRequestParamsRef, argument: CompleteRequestParamsArgument) -> Self {
463        Self { argument, ref_: r }
464    }
465}
466impl CompleteRequestParamsArgument {
467    pub fn new(name: &str, value: &str) -> Self {
468        Self {
469            name: name.to_string(),
470            value: value.to_string(),
471        }
472    }
473}
474
475impl CompleteRequestParamsRef {
476    pub fn new_prompt(name: &str) -> Self {
477        CompleteRequestParamsRef::PromptReference(PromptReference::new(name))
478    }
479    pub fn new_resource(uri: &str) -> Self {
480        CompleteRequestParamsRef::ResourceReference(ResourceReference::new(uri))
481    }
482}
483impl PromptReference {
484    pub fn new(name: &str) -> Self {
485        Self {
486            name: name.to_string(),
487            type_: "ref/prompt".to_string(),
488        }
489    }
490}
491impl ResourceReference {
492    pub fn new(uri: &str) -> Self {
493        Self {
494            uri: uri.to_string(),
495            type_: "ref/resource".to_string(),
496        }
497    }
498}