Skip to main content

xz_mcp_engine/
http.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use reqwest::Client;
5use serde_json::Value;
6use tokio::sync::Mutex;
7use xz_mcp_core::{McpClient, McpError, McpTool, McpToolResult};
8
9static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
10
11/// An MCP client that communicates with a remote server via Streamable HTTP.
12pub struct HttpMcpClient {
13    url: String,
14    http: Client,
15    headers: HashMap<String, String>,
16    session_id: Mutex<Option<String>>,
17    connected: Mutex<bool>,
18}
19
20impl HttpMcpClient {
21    pub fn new(url: impl Into<String>, headers: HashMap<String, String>) -> Self {
22        Self {
23            url: url.into(),
24            http: Client::new(),
25            headers,
26            session_id: Mutex::new(None),
27            connected: Mutex::new(false),
28        }
29    }
30
31    fn next_id() -> u64 {
32        NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
33    }
34
35    fn rpc(method: &str, params: Value) -> Value {
36        serde_json::json!({
37            "jsonrpc": "2.0",
38            "id": Self::next_id(),
39            "method": method,
40            "params": params,
41        })
42    }
43
44    fn notification(method: &str, params: Value) -> Value {
45        serde_json::json!({
46            "jsonrpc": "2.0",
47            "method": method,
48            "params": params,
49        })
50    }
51
52    async fn apply_headers(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
53        let mut b = builder
54            .header("Accept", "application/json, text/event-stream");
55        for (k, v) in &self.headers {
56            b = b.header(k.as_str(), v.as_str());
57        }
58        if let Some(sid) = self.session_id.lock().await.as_ref() {
59            b = b.header("Mcp-Session-Id", sid.as_str());
60        }
61        b
62    }
63
64    async fn capture_session_id(&self, resp: &reqwest::Response) {
65        if let Some(sid) = resp.headers().get("Mcp-Session-Id")
66            .or_else(|| resp.headers().get("mcp-session-id"))
67        {
68            if let Ok(v) = sid.to_str() {
69                *self.session_id.lock().await = Some(v.to_string());
70            }
71        }
72    }
73
74    fn parse_sse(body: &str, request_id: u64) -> Result<Value, McpError> {
75        let mut last_data: Option<&str> = None;
76
77        for block in body.split("\n\n") {
78            let block = block.trim();
79            if block.is_empty() || block.starts_with(':') {
80                continue;
81            }
82
83            let mut event_data: Option<&str> = None;
84            for line in block.lines() {
85                if let Some(data) = line.strip_prefix("data:") {
86                    event_data = Some(data.trim());
87                }
88            }
89
90            if let Some(data) = event_data {
91                if data.is_empty() {
92                    continue;
93                }
94                if let Ok(parsed) = serde_json::from_str::<Value>(data) {
95                    if parsed.get("id").and_then(|i| i.as_u64()) == Some(request_id) {
96                        return Ok(parsed);
97                    }
98                }
99                last_data = Some(data);
100            }
101        }
102
103        if let Some(data) = last_data {
104            serde_json::from_str(data)
105                .map_err(|e| McpError::Protocol(format!("SSE JSON parse: {e}. data: {:.200}", data)))
106        } else {
107            Err(McpError::Protocol(format!(
108                "SSE stream contained no response for request id {request_id}. body: {:.200}",
109                body
110            )))
111        }
112    }
113
114    async fn send(&self, req: &Value) -> Result<Value, McpError> {
115        let request_id = req["id"].as_u64();
116
117        let builder = self.apply_headers(self.http.post(&self.url).json(req)).await;
118        let resp = builder
119            .send()
120            .await
121            .map_err(|e| McpError::Connection(format!("HTTP request failed: {e}")))?;
122
123        self.capture_session_id(&resp).await;
124
125        let status = resp.status();
126        let content_type = resp
127            .headers()
128            .get("content-type")
129            .and_then(|v| v.to_str().ok())
130            .unwrap_or("")
131            .to_lowercase();
132
133        let raw = resp.text().await
134            .map_err(|e| McpError::Protocol(format!("read body: {e}")))?;
135
136        let body: Value = if content_type.contains("text/event-stream") {
137            Self::parse_sse(&raw, request_id.unwrap_or(0))?
138        } else {
139            match serde_json::from_str(&raw) {
140                Ok(v) => v,
141                Err(_) if raw.contains("event:") || raw.contains("data:") => {
142                    Self::parse_sse(&raw, request_id.unwrap_or(0))?
143                }
144                Err(e) => return Err(McpError::Protocol(format!(
145                    "JSON parse: {e}. body ({status}): {:.200}", raw
146                ))),
147            }
148        };
149
150        if !status.is_success() {
151            return Err(McpError::Server(format!("HTTP {status}: {body}")));
152        }
153
154        if let Some(err) = body.get("error") {
155            return Err(McpError::Server(err.to_string()));
156        }
157
158        Ok(body)
159    }
160}
161
162#[async_trait]
163impl McpClient for HttpMcpClient {
164    async fn connect(&mut self) -> Result<(), McpError> {
165        let req = Self::rpc("initialize", serde_json::json!({
166            "protocolVersion": "2024-11-05",
167            "capabilities": {},
168            "clientInfo": { "name": "xz-writer", "version": "1.0" }
169        }));
170        self.send(&req).await?;
171
172        let notif = Self::notification("notifications/initialized", serde_json::json!({}));
173        let resp = self.apply_headers(self.http.post(&self.url).json(&notif))
174            .await
175            .send()
176            .await
177            .map_err(|e| McpError::Connection(format!("initialized notification failed: {e}")))?;
178
179        let status = resp.status();
180        if !status.is_success() {
181            let body = resp.text().await.unwrap_or_default();
182            return Err(McpError::Server(format!(
183                "initialized notification rejected: HTTP {status}: {:.200}", body
184            )));
185        }
186
187        *self.connected.lock().await = true;
188        Ok(())
189    }
190
191    async fn list_tools(&self) -> Result<Vec<McpTool>, McpError> {
192        let req = Self::rpc("tools/list", serde_json::json!({}));
193        let resp = self.send(&req).await?;
194        let tools = resp["result"]["tools"].as_array()
195            .ok_or_else(|| McpError::Protocol("missing tools array".into()))?;
196        tools.iter().map(|t| Ok(McpTool {
197            name: t["name"].as_str().unwrap_or("").into(),
198            description: t["description"].as_str().unwrap_or("").into(),
199            input_schema: t.get("inputSchema").cloned().unwrap_or(serde_json::json!({})),
200        })).collect()
201    }
202
203    async fn call_tool(&self, name: &str, args: Value) -> Result<McpToolResult, McpError> {
204        let req = Self::rpc("tools/call", serde_json::json!({"name":name,"arguments":args}));
205        let resp = self.send(&req).await?;
206        let result = &resp["result"];
207        let content: Vec<Value> = result["content"].as_array().cloned().unwrap_or_default();
208        Ok(McpToolResult {
209            content: serde_json::from_value(Value::Array(content)).unwrap_or_default(),
210            is_error: result["isError"].as_bool().unwrap_or(false),
211        })
212    }
213
214    async fn is_alive(&self) -> bool {
215        *self.connected.lock().await
216    }
217}