trip-test 0.1.1

//! MCP protocol client implementation.
//!
//! Handles communication with MCP servers over stdio transport using JSON-RPC.
//! Manages the initialization handshake and supports tool listing and calling.

use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::process::Stdio;
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};

/// Represents a tool exposed by an MCP server.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tool {
    pub name: String,
    pub description: String,
    #[serde(rename = "inputSchema")]
    pub input_schema: Option<Value>,
}

/// MCP client for communicating with servers over stdio transport.
pub struct MCPClient {
    #[allow(dead_code)]
    child: Child,
    reader: BufReader<ChildStdout>,
    writer: ChildStdin,
    request_id: u64,
}

impl MCPClient {
    /// Create a new MCP client and connect to a server via stdio.
    ///
    /// # Arguments
    /// * `server_cmd` - Command to spawn the server (e.g., "python -m my_server")
    pub async fn new(server_cmd: &str) -> Result<Self> {
        let parts: Vec<&str> = server_cmd.split_whitespace().collect();
        
        if parts.is_empty() {
            return Err(anyhow!("Invalid server command"));
        }

        let mut cmd = Command::new(parts[0]);
        for arg in &parts[1..] {
            cmd.arg(arg);
        }

        let mut child = cmd
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;

        let stdout = child.stdout.take().ok_or(anyhow!("Failed to capture stdout"))?;
        let stdin = child.stdin.take().ok_or(anyhow!("Failed to capture stdin"))?;

        let reader = BufReader::new(stdout);

        let mut client = MCPClient {
            child,
            reader,
            writer: stdin,
            request_id: 1,
        };

        // perform async handshake
        client.initialize().await?;
        Ok(client)
    }

    async fn initialize(&mut self) -> Result<()> {
        let request = json!({
            "jsonrpc": "2.0",
            "id": self.request_id,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {
                    "name": "trip-test",
                    "version": "0.1.0"
                }
            }
        });

        self.request_id += 1;
        self.send_request(&request).await?;
        let response = self.read_response().await?;

        if let Some(error) = response.get("error") {
            return Err(anyhow!("Handshake failed: {}", error));
        }

        Ok(())
    }
    pub async fn call_tool(&mut self, tool_name: &str, args: Value) -> Result<Value> {
            let request = json!({
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "tools/call",
                "params": {
                    "name": tool_name,
                    "arguments": args
                }
            });
    
            self.request_id += 1;
            self.send_request(&request).await?;
            let response = self.read_response().await?;
    
            if let Some(error) = response.get("error") {
                return Err(anyhow!("Tool call failed: {}", error));
            }
    
            Ok(response)
        }

    pub async fn list_tools(&mut self) -> Result<Vec<Tool>> {
        let request = json!({
            "jsonrpc": "2.0",
            "id": self.request_id,
            "method": "tools/list",
            "params": {}
        });

        self.request_id += 1;
        self.send_request(&request).await?;
        let response = self.read_response().await?;

        if let Some(error) = response.get("error") {
            return Err(anyhow!("Failed to list tools: {}", error));
        }

        let tools = response
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(|t| t.as_array())
            .ok_or(anyhow!("Invalid tools response"))?;

        let mut tool_list = Vec::new();
        for tool in tools {
            if let Ok(t) = serde_json::from_value::<Tool>(tool.clone()) {
                tool_list.push(t);
            }
        }

        Ok(tool_list)
    }

    async fn send_request(&mut self, request: &Value) -> Result<()> {
        let json_str = serde_json::to_string(request)?;
        self.writer.write_all(json_str.as_bytes()).await?;
        self.writer.write_all(b"\n").await?;
        self.writer.flush().await?;
        Ok(())
    }

    async fn read_response(&mut self) -> Result<Value> {
        let mut line = String::new();
        let n = self.reader.read_line(&mut line).await?;

        if n == 0 {
            return Err(anyhow!("Connection closed"));
        }

        let response: Value = serde_json::from_str(line.trim_end())?;
        Ok(response)
    }

    pub fn get_server_name(&self) -> String {
        "test-server".to_string()
    }

    pub fn get_server_version(&self) -> String {
        "0.1.0".to_string()
    }
}