agent-harness 0.4.0-alpha.4

//! The transport-agnostic MCP client: the `initialize` handshake, `tools/list`
//! (with cursor pagination), and `tools/call` (flattening result content to
//! text). The wire is one of two [`McpConnection`] transports — [`stdio`] (a
//! launched process) or [`http`] (a remote endpoint) — which differ only in how
//! a single JSON-RPC request/notification is delivered. Grounded in the MCP
//! 2025-11-25 spec.
//!
//! [`stdio`]: super::stdio
//! [`http`]: super::http

use std::path::Path;

use serde_json::{json, Value};

use super::{McpPromptArg, McpServer, McpTransport, PromptMessage};

/// Protocol version we negotiate (the latest as of writing). A server speaking
/// an older version still responds; we don't gate on the echoed value.
pub(super) const PROTOCOL_VERSION: &str = "2025-11-25";

/// One way to deliver a JSON-RPC message to an MCP server. The lifecycle in
/// [`McpClient`] is written once against this; stdio and HTTP each implement it.
pub(crate) trait McpConnection: Send + Sync {
    /// Send a request and return its `result` (or an `Err` for a JSON-RPC error
    /// / transport failure).
    fn request(&self, method: &str, params: Value) -> Result<Value, String>;
    /// Send a fire-and-forget notification (no response expected).
    fn notify(&self, method: &str, params: Value) -> Result<(), String>;
}

/// One tool advertised by an MCP server (`tools/list` entry).
pub(crate) struct McpToolDef {
    pub name: String,
    pub description: String,
    pub input_schema: Value,
}

/// One resource advertised by an MCP server (`resources/list` entry).
pub(crate) struct McpResource {
    pub uri: String,
    pub name: String,
    pub description: String,
}

/// A prompt template from `prompts/list` (the server name is added by the
/// caller, which knows it).
pub(crate) struct McpPromptDef {
    pub name: String,
    pub description: String,
    pub arguments: Vec<McpPromptArg>,
}

/// A connected MCP server: a transport plus the protocol lifecycle. Shared by
/// that server's [`McpTool`]s via `Arc`.
///
/// [`McpTool`]: super::tool::McpTool
pub(crate) struct McpClient {
    conn: Box<dyn McpConnection>,
}

impl McpClient {
    /// Open the transport for `server`, perform the handshake, and return the
    /// client together with the tools it advertises. `Err` on connect / handshake
    /// failure. `cwd` is the working directory for a launched stdio server.
    pub(crate) fn connect(server: &McpServer, cwd: &Path) -> Result<(McpClient, Vec<McpToolDef>), String> {
        let conn: Box<dyn McpConnection> = match &server.transport {
            McpTransport::Stdio { command, args, env } => {
                Box::new(super::stdio::StdioConnection::spawn(&server.name, command, args, env, cwd)?)
            }
            McpTransport::Http { url, headers } => Box::new(super::http::HttpConnection::new(&server.name, url, headers)),
        };
        let client = McpClient { conn };
        client.initialize()?;
        let tools = client.list_tools()?;
        Ok((client, tools))
    }

    /// `initialize` request → `notifications/initialized` (the spec's handshake).
    fn initialize(&self) -> Result<(), String> {
        let params = json!({
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": {},
            "clientInfo": { "name": "openai-compatible", "version": env!("CARGO_PKG_VERSION") }
        });
        self.conn.request("initialize", params)?;
        self.conn.notify("notifications/initialized", json!({}))
    }

    /// `tools/list`, following `nextCursor` pagination to the end.
    fn list_tools(&self) -> Result<Vec<McpToolDef>, String> {
        let mut out = Vec::new();
        let mut cursor: Option<String> = None;
        loop {
            let params = cursor.as_ref().map_or_else(|| json!({}), |c| json!({ "cursor": c }));
            let result = self.conn.request("tools/list", params)?;
            if let Some(arr) = result.get("tools").and_then(Value::as_array) {
                for t in arr {
                    let Some(name) = t.get("name").and_then(Value::as_str) else { continue };
                    out.push(McpToolDef {
                        name: name.to_owned(),
                        description: t.get("description").and_then(Value::as_str).unwrap_or_default().to_owned(),
                        input_schema: t.get("inputSchema").cloned().unwrap_or_else(|| json!({ "type": "object" })),
                    });
                }
            }
            match result.get("nextCursor").and_then(Value::as_str) {
                Some(c) => cursor = Some(c.to_owned()),
                None => return Ok(out),
            }
        }
    }

    /// `tools/call`, flattening the result's content blocks to text. `Err` when
    /// the result is flagged `isError` (the tool itself failed) or on transport
    /// failure.
    pub(crate) fn call(&self, name: &str, arguments: &Value) -> Result<String, String> {
        let result = self.conn.request("tools/call", json!({ "name": name, "arguments": arguments }))?;
        let text = match result.get("content").and_then(Value::as_array) {
            Some(blocks) => flatten_content(blocks),
            None => String::new(),
        };
        if result.get("isError").and_then(Value::as_bool).unwrap_or(false) {
            return Err(if text.is_empty() { "the tool reported an error".to_owned() } else { text });
        }
        Ok(text)
    }

    /// `resources/list` — the server's readable resources. A server without
    /// resource support answers "method not found"; we map that (and any error)
    /// to an empty list, so it's simply skipped — resources are a bonus, not
    /// required.
    pub(crate) fn list_resources(&self) -> Vec<McpResource> {
        let Ok(result) = self.conn.request("resources/list", json!({})) else {
            return Vec::new();
        };
        result
            .get("resources")
            .and_then(Value::as_array)
            .map(|arr| {
                arr.iter()
                    .filter_map(|r| {
                        let uri = r.get("uri").and_then(Value::as_str)?;
                        Some(McpResource {
                            uri: uri.to_owned(),
                            name: r.get("name").and_then(Value::as_str).unwrap_or(uri).to_owned(),
                            description: r.get("description").and_then(Value::as_str).unwrap_or_default().to_owned(),
                        })
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    /// `resources/read` — read a resource's text contents by uri.
    pub(crate) fn read_resource(&self, uri: &str) -> Result<String, String> {
        let result = self.conn.request("resources/read", json!({ "uri": uri }))?;
        Ok(match result.get("contents").and_then(Value::as_array) {
            Some(items) => flatten_resource_contents(items),
            None => String::new(),
        })
    }

    /// `prompts/list` — the server's prompt templates. Unsupported / error → empty.
    pub(crate) fn list_prompts(&self) -> Vec<McpPromptDef> {
        let Ok(result) = self.conn.request("prompts/list", json!({})) else {
            return Vec::new();
        };
        result
            .get("prompts")
            .and_then(Value::as_array)
            .map(|arr| arr.iter().filter_map(parse_prompt_def).collect())
            .unwrap_or_default()
    }

    /// `prompts/get` — resolve a prompt template to its messages (text flattened).
    pub(crate) fn get_prompt(&self, name: &str, arguments: &[(String, String)]) -> Result<Vec<PromptMessage>, String> {
        let args: serde_json::Map<String, Value> =
            arguments.iter().map(|(k, v)| (k.clone(), Value::String(v.clone()))).collect();
        let result = self.conn.request("prompts/get", json!({ "name": name, "arguments": args }))?;
        Ok(result
            .get("messages")
            .and_then(Value::as_array)
            .map(|arr| {
                arr.iter()
                    .filter_map(|m| {
                        let role = m.get("role").and_then(Value::as_str)?;
                        Some(PromptMessage { role: role.to_owned(), content: prompt_content_text(m.get("content")) })
                    })
                    .collect()
            })
            .unwrap_or_default())
    }
}

/// Flatten MCP content blocks to text — text blocks verbatim, other block types
/// noted by kind so the model knows non-text content came back.
pub(super) fn flatten_content(blocks: &[Value]) -> String {
    blocks
        .iter()
        .filter_map(|b| match b.get("type").and_then(Value::as_str) {
            Some("text") => b.get("text").and_then(Value::as_str).map(str::to_owned),
            Some(other) => Some(format!("[{other} content omitted]")),
            None => None,
        })
        .collect::<Vec<_>>()
        .join("\n")
}

/// Flatten `resources/read` content items to text (text items verbatim; a binary
/// blob noted by placeholder).
fn flatten_resource_contents(items: &[Value]) -> String {
    items
        .iter()
        .filter_map(|c| {
            c.get("text")
                .and_then(Value::as_str)
                .map(str::to_owned)
                .or_else(|| c.get("blob").map(|_| "[binary resource omitted]".to_owned()))
        })
        .collect::<Vec<_>>()
        .join("\n")
}

/// Parse one `prompts/list` entry into a [`McpPromptDef`].
fn parse_prompt_def(p: &Value) -> Option<McpPromptDef> {
    let name = p.get("name").and_then(Value::as_str)?;
    let arguments = p
        .get("arguments")
        .and_then(Value::as_array)
        .map(|a| {
            a.iter()
                .filter_map(|arg| {
                    let n = arg.get("name").and_then(Value::as_str)?;
                    Some(McpPromptArg {
                        name: n.to_owned(),
                        description: arg.get("description").and_then(Value::as_str).unwrap_or_default().to_owned(),
                        required: arg.get("required").and_then(Value::as_bool).unwrap_or(false),
                    })
                })
                .collect()
        })
        .unwrap_or_default();
    Some(McpPromptDef {
        name: name.to_owned(),
        description: p.get("description").and_then(Value::as_str).unwrap_or_default().to_owned(),
        arguments,
    })
}

/// Extract text from a prompt message's `content` — a single `{type,text}` block,
/// an array of blocks, or a bare string.
fn prompt_content_text(content: Option<&Value>) -> String {
    match content {
        Some(Value::Array(arr)) => flatten_content(arr),
        Some(Value::Object(_)) => {
            content.and_then(|c| c.get("text")).and_then(Value::as_str).unwrap_or_default().to_owned()
        }
        Some(Value::String(s)) => s.clone(),
        _ => String::new(),
    }
}

/// Extract a JSON-RPC `result` for `id` from a response body that is either a
/// plain JSON object or an SSE (`data:`-prefixed) stream — the two shapes the
/// Streamable-HTTP transport returns. Shared by the HTTP connection.
pub(super) fn parse_rpc_result(text: &str, id: i64) -> Result<Value, String> {
    // Collect every JSON object in the body: the whole thing if it parses, else
    // each SSE `data:` line that parses.
    let mut messages: Vec<Value> = Vec::new();
    if let Ok(v) = serde_json::from_str::<Value>(text) {
        messages.push(v);
    } else {
        for data in text.lines().filter_map(|l| l.strip_prefix("data:").map(str::trim)) {
            if let Ok(v) = serde_json::from_str::<Value>(data) {
                messages.push(v);
            }
        }
    }
    // Prefer the response carrying our id; fall back to any result/error message.
    let response = messages
        .iter()
        .find(|m| m.get("id").and_then(Value::as_i64) == Some(id))
        .or_else(|| messages.iter().find(|m| m.get("result").is_some() || m.get("error").is_some()));
    match response {
        Some(m) if m.get("error").is_some() => Err(format!("{}", m["error"])),
        Some(m) => Ok(m.get("result").cloned().unwrap_or(Value::Null)),
        None => Err("no JSON-RPC response in the server reply".to_owned()),
    }
}

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

    #[test]
    fn flatten_keeps_text_and_notes_other_blocks() {
        let blocks = vec![
            json!({ "type": "text", "text": "hello" }),
            json!({ "type": "image", "data": "..." }),
            json!({ "type": "text", "text": "world" }),
        ];
        assert_eq!(flatten_content(&blocks), "hello\n[image content omitted]\nworld");
    }

    #[test]
    fn parse_rpc_result_reads_json_sse_and_errors() {
        // Plain JSON response.
        let json = r#"{"jsonrpc":"2.0","id":2,"result":{"ok":true}}"#;
        assert_eq!(parse_rpc_result(json, 2).unwrap(), json!({ "ok": true }));
        // SSE-framed response (the matching id among several data lines).
        let sse = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":5,\"result\":{\"v\":1}}\n\n";
        assert_eq!(parse_rpc_result(sse, 5).unwrap(), json!({ "v": 1 }));
        // JSON-RPC error surfaces as Err.
        let err = r#"{"jsonrpc":"2.0","id":3,"error":{"code":-32601,"message":"nope"}}"#;
        assert!(parse_rpc_result(err, 3).unwrap_err().contains("nope"));
    }

    #[test]
    fn connect_reports_a_spawn_failure() {
        let server = McpServer::stdio("missing", "definitely-not-a-real-binary-xyz", vec![]);
        match McpClient::connect(&server, Path::new(".")) {
            Err(e) => assert!(e.contains("spawning"), "got: {e}"),
            Ok(_) => panic!("expected a spawn failure"),
        }
    }

    /// End-to-end against a tiny `sh` server that speaks the protocol: it replies
    /// to `initialize` (id 1), `tools/list` (id 2), and one `tools/call` (id 3) —
    /// the ids our client assigns in order — verifying the full handshake.
    #[test]
    fn connect_handshakes_lists_and_calls() {
        let script = r#"
            read _initialize
            printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-11-25","capabilities":{},"serverInfo":{"name":"fake","version":"0"}}}'
            read _initialized
            read _list
            printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{"tools":[{"name":"echo","description":"echoes input","inputSchema":{"type":"object"}}]}}'
            read _call
            printf '%s\n' '{"jsonrpc":"2.0","id":3,"result":{"content":[{"type":"text","text":"pong"}],"isError":false}}'
            read _reslist
            printf '%s\n' '{"jsonrpc":"2.0","id":4,"result":{"resources":[{"uri":"file:///doc","name":"Doc","description":"a doc"}]}}'
            read _resread
            printf '%s\n' '{"jsonrpc":"2.0","id":5,"result":{"contents":[{"uri":"file:///doc","text":"doc body"}]}}'
            read _promptslist
            printf '%s\n' '{"jsonrpc":"2.0","id":6,"result":{"prompts":[{"name":"greet","description":"greeting","arguments":[{"name":"who","required":true}]}]}}'
            read _promptsget
            printf '%s\n' '{"jsonrpc":"2.0","id":7,"result":{"messages":[{"role":"user","content":{"type":"text","text":"Hello there"}}]}}'
        "#;
        let server = McpServer::stdio("fake", "sh", vec!["-c".to_owned(), script.to_owned()]);
        let (client, tools) = McpClient::connect(&server, Path::new(".")).expect("handshake succeeds");
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0].name, "echo");
        assert_eq!(tools[0].description, "echoes input");
        assert_eq!(client.call("echo", &json!({ "x": 1 })).expect("call succeeds"), "pong");
        // resources/list then resources/read.
        let resources = client.list_resources();
        assert_eq!(resources.len(), 1);
        assert_eq!(resources[0].uri, "file:///doc");
        assert_eq!(client.read_resource("file:///doc").expect("read succeeds"), "doc body");
        // prompts/list then prompts/get.
        let prompts = client.list_prompts();
        assert_eq!(prompts.len(), 1);
        assert_eq!(prompts[0].name, "greet");
        assert!(prompts[0].arguments[0].required && prompts[0].arguments[0].name == "who");
        let msgs = client.get_prompt("greet", &[("who".to_owned(), "world".to_owned())]).expect("get prompt");
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].role, "user");
        assert_eq!(msgs[0].content, "Hello there");
    }
}