Skip to main content

sim_codec_mcp/
canonical.rs

1//! The `codec:mcp` decoder/encoder and its host-registered lib. Decodes one MCP
2//! JSON-RPC envelope per frame into an envelope `Expr` and encodes envelopes
3//! back to JSON-RPC text, validating the envelope on both sides.
4
5use std::{str::FromStr, sync::Arc};
6
7use serde_json::{Map as JsonMap, Number as JsonNumber, Value as JsonValue};
8use sim_codec::{DecodeBudget, Decoder, DomainCodecLib, Encoder, Input, Output, ReadCx};
9use sim_kernel::{
10    CodecId, Error, Expr, Lib, LibManifest, Linker, LoadCx, NumberLiteral, Result, Symbol, WriteCx,
11};
12
13use crate::{
14    envelope::{
15        McpEnvelope, McpError, McpErrorEnvelope, McpNotification, McpRequest, McpResponse,
16        is_jsonrpc_id,
17    },
18    error::codec_error,
19    expr::{envelope_to_expr, expr_to_envelope},
20    wire_keys::reject_duplicate_mcp_wire_keys,
21};
22
23const JSONRPC_VERSION: &str = "2.0";
24
25/// The `codec:mcp` decoder/encoder.
26///
27/// As a [`Decoder`] it parses one MCP JSON-RPC envelope per frame into a
28/// canonical envelope `Expr`; as an [`Encoder`] it validates an envelope `Expr`
29/// and writes it back to JSON-RPC text. Non-MCP JSON is rejected.
30pub struct McpCodec;
31
32impl Decoder for McpCodec {
33    fn decode(&self, cx: &mut ReadCx<'_>, input: Input) -> Result<Expr> {
34        let source = input_text(cx.codec, input)?;
35        let mut budget = DecodeBudget::new(cx.limits);
36        budget.check_input_bytes(cx.codec, source.len())?;
37        reject_duplicate_mcp_wire_keys(cx.codec, &source)?;
38        let value = serde_json::from_str::<JsonValue>(&source)
39            .map_err(|err| codec_error(cx.codec, format!("MCP JSON parse error: {err}")))?;
40        let envelope = json_to_envelope(cx.codec, &value, &mut budget)?;
41        Ok(envelope_to_expr(&envelope))
42    }
43}
44
45impl Encoder for McpCodec {
46    fn encode(&self, cx: &mut WriteCx<'_>, expr: &Expr) -> Result<Output> {
47        let envelope = expr_to_envelope(expr).map_err(|err| Error::CodecError {
48            codec: cx.codec,
49            message: err.to_string(),
50        })?;
51        let value = envelope_to_json(cx.codec, &envelope)?;
52        serde_json::to_string(&value)
53            .map(Output::Text)
54            .map_err(|err| codec_error(cx.codec, err.to_string()))
55    }
56}
57
58fn input_text(codec: CodecId, input: Input) -> Result<String> {
59    match input {
60        Input::Text(text) => Ok(text),
61        Input::Bytes(bytes) => String::from_utf8(bytes)
62            .map_err(|err| codec_error(codec, format!("MCP input is not valid UTF-8: {err}"))),
63    }
64}
65
66fn json_to_envelope(
67    codec: CodecId,
68    value: &JsonValue,
69    budget: &mut DecodeBudget,
70) -> Result<McpEnvelope> {
71    match value {
72        JsonValue::Array(_) => Err(codec_error(
73            codec,
74            "MCP batch arrays are not supported by codec:mcp",
75        )),
76        JsonValue::Object(map) => json_object_to_envelope(codec, map, budget),
77        _ => Err(codec_error(codec, "MCP envelope must be a JSON object")),
78    }
79}
80
81fn json_object_to_envelope(
82    codec: CodecId,
83    map: &JsonMap<String, JsonValue>,
84    budget: &mut DecodeBudget,
85) -> Result<McpEnvelope> {
86    require_jsonrpc(codec, map)?;
87    let has_id = map.contains_key("id");
88    let has_method = map.contains_key("method");
89    let has_result = map.contains_key("result");
90    let has_error = map.contains_key("error");
91
92    match (has_method, has_id, has_result, has_error) {
93        (true, true, false, false) => json_request(codec, map, budget),
94        (true, false, false, false) => json_notification(codec, map, budget),
95        (false, true, true, false) => json_response(codec, map, budget),
96        (false, true, false, true) => json_error_response(codec, map, budget),
97        _ => Err(codec_error(
98            codec,
99            "invalid MCP JSON-RPC envelope field combination",
100        )),
101    }
102}
103
104fn json_request(
105    codec: CodecId,
106    map: &JsonMap<String, JsonValue>,
107    budget: &mut DecodeBudget,
108) -> Result<McpEnvelope> {
109    reject_unknown_json(codec, map, &["jsonrpc", "id", "method", "params"])?;
110    Ok(McpEnvelope::Request(McpRequest {
111        id: json_id(codec, required_json(codec, map, "id")?)?,
112        method: required_json_string(codec, map, "method")?.to_owned(),
113        params: json_value_expr(codec, map.get("params"), budget)?,
114    }))
115}
116
117fn json_notification(
118    codec: CodecId,
119    map: &JsonMap<String, JsonValue>,
120    budget: &mut DecodeBudget,
121) -> Result<McpEnvelope> {
122    reject_unknown_json(codec, map, &["jsonrpc", "method", "params"])?;
123    Ok(McpEnvelope::Notification(McpNotification {
124        method: required_json_string(codec, map, "method")?.to_owned(),
125        params: json_value_expr(codec, map.get("params"), budget)?,
126    }))
127}
128
129fn json_response(
130    codec: CodecId,
131    map: &JsonMap<String, JsonValue>,
132    budget: &mut DecodeBudget,
133) -> Result<McpEnvelope> {
134    reject_unknown_json(codec, map, &["jsonrpc", "id", "result"])?;
135    Ok(McpEnvelope::Response(McpResponse {
136        id: json_id(codec, required_json(codec, map, "id")?)?,
137        result: json_value_expr(codec, map.get("result"), budget)?,
138    }))
139}
140
141fn json_error_response(
142    codec: CodecId,
143    map: &JsonMap<String, JsonValue>,
144    budget: &mut DecodeBudget,
145) -> Result<McpEnvelope> {
146    reject_unknown_json(codec, map, &["jsonrpc", "id", "error"])?;
147    Ok(McpEnvelope::Error(McpErrorEnvelope {
148        id: json_id(codec, required_json(codec, map, "id")?)?,
149        error: json_error_object(codec, required_json(codec, map, "error")?, budget)?,
150    }))
151}
152
153fn json_error_object(
154    codec: CodecId,
155    value: &JsonValue,
156    budget: &mut DecodeBudget,
157) -> Result<McpError> {
158    let JsonValue::Object(map) = value else {
159        return Err(codec_error(codec, "MCP error must be an object"));
160    };
161    reject_unknown_json(codec, map, &["code", "message", "data"])?;
162    let Some(code) = required_json(codec, map, "code")?.as_i64() else {
163        return Err(codec_error(codec, "MCP error code must be an integer"));
164    };
165    Ok(McpError {
166        code,
167        message: required_json_string(codec, map, "message")?.to_owned(),
168        data: json_value_expr(codec, map.get("data"), budget)?,
169    })
170}
171
172fn json_value_expr(
173    codec: CodecId,
174    value: Option<&JsonValue>,
175    budget: &mut DecodeBudget,
176) -> Result<Expr> {
177    match value {
178        Some(value) => sim_codec_json::json_to_expr(codec, value, budget, 1),
179        None => Ok(Expr::Nil),
180    }
181}
182
183fn require_jsonrpc(codec: CodecId, map: &JsonMap<String, JsonValue>) -> Result<()> {
184    match map.get("jsonrpc") {
185        Some(JsonValue::String(version)) if version == JSONRPC_VERSION => Ok(()),
186        _ => Err(codec_error(
187            codec,
188            "MCP JSON-RPC envelope must declare jsonrpc \"2.0\"",
189        )),
190    }
191}
192
193fn required_json<'a>(
194    codec: CodecId,
195    map: &'a JsonMap<String, JsonValue>,
196    key: &str,
197) -> Result<&'a JsonValue> {
198    map.get(key)
199        .ok_or_else(|| codec_error(codec, format!("MCP envelope is missing {key}")))
200}
201
202fn required_json_string<'a>(
203    codec: CodecId,
204    map: &'a JsonMap<String, JsonValue>,
205    key: &str,
206) -> Result<&'a str> {
207    required_json(codec, map, key)?
208        .as_str()
209        .ok_or_else(|| codec_error(codec, format!("MCP envelope {key} must be a string")))
210}
211
212fn reject_unknown_json(
213    codec: CodecId,
214    map: &JsonMap<String, JsonValue>,
215    allowed: &[&str],
216) -> Result<()> {
217    for key in map.keys() {
218        if !allowed.contains(&key.as_str()) {
219            return Err(codec_error(
220                codec,
221                format!("unknown MCP JSON-RPC field {key}"),
222            ));
223        }
224    }
225    Ok(())
226}
227
228fn json_id(codec: CodecId, value: &JsonValue) -> Result<Expr> {
229    match value {
230        JsonValue::String(text) => Ok(Expr::String(text.clone())),
231        JsonValue::Number(number) => Ok(Expr::Number(NumberLiteral {
232            domain: Symbol::qualified("numbers", "f64"),
233            canonical: number.to_string(),
234        })),
235        JsonValue::Null => Ok(Expr::Nil),
236        _ => Err(codec_error(
237            codec,
238            "MCP JSON-RPC id must be a string, number, or null",
239        )),
240    }
241}
242
243fn envelope_to_json(codec: CodecId, envelope: &McpEnvelope) -> Result<JsonValue> {
244    let mut map = JsonMap::new();
245    map.insert(
246        "jsonrpc".to_owned(),
247        JsonValue::String(JSONRPC_VERSION.to_owned()),
248    );
249    match envelope {
250        McpEnvelope::Request(request) => {
251            map.insert("id".to_owned(), id_to_json(codec, &request.id)?);
252            map.insert(
253                "method".to_owned(),
254                JsonValue::String(request.method.clone()),
255            );
256            map.insert(
257                "params".to_owned(),
258                sim_codec_json::expr_to_json(&request.params),
259            );
260        }
261        McpEnvelope::Notification(notification) => {
262            map.insert(
263                "method".to_owned(),
264                JsonValue::String(notification.method.clone()),
265            );
266            map.insert(
267                "params".to_owned(),
268                sim_codec_json::expr_to_json(&notification.params),
269            );
270        }
271        McpEnvelope::Response(response) => {
272            map.insert("id".to_owned(), id_to_json(codec, &response.id)?);
273            map.insert(
274                "result".to_owned(),
275                sim_codec_json::expr_to_json(&response.result),
276            );
277        }
278        McpEnvelope::Error(error) => {
279            map.insert("id".to_owned(), id_to_json(codec, &error.id)?);
280            map.insert(
281                "error".to_owned(),
282                JsonValue::Object(error_to_json(&error.error)),
283            );
284        }
285    }
286    Ok(JsonValue::Object(map))
287}
288
289fn error_to_json(error: &McpError) -> JsonMap<String, JsonValue> {
290    let mut map = JsonMap::new();
291    map.insert(
292        "code".to_owned(),
293        JsonValue::Number(JsonNumber::from(error.code)),
294    );
295    map.insert(
296        "message".to_owned(),
297        JsonValue::String(error.message.clone()),
298    );
299    map.insert("data".to_owned(), sim_codec_json::expr_to_json(&error.data));
300    map
301}
302
303fn id_to_json(codec: CodecId, id: &Expr) -> Result<JsonValue> {
304    if !is_jsonrpc_id(id) {
305        return Err(codec_error(
306            codec,
307            "MCP JSON-RPC id must be a string, number, or nil",
308        ));
309    }
310    match id {
311        Expr::String(text) => Ok(JsonValue::String(text.clone())),
312        Expr::Number(number) => JsonNumber::from_str(&number.canonical)
313            .map(JsonValue::Number)
314            .map_err(|err| codec_error(codec, format!("invalid MCP numeric id: {err}"))),
315        Expr::Nil => Ok(JsonValue::Null),
316        _ => unreachable!("validated MCP id variants above"),
317    }
318}
319
320/// The host-registered [`Lib`] that installs [`McpCodec`] as the domain codec
321/// `codec:mcp`.
322pub struct McpCodecLib {
323    symbol: Symbol,
324    codec_id: CodecId,
325}
326
327impl McpCodecLib {
328    /// Create the lib bound to the given codec id (obtained from
329    /// [`Registry::fresh_codec_id`](sim_kernel::Registry::fresh_codec_id)).
330    pub fn new(id: CodecId) -> Self {
331        Self {
332            symbol: Symbol::qualified("codec", "mcp"),
333            codec_id: id,
334        }
335    }
336
337    fn domain_lib(&self) -> DomainCodecLib {
338        DomainCodecLib::new(
339            self.symbol.clone(),
340            self.codec_id,
341            Arc::new(McpCodec),
342            Arc::new(McpCodec),
343            Symbol::qualified("codec", "McpEnvelope"),
344        )
345    }
346}
347
348impl Lib for McpCodecLib {
349    fn manifest(&self) -> LibManifest {
350        self.domain_lib().manifest()
351    }
352
353    fn load(&self, cx: &mut LoadCx, linker: &mut Linker<'_>) -> Result<()> {
354        self.domain_lib().load(cx, linker)
355    }
356}