slave 0.13.1

//! A small client for the project's MCP backend (Streamable HTTP).
//!
//! The backend is stateless, so each tool call is one JSON-RPC `tools/call`
//! POST. The signed-in user's bearer token is passed per call. In the single
//! runtime the POST goes straight to the backend's router in the same process
//! (no socket); otherwise over HTTP.

use std::fmt;

use axum::body::Body;
use axum::http::{Request, header};
use serde::de::DeserializeOwned;
use serde_json::{Value, json};
use tower::ServiceExt;

/// Largest MCP response read, in bytes.
const MAX_RESPONSE: usize = 16 * 1024 * 1024;

#[derive(Clone)]
pub struct McpClient {
    transport: Transport,
}

#[derive(Clone)]
enum Transport {
    Http {
        http: reqwest::Client,
        url: String,
    },
    /// The MCP backend's router, in this process.
    Local(axum::Router),
}

/// A failed tool call.
#[derive(Debug, Clone)]
pub struct ToolError {
    /// The backend's error code (`unauthenticated`, `forbidden`, `not_found`,
    /// ...), or `unavailable` when the backend could not be reached.
    pub code: String,
    pub message: String,
}

impl fmt::Display for ToolError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.message)
    }
}

impl ToolError {
    pub(crate) fn new(code: &str, message: impl Into<String>) -> Self {
        Self {
            code: code.to_owned(),
            message: message.into(),
        }
    }
}

impl McpClient {
    /// A client for the backend at `url` (`http://…/mcp`).
    pub fn new(url: String) -> Self {
        let http = reqwest::Client::new();
        Self {
            transport: Transport::Http { http, url },
        }
    }

    /// A client for `backend`, the MCP backend's router in this process.
    pub fn local(backend: axum::Router) -> Self {
        Self {
            transport: Transport::Local(backend),
        }
    }

    /// Where the backend is, for logs.
    pub fn url(&self) -> &str {
        match &self.transport {
            Transport::Http { url, .. } => url,
            Transport::Local(_) => "(in-process)/mcp",
        }
    }

    /// Calls `tool` with `arguments`, as the user owning `token` if given, and
    /// decodes its structured result.
    pub async fn call<T: DeserializeOwned>(
        &self,
        tool: &str,
        arguments: Value,
        token: Option<&str>,
    ) -> Result<T, ToolError> {
        let message = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {"name": tool, "arguments": arguments},
        });
        let response = match &self.transport {
            Transport::Http { http, url } => {
                let mut request = http
                    .post(url)
                    .header("accept", "application/json, text/event-stream")
                    .json(&message);
                if let Some(token) = token {
                    request = request.bearer_auth(token);
                }
                let unavailable = |err: reqwest::Error| {
                    tracing::warn!("MCP backend at {url}: {err}");
                    ToolError::new("unavailable", "the backend is not reachable right now")
                };
                request
                    .send()
                    .await
                    .and_then(reqwest::Response::error_for_status)
                    .map_err(unavailable)?
                    .json::<Value>()
                    .await
                    .map_err(unavailable)?
            }
            Transport::Local(backend) => call_local(backend, &message, token).await?,
        };

        if let Some(error) = response.get("error") {
            let message = error["message"].as_str().unwrap_or("request failed");
            return Err(ToolError::new("protocol_error", message));
        }
        let result = &response["result"];
        let content = result["structuredContent"].clone();
        if result["isError"] == true {
            let code = content["code"].as_str().unwrap_or("tool_error");
            let message = content["error"]
                .as_str()
                .or_else(|| result["content"][0]["text"].as_str())
                .unwrap_or("the request failed");
            return Err(ToolError::new(code, message));
        }
        serde_json::from_value(content)
            .map_err(|err| ToolError::new("protocol_error", format!("unexpected result: {err}")))
    }
}

/// One JSON-RPC POST to the in-process backend router, as over HTTP.
async fn call_local(
    backend: &axum::Router,
    message: &Value,
    token: Option<&str>,
) -> Result<Value, ToolError> {
    let failed = |err: String| {
        tracing::warn!("MCP backend (in-process): {err}");
        ToolError::new("unavailable", "the backend is not reachable right now")
    };
    // `localhost`: the backend only accepts its own host names.
    let mut request = Request::post("/mcp")
        .header(header::HOST, "localhost")
        .header(header::CONTENT_TYPE, "application/json")
        .header(header::ACCEPT, "application/json, text/event-stream");
    if let Some(token) = token {
        request = request.header(header::AUTHORIZATION, format!("Bearer {token}"));
    }
    let request = request
        .body(Body::from(message.to_string()))
        .map_err(|err| failed(err.to_string()))?;
    let response = match backend.clone().oneshot(request).await {
        Ok(response) => response,
        Err(never) => match never {},
    };
    if !response.status().is_success() {
        return Err(failed(format!("status {}", response.status())));
    }
    let body = axum::body::to_bytes(response.into_body(), MAX_RESPONSE)
        .await
        .map_err(|err| failed(err.to_string()))?;
    serde_json::from_slice(&body).map_err(|err| failed(err.to_string()))
}