subako-function-sdk 0.2.4

//! The MCP messages a guest reads and writes over stdio, a JSON-RPC 2.0
//! message per line.

use serde::{Deserialize, Serialize};

/// The MCP revision a guest speaks.
// CODESYNC(mcp-protocol-version)
pub(crate) const MCP_PROTOCOL_VERSION: &str = "2026-07-28";

// CODESYNC(mcp-wire)
const JSONRPC_VERSION: &str = "2.0";
const DISCOVER: &str = "server/discover";
const TOOLS_LIST: &str = "tools/list";
const TOOLS_CALL: &str = "tools/call";
const RESULT_COMPLETE: &str = "complete";
const PROTOCOL_VERSION_KEY: &str = "io.modelcontextprotocol/protocolVersion";
const CLIENT_CAPABILITIES_KEY: &str = "io.modelcontextprotocol/clientCapabilities";
const SERVER_INFO_KEY: &str = "io.modelcontextprotocol/serverInfo";

const PARSE_ERROR: i64 = -32700;
const INVALID_REQUEST: i64 = -32600;
const METHOD_NOT_FOUND: i64 = -32601;
const INVALID_PARAMS: i64 = -32602;
const INTERNAL_ERROR: i64 = -32603;
const UNSUPPORTED_PROTOCOL_VERSION: i64 = -32022;

/// How long a client may keep a `tools/list` result: a module's tools never
/// change.
const TOOLS_TTL_MS: u64 = 24 * 60 * 60 * 1000;

/// A request's id, which its response echoes.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum RequestId {
    Number(i64),
    String(String),
}

/// What one line of stdin asks of the guest.
pub(crate) enum Incoming {
    Discover {
        id: RequestId,
    },
    ListTools {
        id: RequestId,
    },
    Call(ToolsCall),
    /// A notification or a response, which the guest reads past.
    Ignored,
    /// A line the guest answers with a JSON-RPC error.
    Refused(Refusal),
}

/// A `tools/call` request.
pub(crate) struct ToolsCall {
    pub(crate) id: RequestId,
    pub(crate) name: String,
    pub(crate) arguments: serde_json::Value,
}

/// A message as JSON-RPC lays it out, before its members say which kind it
/// is.
#[derive(Deserialize)]
struct RawMessage {
    jsonrpc: String,
    #[serde(default)]
    id: Option<RequestId>,
    #[serde(default)]
    method: Option<String>,
    #[serde(default)]
    params: Option<serde_json::Value>,
    #[serde(default)]
    result: Option<serde_json::Value>,
    #[serde(default)]
    error: Option<serde_json::Value>,
}

#[derive(Deserialize)]
struct ToolsCallParams {
    name: String,
    /// MCP lets a call leave its arguments out.
    #[serde(default = "no_arguments")]
    arguments: serde_json::Value,
}

fn no_arguments() -> serde_json::Value {
    serde_json::Value::Object(serde_json::Map::new())
}

/// Why a line is answered with a JSON-RPC error rather than served.
pub(crate) struct Refusal {
    /// Absent where no id could be read.
    id: Option<RequestId>,
    code: i64,
    message: String,
    data: Option<Box<serde_json::Value>>,
}

impl Refusal {
    fn new(id: Option<RequestId>, code: i64, message: String) -> Refusal {
        Refusal {
            id,
            code,
            message,
            data: None,
        }
    }

    pub(crate) fn unreadable(error: &std::io::Error) -> Refusal {
        Refusal::new(None, INTERNAL_ERROR, format!("stdin unreadable: {error}"))
    }

    pub(crate) fn line(&self) -> Vec<u8> {
        let error = serde_json::json!({"code": self.code, "message": self.message});
        let error = match &self.data {
            None => error,
            Some(data) => {
                serde_json::json!({"code": self.code, "message": self.message, "data": data})
            }
        };
        line(&serde_json::json!({"jsonrpc": JSONRPC_VERSION, "id": self.id, "error": error}))
    }
}

/// Reads the message on `line`, which may end in its newline.
pub(crate) fn read(line: &[u8]) -> Incoming {
    let refused = |id, code, message| Incoming::Refused(Refusal::new(id, code, message));
    let value: serde_json::Value = match serde_json::from_slice(line) {
        Ok(value) => value,
        Err(e) => return refused(None, PARSE_ERROR, format!("not JSON: {e}")),
    };
    let message: RawMessage = match serde_json::from_value(value) {
        Ok(message) => message,
        Err(e) => return refused(None, INVALID_REQUEST, format!("not a message: {e}")),
    };
    match message.jsonrpc.as_str() {
        JSONRPC_VERSION => {}
        version => {
            return refused(
                message.id,
                INVALID_REQUEST,
                format!("unknown JSON-RPC version: {version}"),
            );
        }
    }
    let (id, method) = match (message.id, message.method) {
        (Some(id), Some(method)) => (id, method),
        (None, Some(_)) => return Incoming::Ignored,
        (Some(_), None) if message.result.is_some() || message.error.is_some() => {
            return Incoming::Ignored;
        }
        (id, None) => {
            return refused(
                id,
                INVALID_REQUEST,
                "neither a request, a notification, nor a response".to_string(),
            );
        }
    };
    let params = message.params.unwrap_or(serde_json::Value::Null);
    if let Err(refusal) = check_meta(&id, &params) {
        return Incoming::Refused(refusal);
    }
    match method.as_str() {
        DISCOVER => Incoming::Discover { id },
        TOOLS_LIST => match params.get("cursor") {
            // Every tool is on the first page, so no cursor names another.
            Some(_) => refused(Some(id), INVALID_PARAMS, "no such cursor".to_string()),
            None => Incoming::ListTools { id },
        },
        TOOLS_CALL => match serde_json::from_value::<ToolsCallParams>(params) {
            Ok(params) => Incoming::Call(ToolsCall {
                id,
                name: params.name,
                arguments: params.arguments,
            }),
            Err(e) => refused(
                Some(id),
                INVALID_PARAMS,
                format!("invalid tools/call params: {e}"),
            ),
        },
        other => refused(
            Some(id),
            METHOD_NOT_FOUND,
            format!("unknown method: {other}"),
        ),
    }
}

/// Checks the `_meta` every request carries: the revision it speaks, which
/// must be this guest's, and the client's capabilities.
fn check_meta(id: &RequestId, params: &serde_json::Value) -> Result<(), Refusal> {
    let meta = params.get("_meta");
    let field = |key| meta.and_then(|meta| meta.get(key));
    let missing = |key| {
        Refusal::new(
            Some(id.clone()),
            INVALID_PARAMS,
            format!("missing _meta {key}; this server speaks MCP {MCP_PROTOCOL_VERSION}"),
        )
    };
    let version = match field(PROTOCOL_VERSION_KEY) {
        Some(serde_json::Value::String(version)) => version,
        Some(_) | None => return Err(missing(PROTOCOL_VERSION_KEY)),
    };
    match version.as_str() {
        MCP_PROTOCOL_VERSION => {}
        other => {
            return Err(Refusal {
                id: Some(id.clone()),
                code: UNSUPPORTED_PROTOCOL_VERSION,
                message: "Unsupported protocol version".to_string(),
                data: Some(Box::new(serde_json::json!({
                    "supported": [MCP_PROTOCOL_VERSION],
                    "requested": other,
                }))),
            });
        }
    }
    match field(CLIENT_CAPABILITIES_KEY) {
        Some(serde_json::Value::Object(_)) => Ok(()),
        Some(_) | None => Err(missing(CLIENT_CAPABILITIES_KEY)),
    }
}

/// Who answers: MCP's `serverInfo`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ServerInfo<'a> {
    pub name: &'a str,
    pub version: &'a str,
}

/// The answer to `server/discover`: the revision this guest speaks, and that
/// it serves tools.
pub(crate) fn discover_line(id: &RequestId, server: &ServerInfo<'_>) -> Vec<u8> {
    result_line(
        id,
        server,
        [
            (
                "supportedVersions",
                serde_json::json!([MCP_PROTOCOL_VERSION]),
            ),
            ("capabilities", serde_json::json!({"tools": {}})),
        ]
        .into_iter(),
    )
}

pub(crate) fn tools_line(
    id: &RequestId,
    server: &ServerInfo<'_>,
    tools: Vec<serde_json::Value>,
) -> Vec<u8> {
    result_line(
        id,
        server,
        [
            ("tools", serde_json::json!(tools)),
            ("ttlMs", serde_json::json!(TOOLS_TTL_MS)),
            ("cacheScope", serde_json::json!("public")),
        ]
        .into_iter(),
    )
}

/// The response carrying a result of `fields`, marked complete and signed
/// with `server`.
pub(crate) fn result_line<'k>(
    id: &RequestId,
    server: &ServerInfo<'_>,
    fields: impl Iterator<Item = (&'k str, serde_json::Value)>,
) -> Vec<u8> {
    let signed = [
        ("resultType", serde_json::json!(RESULT_COMPLETE)),
        (
            "_meta",
            serde_json::json!({
                SERVER_INFO_KEY: {"name": server.name, "version": server.version},
            }),
        ),
    ];
    let result: serde_json::Map<String, serde_json::Value> = fields
        .chain(signed)
        .map(|(key, value)| (key.to_string(), value))
        .collect();
    line(&serde_json::json!({"jsonrpc": JSONRPC_VERSION, "id": id, "result": result}))
}

/// `message` as one line: compact JSON escapes every newline inside it.
fn line(message: &serde_json::Value) -> Vec<u8> {
    // A `Value` always serializes.
    let mut line = serde_json::to_vec(message).unwrap_or_default();
    line.push(b'\n');
    line
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn an_unreadable_stdin_is_an_internal_error_answering_no_id() {
        let line = Refusal::unreadable(&std::io::Error::other("gone")).line();
        let message: serde_json::Value = serde_json::from_slice(&line).unwrap();
        assert_eq!(
            message,
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": null,
                "error": {"code": INTERNAL_ERROR, "message": "stdin unreadable: gone"},
            })
        );
    }
}