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