mcp_attr/
schema_ext.rs

1#![allow(missing_docs)]
2
3use jsoncall::{ErrorCode, bail_public};
4use schemars::{JsonSchema, r#gen::SchemaSettings, schema::Metadata};
5use serde::Serialize;
6use serde_json::{Value, to_value};
7use url::Url;
8
9use crate::{
10    Result,
11    schema::{
12        Annotations, BlobResourceContents, CallToolRequestParams, CallToolResult,
13        CallToolResultContentItem, CompleteRequestParams, CompleteRequestParamsArgument,
14        CompleteRequestParamsRef, CompleteResult, CompleteResultCompletion, EmbeddedResource,
15        EmbeddedResourceResource, GetPromptRequestParams, GetPromptResult, ImageContent,
16        Implementation, ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult,
17        ListRootsResult, ListToolsResult, Prompt, PromptArgument, PromptMessage,
18        PromptMessageContent, PromptReference, ReadResourceRequestParams, ReadResourceResult,
19        ReadResourceResultContentsItem, Resource, ResourceReference, ResourceTemplate, Role, Root,
20        TextContent, TextResourceContents, Tool, ToolAnnotations, ToolInputSchema,
21    },
22    utils::{Base64Bytes, Json},
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<Annotations>) -> 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<Annotations>) -> 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            annotations: None,
246        }
247    }
248    pub fn with_description(mut self, description: &str) -> Self {
249        self.description = Some(description.to_string());
250        self
251    }
252    pub fn with_annotations(mut self, annotations: impl Into<ToolAnnotations>) -> Self {
253        self.annotations = Some(annotations.into());
254        self
255    }
256}
257
258impl ToolInputSchema {
259    pub fn new() -> Self {
260        Self {
261            properties: BTreeMap::new(),
262            required: vec![],
263            type_: "object".to_string(),
264        }
265    }
266    pub fn insert_property<T: JsonSchema>(
267        &mut self,
268        name: &str,
269        description: &str,
270        required: bool,
271    ) -> Result<()> {
272        let mut settings = SchemaSettings::default();
273        settings.inline_subschemas = true;
274        let g = settings.into_generator();
275        let mut root = g.into_root_schema_for::<T>();
276
277        if !description.is_empty() {
278            let metadata = root
279                .schema
280                .metadata
281                .get_or_insert(Box::new(Metadata::default()));
282            metadata.description = Some(description.to_string());
283        }
284        let value = to_value(root.schema)?;
285        let Value::Object(obj) = value else {
286            bail_public!(
287                ErrorCode::INVALID_PARAMS,
288                "schema for `{name}` is not an object"
289            );
290        };
291        self.properties.insert(name.to_string(), obj);
292        if required {
293            self.required.push(name.to_string());
294        }
295        Ok(())
296    }
297    pub fn with_property<T: JsonSchema>(
298        mut self,
299        name: &str,
300        description: &str,
301        required: bool,
302    ) -> Result<Self> {
303        self.insert_property::<T>(name, description, required)?;
304        Ok(self)
305    }
306}
307impl Default for ToolInputSchema {
308    fn default() -> Self {
309        Self::new()
310    }
311}
312impl CallToolRequestParams {
313    pub fn new(name: &str) -> Self {
314        CallToolRequestParams {
315            name: name.to_string(),
316            arguments: None,
317        }
318    }
319    pub fn with_argument(mut self, name: &str, value: impl Serialize) -> Result<Self> {
320        let mut arguments = self.arguments.unwrap_or_default();
321        arguments.insert(name.to_string(), to_value(value)?);
322        self.arguments = Some(arguments);
323        Ok(self)
324    }
325}
326
327impl TextContent {
328    pub fn new(text: impl std::fmt::Display) -> Self {
329        Self {
330            text: text.to_string(),
331            annotations: None,
332            type_: "text".to_string(),
333        }
334    }
335}
336impl From<String> for TextContent {
337    fn from(text: String) -> Self {
338        Self::new(text)
339    }
340}
341impl From<&str> for TextContent {
342    fn from(text: &str) -> Self {
343        Self::new(text)
344    }
345}
346
347impl ImageContent {
348    pub fn new(data: Base64Bytes, mime_type: &str) -> Self {
349        Self {
350            data,
351            mime_type: mime_type.to_string(),
352            annotations: None,
353            type_: "image".to_string(),
354        }
355    }
356}
357
358impl EmbeddedResource {
359    pub fn new(resource: impl Into<EmbeddedResourceResource>) -> Self {
360        Self {
361            annotations: None,
362            resource: resource.into(),
363            type_: "resource".to_string(),
364        }
365    }
366}
367
368impl From<String> for TextResourceContents {
369    fn from(text: String) -> Self {
370        TextResourceContents {
371            text,
372            ..Default::default()
373        }
374    }
375}
376impl From<&str> for TextResourceContents {
377    fn from(text: &str) -> Self {
378        text.to_string().into()
379    }
380}
381
382impl From<Base64Bytes> for BlobResourceContents {
383    fn from(blob: Base64Bytes) -> Self {
384        BlobResourceContents {
385            blob,
386            ..Default::default()
387        }
388    }
389}
390
391impl Implementation {
392    pub fn new(name: &str, version: &str) -> Self {
393        Self {
394            name: name.to_string(),
395            version: version.to_string(),
396        }
397    }
398    pub fn from_compile_time_env() -> Self {
399        Self::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"))
400    }
401}
402
403impl Root {
404    pub fn new(uri: &str) -> Self {
405        Self {
406            uri: uri.to_string(),
407            name: None,
408        }
409    }
410    pub fn with_name(mut self, name: impl Display) -> Self {
411        self.name = Some(name.to_string());
412        self
413    }
414
415    pub fn from_file_path(path: impl AsRef<Path>) -> Option<Self> {
416        Some(Self::new(Url::from_file_path(path).ok()?.as_str()))
417    }
418    pub fn to_file_path(&self) -> Option<PathBuf> {
419        Url::from_str(&self.uri).ok()?.to_file_path().ok()
420    }
421}
422
423impl From<Vec<Root>> for ListRootsResult {
424    fn from(roots: Vec<Root>) -> Self {
425        ListRootsResult {
426            roots,
427            meta: Default::default(),
428        }
429    }
430}
431
432impl From<CompleteResultCompletion> for CompleteResult {
433    fn from(completion: CompleteResultCompletion) -> Self {
434        Self {
435            completion,
436            meta: Default::default(),
437        }
438    }
439}
440impl CompleteResultCompletion {
441    pub const MAX_VALUES: usize = 100;
442}
443
444impl From<Vec<String>> for CompleteResultCompletion {
445    fn from(mut values: Vec<String>) -> Self {
446        let total = Some(values.len() as i64);
447        let has_more = if values.len() > Self::MAX_VALUES {
448            values.truncate(Self::MAX_VALUES);
449            Some(true)
450        } else {
451            None
452        };
453        Self {
454            has_more,
455            total,
456            values,
457        }
458    }
459}
460impl From<&[&str]> for CompleteResultCompletion {
461    fn from(values: &[&str]) -> Self {
462        values
463            .iter()
464            .map(|s| s.to_string())
465            .collect::<Vec<String>>()
466            .into()
467    }
468}
469
470impl CompleteRequestParams {
471    pub fn new(r: CompleteRequestParamsRef, argument: CompleteRequestParamsArgument) -> Self {
472        Self { argument, ref_: r }
473    }
474}
475impl CompleteRequestParamsArgument {
476    pub fn new(name: &str, value: &str) -> Self {
477        Self {
478            name: name.to_string(),
479            value: value.to_string(),
480        }
481    }
482}
483
484impl CompleteRequestParamsRef {
485    pub fn new_prompt(name: &str) -> Self {
486        CompleteRequestParamsRef::PromptReference(PromptReference::new(name))
487    }
488    pub fn new_resource(uri: &str) -> Self {
489        CompleteRequestParamsRef::ResourceReference(ResourceReference::new(uri))
490    }
491}
492impl PromptReference {
493    pub fn new(name: &str) -> Self {
494        Self {
495            name: name.to_string(),
496            type_: "ref/prompt".to_string(),
497        }
498    }
499}
500impl ResourceReference {
501    pub fn new(uri: &str) -> Self {
502        Self {
503            uri: uri.to_string(),
504            type_: "ref/resource".to_string(),
505        }
506    }
507}
508
509impl<T> From<Json<T>> for TextContent {
510    fn from(value: Json<T>) -> Self {
511        TextContent::new(value.into_string())
512    }
513}
514impl<T> From<Json<T>> for CallToolResult {
515    fn from(value: Json<T>) -> Self {
516        CallToolResultContentItem::from(value).into()
517    }
518}
519impl<T> From<Json<T>> for CallToolResultContentItem {
520    fn from(value: Json<T>) -> Self {
521        TextContent::from(value).into()
522    }
523}