context_mcp/
protocol.rs

1//! MCP protocol types and message handling
2//!
3//! Implements the Model Context Protocol specification for
4//! communication between AI assistants and context servers.
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::collections::HashMap;
9
10/// JSON-RPC version constant
11pub const JSONRPC_VERSION: &str = "2.0";
12
13/// MCP protocol version
14pub const MCP_VERSION: &str = "2024-11-05";
15
16/// JSON-RPC request
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct JsonRpcRequest {
19    pub jsonrpc: String,
20    pub id: RequestId,
21    pub method: String,
22    #[serde(default, skip_serializing_if = "Option::is_none")]
23    pub params: Option<Value>,
24}
25
26impl JsonRpcRequest {
27    /// Create a new request
28    pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
29        Self {
30            jsonrpc: JSONRPC_VERSION.to_string(),
31            id: RequestId::Number(rand::random()),
32            method: method.into(),
33            params,
34        }
35    }
36}
37
38/// JSON-RPC response
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct JsonRpcResponse {
41    pub jsonrpc: String,
42    pub id: RequestId,
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub result: Option<Value>,
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub error: Option<JsonRpcError>,
47}
48
49impl JsonRpcResponse {
50    /// Create a success response
51    pub fn success(id: RequestId, result: Value) -> Self {
52        Self {
53            jsonrpc: JSONRPC_VERSION.to_string(),
54            id,
55            result: Some(result),
56            error: None,
57        }
58    }
59
60    /// Create an error response
61    pub fn error(id: RequestId, error: JsonRpcError) -> Self {
62        Self {
63            jsonrpc: JSONRPC_VERSION.to_string(),
64            id,
65            result: None,
66            error: Some(error),
67        }
68    }
69}
70
71/// JSON-RPC request ID
72#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
73#[serde(untagged)]
74pub enum RequestId {
75    Number(i64),
76    String(String),
77}
78
79/// JSON-RPC error
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct JsonRpcError {
82    pub code: i32,
83    pub message: String,
84    #[serde(skip_serializing_if = "Option::is_none")]
85    pub data: Option<Value>,
86}
87
88/// Standard JSON-RPC error codes
89pub mod error_codes {
90    pub const PARSE_ERROR: i32 = -32700;
91    pub const INVALID_REQUEST: i32 = -32600;
92    pub const METHOD_NOT_FOUND: i32 = -32601;
93    pub const INVALID_PARAMS: i32 = -32602;
94    pub const INTERNAL_ERROR: i32 = -32603;
95}
96
97impl JsonRpcError {
98    pub fn parse_error() -> Self {
99        Self {
100            code: error_codes::PARSE_ERROR,
101            message: "Parse error".to_string(),
102            data: None,
103        }
104    }
105
106    pub fn invalid_request(msg: impl Into<String>) -> Self {
107        Self {
108            code: error_codes::INVALID_REQUEST,
109            message: msg.into(),
110            data: None,
111        }
112    }
113
114    pub fn method_not_found(method: &str) -> Self {
115        Self {
116            code: error_codes::METHOD_NOT_FOUND,
117            message: format!("Method not found: {}", method),
118            data: None,
119        }
120    }
121
122    pub fn invalid_params(msg: impl Into<String>) -> Self {
123        Self {
124            code: error_codes::INVALID_PARAMS,
125            message: msg.into(),
126            data: None,
127        }
128    }
129
130    pub fn internal_error(msg: impl Into<String>) -> Self {
131        Self {
132            code: error_codes::INTERNAL_ERROR,
133            message: msg.into(),
134            data: None,
135        }
136    }
137}
138
139/// MCP server capabilities
140#[derive(Debug, Clone, Default, Serialize, Deserialize)]
141pub struct ServerCapabilities {
142    #[serde(skip_serializing_if = "Option::is_none")]
143    pub tools: Option<ToolsCapability>,
144    #[serde(skip_serializing_if = "Option::is_none")]
145    pub resources: Option<ResourcesCapability>,
146    #[serde(skip_serializing_if = "Option::is_none")]
147    pub prompts: Option<PromptsCapability>,
148}
149
150#[derive(Debug, Clone, Default, Serialize, Deserialize)]
151pub struct ToolsCapability {
152    #[serde(default)]
153    pub list_changed: bool,
154}
155
156#[derive(Debug, Clone, Default, Serialize, Deserialize)]
157pub struct ResourcesCapability {
158    #[serde(default)]
159    pub subscribe: bool,
160    #[serde(default)]
161    pub list_changed: bool,
162}
163
164#[derive(Debug, Clone, Default, Serialize, Deserialize)]
165pub struct PromptsCapability {
166    #[serde(default)]
167    pub list_changed: bool,
168}
169
170/// MCP server info
171#[derive(Debug, Clone, Serialize, Deserialize)]
172pub struct ServerInfo {
173    pub name: String,
174    pub version: String,
175}
176
177/// MCP initialize result
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct InitializeResult {
180    pub protocol_version: String,
181    pub capabilities: ServerCapabilities,
182    pub server_info: ServerInfo,
183}
184
185/// MCP tool definition
186#[derive(Debug, Clone, Serialize, Deserialize)]
187pub struct Tool {
188    pub name: String,
189    #[serde(skip_serializing_if = "Option::is_none")]
190    pub description: Option<String>,
191    pub input_schema: InputSchema,
192}
193
194/// JSON Schema for tool input
195#[derive(Debug, Clone, Serialize, Deserialize)]
196pub struct InputSchema {
197    #[serde(rename = "type")]
198    pub schema_type: String,
199    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
200    pub properties: HashMap<String, PropertySchema>,
201    #[serde(default, skip_serializing_if = "Vec::is_empty")]
202    pub required: Vec<String>,
203}
204
205impl InputSchema {
206    /// Create an object schema
207    pub fn object() -> Self {
208        Self {
209            schema_type: "object".to_string(),
210            properties: HashMap::new(),
211            required: Vec::new(),
212        }
213    }
214
215    /// Add a property
216    pub fn with_property(mut self, name: impl Into<String>, schema: PropertySchema) -> Self {
217        self.properties.insert(name.into(), schema);
218        self
219    }
220
221    /// Add a required property
222    pub fn with_required(mut self, name: impl Into<String>, schema: PropertySchema) -> Self {
223        let name = name.into();
224        self.required.push(name.clone());
225        self.properties.insert(name, schema);
226        self
227    }
228}
229
230/// Property schema definition
231#[derive(Debug, Clone, Serialize, Deserialize)]
232pub struct PropertySchema {
233    #[serde(rename = "type")]
234    pub schema_type: String,
235    #[serde(skip_serializing_if = "Option::is_none")]
236    pub description: Option<String>,
237    #[serde(skip_serializing_if = "Option::is_none")]
238    pub default: Option<Value>,
239    #[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
240    pub enum_values: Option<Vec<String>>,
241}
242
243impl PropertySchema {
244    pub fn string(description: impl Into<String>) -> Self {
245        Self {
246            schema_type: "string".to_string(),
247            description: Some(description.into()),
248            default: None,
249            enum_values: None,
250        }
251    }
252
253    pub fn number(description: impl Into<String>) -> Self {
254        Self {
255            schema_type: "number".to_string(),
256            description: Some(description.into()),
257            default: None,
258            enum_values: None,
259        }
260    }
261
262    pub fn boolean(description: impl Into<String>) -> Self {
263        Self {
264            schema_type: "boolean".to_string(),
265            description: Some(description.into()),
266            default: None,
267            enum_values: None,
268        }
269    }
270
271    pub fn array(description: impl Into<String>) -> Self {
272        Self {
273            schema_type: "array".to_string(),
274            description: Some(description.into()),
275            default: None,
276            enum_values: None,
277        }
278    }
279
280    pub fn with_default(mut self, value: Value) -> Self {
281        self.default = Some(value);
282        self
283    }
284
285    pub fn with_enum(mut self, values: Vec<&str>) -> Self {
286        self.enum_values = Some(values.into_iter().map(|s| s.to_string()).collect());
287        self
288    }
289}
290
291/// MCP tool call request
292#[derive(Debug, Clone, Serialize, Deserialize)]
293pub struct CallToolRequest {
294    pub name: String,
295    #[serde(default)]
296    pub arguments: HashMap<String, Value>,
297}
298
299/// MCP tool call result
300#[derive(Debug, Clone, Serialize, Deserialize)]
301pub struct CallToolResult {
302    pub content: Vec<Content>,
303    #[serde(default)]
304    pub is_error: bool,
305}
306
307impl CallToolResult {
308    /// Create a text result
309    pub fn text(text: impl Into<String>) -> Self {
310        Self {
311            content: vec![Content::text(text)],
312            is_error: false,
313        }
314    }
315
316    /// Create an error result
317    pub fn error(message: impl Into<String>) -> Self {
318        Self {
319            content: vec![Content::text(message)],
320            is_error: true,
321        }
322    }
323
324    /// Create a JSON result
325    pub fn json(value: Value) -> Self {
326        Self {
327            content: vec![Content::text(
328                serde_json::to_string_pretty(&value).unwrap_or_default(),
329            )],
330            is_error: false,
331        }
332    }
333}
334
335/// Content item in tool result
336#[derive(Debug, Clone, Serialize, Deserialize)]
337#[serde(tag = "type")]
338pub enum Content {
339    #[serde(rename = "text")]
340    Text { text: String },
341    #[serde(rename = "image")]
342    Image { data: String, mime_type: String },
343    #[serde(rename = "resource")]
344    Resource { resource: ResourceContent },
345}
346
347impl Content {
348    pub fn text(text: impl Into<String>) -> Self {
349        Self::Text { text: text.into() }
350    }
351
352    pub fn image(data: String, mime_type: impl Into<String>) -> Self {
353        Self::Image {
354            data,
355            mime_type: mime_type.into(),
356        }
357    }
358}
359
360/// Resource content
361#[derive(Debug, Clone, Serialize, Deserialize)]
362pub struct ResourceContent {
363    pub uri: String,
364    pub mime_type: Option<String>,
365    pub text: Option<String>,
366    pub blob: Option<String>,
367}
368
369/// MCP resource definition
370#[derive(Debug, Clone, Serialize, Deserialize)]
371pub struct Resource {
372    pub uri: String,
373    pub name: String,
374    #[serde(skip_serializing_if = "Option::is_none")]
375    pub description: Option<String>,
376    #[serde(skip_serializing_if = "Option::is_none")]
377    pub mime_type: Option<String>,
378}
379
380/// MCP prompt definition
381#[derive(Debug, Clone, Serialize, Deserialize)]
382pub struct Prompt {
383    pub name: String,
384    #[serde(skip_serializing_if = "Option::is_none")]
385    pub description: Option<String>,
386    #[serde(default, skip_serializing_if = "Vec::is_empty")]
387    pub arguments: Vec<PromptArgument>,
388}
389
390/// Prompt argument definition
391#[derive(Debug, Clone, Serialize, Deserialize)]
392pub struct PromptArgument {
393    pub name: String,
394    #[serde(skip_serializing_if = "Option::is_none")]
395    pub description: Option<String>,
396    #[serde(default)]
397    pub required: bool,
398}
399
400/// MCP notification
401#[derive(Debug, Clone, Serialize, Deserialize)]
402pub struct Notification {
403    pub jsonrpc: String,
404    pub method: String,
405    #[serde(skip_serializing_if = "Option::is_none")]
406    pub params: Option<Value>,
407}
408
409impl Notification {
410    /// Create a new notification
411    pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
412        Self {
413            jsonrpc: JSONRPC_VERSION.to_string(),
414            method: method.into(),
415            params,
416        }
417    }
418
419    /// Tools list changed notification
420    pub fn tools_list_changed() -> Self {
421        Self::new("notifications/tools/list_changed", None)
422    }
423
424    /// Resources list changed notification
425    pub fn resources_list_changed() -> Self {
426        Self::new("notifications/resources/list_changed", None)
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use super::*;
433
434    #[test]
435    fn test_json_rpc_request() {
436        let req = JsonRpcRequest::new("test_method", None);
437        assert_eq!(req.jsonrpc, JSONRPC_VERSION);
438        assert_eq!(req.method, "test_method");
439    }
440
441    #[test]
442    fn test_input_schema() {
443        let schema = InputSchema::object()
444            .with_required("content", PropertySchema::string("The content"))
445            .with_property(
446                "domain",
447                PropertySchema::string("Domain").with_enum(vec!["Code", "Docs"]),
448            );
449
450        assert_eq!(schema.schema_type, "object");
451        assert!(schema.required.contains(&"content".to_string()));
452        assert!(schema.properties.contains_key("domain"));
453    }
454
455    #[test]
456    fn test_tool_result() {
457        let result = CallToolResult::text("Success");
458        assert!(!result.is_error);
459        assert_eq!(result.content.len(), 1);
460    }
461}
context_mcp/protocol.rs

context_mcp/
protocol.rs
