kvlar_proxy/
mcp.rs

1//! MCP JSON-RPC message parsing.
2//!
3//! Parses Model Context Protocol messages according to the JSON-RPC 2.0
4//! specification. Extracts tool call names and parameters for policy evaluation.
5
6use serde::{Deserialize, Serialize};
7
8/// A parsed MCP message — either a request or a response.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10#[serde(untagged)]
11pub enum McpMessage {
12    /// A JSON-RPC request (has a "method" field).
13    Request(McpRequest),
14    /// A JSON-RPC response (has a "result" or "error" field).
15    Response(McpResponse),
16}
17
18/// A JSON-RPC 2.0 request from an agent.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct McpRequest {
21    /// JSON-RPC version — must be "2.0".
22    pub jsonrpc: String,
23
24    /// Request ID.
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub id: Option<serde_json::Value>,
27
28    /// The method being called (e.g., "tools/call", "resources/read").
29    pub method: String,
30
31    /// Request parameters.
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub params: Option<serde_json::Value>,
34}
35
36/// A JSON-RPC 2.0 response from a tool server.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct McpResponse {
39    /// JSON-RPC version — must be "2.0".
40    pub jsonrpc: String,
41
42    /// Response ID (matches the request ID).
43    pub id: serde_json::Value,
44
45    /// The result (present on success).
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub result: Option<serde_json::Value>,
48
49    /// The error (present on failure).
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub error: Option<McpError>,
52}
53
54/// A JSON-RPC error object.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct McpError {
57    /// Error code.
58    pub code: i64,
59    /// Error message.
60    pub message: String,
61    /// Optional additional data.
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub data: Option<serde_json::Value>,
64}
65
66/// Extracted tool call parameters from an MCP `tools/call` request.
67#[derive(Debug, Clone)]
68pub struct ToolCallParams {
69    /// The name of the tool being called.
70    pub tool_name: String,
71    /// The arguments passed to the tool.
72    pub arguments: serde_json::Value,
73}
74
75impl McpRequest {
76    /// Returns true if this is a tool call request.
77    pub fn is_tool_call(&self) -> bool {
78        self.method == "tools/call"
79    }
80
81    /// Extracts tool call parameters from this request.
82    ///
83    /// Returns `None` if this is not a `tools/call` request or if the
84    /// params don't contain the expected `name` and `arguments` fields.
85    pub fn extract_tool_call(&self) -> Option<ToolCallParams> {
86        if !self.is_tool_call() {
87            return None;
88        }
89
90        let params = self.params.as_ref()?;
91        let name = params.get("name")?.as_str()?.to_string();
92        let arguments = params
93            .get("arguments")
94            .cloned()
95            .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
96
97        Some(ToolCallParams {
98            tool_name: name,
99            arguments,
100        })
101    }
102}
103
104impl McpMessage {
105    /// Parses a JSON string into an MCP message.
106    pub fn parse(json: &str) -> Result<Self, serde_json::Error> {
107        serde_json::from_str(json)
108    }
109
110    /// Returns the request if this is a request message.
111    pub fn as_request(&self) -> Option<&McpRequest> {
112        match self {
113            McpMessage::Request(req) => Some(req),
114            _ => None,
115        }
116    }
117
118    /// Returns the response if this is a response message.
119    pub fn as_response(&self) -> Option<&McpResponse> {
120        match self {
121            McpMessage::Response(resp) => Some(resp),
122            _ => None,
123        }
124    }
125
126    /// Returns true if this message is a tool call request.
127    pub fn is_tool_call(&self) -> bool {
128        self.as_request().is_some_and(|r| r.is_tool_call())
129    }
130
131    /// Serializes this message back to JSON.
132    pub fn to_json(&self) -> Result<String, serde_json::Error> {
133        serde_json::to_string(self)
134    }
135}
136
137/// Creates a JSON-RPC error response for a denied tool call.
138///
139/// Returns the denial as a tool result with `isError: true` so the LLM
140/// receives the message as conversation content (not a transport error)
141/// and can relay the policy decision to the user.
142///
143/// Includes a `_kvlar` metadata field with machine-readable error detail
144/// for programmatic consumers.
145pub fn deny_response(
146    request_id: serde_json::Value,
147    reason: &str,
148    tool_name: &str,
149    rule_id: &str,
150) -> McpResponse {
151    let code = if rule_id == "_default_deny" {
152        "POLICY_DEFAULT_DENY"
153    } else {
154        "POLICY_DENY"
155    };
156    let message = format!(
157        "[BLOCKED BY KVLAR]\n\
158         Tool: {tool_name}\n\
159         Policy rule: {rule_id}\n\
160         Reason: {reason}\n\n\
161         This action was blocked by the Kvlar security policy. \
162         Tell the user exactly what was blocked and why.",
163    );
164    McpResponse {
165        jsonrpc: "2.0".into(),
166        id: request_id,
167        result: Some(serde_json::json!({
168            "content": [{"type": "text", "text": message}],
169            "isError": true,
170            "_kvlar": {
171                "code": code,
172                "decision": "deny",
173                "rule_id": rule_id,
174                "reason": reason,
175                "tool": tool_name
176            }
177        })),
178        error: None,
179    }
180}
181
182/// Creates a JSON-RPC error response for an action requiring approval.
183///
184/// Returns the denial as a tool result with `isError: true` so the LLM
185/// receives the message as conversation content and can inform the user.
186///
187/// Includes a `_kvlar` metadata field with machine-readable error detail
188/// for programmatic consumers.
189pub fn approval_required_response(
190    request_id: serde_json::Value,
191    reason: &str,
192    tool_name: &str,
193    rule_id: &str,
194) -> McpResponse {
195    let message = format!(
196        "[KVLAR — APPROVAL REQUIRED]\n\
197         Tool: {tool_name}\n\
198         Policy rule: {rule_id}\n\
199         Reason: {reason}\n\n\
200         This action requires explicit human approval before it can proceed. \
201         Tell the user what action needs their approval and why.",
202    );
203    McpResponse {
204        jsonrpc: "2.0".into(),
205        id: request_id,
206        result: Some(serde_json::json!({
207            "content": [{"type": "text", "text": message}],
208            "isError": true,
209            "_kvlar": {
210                "code": "POLICY_APPROVAL_REQUIRED",
211                "decision": "require_approval",
212                "rule_id": rule_id,
213                "reason": reason,
214                "tool": tool_name
215            }
216        })),
217        error: None,
218    }
219}
220
221/// Creates a JSON-RPC error response for upstream errors
222/// (disconnect, timeout, communication failure).
223///
224/// Uses standard JSON-RPC error codes:
225/// - -32000: Server error (upstream disconnect/timeout)
226pub fn upstream_error_response(request_id: serde_json::Value, message: &str) -> McpResponse {
227    McpResponse {
228        jsonrpc: "2.0".into(),
229        id: request_id,
230        result: None,
231        error: Some(McpError {
232            code: -32000,
233            message: format!("[KVLAR] Upstream error: {}", message),
234            data: Some(serde_json::json!({
235                "_kvlar": {
236                    "code": "UPSTREAM_ERROR",
237                    "reason": message
238                }
239            })),
240        }),
241    }
242}
243
244/// Creates a JSON-RPC error response for malformed/unparseable messages.
245///
246/// Uses standard JSON-RPC error code -32700 (Parse error).
247pub fn parse_error_response(message: &str) -> McpResponse {
248    McpResponse {
249        jsonrpc: "2.0".into(),
250        id: serde_json::json!(null),
251        result: None,
252        error: Some(McpError {
253            code: -32700,
254            message: format!("[KVLAR] Parse error: {}", message),
255            data: Some(serde_json::json!({
256                "_kvlar": {
257                    "code": "PARSE_ERROR",
258                    "reason": message
259                }
260            })),
261        }),
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn test_parse_tool_call_request() {
271        let json = r#"{
272            "jsonrpc": "2.0",
273            "id": 1,
274            "method": "tools/call",
275            "params": {
276                "name": "bash",
277                "arguments": {
278                    "command": "ls -la"
279                }
280            }
281        }"#;
282
283        let msg = McpMessage::parse(json).unwrap();
284        assert!(msg.is_tool_call());
285
286        let req = msg.as_request().unwrap();
287        let tool_call = req.extract_tool_call().unwrap();
288        assert_eq!(tool_call.tool_name, "bash");
289        assert_eq!(
290            tool_call
291                .arguments
292                .get("command")
293                .unwrap()
294                .as_str()
295                .unwrap(),
296            "ls -la"
297        );
298    }
299
300    #[test]
301    fn test_parse_non_tool_call_request() {
302        let json = r#"{
303            "jsonrpc": "2.0",
304            "id": 2,
305            "method": "resources/read",
306            "params": {
307                "uri": "file:///tmp/test.txt"
308            }
309        }"#;
310
311        let msg = McpMessage::parse(json).unwrap();
312        assert!(!msg.is_tool_call());
313
314        let req = msg.as_request().unwrap();
315        assert!(req.extract_tool_call().is_none());
316    }
317
318    #[test]
319    fn test_parse_response() {
320        let json = r#"{
321            "jsonrpc": "2.0",
322            "id": 1,
323            "result": {
324                "content": [{"type": "text", "text": "hello"}]
325            }
326        }"#;
327
328        let msg = McpMessage::parse(json).unwrap();
329        assert!(!msg.is_tool_call());
330        assert!(msg.as_response().is_some());
331        assert!(msg.as_request().is_none());
332    }
333
334    #[test]
335    fn test_parse_error_response() {
336        let json = r#"{
337            "jsonrpc": "2.0",
338            "id": 1,
339            "error": {
340                "code": -32600,
341                "message": "Invalid request"
342            }
343        }"#;
344
345        let msg = McpMessage::parse(json).unwrap();
346        let resp = msg.as_response().unwrap();
347        assert!(resp.result.is_none());
348        assert!(resp.error.is_some());
349        assert_eq!(resp.error.as_ref().unwrap().code, -32600);
350    }
351
352    #[test]
353    fn test_deny_response() {
354        let resp = deny_response(
355            serde_json::json!(42),
356            "bash is not allowed",
357            "bash",
358            "deny-shell",
359        );
360        assert_eq!(resp.id, serde_json::json!(42));
361        assert!(resp.error.is_none());
362        let result = resp.result.unwrap();
363        assert_eq!(result["isError"], true);
364        let text = result["content"][0]["text"].as_str().unwrap();
365        assert!(text.contains("BLOCKED BY KVLAR"));
366        assert!(text.contains("bash is not allowed"));
367        assert!(text.contains("deny-shell"));
368        assert!(text.contains("Tool: bash"));
369
370        // Verify structured metadata
371        let kvlar = &result["_kvlar"];
372        assert_eq!(kvlar["code"], "POLICY_DENY");
373        assert_eq!(kvlar["decision"], "deny");
374        assert_eq!(kvlar["rule_id"], "deny-shell");
375        assert_eq!(kvlar["reason"], "bash is not allowed");
376        assert_eq!(kvlar["tool"], "bash");
377    }
378
379    #[test]
380    fn test_deny_response_default_deny() {
381        let resp = deny_response(
382            serde_json::json!(1),
383            "no matching rule",
384            "dangerous_tool",
385            "_default_deny",
386        );
387        let result = resp.result.unwrap();
388        let kvlar = &result["_kvlar"];
389        assert_eq!(kvlar["code"], "POLICY_DEFAULT_DENY");
390        assert_eq!(kvlar["rule_id"], "_default_deny");
391    }
392
393    #[test]
394    fn test_approval_required_response() {
395        let resp = approval_required_response(
396            serde_json::json!(7),
397            "email requires approval",
398            "send_email",
399            "approve-email",
400        );
401        assert!(resp.error.is_none());
402        let result = resp.result.unwrap();
403        assert_eq!(result["isError"], true);
404        let text = result["content"][0]["text"].as_str().unwrap();
405        assert!(text.contains("APPROVAL REQUIRED"));
406        assert!(text.contains("email requires approval"));
407        assert!(text.contains("approve-email"));
408
409        // Verify structured metadata
410        let kvlar = &result["_kvlar"];
411        assert_eq!(kvlar["code"], "POLICY_APPROVAL_REQUIRED");
412        assert_eq!(kvlar["decision"], "require_approval");
413        assert_eq!(kvlar["rule_id"], "approve-email");
414        assert_eq!(kvlar["reason"], "email requires approval");
415        assert_eq!(kvlar["tool"], "send_email");
416    }
417
418    #[test]
419    fn test_tool_call_no_arguments() {
420        let json = r#"{
421            "jsonrpc": "2.0",
422            "id": 1,
423            "method": "tools/call",
424            "params": {
425                "name": "list_files"
426            }
427        }"#;
428
429        let msg = McpMessage::parse(json).unwrap();
430        let req = msg.as_request().unwrap();
431        let tool_call = req.extract_tool_call().unwrap();
432        assert_eq!(tool_call.tool_name, "list_files");
433        assert!(tool_call.arguments.is_object());
434    }
435
436    #[test]
437    fn test_message_roundtrip() {
438        let json = r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"bash","arguments":{"cmd":"ls"}}}"#;
439        let msg = McpMessage::parse(json).unwrap();
440        let back = msg.to_json().unwrap();
441        let msg2 = McpMessage::parse(&back).unwrap();
442        assert!(msg2.is_tool_call());
443    }
444}
kvlar_proxy/mcp.rs

kvlar_proxy/
mcp.rs
