1use 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
24pub 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(¬ification.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
318pub struct McpCodecLib {
321 symbol: Symbol,
322 codec_id: CodecId,
323}
324
325impl McpCodecLib {
326 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}