//! MCP server 主循环 — stdio JSON-RPC 2.0 //! //! 协议流程: //! 1. 客户端发 initialize → 回 serverInfo + capabilities.tools //! 2. 客户端发 notifications/initialized(通知,不回) //! 3. 客户端发 tools/list → 回工具清单(read-only 模式过滤) //! 4. 客户端发 tools/call → dispatch 到 handler //! //! 高风险(High)工具:tools/list 不暴露(从清单剔除),tools/call 即便绕过也由 handler 兜底拒绝。 //! read-only:tools/list 仅留 Low,tools/call Medium/High 一律拒绝。 use std::path::Path; use std::sync::Arc; use df_storage::db::Database; use serde_json::{json, Value}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use crate::protocol::{ CallToolResult, Capabilities, InitializeResult, McpMethod, Request, Response, ServerInfo, ToolsCapability, }; use crate::tools::{self, Ctx, RiskLevel}; /// 协议版本(MCP 2025-06-18) const PROTOCOL_VERSION: &str = "2025-06-18"; const SERVER_NAME: &str = "devflow-mcp"; const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION"); /// 启动 MCP server。 /// /// - `db_path`:SQLite 数据库文件路径(应用同库,开 WAL 模式) /// - `read_only`:true 则仅暴露 list/get 工具 /// /// 复用 [`Database::open`](df_storage::Database::open)(已含 `PRAGMA journal_mode=WAL`), /// 故 SQLite WAL 状态:随 df-storage 一起已启用,无需额外处理。 pub async fn run_server(db_path: &Path, read_only: bool) -> anyhow::Result<()> { let db = Arc::new(Database::open(db_path).await?); let ctx = Ctx::new(db); let stdin = tokio::io::stdin(); let stdout = tokio::io::stdout(); main_loop(stdin, stdout, &ctx, read_only).await } /// 可单测的主循环(参数化 stdin/stdout)。 /// /// 协议正确性:每行一个 JSON-RPC 消息,Response 单行写回(末尾 \n)。 /// Notification(id=None)不回响应。 pub async fn main_loop( stdin: R, stdout: W, ctx: &Ctx, read_only: bool, ) -> anyhow::Result<()> where R: tokio::io::AsyncRead + Unpin, W: tokio::io::AsyncWrite + Unpin, { let mut reader = BufReader::new(stdin); let mut writer = stdout; let mut line = String::new(); loop { line.clear(); let n = reader.read_line(&mut line).await?; if n == 0 { // EOF(stdin 关闭),优雅退出 break; } let trimmed = line.trim(); if trimmed.is_empty() { continue; } // 解析 JSON-RPC 请求 let req: Request = match serde_json::from_str(trimmed) { Ok(r) => r, Err(e) => { // 解析失败:无 id 时无法回响应,只能 log;有 id(尽力猜)回 PARSE_ERROR tracing::warn!(target: "df_mcp", line = %trimmed, err = %e, "解析 JSON-RPC 行失败"); let resp = Response::err(None, crate::protocol::PARSE_ERROR, "Parse error", None); write_response(&mut writer, &resp).await?; continue; } }; let is_notification = req.id.is_none(); let method = McpMethod::from_request(&req); // Notification 不回响应(除解析错外) if is_notification { match method { McpMethod::Initialized => { tracing::debug!(target: "df_mcp", "客户端 initialized 通知已收"); } _ => { tracing::debug!(target: "df_mcp", m = ?method, "忽略未识别通知"); } } continue; } let resp = dispatch(ctx, read_only, req.id.clone(), method).await; write_response(&mut writer, &resp).await?; } Ok(()) } /// 方法分发 → 构造 Response。 /// /// `id`:JSON-RPC 请求 id(回响应时原样回填;通知由 main_loop 已过滤)。 async fn dispatch(ctx: &Ctx, read_only: bool, id: Option, method: McpMethod) -> Response { match method { McpMethod::Initialize { .. } => { let result = InitializeResult { protocol_version: PROTOCOL_VERSION, capabilities: Capabilities { tools: ToolsCapability { list_changed: false }, }, server_info: ServerInfo { name: SERVER_NAME, version: SERVER_VERSION, }, }; Response::ok( id, serde_json::to_value(result).unwrap_or(Value::Null), ) } McpMethod::Initialized => { // 通知分支(由 main_loop 已过滤);到这里说明带 id 的 initialized,回空 result Response::ok(id, Value::Object(Default::default())) } McpMethod::ToolsList => { let tools: Vec<_> = tools::all_tools() .into_iter() .filter(|t| visible(read_only, t.risk)) .map(|t| serde_json::to_value(&t.tool).unwrap_or(Value::Null)) .collect(); Response::ok(id, json!({ "tools": tools })) } McpMethod::Ping => Response::ok(id, Value::Object(Default::default())), McpMethod::ToolsCall { name, arguments } => { // 查工具 let Some(spec) = tools::find(&name) else { return Response::ok( id, serde_json::to_value(CallToolResult::error(format!("未知工具: {name}"))) .unwrap_or(Value::Null), ); }; // read-only 模式:Medium/High 全拒 if read_only && spec.risk != RiskLevel::Low { let r = CallToolResult::error(format!( "只读模式拒绝执行 {name}(风险等级 {:?})", spec.risk )); return Response::ok(id, serde_json::to_value(r).unwrap_or(Value::Null)); } // 非 read-only:High 兜底拒绝(handler 内也会拒,双保险) if spec.risk == RiskLevel::High { let r = CallToolResult::error(format!( "High 风险操作 {name} 默认拒绝,请在 DevFlow 应用内执行。" )); return Response::ok(id, serde_json::to_value(r).unwrap_or(Value::Null)); } // Low / Medium:执行 let result = (spec.handler)(ctx, arguments).await; Response::ok(id, serde_json::to_value(result).unwrap_or(Value::Null)) } McpMethod::Unknown(m) => Response::err( id, crate::protocol::METHOD_NOT_FOUND, format!("Method not found: {m}"), None, ), } } /// 工具可见性:read-only 仅 Low,否则 Low + Medium(High 永不可见) fn visible(read_only: bool, risk: RiskLevel) -> bool { if read_only { risk == RiskLevel::Low } else { risk != RiskLevel::High } } /// 单行写 Response(末尾 \n,MCP 客户端按行 framing) async fn write_response( writer: &mut W, resp: &Response, ) -> anyhow::Result<()> { let mut buf = serde_json::to_string(resp)?; buf.push('\n'); writer.write_all(buf.as_bytes()).await?; writer.flush().await?; Ok(()) } // ============================================================ // 单测:协议 framing / initialize / tools/list / tools/call / High 拒绝 / read-only 过滤 // ============================================================ // // 测试驱动 dispatch()(协议层逻辑核心),不跑真实 stdio(真实 stdio 需集成测试)。 // main_loop 的 framing(write_response 单行 +\n、notification 不回)由 dispatch 的输入输出间接覆盖。 #[cfg(test)] mod tests { use super::*; use crate::tools::RiskLevel; /// 构造内存 DB + Ctx async fn test_ctx() -> Ctx { let db = Arc::new(Database::open_in_memory().await.unwrap()); Ctx::new(db) } /// 直接驱动 dispatch:喂一行 → 收一行(逐行,notification 跳过) async fn run_io_lines(lines: &[&str], read_only: bool) -> Vec { let ctx = test_ctx().await; let mut out = Vec::new(); for line in lines { let req: Request = serde_json::from_str(line).unwrap(); let is_notif = req.id.is_none(); let method = McpMethod::from_request(&req); if is_notif { continue; } let resp = dispatch(&ctx, read_only, req.id.clone(), method).await; out.push(serde_json::to_string(&resp).unwrap()); } out } #[tokio::test] async fn initialize_returns_server_info_and_capabilities() { let input = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#; let out = run_io_lines(&[input], false).await; assert_eq!(out.len(), 1); let v: Value = serde_json::from_str(&out[0]).unwrap(); assert_eq!(v["jsonrpc"], "2.0"); assert_eq!(v["id"], 1); assert_eq!(v["result"]["protocolVersion"], PROTOCOL_VERSION); assert_eq!(v["result"]["serverInfo"]["name"], SERVER_NAME); // capabilities.tools 存在 assert!(v["result"]["capabilities"]["tools"].is_object()); } #[tokio::test] async fn tools_list_excludes_high_risk_by_default() { let input = r#"{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}"#; let out = run_io_lines(&[input], false).await; let v: Value = serde_json::from_str(&out[0]).unwrap(); let tools = v["result"]["tools"].as_array().unwrap(); let names: Vec<&str> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect(); // 默认(非 read-only):Low + Medium 可见,High 不可见 assert!(names.contains(&"list_projects")); assert!(names.contains(&"create_project")); // Medium assert!(!names.contains(&"delete_project")); // High assert!(!names.contains(&"run_workflow")); // High // MCP Tool schema:name + inputSchema let first = &tools[0]; assert!(first["name"].is_string()); assert!(first["inputSchema"]["type"] == "object"); } #[tokio::test] async fn tools_list_read_only_keeps_only_low() { let input = r#"{"jsonrpc":"2.0","id":3,"method":"tools/list","params":{}}"#; let out = run_io_lines(&[input], true).await; let v: Value = serde_json::from_str(&out[0]).unwrap(); let tools = v["result"]["tools"].as_array().unwrap(); let names: Vec<&str> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect(); assert!(names.contains(&"list_projects")); assert!(!names.contains(&"create_project")); // Medium 被过滤 assert!(!names.contains(&"delete_project")); // High } #[tokio::test] async fn tools_call_unknown_method_returns_method_not_found() { let input = r#"{"jsonrpc":"2.0","id":4,"method":"bogus/method","params":{}}"#; let out = run_io_lines(&[input], false).await; let v: Value = serde_json::from_str(&out[0]).unwrap(); assert_eq!(v["error"]["code"], crate::protocol::METHOD_NOT_FOUND); } #[tokio::test] async fn tools_call_high_risk_is_rejected() { let input = r#"{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"delete_project","arguments":{"id":"x"}}}"#; let out = run_io_lines(&[input], false).await; let v: Value = serde_json::from_str(&out[0]).unwrap(); let content = v["result"]["content"].as_array().unwrap(); assert_eq!(content[0]["type"], "text"); let text = content[0]["text"].as_str().unwrap(); assert!(text.contains("拒绝") || text.contains("High")); assert_eq!(v["result"]["isError"], true); } #[tokio::test] async fn tools_call_read_only_rejects_medium() { let input = r#"{"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"name":"create_project","arguments":{"name":"t","description":"d"}}}"#; let out = run_io_lines(&[input], true).await; let v: Value = serde_json::from_str(&out[0]).unwrap(); let text = v["result"]["content"][0]["text"].as_str().unwrap(); assert!(text.contains("只读模式拒绝")); } #[tokio::test] async fn tools_call_create_then_list_project_roundtrip() { let create = r#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"create_project","arguments":{"name":"McpProj","description":"via mcp"}}}"#; let list = r#"{"jsonrpc":"2.0","id":8,"method":"tools/call","params":{"name":"list_projects","arguments":{}}}"#; let out = run_io_lines(&[create, list], false).await; let created: Value = serde_json::from_str(&out[0]).unwrap(); let id = created["result"]["content"][0]["text"] .as_str() .unwrap(); let created_v: Value = serde_json::from_str(id).unwrap(); assert!(created_v["id"].is_string()); let listed: Value = serde_json::from_str(&out[1]).unwrap(); let list_text = listed["result"]["content"][0]["text"].as_str().unwrap(); let list_v: Value = serde_json::from_str(list_text).unwrap(); assert_eq!(list_v["count"], 1); assert_eq!(list_v["projects"][0]["name"], "McpProj"); } #[tokio::test] async fn visible_predicate() { assert!(visible(false, RiskLevel::Low)); assert!(visible(false, RiskLevel::Medium)); assert!(!visible(false, RiskLevel::High)); assert!(visible(true, RiskLevel::Low)); assert!(!visible(true, RiskLevel::Medium)); assert!(!visible(true, RiskLevel::High)); } }