trip_test/
mcp.rs

1//! MCP protocol client implementation.
2//!
3//! Handles communication with MCP servers over stdio transport using JSON-RPC.
4//! Manages the initialization handshake and supports tool listing and calling.
5
6use anyhow::{anyhow, Result};
7use serde::{Deserialize, Serialize};
8use serde_json::{json, Value};
9use std::process::Stdio;
10use tokio::process::{Child, ChildStdin, ChildStdout, Command};
11use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
12
13/// Represents a tool exposed by an MCP server.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Tool {
16    pub name: String,
17    pub description: String,
18    #[serde(rename = "inputSchema")]
19    pub input_schema: Option<Value>,
20}
21
22/// MCP client for communicating with servers over stdio transport.
23pub struct MCPClient {
24    #[allow(dead_code)]
25    child: Child,
26    reader: BufReader<ChildStdout>,
27    writer: ChildStdin,
28    request_id: u64,
29}
30
31impl MCPClient {
32    /// Create a new MCP client and connect to a server via stdio.
33    ///
34    /// # Arguments
35    /// * `server_cmd` - Command to spawn the server (e.g., "python -m my_server")
36    pub async fn new(server_cmd: &str) -> Result<Self> {
37        let parts: Vec<&str> = server_cmd.split_whitespace().collect();
38        
39        if parts.is_empty() {
40            return Err(anyhow!("Invalid server command"));
41        }
42
43        let mut cmd = Command::new(parts[0]);
44        for arg in &parts[1..] {
45            cmd.arg(arg);
46        }
47
48        let mut child = cmd
49            .stdin(Stdio::piped())
50            .stdout(Stdio::piped())
51            .stderr(Stdio::piped())
52            .spawn()?;
53
54        let stdout = child.stdout.take().ok_or(anyhow!("Failed to capture stdout"))?;
55        let stdin = child.stdin.take().ok_or(anyhow!("Failed to capture stdin"))?;
56
57        let reader = BufReader::new(stdout);
58
59        let mut client = MCPClient {
60            child,
61            reader,
62            writer: stdin,
63            request_id: 1,
64        };
65
66        // perform async handshake
67        client.initialize().await?;
68        Ok(client)
69    }
70
71    async fn initialize(&mut self) -> Result<()> {
72        let request = json!({
73            "jsonrpc": "2.0",
74            "id": self.request_id,
75            "method": "initialize",
76            "params": {
77                "protocolVersion": "2024-11-05",
78                "capabilities": {},
79                "clientInfo": {
80                    "name": "tripwire",
81                    "version": "0.1.0"
82                }
83            }
84        });
85
86        self.request_id += 1;
87        self.send_request(&request).await?;
88        let response = self.read_response().await?;
89
90        if let Some(error) = response.get("error") {
91            return Err(anyhow!("Handshake failed: {}", error));
92        }
93
94        Ok(())
95    }
96    pub async fn call_tool(&mut self, tool_name: &str, args: Value) -> Result<Value> {
97            let request = json!({
98                "jsonrpc": "2.0",
99                "id": self.request_id,
100                "method": "tools/call",
101                "params": {
102                    "name": tool_name,
103                    "arguments": args
104                }
105            });
106    
107            self.request_id += 1;
108            self.send_request(&request).await?;
109            let response = self.read_response().await?;
110    
111            if let Some(error) = response.get("error") {
112                return Err(anyhow!("Tool call failed: {}", error));
113            }
114    
115            Ok(response)
116        }
117
118    pub async fn list_tools(&mut self) -> Result<Vec<Tool>> {
119        let request = json!({
120            "jsonrpc": "2.0",
121            "id": self.request_id,
122            "method": "tools/list",
123            "params": {}
124        });
125
126        self.request_id += 1;
127        self.send_request(&request).await?;
128        let response = self.read_response().await?;
129
130        if let Some(error) = response.get("error") {
131            return Err(anyhow!("Failed to list tools: {}", error));
132        }
133
134        let tools = response
135            .get("result")
136            .and_then(|r| r.get("tools"))
137            .and_then(|t| t.as_array())
138            .ok_or(anyhow!("Invalid tools response"))?;
139
140        let mut tool_list = Vec::new();
141        for tool in tools {
142            if let Ok(t) = serde_json::from_value::<Tool>(tool.clone()) {
143                tool_list.push(t);
144            }
145        }
146
147        Ok(tool_list)
148    }
149
150    async fn send_request(&mut self, request: &Value) -> Result<()> {
151        let json_str = serde_json::to_string(request)?;
152        self.writer.write_all(json_str.as_bytes()).await?;
153        self.writer.write_all(b"\n").await?;
154        self.writer.flush().await?;
155        Ok(())
156    }
157
158    async fn read_response(&mut self) -> Result<Value> {
159        let mut line = String::new();
160        let n = self.reader.read_line(&mut line).await?;
161
162        if n == 0 {
163            return Err(anyhow!("Connection closed"));
164        }
165
166        let response: Value = serde_json::from_str(line.trim_end())?;
167        Ok(response)
168    }
169
170    pub fn get_server_name(&self) -> String {
171        "test-server".to_string()
172    }
173
174    pub fn get_server_version(&self) -> String {
175        "0.1.0".to_string()
176    }
177}
trip_test/mcp.rs

trip_test/
mcp.rs
