haematite_cli/mcp/
rpc.rs

1use std::collections::BTreeMap;
2use std::fmt;
3
4use serde::de::DeserializeOwned;
5use serde::{Deserialize, Deserializer, Serialize};
6use serde_json::value::{RawValue, to_raw_value};
7
8use super::schema::{self, ToolSchema};
9
10pub const SUPPORTED_PROTOCOL_VERSION: &str = "2025-11-25";
11const JSON_RPC_VERSION: &str = "2.0";
12
13#[derive(Debug, Clone, Serialize)]
14#[serde(transparent)]
15pub struct RequestId(Box<RawValue>);
16
17impl RequestId {
18    /// Validate a JSON-RPC request id.
19    ///
20    /// # Errors
21    /// Returns `rpc_invalid_id` unless the value is a string or an integer.
22    pub fn from_raw(raw: Box<RawValue>) -> Result<Self, Refusal> {
23        let text = raw.get();
24        if text.starts_with('"') || is_integer_number(text) {
25            return Ok(Self(raw));
26        }
27        Err(Refusal::new(
28            None,
29            true,
30            -32600,
31            "rpc_invalid_id",
32            "id must be a string or integer",
33        ))
34    }
35}
36
37fn is_integer_number(text: &str) -> bool {
38    let text = text.strip_prefix('-').unwrap_or(text);
39    if !text.as_bytes().first().is_some_and(u8::is_ascii_digit) {
40        return false;
41    }
42    let (mantissa, exponent) = text.split_once(['e', 'E']).unwrap_or((text, "0"));
43    let zero = mantissa.bytes().all(|byte| byte == b'0' || byte == b'.');
44    if zero {
45        return true;
46    }
47    let fractional_places = mantissa
48        .split_once('.')
49        .map_or(0, |(_, fraction)| fraction.len());
50    let trailing_zeroes = mantissa
51        .bytes()
52        .rev()
53        .filter(|byte| *byte != b'.')
54        .take_while(|byte| *byte == b'0')
55        .count();
56    let negative_exponent = exponent.starts_with('-');
57    let magnitude = exponent.trim_start_matches(['+', '-']).parse::<usize>();
58    // Count decimal places rather than round through a floating-point value.
59    match magnitude {
60        Ok(magnitude) if negative_exponent => fractional_places
61            .checked_add(magnitude)
62            .is_some_and(|needed| needed <= trailing_zeroes),
63        Ok(magnitude) => fractional_places.saturating_sub(magnitude) <= trailing_zeroes,
64        Err(error) if matches!(error.kind(), std::num::IntErrorKind::PosOverflow) => {
65            !negative_exponent
66        }
67        Err(_) => false,
68    }
69}
70
71#[derive(Debug)]
72pub struct Refusal {
73    id: Option<RequestId>,
74    respond: bool,
75    number: i32,
76    name: &'static str,
77    message: String,
78}
79
80impl Refusal {
81    fn new(
82        id: Option<RequestId>,
83        respond: bool,
84        number: i32,
85        name: &'static str,
86        message: impl Into<String>,
87    ) -> Self {
88        Self {
89            id,
90            respond,
91            number,
92            name,
93            message: message.into(),
94        }
95    }
96
97    /// Encode this refusal when the peer expects a response.
98    ///
99    /// # Errors
100    /// Propagates JSON serialization errors; notifications produce no bytes.
101    pub fn encode(&self) -> Result<Option<Vec<u8>>, serde_json::Error> {
102        if !self.respond {
103            return Ok(None);
104        }
105        serde_json::to_vec(&ErrorResponse {
106            jsonrpc: JSON_RPC_VERSION,
107            id: self.id.as_ref(),
108            error: ErrorBody {
109                code: self.number,
110                message: &self.message,
111                data: ErrorData { code: self.name },
112            },
113        })
114        .map(Some)
115    }
116}
117
118impl fmt::Display for Refusal {
119    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
120        write!(formatter, "{}: {}", self.name, self.message)
121    }
122}
123
124impl std::error::Error for Refusal {}
125
126#[derive(Serialize)]
127struct ErrorResponse<'a> {
128    jsonrpc: &'static str,
129    id: Option<&'a RequestId>,
130    error: ErrorBody<'a>,
131}
132
133#[derive(Serialize)]
134struct ErrorBody<'a> {
135    code: i32,
136    message: &'a str,
137    data: ErrorData,
138}
139
140#[derive(Serialize)]
141struct ErrorData {
142    code: &'static str,
143}
144
145#[derive(Serialize)]
146struct Success<'a> {
147    jsonrpc: &'static str,
148    id: &'a RequestId,
149    result: &'a RawValue,
150}
151
152/// Encode an object result under the request id.
153///
154/// # Errors
155/// Returns `rpc_result_not_object` for another JSON kind, or `rpc_result_encoding_failed` when serialization fails.
156pub fn encode_result<T: Serialize + ?Sized>(
157    id: &RequestId,
158    result: &T,
159) -> Result<Vec<u8>, Refusal> {
160    let failure = |error: serde_json::Error| {
161        Refusal::new(
162            Some(id.clone()),
163            true,
164            -32603,
165            "rpc_result_encoding_failed",
166            error.to_string(),
167        )
168    };
169    let raw = to_raw_value(result).map_err(failure)?;
170    if !raw.get().starts_with('{') {
171        return Err(Refusal::new(
172            Some(id.clone()),
173            true,
174            -32603,
175            "rpc_result_not_object",
176            "result must be an object",
177        ));
178    }
179    serde_json::to_vec(&Success {
180        jsonrpc: JSON_RPC_VERSION,
181        id,
182        result: &raw,
183    })
184    .map_err(failure)
185}
186
187fn present_raw<'de, D: Deserializer<'de>>(
188    deserializer: D,
189) -> Result<Option<Box<RawValue>>, D::Error> {
190    Box::<RawValue>::deserialize(deserializer).map(Some)
191}
192
193#[derive(Deserialize)]
194#[serde(deny_unknown_fields)]
195struct Envelope {
196    #[serde(default, deserialize_with = "present_raw")]
197    jsonrpc: Option<Box<RawValue>>,
198    #[serde(default, deserialize_with = "present_raw")]
199    id: Option<Box<RawValue>>,
200    #[serde(default, deserialize_with = "present_raw")]
201    method: Option<Box<RawValue>>,
202    #[serde(default, deserialize_with = "present_raw")]
203    params: Option<Box<RawValue>>,
204}
205
206#[derive(Debug)]
207pub struct Request {
208    id: Option<RequestId>,
209    method: Method,
210}
211
212impl Request {
213    #[must_use]
214    pub const fn id(&self) -> Option<&RequestId> {
215        self.id.as_ref()
216    }
217
218    #[must_use]
219    pub const fn method(&self) -> &Method {
220        &self.method
221    }
222}
223
224#[derive(Debug)]
225pub enum Method {
226    Initialize(InitializeParams),
227    Initialized,
228    Ping,
229    ToolsList,
230    ToolsCall(ToolCall),
231}
232
233type RawObject = BTreeMap<String, Box<RawValue>>;
234
235#[derive(Debug, Deserialize)]
236#[serde(rename_all = "camelCase")]
237pub struct InitializeParams {
238    pub protocol_version: String,
239    pub capabilities: RawObject,
240    pub client_info: Implementation,
241    #[serde(rename = "_meta")]
242    pub metadata: Option<RawObject>,
243}
244
245#[derive(Debug, Deserialize)]
246#[serde(rename_all = "camelCase")]
247pub struct Implementation {
248    pub name: String,
249    pub version: String,
250    pub title: Option<String>,
251    pub description: Option<String>,
252    pub icons: Option<Vec<Icon>>,
253    pub website_url: Option<String>,
254}
255
256#[derive(Debug, Deserialize)]
257#[serde(rename_all = "camelCase")]
258pub struct Icon {
259    pub src: String,
260    pub mime_type: Option<String>,
261    pub sizes: Option<Vec<String>>,
262    pub theme: Option<String>,
263}
264
265#[derive(Deserialize)]
266struct EmptyParams {
267    #[serde(rename = "_meta")]
268    metadata: Option<RawObject>,
269}
270
271#[derive(Deserialize)]
272struct ListParams {
273    cursor: Option<String>,
274    #[serde(rename = "_meta")]
275    metadata: Option<RawObject>,
276}
277
278#[derive(Deserialize)]
279struct CallParams {
280    name: String,
281    #[serde(default, deserialize_with = "present_raw")]
282    arguments: Option<Box<RawValue>>,
283    #[serde(default, deserialize_with = "present_raw")]
284    task: Option<Box<RawValue>>,
285    #[serde(rename = "_meta")]
286    metadata: Option<RawObject>,
287}
288
289#[derive(Debug)]
290pub struct ToolCall {
291    pub name: String,
292    pub arguments: Box<RawValue>,
293    pub metadata: Option<RawObject>,
294}
295
296#[derive(Debug)]
297pub enum Prepared {
298    NoResponse,
299    Response(Vec<u8>),
300    ToolCall { id: RequestId, call: ToolCall },
301}
302
303#[derive(Debug)]
304pub struct RpcCodec {
305    tools: [ToolSchema; 11],
306}
307
308impl Default for RpcCodec {
309    fn default() -> Self {
310        Self::new()
311    }
312}
313
314impl RpcCodec {
315    #[must_use]
316    pub fn new() -> Self {
317        Self::with_registry(schema::tools)
318    }
319
320    fn with_registry(build: impl FnOnce() -> [ToolSchema; 11]) -> Self {
321        Self { tools: build() }
322    }
323
324    /// Decode one bounded request against the served tool registry.
325    ///
326    /// # Errors
327    /// Refuses oversized frames, malformed JSON or envelopes, unsupported methods, and invalid method parameters by name.
328    pub fn decode(&self, frame: &[u8], byte_limit: usize) -> Result<Request, Refusal> {
329        if frame.len() > byte_limit {
330            return Err(Refusal::new(
331                None,
332                true,
333                -32000,
334                "rpc_frame_too_large",
335                "frame exceeds byte limit",
336            ));
337        }
338        let raw: Box<RawValue> = serde_json::from_slice(frame).map_err(|error| {
339            Refusal::new(None, true, -32700, "rpc_parse_error", error.to_string())
340        })?;
341        match raw.get().as_bytes().first() {
342            Some(b'{') => {}
343            Some(b'[') => {
344                return Err(Refusal::new(
345                    None,
346                    true,
347                    -32600,
348                    "rpc_batch_unsupported",
349                    "batch requests are unsupported",
350                ));
351            }
352            _ => {
353                return Err(Refusal::new(
354                    None,
355                    true,
356                    -32600,
357                    "rpc_not_object",
358                    "request must be an object",
359                ));
360            }
361        }
362        let envelope: Envelope = serde_json::from_str(raw.get()).map_err(|error| {
363            Refusal::new(None, true, -32600, "rpc_invalid_request", error.to_string())
364        })?;
365        self.decode_envelope(envelope)
366    }
367
368    fn decode_envelope(&self, envelope: Envelope) -> Result<Request, Refusal> {
369        let id = envelope.id.map(RequestId::from_raw).transpose()?;
370        let respond = id.is_some();
371        let invalid = |message: String| {
372            Refusal::new(id.clone(), true, -32600, "rpc_invalid_request", message)
373        };
374        let version = envelope
375            .jsonrpc
376            .as_ref()
377            .ok_or_else(|| invalid("jsonrpc is required".to_owned()))?;
378        let version: String =
379            serde_json::from_str(version.get()).map_err(|error| invalid(error.to_string()))?;
380        if version != JSON_RPC_VERSION {
381            return Err(invalid("jsonrpc must be 2.0".to_owned()));
382        }
383        let method = envelope
384            .method
385            .as_ref()
386            .ok_or_else(|| invalid("method is required".to_owned()))?;
387        let method: String =
388            serde_json::from_str(method.get()).map_err(|error| invalid(error.to_string()))?;
389        let params = envelope.params.as_deref();
390        let invalid_params = |message: String| {
391            Refusal::new(id.clone(), respond, -32602, "rpc_invalid_params", message)
392        };
393        if params.is_some_and(|raw| !raw.get().starts_with('{')) {
394            return Err(invalid_params("params must be an object".to_owned()));
395        }
396        let decoded = match method.as_str() {
397            "initialize" => Method::Initialize(read_initialize(params, id.as_ref())?),
398            "notifications/initialized" => {
399                if id.is_some() {
400                    return Err(invalid(
401                        "notifications/initialized must not have an id".to_owned(),
402                    ));
403                }
404                let input: EmptyParams = read_params(params).map_err(invalid_params)?;
405                drop(input.metadata);
406                Method::Initialized
407            }
408            "ping" => {
409                let input: EmptyParams = read_params(params).map_err(invalid_params)?;
410                drop(input.metadata);
411                Method::Ping
412            }
413            "tools/list" => {
414                let input: ListParams = read_params(params).map_err(invalid_params)?;
415                if input.cursor.is_some() {
416                    return Err(Refusal::new(
417                        id,
418                        respond,
419                        -32602,
420                        "mcp_cursor_unsupported",
421                        "the fixed tool list has no continuation cursor",
422                    ));
423                }
424                drop(input.metadata);
425                Method::ToolsList
426            }
427            "tools/call" => Method::ToolsCall(self.read_call(params, id.as_ref())?),
428            _ => {
429                return Err(Refusal::new(
430                    id,
431                    respond,
432                    -32601,
433                    "rpc_unknown_method",
434                    format!("unknown method {method}"),
435                ));
436            }
437        };
438        Ok(Request {
439            id,
440            method: decoded,
441        })
442    }
443
444    fn read_call(
445        &self,
446        params: Option<&RawValue>,
447        id: Option<&RequestId>,
448    ) -> Result<ToolCall, Refusal> {
449        let input: CallParams =
450            read_params(params).map_err(|message| parameter_refusal(id, message))?;
451        if input.task.is_some() {
452            return Err(Refusal::new(
453                id.cloned(),
454                id.is_some(),
455                -32602,
456                "mcp_tasks_unsupported",
457                "task execution is unsupported",
458            ));
459        }
460        if !self.tools.iter().any(|tool| tool.name == input.name) {
461            return Err(Refusal::new(
462                id.cloned(),
463                id.is_some(),
464                -32602,
465                "mcp_tool_unknown",
466                format!("unknown tool {}", input.name),
467            ));
468        }
469        let arguments = match input.arguments {
470            Some(arguments) => arguments,
471            None => RawValue::from_string("{}".to_owned())
472                .map_err(|error| parameter_refusal(id, error.to_string()))?,
473        };
474        if !arguments.get().starts_with('{') {
475            return Err(parameter_refusal(id, "arguments must be an object"));
476        }
477        Ok(ToolCall {
478            name: input.name,
479            arguments,
480            metadata: input.metadata,
481        })
482    }
483
484    /// Prepare the response or tool call for a decoded request.
485    ///
486    /// # Errors
487    /// Propagates result-encoding refusals and rejects an initialized notification carrying an id.
488    pub fn prepare(&self, request: Request) -> Result<Prepared, Refusal> {
489        let Some(id) = request.id else {
490            return Ok(Prepared::NoResponse);
491        };
492        let bytes = match request.method {
493            Method::Initialize(_) => encode_result(
494                &id,
495                &InitializeResult {
496                    protocol_version: SUPPORTED_PROTOCOL_VERSION,
497                    capabilities: Capabilities {
498                        tools: ToolsCapability {
499                            list_changed: false,
500                        },
501                    },
502                    server_info: ServerInfo {
503                        name: "haem",
504                        version: env!("CARGO_PKG_VERSION"),
505                    },
506                },
507            )?,
508            Method::Ping => encode_result(&id, &EmptyResult {})?,
509            Method::ToolsList => encode_result(&id, &ToolsListResult { tools: &self.tools })?,
510            Method::ToolsCall(call) => return Ok(Prepared::ToolCall { id, call }),
511            Method::Initialized => {
512                return Err(Refusal::new(
513                    Some(id),
514                    true,
515                    -32600,
516                    "rpc_invalid_request",
517                    "notifications/initialized must not have an id",
518                ));
519            }
520        };
521        Ok(Prepared::Response(bytes))
522    }
523}
524
525fn read_initialize(
526    params: Option<&RawValue>,
527    id: Option<&RequestId>,
528) -> Result<InitializeParams, Refusal> {
529    if id.is_none() {
530        return Err(Refusal::new(
531            None,
532            true,
533            -32600,
534            "rpc_invalid_request",
535            "initialize requires an id",
536        ));
537    }
538    let input: InitializeParams =
539        read_params(params).map_err(|message| parameter_refusal(id, message))?;
540    if input.protocol_version.is_empty() {
541        return Err(parameter_refusal(
542            id,
543            "protocolVersion must be a nonempty string",
544        ));
545    }
546    if input.client_info.name.is_empty() || input.client_info.version.is_empty() {
547        return Err(parameter_refusal(
548            id,
549            "clientInfo name and version must be nonempty",
550        ));
551    }
552    Ok(input)
553}
554
555fn parameter_refusal(id: Option<&RequestId>, message: impl Into<String>) -> Refusal {
556    Refusal::new(
557        id.cloned(),
558        id.is_some(),
559        -32602,
560        "rpc_invalid_params",
561        message,
562    )
563}
564
565fn read_params<T: DeserializeOwned>(params: Option<&RawValue>) -> Result<T, String> {
566    serde_json::from_str(params.map_or("{}", RawValue::get)).map_err(|error| error.to_string())
567}
568
569#[derive(Serialize)]
570#[serde(rename_all = "camelCase")]
571struct InitializeResult {
572    protocol_version: &'static str,
573    capabilities: Capabilities,
574    server_info: ServerInfo,
575}
576
577#[derive(Serialize)]
578struct Capabilities {
579    tools: ToolsCapability,
580}
581
582#[derive(Serialize)]
583#[serde(rename_all = "camelCase")]
584struct ToolsCapability {
585    list_changed: bool,
586}
587
588#[derive(Serialize)]
589struct ServerInfo {
590    name: &'static str,
591    version: &'static str,
592}
593
594#[derive(Serialize)]
595struct EmptyResult {}
596
597#[derive(Serialize)]
598struct ToolsListResult<'a> {
599    tools: &'a [ToolSchema; 11],
600}
601
602#[cfg(test)]
603mod tests {
604    use std::error::Error;
605
606    use serde_json::value::RawValue;
607    use serde_json::{Value, json};
608
609    use super::{
610        Method, Prepared, Refusal, RequestId, RpcCodec, SUPPORTED_PROTOCOL_VERSION, encode_result,
611    };
612
613    type TestResult = Result<(), Box<dyn Error>>;
614
615    #[derive(serde::Deserialize)]
616    struct IdAnswer {
617        id: Box<RawValue>,
618    }
619
620    #[derive(serde::Deserialize)]
621    struct ResultAnswer {
622        result: Box<RawValue>,
623    }
624
625    fn refusal(frame: &[u8], limit: usize) -> Result<Refusal, Box<dyn Error>> {
626        RpcCodec::new()
627            .decode(frame, limit)
628            .err()
629            .ok_or_else(|| "frame accepted".into())
630    }
631
632    fn error_name(error: &Refusal) -> Result<String, Box<dyn Error>> {
633        let bytes = error.encode()?.ok_or("refusal had no answer")?;
634        let body: Value = serde_json::from_slice(&bytes)?;
635        let name = body
636            .get("error")
637            .and_then(|v| v.get("data"))
638            .and_then(|v| v.get("code"))
639            .and_then(Value::as_str)
640            .ok_or("refusal has no code name")?;
641        Ok(name.to_owned())
642    }
643
644    fn response(codec: &RpcCodec, frame: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
645        match codec.prepare(codec.decode(frame, frame.len())?)? {
646            Prepared::Response(bytes) => Ok(bytes),
647            _ => Err("request did not return a response".into()),
648        }
649    }
650
651    #[test]
652    fn malformed_json_is_a_named_parse_error() -> TestResult {
653        let error = refusal(br"{", 4096)?;
654        assert_eq!(error_name(&error)?, "rpc_parse_error");
655        Ok(())
656    }
657
658    #[test]
659    fn nonobjects_and_batch_arrays_are_refused_by_name() -> TestResult {
660        for text in [r#""text""#, "null", "1", "true"] {
661            assert_eq!(
662                error_name(&refusal(text.as_bytes(), 4096)?)?,
663                "rpc_not_object"
664            );
665        }
666        assert_eq!(error_name(&refusal(b"[]", 4096)?)?, "rpc_batch_unsupported");
667        Ok(())
668    }
669
670    #[test]
671    fn unknown_method_keeps_the_request_id() -> TestResult {
672        let error = refusal(br#"{"jsonrpc":"2.0","id":"a","method":"unknown"}"#, 4096)?;
673        assert_eq!(error_name(&error)?, "rpc_unknown_method");
674        let bytes = error.encode()?.ok_or("missing error")?;
675        let text = String::from_utf8(bytes)?;
676        assert!(text.contains(r#""id":"a""#));
677        Ok(())
678    }
679
680    #[test]
681    fn null_and_noninteger_ids_are_refused() -> TestResult {
682        for id in ["null", "true", "{}", "[]", "1.5", "1e-2", "1.2e0"] {
683            let text = format!(r#"{{"jsonrpc":"2.0","id":{id},"method":"ping"}}"#);
684            assert_eq!(
685                error_name(&refusal(text.as_bytes(), 4096)?)?,
686                "rpc_invalid_id"
687            );
688        }
689        Ok(())
690    }
691
692    #[test]
693    fn number_string_and_large_ids_keep_their_exact_spelling() -> TestResult {
694        let codec = RpcCodec::new();
695        for spelling in [
696            "1",
697            "-0",
698            "1e2",
699            "1.0",
700            "1.2e1",
701            "18446744073709551616001",
702            r#""a\u0062""#,
703        ] {
704            let text = format!(r#"{{"jsonrpc":"2.0","id":{spelling},"method":"ping"}}"#);
705            let bytes = response(&codec, text.as_bytes())?;
706            let answer: IdAnswer = serde_json::from_slice(&bytes)?;
707            assert_eq!(answer.id.get(), spelling);
708        }
709        Ok(())
710    }
711
712    #[test]
713    fn notifications_return_no_answer_even_when_the_method_is_unknown() -> TestResult {
714        let codec = RpcCodec::new();
715        for method in ["notifications/initialized", "ping", "tools/list"] {
716            let text = format!(r#"{{"jsonrpc":"2.0","method":"{method}"}}"#);
717            assert!(matches!(
718                codec.prepare(codec.decode(text.as_bytes(), 4096)?)?,
719                Prepared::NoResponse
720            ));
721        }
722        let error = refusal(br#"{"jsonrpc":"2.0","method":"unknown"}"#, 4096)?;
723        assert!(error.encode()?.is_none());
724        Ok(())
725    }
726
727    fn initialize(version: &str) -> String {
728        format!(
729            r#"{{"jsonrpc":"2.0","id":1,"method":"initialize","params":{{"protocolVersion":"{version}","capabilities":{{}},"clientInfo":{{"name":"test","version":"1"}}}}}}"#
730        )
731    }
732
733    #[test]
734    fn initialize_refuses_an_empty_version_by_name() -> TestResult {
735        let text = initialize("");
736        assert_eq!(
737            error_name(&refusal(text.as_bytes(), 4096)?)?,
738            "rpc_invalid_params"
739        );
740        Ok(())
741    }
742
743    #[test]
744    fn initialize_names_one_version_and_only_static_tools() -> TestResult {
745        let codec = RpcCodec::new();
746        let bytes = response(&codec, initialize(SUPPORTED_PROTOCOL_VERSION).as_bytes())?;
747        let answer: Value = serde_json::from_slice(&bytes)?;
748        let result = answer.get("result").ok_or("no result")?;
749        assert_eq!(
750            result.get("protocolVersion"),
751            Some(&json!(SUPPORTED_PROTOCOL_VERSION))
752        );
753        let capabilities = result
754            .get("capabilities")
755            .and_then(Value::as_object)
756            .ok_or("no capabilities")?;
757        assert_eq!(capabilities.len(), 1);
758        assert_eq!(
759            capabilities.get("tools"),
760            Some(&json!({"listChanged":false}))
761        );
762        assert_eq!(
763            result.get("serverInfo").and_then(|v| v.get("name")),
764            Some(&json!("haem"))
765        );
766        Ok(())
767    }
768
769    #[test]
770    fn ping_returns_an_empty_object() -> TestResult {
771        let codec = RpcCodec::new();
772        let answer: Value = serde_json::from_slice(&response(
773            &codec,
774            br#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#,
775        )?)?;
776        assert_eq!(answer.get("result"), Some(&json!({})));
777        Ok(())
778    }
779
780    #[test]
781    fn tools_list_has_eleven_pairs_and_reuses_the_held_registry() -> TestResult {
782        let builds = std::cell::Cell::new(0);
783        let codec = RpcCodec::with_registry(|| {
784            builds.set(builds.get() + 1);
785            crate::mcp::schema::tools()
786        });
787        let held = codec.tools.as_ptr();
788        assert_eq!(builds.get(), 1);
789        for _ in 0..2 {
790            let answer: Value = serde_json::from_slice(&response(
791                &codec,
792                br#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
793            )?)?;
794            let tools = answer
795                .get("result")
796                .and_then(|v| v.get("tools"))
797                .and_then(Value::as_array)
798                .ok_or("no tools")?;
799            let names = tools
800                .iter()
801                .map(|tool| {
802                    tool.get("name")
803                        .and_then(Value::as_str)
804                        .ok_or("missing name")
805                })
806                .collect::<Result<Vec<_>, _>>()?;
807            assert_eq!(
808                names,
809                [
810                    "describe_store",
811                    "run_query",
812                    "read_history",
813                    "define_collection",
814                    "put_records",
815                    "delete_record",
816                    "link_records",
817                    "unlink_records",
818                    "create_branch",
819                    "review_branch",
820                    "merge_branch"
821                ]
822            );
823            for tool in tools {
824                assert!(tool.get("inputSchema").is_some_and(Value::is_object));
825                assert!(tool.get("outputSchema").is_some_and(Value::is_object));
826            }
827            assert_eq!(codec.tools.as_ptr(), held);
828            assert_eq!(builds.get(), 1);
829        }
830        Ok(())
831    }
832
833    #[test]
834    fn an_unknown_tool_is_refused_without_dispatch() -> TestResult {
835        let error = refusal(br#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"missing","arguments":{}}}"#, 4096)?;
836        assert_eq!(error_name(&error)?, "mcp_tool_unknown");
837        Ok(())
838    }
839
840    #[test]
841    fn tool_call_stops_at_a_typed_name_and_raw_arguments() -> TestResult {
842        let codec = RpcCodec::new();
843        let body = r#"{ "collection":"records", "n":1e2 }"#;
844        let text = format!(
845            r#"{{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{{"name":"put_records","arguments":{body}}}}}"#
846        );
847        let request = codec.decode(text.as_bytes(), 4096)?;
848        assert!(matches!(&request.method, Method::ToolsCall(_)));
849        match codec.prepare(request)? {
850            Prepared::ToolCall { call, .. } => {
851                assert_eq!(call.name, "put_records");
852                assert_eq!(call.arguments.get(), body);
853            }
854            _ => return Err("tool call was answered or discarded".into()),
855        }
856        Ok(())
857    }
858
859    #[test]
860    fn a_frame_one_byte_over_limit_is_refused_before_parsing() -> TestResult {
861        let malformed = br#"{"jsonrpc":"#;
862        assert_eq!(
863            error_name(&refusal(malformed, malformed.len() - 1)?)?,
864            "rpc_frame_too_large"
865        );
866        assert_eq!(
867            error_name(&refusal(malformed, malformed.len())?)?,
868            "rpc_parse_error"
869        );
870        Ok(())
871    }
872
873    #[test]
874    fn result_body_numbers_leave_with_their_original_spelling() -> TestResult {
875        let id = RequestId::from_raw(RawValue::from_string(r#""r1""#.to_owned())?)?;
876        let body =
877            RawValue::from_string(r#"{ "n":1e2, "big":18446744073709551616001 }"#.to_owned())?;
878        let bytes = encode_result(&id, body.as_ref())?;
879        let answer: ResultAnswer = serde_json::from_slice(&bytes)?;
880        assert_eq!(answer.result.get(), body.get());
881        Ok(())
882    }
883
884    #[test]
885    fn invalid_params_and_wrong_rpc_version_are_named() -> TestResult {
886        for text in [
887            r#"{"jsonrpc":"2.0","id":1,"method":"ping","params":null}"#,
888            r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"put_records","arguments":[]}}"#,
889            r#"{"jsonrpc":"2.0","id":1,"method":"ping","params":{"_meta":[]}}"#,
890        ] {
891            assert_eq!(
892                error_name(&refusal(text.as_bytes(), 4096)?)?,
893                "rpc_invalid_params"
894            );
895        }
896        assert_eq!(
897            error_name(&refusal(
898                br#"{"jsonrpc":"1.0","id":1,"method":"ping"}"#,
899                4096
900            )?)?,
901            "rpc_invalid_request"
902        );
903        Ok(())
904    }
905    #[test]
906    fn older_protocol_versions_negotiate_the_supported_version() -> TestResult {
907        let codec = RpcCodec::new();
908        for version in ["2025-06-18", "2024-11-05"] {
909            let answer: Value =
910                serde_json::from_slice(&response(&codec, initialize(version).as_bytes())?)?;
911            assert_eq!(
912                answer
913                    .get("result")
914                    .and_then(|value| value.get("protocolVersion")),
915                Some(&json!(SUPPORTED_PROTOCOL_VERSION)),
916            );
917        }
918        Ok(())
919    }
920
921    #[test]
922    fn initialization_ignores_unconsumed_protocol_members() -> TestResult {
923        let codec = RpcCodec::new();
924        let text = format!(
925            r#"{{"jsonrpc":"2.0","id":1,"method":"initialize","params":{{"protocolVersion":"{SUPPORTED_PROTOCOL_VERSION}","capabilities":{{"future":{{"enabled":true}}}},"future":true,"clientInfo":{{"name":"test","version":"1","future":true,"icons":[{{"src":"https://example.com/icon.png","future":true}}]}}}}}}"#
926        );
927        let answer: Value = serde_json::from_slice(&response(&codec, text.as_bytes())?)?;
928        assert_eq!(
929            answer
930                .get("result")
931                .and_then(|value| value.get("protocolVersion")),
932            Some(&json!(SUPPORTED_PROTOCOL_VERSION)),
933        );
934        Ok(())
935    }
936
937    #[test]
938    fn tools_list_ignores_unconsumed_protocol_params() -> TestResult {
939        let codec = RpcCodec::new();
940        let frame = br#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"future":true}}"#;
941        let answer: Value = serde_json::from_slice(&response(&codec, frame)?)?;
942        let tools = answer
943            .get("result")
944            .and_then(|value| value.get("tools"))
945            .and_then(Value::as_array)
946            .ok_or("tools absent")?;
947        assert_eq!(tools.len(), 11);
948        Ok(())
949    }
950
951    #[test]
952    fn tools_call_ignores_unconsumed_protocol_members() -> TestResult {
953        let codec = RpcCodec::new();
954        let frame = br#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"put_records","arguments":{},"future":{"n":1e2}}}"#;
955        match codec.prepare(codec.decode(frame, 4096)?)? {
956            Prepared::ToolCall { call, .. } => {
957                assert_eq!(call.name, "put_records");
958                assert_eq!(call.arguments.get(), "{}");
959            }
960            _ => return Err("tools/call was not decoded".into()),
961        }
962        Ok(())
963    }
964
965    #[test]
966    fn any_present_task_is_refused_by_name_without_dispatch() -> TestResult {
967        for task in ["{}", "null", "true", r#"{"ttl":1000}"#] {
968            let frame = format!(
969                r#"{{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{{"name":"put_records","arguments":{{}},"task":{task}}}}}"#
970            );
971            assert_eq!(
972                error_name(&refusal(frame.as_bytes(), 4096)?)?,
973                "mcp_tasks_unsupported"
974            );
975        }
976        Ok(())
977    }
978
979    #[test]
980    fn tool_argument_members_remain_strict_after_protocol_decoding() -> TestResult {
981        use crate::mcp::types::PutRecordsInput;
982
983        let codec = RpcCodec::new();
984        let arguments = r#"{"request_id":"r1","branch":"feature","expected":{"generation":"1","sequence":"2"},"collection":"records","records":[{"document_id":"a","document":{"n":1e2}}],"future":true}"#;
985        let frame = format!(
986            r#"{{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{{"name":"put_records","arguments":{arguments},"future":true}}}}"#
987        );
988        let Prepared::ToolCall { call, .. } =
989            codec.prepare(codec.decode(frame.as_bytes(), 4096)?)?
990        else {
991            return Err("tools/call was not decoded".into());
992        };
993        assert_eq!(call.arguments.get(), arguments);
994        let error = serde_json::from_str::<PutRecordsInput>(call.arguments.get())
995            .err()
996            .ok_or("unknown tool argument was accepted")?;
997        assert!(error.to_string().contains("unknown field"), "{error}");
998        Ok(())
999    }
1000
1001    #[test]
1002    fn ping_and_initialized_ignore_unconsumed_protocol_params() -> TestResult {
1003        let codec = RpcCodec::new();
1004        let ping = br#"{"jsonrpc":"2.0","id":1,"method":"ping","params":{"future":true}}"#;
1005        let answer: Value = serde_json::from_slice(&response(&codec, ping)?)?;
1006        assert_eq!(answer.get("result"), Some(&json!({})));
1007        let initialized =
1008            br#"{"jsonrpc":"2.0","method":"notifications/initialized","params":{"future":true}}"#;
1009        assert!(matches!(
1010            codec.prepare(codec.decode(initialized, 4096)?)?,
1011            Prepared::NoResponse
1012        ));
1013        Ok(())
1014    }
1015
1016    #[test]
1017    fn envelope_members_remain_strict() -> TestResult {
1018        let frame = br#"{"jsonrpc":"2.0","id":1,"method":"ping","future":true}"#;
1019        assert_eq!(error_name(&refusal(frame, 4096)?)?, "rpc_invalid_request");
1020        Ok(())
1021    }
1022}
haematite_cli/mcp/rpc.rs

haematite_cli/mcp/
rpc.rs
