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, HashSet},
    convert::Infallible,
    net::SocketAddr,
    sync::{Arc, OnceLock, RwLock},
};
use tokio::sync::{mpsc, Mutex};
use tokio_stream::{wrappers::ReceiverStream, StreamExt as _};

use crate::db;

struct SessionEntry {
    sender: mpsc::Sender<(String, String)>,
    project_scope: Option<String>,
}

/// SSE 会话表：session_id → 发送端与连接作用域
type Sessions = Arc<Mutex<HashMap<String, SessionEntry>>>;

const MCP_LOOPBACK_HOST: &str = "127.0.0.1";
const MCP_SERVER_INSTRUCTIONS: &str = "先读取 get_diagnostics、verify_onboarding 或项目状态，再执行后续操作。项目级工具必须使用当前项目的 canonical project_id，不得跨项目猜测或改写身份。写操作仅在用户明确确认目标与影响后调用；状态 unavailable/unknown 不得当作成功。";
static INITIALIZED_PROJECT_CLIENTS: OnceLock<RwLock<HashSet<String>>> = OnceLock::new();

fn initialized_project_clients() -> &'static RwLock<HashSet<String>> {
    INITIALIZED_PROJECT_CLIENTS.get_or_init(|| RwLock::new(HashSet::new()))
}

fn record_project_client_initialize(req: &RpcRequest, project_scope: Option<&str>) {
    let Some(project_id) = project_scope else {
        return;
    };
    let client_name = req
        .params
        .as_ref()
        .and_then(|params| params.get("clientInfo"))
        .and_then(|info| info.get("name"))
        .and_then(Value::as_str)
        .unwrap_or_default();
    if client_name == "lian-jing-readiness" {
        return;
    }
    if let Ok(mut clients) = initialized_project_clients().write() {
        clients.insert(project_id.to_string());
    }
}

pub(crate) fn has_initialized_project_client(project_id: &str) -> bool {
    initialized_project_clients()
        .read()
        .map(|clients| clients.contains(project_id))
        .unwrap_or(false)
}

/// MCP 仅监听 IPv4 loopback，避免把本机 agent 控制面暴露到局域网。
fn mcp_bind_addr(port: u16) -> SocketAddr {
    SocketAddr::from(([127, 0, 0, 1], port))
}

/// 所有 MCP 对外 URL 的单一基址事实。
fn mcp_base_url(port: u16) -> String {
    format!("http://{MCP_LOOPBACK_HOST}:{port}")
}

pub(crate) fn mcp_health_url(port: u16) -> String {
    format!("{}/health", mcp_base_url(port))
}

fn sse_messages_url(port: u16, session_id: &str) -> String {
    format!(
        "{}/mcp/messages?session_id={session_id}",
        mcp_base_url(port)
    )
}

pub(crate) fn scoped_mcp_url(port: u16, project_id: &str) -> String {
    format!(
        "{}/mcp/projects/{}",
        mcp_base_url(port),
        urlencoding::encode(project_id)
    )
}

fn scoped_sse_messages_url(port: u16, project_id: &str, session_id: &str) -> String {
    format!(
        "{}/messages?session_id={session_id}",
        scoped_mcp_url(port, project_id)
    )
}

/// 从 settings 表读取 MCP 端口，默认 27190
pub(crate) fn read_port_from_conn(conn: &rusqlite::Connection) -> u16 {
    conn.query_row(
        "SELECT value FROM settings WHERE key = 'mcp_port'",
        [],
        |row| row.get::<_, String>(0),
    )
    .ok()
    .and_then(|value| value.parse().ok())
    .unwrap_or(27190)
}

pub fn read_port() -> u16 {
    db::pool()
        .get()
        .ok()
        .map(|conn| read_port_from_conn(&conn))
        .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 = mcp_router(sessions);

    let addr = mcp_bind_addr(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,
        ),
    }
}

fn mcp_router(sessions: Sessions) -> Router {
    Router::new()
        .route("/health", get(health))
        // 单一全局服务器路由（lian-jing）：寻址不依赖 group_id，工具显式传 project_id
        .route("/mcp", post(handle_mcp).get(handle_sse_connect))
        .route(
            "/mcp/projects/:project_id",
            post(handle_scoped_mcp).get(handle_scoped_sse_connect),
        )
        // SSE 消息端点：客户端通过此 URL 发送 JSON-RPC，响应经 SSE 推回
        .route("/mcp/messages", post(handle_sse_message))
        .route(
            "/mcp/projects/:project_id/messages",
            post(handle_scoped_sse_message),
        )
        // Gitea webhook：接收 push/pull_request 事件写入飞轮
        .route(
            "/webhook/gitea/:project_id",
            post(crate::mcp_webhook::handle_gitea_webhook),
        )
        .with_state(sessions)
}

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(req: &RpcRequest, project_scope: Option<&str>) -> Result<Value, String> {
    if req.method == "initialize" {
        record_project_client_initialize(req, project_scope);
    }
    match req.method.as_str() {
        "initialize" => Ok(json!({
            "protocolVersion": "2024-11-05",
            "capabilities": { "tools": {} },
            "serverInfo": { "name": "炼境", "version": env!("CARGO_PKG_VERSION") },
            "instructions": MCP_SERVER_INSTRUCTIONS,
            "lianJing": match project_scope {
                Some(project_id) => json!({
                    "scope_mode": "project",
                    "project_id": project_id
                }),
                None => json!({
                    "scope_mode": "legacy",
                    "upgrade_hint": "重跑 inject_mcp，将项目配置升级到 /mcp/projects/:project_id"
                })
            }
        })),
        "tools/list" => Ok(match project_scope {
            Some(project_id) => crate::mcp_tools::tools_list_result_for_project(project_id),
            None => crate::mcp_tools::tools_list_result(),
        }),
        "tools/call" => Ok(match project_scope {
            Some(project_id) => {
                crate::mcp_tools::tools_call_for_project(req.params.as_ref(), project_id).await
            }
            None => crate::mcp_tools::tools_call(req.params.as_ref()).await,
        }),
        _ => Err(format!("Method not found: {}", req.method)),
    }
}

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

async fn handle_mcp(Json(req): Json<RpcRequest>) -> Response {
    if req.id.is_none() {
        return StatusCode::ACCEPTED.into_response();
    }
    let id = req.id.clone();
    match process_rpc(&req, None).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(),
    }
}

async fn handle_scoped_mcp(
    Path(project_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(&req, Some(&project_id)).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 — Claude Code SSE 握手，首条事件告知后续 POST URL
async fn handle_sse_connect(
    State(sessions): State<Sessions>,
) -> 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 = sse_messages_url(port, &session_id);
    let _ = tx.send(("endpoint".to_string(), endpoint_uri)).await;
    sessions.lock().await.insert(
        session_id,
        SessionEntry {
            sender: tx,
            project_scope: None,
        },
    );

    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())
}

async fn handle_scoped_sse_connect(
    Path(project_id): Path<String>,
    State(sessions): State<Sessions>,
) -> 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 endpoint_uri = scoped_sse_messages_url(read_port(), &project_id, &session_id);
    let _ = tx.send(("endpoint".to_string(), endpoint_uri)).await;
    sessions.lock().await.insert(
        session_id,
        SessionEntry {
            sender: tx,
            project_scope: Some(project_id),
        },
    );
    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())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn rpc(method: &str) -> RpcRequest {
        RpcRequest {
            jsonrpc: "2.0".to_string(),
            method: method.to_string(),
            id: Some(json!(1)),
            params: None,
        }
    }

    /// 单一 lian-jing 服务器的握手路径：initialize / tools/list 必须可用
    #[tokio::test]
    async fn process_rpc_handshake() {
        let init = process_rpc(&rpc("initialize"), None).await.unwrap();
        assert_eq!(init["serverInfo"]["name"], "炼境");
        let instructions = init["instructions"].as_str().unwrap();
        assert!(instructions.len() <= 512);
        assert!(instructions.contains("先读取"));
        assert!(instructions.contains("项目"));
        assert!(instructions.contains("用户明确确认"));
        assert_eq!(init["lianJing"]["scope_mode"], "legacy");
        assert!(init["lianJing"]["upgrade_hint"].is_string());

        let tools = process_rpc(&rpc("tools/list"), None).await.unwrap();
        let list = tools["tools"].as_array().unwrap();
        assert!(
            list.len() >= 20,
            "工具清单不应因去组化而缺失，实际 {}",
            list.len()
        );
        assert!(
            list.iter()
                .any(|tool| tool["name"] == "get_delivery_status"),
            "MCP 握手后的工具清单应暴露 get_delivery_status"
        );
    }

    #[tokio::test]
    async fn scoped_handshake_binds_project_and_filters_global_tools() {
        let init = process_rpc(&rpc("initialize"), Some("p1")).await.unwrap();
        assert_eq!(init["lianJing"]["scope_mode"], "project");
        assert_eq!(init["lianJing"]["project_id"], "p1");

        let tools = process_rpc(&rpc("tools/list"), Some("p1")).await.unwrap();
        let list = tools["tools"].as_array().unwrap();
        assert!(list
            .iter()
            .any(|tool| tool["name"] == "get_delivery_status"));
        assert!(!list.iter().any(|tool| tool["name"] == "unregister_project"));
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn scoped_enterprise_fixture_reads_non_secret_resource_status_over_mcp() {
        db::init_test_pool();
        let listener = tokio::net::TcpListener::bind(mcp_bind_addr(0))
            .await
            .unwrap();
        let port = listener.local_addr().unwrap().port();
        {
            let conn = db::pool().get().unwrap();
            conn.execute("INSERT OR IGNORE INTO project_profiles (id, name) VALUES ('enterprise-fixture-profile', 'Enterprise fixture')", []).unwrap();
            conn.execute("INSERT OR IGNORE INTO project_workspaces (id, profile_id) VALUES ('enterprise-fixture', 'enterprise-fixture-profile')", []).unwrap();
            conn.execute(
                "UPDATE settings SET value=?1 WHERE key='mcp_port'",
                [port.to_string()],
            )
            .unwrap();
        }
        let sessions: Sessions = Arc::new(Mutex::new(HashMap::new()));
        let server = tokio::spawn(async move {
            axum::serve(listener, mcp_router(sessions)).await.unwrap();
        });
        let url = format!("http://127.0.0.1:{port}/mcp/projects/enterprise-fixture");
        let call = |method: &str, params: Value| -> Value {
            ureq::post(&url)
                .send_json(json!({"jsonrpc":"2.0","id":1,"method":method,"params":params}))
                .unwrap()
                .into_json()
                .unwrap()
        };
        let init = call("initialize", json!({}));
        assert_eq!(
            init["result"]["lianJing"]["project_id"],
            "enterprise-fixture"
        );
        let catalog = call("tools/list", json!({}));
        assert!(catalog["result"]["tools"]
            .as_array()
            .unwrap()
            .iter()
            .any(|tool| tool["name"] == "get_project_resource_status"));
        let result = call(
            "tools/call",
            json!({"name":"get_project_resource_status","arguments":{}}),
        );
        assert_eq!(
            result["result"]["structuredContent"]["project_id"],
            "enterprise-fixture"
        );
        assert_eq!(
            result["result"]["structuredContent"]["gitea_pr"]["state"],
            "not_configured"
        );
        assert!(!result.to_string().contains("access_token"));
        let denied = call(
            "tools/call",
            json!({"name":"get_project_resource_status","arguments":{"project_id":"other-project"}}),
        );
        assert_eq!(
            denied["result"]["structuredContent"]["code"],
            "scope_violation"
        );

        use std::io::BufRead;
        let sse = ureq::get(&url).call().unwrap();
        assert!(sse
            .header("Content-Type")
            .unwrap_or("")
            .contains("text/event-stream"));
        let mut line = String::new();
        std::io::BufReader::new(sse.into_reader())
            .read_line(&mut line)
            .unwrap();
        assert!(line.contains("endpoint"));
        server.abort();
    }

    #[tokio::test]
    async fn readiness_probe_does_not_satisfy_client_reload_evidence() {
        let project_id = "reload-evidence-fixture";
        let mut readiness = rpc("initialize");
        readiness.params = Some(json!({
            "clientInfo": { "name": "lian-jing-readiness", "version": "test" }
        }));
        process_rpc(&readiness, Some(project_id)).await.unwrap();
        assert!(!has_initialized_project_client(project_id));

        let mut codex = rpc("initialize");
        codex.params = Some(json!({
            "clientInfo": { "name": "codex", "version": "test" }
        }));
        process_rpc(&codex, Some(project_id)).await.unwrap();
        assert!(has_initialized_project_client(project_id));
    }

    #[test]
    fn loopback_urls_share_one_ipv4_base() {
        let port = 27190;
        let session_id = "test-session";

        assert_eq!(mcp_bind_addr(port).to_string(), "127.0.0.1:27190");
        assert_eq!(mcp_health_url(port), "http://127.0.0.1:27190/health");
        assert_eq!(
            scoped_mcp_url(port, "project one"),
            "http://127.0.0.1:27190/mcp/projects/project%20one"
        );
        assert_eq!(
            sse_messages_url(port, session_id),
            "http://127.0.0.1:27190/mcp/messages?session_id=test-session"
        );
    }
}

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

    let session_id = query.session_id.clone();
    tokio::spawn(async move {
        let id = req.id.clone();
        let response_json = match process_rpc(&req, None).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.lock().await;
        if let Some(entry) = guard.get(&session_id) {
            if entry.project_scope.is_some() {
                return;
            }
            if entry
                .sender
                .send(("message".to_string(), response_json.to_string()))
                .await
                .is_err()
            {
                guard.remove(&session_id);
            }
        }
    });

    StatusCode::ACCEPTED.into_response()
}

async fn handle_scoped_sse_message(
    Path(project_id): Path<String>,
    State(sessions): State<Sessions>,
    Query(query): Query<SessionQuery>,
    Json(req): Json<RpcRequest>,
) -> Response {
    if req.id.is_none() {
        return StatusCode::ACCEPTED.into_response();
    }
    let session_id = query.session_id.clone();
    let scope_matches = sessions
        .lock()
        .await
        .get(&session_id)
        .and_then(|entry| entry.project_scope.as_deref())
        == Some(project_id.as_str());
    if !scope_matches {
        return StatusCode::FORBIDDEN.into_response();
    }
    tokio::spawn(async move {
        let id = req.id.clone();
        let response_json = match process_rpc(&req, Some(&project_id)).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.lock().await;
        if let Some(entry) = guard.get(&session_id) {
            if entry
                .sender
                .send(("message".to_string(), response_json.to_string()))
                .await
                .is_err()
            {
                guard.remove(&session_id);
            }
        }
    });
    StatusCode::ACCEPTED.into_response()
}