dev-manager-tauri/src-tauri/src/mcp_server.rs

use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::{
        sse::{Event, KeepAlive, Sse},
        IntoResponse, Response,
    },
    routing::{get, post},
    Json, Router,
};
use serde::Deserialize;
use serde_json::{json, Value};
use std::{collections::HashMap, convert::Infallible, net::SocketAddr, sync::Arc};
use tokio::sync::{mpsc, Mutex};
use tokio_stream::{wrappers::ReceiverStream, StreamExt as _};

use crate::db;

/// SSE 会话表：session_id → 发送端
type Sessions = Arc<Mutex<HashMap<String, mpsc::Sender<(String, String)>>>>;

/// 从 settings 表读取 MCP 端口，默认 27190
pub fn read_port() -> u16 {
    db::pool()
        .get()
        .ok()
        .and_then(|c| {
            c.query_row(
                "SELECT value FROM settings WHERE key = 'mcp_port'",
                [],
                |r| r.get::<_, String>(0),
            )
            .ok()
        })
        .and_then(|v| v.parse().ok())
        .unwrap_or(27190)
}

/// 启动 MCP HTTP 服务（同时支持 Streamable HTTP 和 SSE 两种传输）
pub async fn start(port: u16) {
    let sessions: Sessions = Arc::new(Mutex::new(HashMap::new()));

    let app = Router::new()
        .route("/health", get(health))
        // Streamable HTTP（POST）— 向后兼容，Python 直调仍可用
        .route("/mcp/group/:group_id", post(handle_mcp).get(handle_sse_connect))
        // SSE 消息端点：客户端通过此 URL 发送 JSON-RPC，响应经 SSE 推回
        .route("/mcp/group/:group_id/messages", post(handle_sse_message))
        // Gitea webhook：接收 push/pull_request 事件写入飞轮
        .route("/webhook/gitea/:project_id", post(crate::mcp_webhook::handle_gitea_webhook))
        .with_state(sessions);

    let addr = SocketAddr::from(([0, 0, 0, 0], port));
    match tokio::net::TcpListener::bind(addr).await {
        Ok(listener) => {
            db::log_event("mcp_server", "info", &format!("MCP Server 已启动，端口 {port}（HTTP + SSE）"), None);
            if let Err(e) = axum::serve(listener, app).await {
                db::log_event("mcp_server", "error", &format!("Server 异常退出: {e}"), None);
            }
        }
        Err(e) => db::log_event("mcp_server", "error", &format!("绑定端口 {port} 失败: {e}"), None),
    }
}

async fn health() -> &'static str {
    "炼境 MCP Server running"
}

// ── JSON-RPC 类型 ─────────────────────────────────────────────────────────────

#[derive(Deserialize)]
struct RpcRequest {
    #[allow(dead_code)]
    jsonrpc: String,
    method: String,
    id: Option<Value>,
    #[serde(default)]
    params: Option<Value>,
}

#[derive(Deserialize)]
struct SessionQuery {
    session_id: String,
}

// ── 共享 RPC 处理逻辑 ─────────────────────────────────────────────────────────

async fn process_rpc(group_id: &str, req: &RpcRequest) -> Result<Value, String> {
    match req.method.as_str() {
        "initialize" => Ok(json!({
            "protocolVersion": "2024-11-05",
            "capabilities": { "tools": {} },
            "serverInfo": { "name": "炼境", "version": env!("CARGO_PKG_VERSION") }
        })),
        "tools/list" => Ok(crate::mcp_tools::tools_list_result()),
        "tools/call" => Ok(crate::mcp_tools::tools_call(group_id, req.params.as_ref()).await),
        _ => Err(format!("Method not found: {}", req.method)),
    }
}

// ── Streamable HTTP（POST）─────────────────────────────────────────────────────

async fn handle_mcp(Path(group_id): Path<String>, Json(req): Json<RpcRequest>) -> Response {
    if req.id.is_none() {
        return StatusCode::ACCEPTED.into_response();
    }
    let id = req.id.clone();
    match process_rpc(&group_id, &req).await {
        Ok(result) => Json(json!({ "jsonrpc": "2.0", "id": id, "result": result })).into_response(),
        Err(msg) => Json(json!({
            "jsonrpc": "2.0", "id": id,
            "error": { "code": -32601, "message": msg }
        }))
        .into_response(),
    }
}

// ── SSE 传输（GET + POST /messages）─────────────────────────────────────────

/// GET /mcp/group/:group_id — Claude Code SSE 握手，首条事件告知后续 POST URL
async fn handle_sse_connect(
    State(sessions): State<Sessions>,
    Path(group_id): Path<String>,
) -> Sse<impl futures_core::Stream<Item = Result<Event, Infallible>>> {
    let session_id = uuid::Uuid::new_v4().to_string();
    let (tx, rx) = mpsc::channel::<(String, String)>(32);

    let port = read_port();
    let endpoint_uri = format!(
        "http://127.0.0.1:{}/mcp/group/{}/messages?session_id={}",
        port, group_id, session_id
    );
    let _ = tx.send(("endpoint".to_string(), endpoint_uri)).await;
    sessions.lock().await.insert(session_id, tx);

    let stream = ReceiverStream::new(rx)
        .map(|(evt_name, data)| Ok::<Event, Infallible>(Event::default().event(evt_name).data(data)));

    Sse::new(stream).keep_alive(KeepAlive::default())
}

/// POST /mcp/group/:group_id/messages?session_id=xxx
/// 立即返回 202，后台异步处理工具调用并经 SSE channel 推回响应
async fn handle_sse_message(
    State(sessions): State<Sessions>,
    Path(group_id): Path<String>,
    Query(query): Query<SessionQuery>,
    Json(req): Json<RpcRequest>,
) -> Response {
    if req.id.is_none() {
        return StatusCode::ACCEPTED.into_response();
    }

    let sessions_clone = sessions.clone();
    let session_id = query.session_id.clone();
    tokio::spawn(async move {
        let id = req.id.clone();
        let response_json = match process_rpc(&group_id, &req).await {
            Ok(result) => json!({ "jsonrpc": "2.0", "id": id, "result": result }),
            Err(msg) => json!({
                "jsonrpc": "2.0", "id": id,
                "error": { "code": -32601, "message": msg }
            }),
        };
        let mut guard = sessions_clone.lock().await;
        if let Some(tx) = guard.get(&session_id) {
            if tx
                .send(("message".to_string(), response_json.to_string()))
                .await
                .is_err()
            {
                guard.remove(&session_id);
            }
        }
    });

    StatusCode::ACCEPTED.into_response()
}