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 wire_keys::reject_duplicate_mcp_wire_keys,
21};
22
23const JSONRPC_VERSION: &str = "2.0";
24
25pub 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(¬ification.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
320pub struct McpCodecLib {
323 symbol: Symbol,
324 codec_id: CodecId,
325}
326
327impl McpCodecLib {
328 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}