xtap-core-lib 0.5.0

/*
 * mcp/mod.rs
 * Project: core_lib
 * Description: 通用 MCP HTTP 服务框架（AI Shell 的“给 AI 用”这张脸）
 *
 * 设计（路由挂载模式抄 sts-x src/server/mod.rs:79-93）：
 * - McpServer 负责自动生成 /health、/tools、/ 三个“壳层”端点（工具发现 + 自解释）
 * - app 只需 .tool(...) 声明工具 Schema，再 .merge(自有业务 Router) 挂载 /search 等 handler
 * - 重新导出 axum / tokio，下游 crate 无需再直接依赖 HTTP 框架（去重，见白皮书 §2）
 * - 全部在 `mcp` feature 下编译；不开该 feature 时 axum/tokio 完全不进依赖图
 */

pub mod instructions;

// 重新导出，供下游 crate（如 sts-x）直接使用，避免各自再声明 axum/tokio 依赖。
pub use axum;
pub use tokio;

use axum::{routing::get, Json, Router};
use std::net::SocketAddr;
use std::sync::Arc;

/// 一个 MCP 工具的自描述（用于 GET /tools 工具发现）。
///
/// `input_schema` 使用 JSON Schema（object/properties/required），格式对齐
/// sts-x src/server/mod.rs:271-354 的 handle_tools 输出。
#[derive(Clone, serde::Serialize)]
pub struct Tool {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
}

impl Tool {
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        input_schema: serde_json::Value,
    ) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            input_schema,
        }
    }
}

/// 通用 MCP 服务框架。
///
/// 典型用法（app 侧）：
/// ```ignore
/// let biz = Router::new()
///     .route("/search", post(handle_search))
///     .route("/file", post(handle_file))
///     .with_state(state);            // 先把 state 解析成 Router<()>
///
/// McpServer::new("sts-x", env!("CARGO_PKG_VERSION"), HINT)
///     .tool(Tool::new("search", "...", schema))
///     .tool(Tool::new("file",   "...", schema))
///     .merge(biz)                     // 挂载业务路由
///     .serve("127.0.0.1", 9876)
///     .await?;
/// ```
pub struct McpServer {
    service: String,
    version: String,
    instructions: String,
    tools: Vec<Tool>,
    extra: Router,
}

impl McpServer {
    pub fn new(
        service: impl Into<String>,
        version: impl Into<String>,
        instructions: impl Into<String>,
    ) -> Self {
        Self {
            service: service.into(),
            version: version.into(),
            instructions: instructions.into(),
            tools: Vec::new(),
            extra: Router::new(),
        }
    }

    /// 声明一个工具（追加到 /tools 工具发现列表）。
    pub fn tool(mut self, tool: Tool) -> Self {
        self.tools.push(tool);
        self
    }

    /// 批量声明工具。
    pub fn tools(mut self, tools: impl IntoIterator<Item = Tool>) -> Self {
        self.tools.extend(tools);
        self
    }

    /// 合并 app 自有业务路由（如 /search /file）。
    ///
    /// 传入的 `Router` 需已 `.with_state(...)` 解析为 `Router<()>`。
    pub fn merge(mut self, router: Router) -> Self {
        self.extra = self.extra.merge(router);
        self
    }

    /// 组装成最终 axum Router（自动挂载 /health、/tools、/ 后再合并业务路由）。
    pub fn into_router(self) -> Router {
        let health = Arc::new(serde_json::json!({
            "status": "ok",
            "version": self.version,
            "service": self.service,
        }));
        let tools_val = Arc::new(serde_json::json!({ "tools": self.tools }));
        let tool_names: Vec<&str> = self.tools.iter().map(|t| t.name.as_str()).collect();
        let root = Arc::new(serde_json::json!({
            "service": self.service,
            "version": self.version,
            "_ai_instructions": self.instructions,
            "tools": tool_names,
            "endpoints": {
                "/health": "GET — service health",
                "/tools": "GET — tool discovery (MCP schema)",
            },
        }));

        Router::new()
            .route(
                "/health",
                get({
                    let v = health.clone();
                    move || {
                        let v = v.clone();
                        async move { Json((*v).clone()) }
                    }
                }),
            )
            .route(
                "/tools",
                get({
                    let v = tools_val.clone();
                    move || {
                        let v = v.clone();
                        async move { Json((*v).clone()) }
                    }
                }),
            )
            .route(
                "/",
                get({
                    let v = root.clone();
                    move || {
                        let v = v.clone();
                        async move { Json((*v).clone()) }
                    }
                }),
            )
            .merge(self.extra)
    }

    /// 启动 HTTP 服务（调用方需已在 tokio runtime 内，如 `#[tokio::main]`）。
    pub async fn serve(self, host: &str, port: u16) -> anyhow::Result<()> {
        let app = self.into_router();
        let addr: SocketAddr = format!("{host}:{port}").parse()?;
        tracing::info!("MCP server listening on http://{}", addr);
        let listener = tokio::net::TcpListener::bind(addr).await?;
        axum::serve(listener, app).await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request;
    use http_body_util::BodyExt;
    use tower::ServiceExt;

    #[test]
    fn test_tool_new() {
        let tool = Tool::new("search", "全库搜索", serde_json::json!({"type": "object"}));
        assert_eq!(tool.name, "search");
        assert_eq!(tool.description, "全库搜索");
        assert_eq!(tool.input_schema["type"], "object");
    }

    #[tokio::test]
    async fn test_into_router_tools_endpoint() {
        let app = McpServer::new("test-svc", "0.1.0", "hint")
            .tool(Tool::new(
                "search",
                "全库搜索",
                serde_json::json!({"type": "object"}),
            ))
            .into_router();

        let resp = app
            .oneshot(
                Request::builder()
                    .uri("/tools")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .expect("oneshot 请求应成功");

        assert_eq!(resp.status(), axum::http::StatusCode::OK);
        let bytes = resp.into_body().collect().await.unwrap().to_bytes();
        let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
        let tools = json["tools"].as_array().expect("tools 应为数组");
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0]["name"], "search");
    }
}