use std::collections::HashMap;
use async_trait::async_trait;
use reqwest::Client;
use serde_json::Value;
use tokio::sync::Mutex;
use xz_mcp_core::{McpClient, McpError, McpTool, McpToolResult};
static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
pub struct HttpMcpClient {
url: String,
http: Client,
headers: HashMap<String, String>,
session_id: Mutex<Option<String>>,
connected: Mutex<bool>,
}
impl HttpMcpClient {
pub fn new(url: impl Into<String>, headers: HashMap<String, String>) -> Self {
Self {
url: url.into(),
http: Client::new(),
headers,
session_id: Mutex::new(None),
connected: Mutex::new(false),
}
}
fn next_id() -> u64 {
NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
}
fn rpc(method: &str, params: Value) -> Value {
serde_json::json!({
"jsonrpc": "2.0",
"id": Self::next_id(),
"method": method,
"params": params,
})
}
fn notification(method: &str, params: Value) -> Value {
serde_json::json!({
"jsonrpc": "2.0",
"method": method,
"params": params,
})
}
async fn apply_headers(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
let mut b = builder
.header("Accept", "application/json, text/event-stream");
for (k, v) in &self.headers {
b = b.header(k.as_str(), v.as_str());
}
if let Some(sid) = self.session_id.lock().await.as_ref() {
b = b.header("Mcp-Session-Id", sid.as_str());
}
b
}
async fn capture_session_id(&self, resp: &reqwest::Response) {
if let Some(sid) = resp.headers().get("Mcp-Session-Id")
.or_else(|| resp.headers().get("mcp-session-id"))
{
if let Ok(v) = sid.to_str() {
*self.session_id.lock().await = Some(v.to_string());
}
}
}
fn parse_sse(body: &str, request_id: u64) -> Result<Value, McpError> {
let mut last_data: Option<&str> = None;
for block in body.split("\n\n") {
let block = block.trim();
if block.is_empty() || block.starts_with(':') {
continue;
}
let mut event_data: Option<&str> = None;
for line in block.lines() {
if let Some(data) = line.strip_prefix("data:") {
event_data = Some(data.trim());
}
}
if let Some(data) = event_data {
if data.is_empty() {
continue;
}
if let Ok(parsed) = serde_json::from_str::<Value>(data) {
if parsed.get("id").and_then(|i| i.as_u64()) == Some(request_id) {
return Ok(parsed);
}
}
last_data = Some(data);
}
}
if let Some(data) = last_data {
serde_json::from_str(data)
.map_err(|e| McpError::Protocol(format!("SSE JSON parse: {e}. data: {:.200}", data)))
} else {
Err(McpError::Protocol(format!(
"SSE stream contained no response for request id {request_id}. body: {:.200}",
body
)))
}
}
async fn send(&self, req: &Value) -> Result<Value, McpError> {
let request_id = req["id"].as_u64();
let builder = self.apply_headers(self.http.post(&self.url).json(req)).await;
let resp = builder
.send()
.await
.map_err(|e| McpError::Connection(format!("HTTP request failed: {e}")))?;
self.capture_session_id(&resp).await;
let status = resp.status();
let content_type = resp
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_lowercase();
let raw = resp.text().await
.map_err(|e| McpError::Protocol(format!("read body: {e}")))?;
let body: Value = if content_type.contains("text/event-stream") {
Self::parse_sse(&raw, request_id.unwrap_or(0))?
} else {
match serde_json::from_str(&raw) {
Ok(v) => v,
Err(_) if raw.contains("event:") || raw.contains("data:") => {
Self::parse_sse(&raw, request_id.unwrap_or(0))?
}
Err(e) => return Err(McpError::Protocol(format!(
"JSON parse: {e}. body ({status}): {:.200}", raw
))),
}
};
if !status.is_success() {
return Err(McpError::Server(format!("HTTP {status}: {body}")));
}
if let Some(err) = body.get("error") {
return Err(McpError::Server(err.to_string()));
}
Ok(body)
}
}
#[async_trait]
impl McpClient for HttpMcpClient {
async fn connect(&mut self) -> Result<(), McpError> {
let req = Self::rpc("initialize", serde_json::json!({
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": { "name": "xz-writer", "version": "1.0" }
}));
self.send(&req).await?;
let notif = Self::notification("notifications/initialized", serde_json::json!({}));
let resp = self.apply_headers(self.http.post(&self.url).json(¬if))
.await
.send()
.await
.map_err(|e| McpError::Connection(format!("initialized notification failed: {e}")))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(McpError::Server(format!(
"initialized notification rejected: HTTP {status}: {:.200}", body
)));
}
*self.connected.lock().await = true;
Ok(())
}
async fn list_tools(&self) -> Result<Vec<McpTool>, McpError> {
let req = Self::rpc("tools/list", serde_json::json!({}));
let resp = self.send(&req).await?;
let tools = resp["result"]["tools"].as_array()
.ok_or_else(|| McpError::Protocol("missing tools array".into()))?;
tools.iter().map(|t| Ok(McpTool {
name: t["name"].as_str().unwrap_or("").into(),
description: t["description"].as_str().unwrap_or("").into(),
input_schema: t.get("inputSchema").cloned().unwrap_or(serde_json::json!({})),
})).collect()
}
async fn call_tool(&self, name: &str, args: Value) -> Result<McpToolResult, McpError> {
let req = Self::rpc("tools/call", serde_json::json!({"name":name,"arguments":args}));
let resp = self.send(&req).await?;
let result = &resp["result"];
let content: Vec<Value> = result["content"].as_array().cloned().unwrap_or_default();
Ok(McpToolResult {
content: serde_json::from_value(Value::Array(content)).unwrap_or_default(),
is_error: result["isError"].as_bool().unwrap_or(false),
})
}
async fn is_alive(&self) -> bool {
*self.connected.lock().await
}
}