yoagent 0.26.6

//! MCP (Model Context Protocol) JSON-RPC 2.0 types.

use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};

static REQUEST_ID_COUNTER: AtomicU64 = AtomicU64::new(1);

/// Generate a unique JSON-RPC request ID.
pub fn next_request_id() -> u64 {
    REQUEST_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
}

// ---------------------------------------------------------------------------
// JSON-RPC 2.0
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcRequest {
    pub jsonrpc: String,
    pub id: u64,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<serde_json::Value>,
}

impl JsonRpcRequest {
    pub fn new(method: impl Into<String>, params: Option<serde_json::Value>) -> Self {
        Self {
            jsonrpc: "2.0".into(),
            id: next_request_id(),
            method: method.into(),
            params,
        }
    }
}

/// A JSON-RPC notification: no `id`, and no response is expected (or sent).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcNotification {
    pub jsonrpc: String,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<serde_json::Value>,
}

impl JsonRpcNotification {
    pub fn new(method: impl Into<String>, params: Option<serde_json::Value>) -> Self {
        Self {
            jsonrpc: "2.0".into(),
            method: method.into(),
            params,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcResponse {
    pub jsonrpc: String,
    pub id: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<JsonRpcError>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcError {
    pub code: i64,
    pub message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<serde_json::Value>,
}

// ---------------------------------------------------------------------------
// MCP Protocol types
// ---------------------------------------------------------------------------

/// Client info sent during initialization.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientInfo {
    pub name: String,
    pub version: String,
}

impl Default for ClientInfo {
    fn default() -> Self {
        Self {
            name: "yoagent".into(),
            version: env!("CARGO_PKG_VERSION").into(),
        }
    }
}

/// Server info received during initialization.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServerInfo {
    pub name: String,
    #[serde(default)]
    pub version: String,
}

/// Server capabilities received during initialization.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ServerCapabilities {
    #[serde(default)]
    pub tools: Option<serde_json::Value>,
    #[serde(default)]
    pub resources: Option<serde_json::Value>,
    #[serde(default)]
    pub prompts: Option<serde_json::Value>,
}

/// Initialize result from the server.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InitializeResult {
    pub protocol_version: String,
    pub capabilities: ServerCapabilities,
    pub server_info: ServerInfo,
}

/// MCP tool as returned by tools/list.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpToolInfo {
    pub name: String,
    #[serde(default)]
    pub description: Option<String>,
    #[serde(default)]
    pub input_schema: serde_json::Value,
}

/// tools/list result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolsListResult {
    pub tools: Vec<McpToolInfo>,
}

/// Content item in a tool call result.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum McpContent {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "image")]
    Image {
        data: String,
        #[serde(rename = "mimeType")]
        mime_type: String,
    },
}

/// tools/call result.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpToolCallResult {
    /// Content kinds this crate does not model (`resource`, `resource_link`,
    /// `audio`, anything newer) arrive as [`McpContent::Text`] holding the
    /// block's JSON, rather than failing the whole call.
    #[serde(deserialize_with = "lenient_content")]
    pub content: Vec<McpContent>,
    #[serde(default)]
    pub is_error: bool,
}

/// Read each content block, keeping one this crate does not model as text
/// the model can use: a `resource`'s own text, or the block's JSON with any
/// large `data` / `blob` payload replaced by its size (base64 audio would
/// otherwise go into the context verbatim).
fn lenient_content<'de, D>(deserializer: D) -> Result<Vec<McpContent>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let blocks = Vec::<serde_json::Value>::deserialize(deserializer)?;
    Ok(blocks
        .into_iter()
        .map(
            |block| match serde_json::from_value::<McpContent>(block.clone()) {
                Ok(content) => content,
                Err(e) => {
                    let kind = block
                        .get("type")
                        .and_then(|t| t.as_str())
                        .unwrap_or("(none)")
                        .to_string();
                    warn_unmodelled(&kind, &e);
                    McpContent::Text {
                        text: unmodelled_text(block),
                    }
                }
            },
        )
        .collect())
}

/// Warn once per content type that had to be passed through as text.
fn warn_unmodelled(kind: &str, error: &serde_json::Error) {
    use std::collections::HashSet;
    use std::sync::{Mutex, OnceLock};
    static SEEN: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
    let first = SEEN
        .get_or_init(|| Mutex::new(HashSet::new()))
        .lock()
        .map(|mut seen| seen.insert(kind.to_string()))
        .unwrap_or(true);
    if first {
        tracing::warn!(
            "MCP content of type '{kind}' is not modelled ({error}); passing it to the model as text"
        );
    }
}

/// Payloads longer than this are replaced by a note of their size.
const MAX_INLINE_PAYLOAD: usize = 1024;

fn unmodelled_text(mut block: serde_json::Value) -> String {
    // A text resource: its own text is what the model wants.
    if let Some(resource) = block.get("resource") {
        if let Some(text) = resource.get("text").and_then(|t| t.as_str()) {
            let uri = resource.get("uri").and_then(|u| u.as_str()).unwrap_or("");
            return format!("[resource {uri}]\n{text}");
        }
    }
    shrink_payloads(&mut block);
    block.to_string()
}

fn shrink_payloads(value: &mut serde_json::Value) {
    match value {
        serde_json::Value::Object(map) => {
            for (key, v) in map.iter_mut() {
                match v {
                    serde_json::Value::String(s)
                        if (key == "data" || key == "blob") && s.len() > MAX_INLINE_PAYLOAD =>
                    {
                        *v = serde_json::Value::String(format!(
                            "[{} bytes of base64 omitted]",
                            s.len()
                        ));
                    }
                    other => shrink_payloads(other),
                }
            }
        }
        serde_json::Value::Array(items) => items.iter_mut().for_each(shrink_payloads),
        _ => {}
    }
}

// ---------------------------------------------------------------------------
// MCP Error
// ---------------------------------------------------------------------------

/// What went wrong talking to an MCP server. `#[non_exhaustive]`: match
/// with a wildcard arm.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum McpError {
    #[error("Transport error: {0}")]
    Transport(String),
    /// The server answered with a non-success HTTP status (the HTTP
    /// transport). `status` is its code: 401 or 403, it wants a credential or
    /// a sign-in; 404, the session expired when one was held (the
    /// message says so), else most likely a wrong URL. `message` is the whole report
    /// as a `Transport` error read before — the server's body (first 200
    /// chars) included when there was one.
    ///
    /// `www_authenticate` is the response's `WWW-Authenticate` header on a
    /// 401 or 403 (the challenge naming how to sign in, or the scope a 403
    /// `insufficient_scope` wants), when it sent one; read it with
    /// [`McpError::www_authenticate`].
    ///
    /// `#[non_exhaustive]`: match with `Http { status, .. }` and build one
    /// with [`McpError::http`], so fields can be added without a breaking
    /// change.
    #[error("Transport error: {message}")]
    #[non_exhaustive]
    Http {
        status: u16,
        message: String,
        www_authenticate: Option<String>,
    },
    #[error("Protocol error: {0}")]
    Protocol(String),
    #[error("JSON-RPC error {code}: {message}")]
    JsonRpc { code: i64, message: String },
    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Connection closed")]
    ConnectionClosed,
}

impl McpError {
    /// An [`Http`](Self::Http) error, for a custom [`McpTransport`](super::McpTransport)
    /// reporting a non-success HTTP status.
    pub fn http(status: u16, message: impl Into<String>) -> Self {
        Self::Http {
            status,
            message: message.into(),
            www_authenticate: None,
        }
    }

    /// This error with the response's `WWW-Authenticate` header, for an
    /// [`Http`](Self::Http) error; any other error is returned unchanged.
    pub(crate) fn with_www_authenticate(mut self, challenge: Option<String>) -> Self {
        if let Self::Http {
            www_authenticate, ..
        } = &mut self
        {
            *www_authenticate = challenge;
        }
        self
    }

    /// The `WWW-Authenticate` header of an [`Http`](Self::Http) error's
    /// response (401 / 403), if it sent one; `None` for any other error.
    pub fn www_authenticate(&self) -> Option<&str> {
        match self {
            Self::Http {
                www_authenticate, ..
            } => www_authenticate.as_deref(),
            _ => None,
        }
    }

    /// The HTTP status of an [`Http`](Self::Http) error; `None` for any other.
    pub fn http_status(&self) -> Option<u16> {
        match self {
            Self::Http { status, .. } => Some(*status),
            _ => None,
        }
    }
}

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

    #[test]
    fn http_errors_carry_their_status() {
        let err = McpError::http(401, "HTTP 401 Unauthorized from server on 'tools/list'");
        assert_eq!(err.http_status(), Some(401));
        assert_eq!(
            err.to_string(),
            "Transport error: HTTP 401 Unauthorized from server on 'tools/list'"
        );
        assert_eq!(McpError::Transport("x".into()).http_status(), None);
        let challenged = McpError::http(401, "sign in")
            .with_www_authenticate(Some("Bearer realm=\"mcp\"".into()));
        assert_eq!(challenged.www_authenticate(), Some("Bearer realm=\"mcp\""));
        assert_eq!(err.www_authenticate(), None);
        assert_eq!(
            McpError::Transport("x".into())
                .with_www_authenticate(Some("y".into()))
                .www_authenticate(),
            None
        );
    }

    #[test]
    fn test_json_rpc_request_serialization() {
        let req = JsonRpcRequest {
            jsonrpc: "2.0".into(),
            id: 1,
            method: "initialize".into(),
            params: Some(serde_json::json!({"protocolVersion": "2024-11-05"})),
        };
        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("\"jsonrpc\":\"2.0\""));
        assert!(json.contains("\"method\":\"initialize\""));

        let parsed: JsonRpcRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.id, 1);
        assert_eq!(parsed.method, "initialize");
    }

    #[test]
    fn test_json_rpc_response_deserialization() {
        let json = r#"{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","capabilities":{},"serverInfo":{"name":"test","version":"1.0"}}}"#;
        let resp: JsonRpcResponse = serde_json::from_str(json).unwrap();
        assert_eq!(resp.id, Some(1));
        assert!(resp.result.is_some());
        assert!(resp.error.is_none());
    }

    #[test]
    fn test_json_rpc_error_response() {
        let json =
            r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"Method not found"}}"#;
        let resp: JsonRpcResponse = serde_json::from_str(json).unwrap();
        assert!(resp.error.is_some());
        let err = resp.error.unwrap();
        assert_eq!(err.code, -32601);
    }

    #[test]
    fn test_initialize_result_deserialization() {
        let json = r#"{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"test-server","version":"0.1.0"}}"#;
        let result: InitializeResult = serde_json::from_str(json).unwrap();
        assert_eq!(result.server_info.name, "test-server");
        assert!(result.capabilities.tools.is_some());
    }

    #[test]
    fn test_mcp_tool_info_deserialization() {
        let json = r#"{"name":"read_file","description":"Read a file","inputSchema":{"type":"object","properties":{"path":{"type":"string"}}}}"#;
        let tool: McpToolInfo = serde_json::from_str(json).unwrap();
        assert_eq!(tool.name, "read_file");
        assert_eq!(tool.description.as_deref(), Some("Read a file"));
    }

    #[test]
    fn test_mcp_tool_call_result() {
        let json = r#"{"content":[{"type":"text","text":"file contents here"}],"isError":false}"#;
        let result: McpToolCallResult = serde_json::from_str(json).unwrap();
        assert_eq!(result.content.len(), 1);
        assert!(!result.is_error);
    }

    #[test]
    fn unknown_content_blocks_are_kept_as_json_text() {
        let json = r#"{"content":[
            {"type":"text","text":"hi"},
            {"type":"resource","resource":{"uri":"file:///a.txt","text":"body"}},
            {"type":"audio","data":"AAAA","mimeType":"audio/wav"}
        ]}"#;
        let result: McpToolCallResult = serde_json::from_str(json).unwrap();
        assert_eq!(result.content.len(), 3);
        assert!(matches!(&result.content[0], McpContent::Text { text } if text == "hi"));
        assert!(
            matches!(&result.content[1], McpContent::Text { text } if text == "[resource file:///a.txt]\nbody")
        );
        assert!(matches!(&result.content[2], McpContent::Text { text } if text.contains("audio")));
    }

    #[test]
    fn large_unmodelled_payloads_are_replaced_by_their_size() {
        let blob = "A".repeat(5000);
        let json = format!(
            r#"{{"content":[{{"type":"audio","data":"{blob}","mimeType":"audio/wav"}},{{"type":"resource","resource":{{"uri":"file:///b.bin","blob":"{blob}"}}}}]}}"#
        );
        let result: McpToolCallResult = serde_json::from_str(&json).unwrap();
        for content in &result.content {
            let McpContent::Text { text } = content else {
                panic!("expected text")
            };
            assert!(text.contains("[5000 bytes of base64 omitted]"), "{text}");
            assert!(text.len() < 500, "{} bytes kept", text.len());
        }
    }

    #[test]
    fn a_notification_has_no_id() {
        let json =
            serde_json::to_string(&JsonRpcNotification::new("notifications/initialized", None))
                .unwrap();
        assert_eq!(
            json,
            r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#
        );
    }

    #[test]
    fn test_unique_request_ids() {
        let id1 = next_request_id();
        let id2 = next_request_id();
        assert_ne!(id1, id2);
    }
}