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
11pub 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(¬if))
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}