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>>>; /// 从 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, #[serde(default)] params: Option, } #[derive(Deserialize)] struct SessionQuery { session_id: String, } // ── 共享 RPC 处理逻辑 ───────────────────────────────────────────────────────── async fn process_rpc(group_id: &str, req: &RpcRequest) -> Result { 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, Json(req): Json) -> 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, Path(group_id): Path, ) -> Sse>> { 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::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, Path(group_id): Path, Query(query): Query, Json(req): Json, ) -> 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() }