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(serde_json::to_string_pretty(&value).unwrap_or_default())],
328            is_error: false,
329        }
330    }
331}
332
333/// Content item in tool result
334#[derive(Debug, Clone, Serialize, Deserialize)]
335#[serde(tag = "type")]
336pub enum Content {
337    #[serde(rename = "text")]
338    Text { text: String },
339    #[serde(rename = "image")]
340    Image { data: String, mime_type: String },
341    #[serde(rename = "resource")]
342    Resource { resource: ResourceContent },
343}
344
345impl Content {
346    pub fn text(text: impl Into<String>) -> Self {
347        Self::Text { text: text.into() }
348    }
349
350    pub fn image(data: String, mime_type: impl Into<String>) -> Self {
351        Self::Image {
352            data,
353            mime_type: mime_type.into(),
354        }
355    }
356}
357
358/// Resource content
359#[derive(Debug, Clone, Serialize, Deserialize)]
360pub struct ResourceContent {
361    pub uri: String,
362    pub mime_type: Option<String>,
363    pub text: Option<String>,
364    pub blob: Option<String>,
365}
366
367/// MCP resource definition
368#[derive(Debug, Clone, Serialize, Deserialize)]
369pub struct Resource {
370    pub uri: String,
371    pub name: String,
372    #[serde(skip_serializing_if = "Option::is_none")]
373    pub description: Option<String>,
374    #[serde(skip_serializing_if = "Option::is_none")]
375    pub mime_type: Option<String>,
376}
377
378/// MCP prompt definition
379#[derive(Debug, Clone, Serialize, Deserialize)]
380pub struct Prompt {
381    pub name: String,
382    #[serde(skip_serializing_if = "Option::is_none")]
383    pub description: Option<String>,
384    #[serde(default, skip_serializing_if = "Vec::is_empty")]
385    pub arguments: Vec<PromptArgument>,
386}
387
388/// Prompt argument definition
389#[derive(Debug, Clone, Serialize, Deserialize)]
390pub struct PromptArgument {
391    pub name: String,
392    #[serde(skip_serializing_if = "Option::is_none")]
393    pub description: Option<String>,
394    #[serde(default)]
395    pub required: bool,
396}
397
398/// MCP notification
399#[derive(Debug, Clone, Serialize, Deserialize)]
400pub struct Notification {
401    pub jsonrpc: String,
402    pub method: String,
403    #[serde(skip_serializing_if = "Option::is_none")]
404    pub params: Option<Value>,
405}
406
407impl Notification {
408    /// Create a new notification
409    pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
410        Self {
411            jsonrpc: JSONRPC_VERSION.to_string(),
412            method: method.into(),
413            params,
414        }
415    }
416
417    /// Tools list changed notification
418    pub fn tools_list_changed() -> Self {
419        Self::new("notifications/tools/list_changed", None)
420    }
421
422    /// Resources list changed notification
423    pub fn resources_list_changed() -> Self {
424        Self::new("notifications/resources/list_changed", None)
425    }
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431
432    #[test]
433    fn test_json_rpc_request() {
434        let req = JsonRpcRequest::new("test_method", None);
435        assert_eq!(req.jsonrpc, JSONRPC_VERSION);
436        assert_eq!(req.method, "test_method");
437    }
438
439    #[test]
440    fn test_input_schema() {
441        let schema = InputSchema::object()
442            .with_required("content", PropertySchema::string("The content"))
443            .with_property("domain", PropertySchema::string("Domain").with_enum(vec!["Code", "Docs"]));
444
445        assert_eq!(schema.schema_type, "object");
446        assert!(schema.required.contains(&"content".to_string()));
447        assert!(schema.properties.contains_key("domain"));
448    }
449
450    #[test]
451    fn test_tool_result() {
452        let result = CallToolResult::text("Success");
453        assert!(!result.is_error);
454        assert_eq!(result.content.len(), 1);
455    }
456}