修复: MCP 多进程缺陷(update CAS乐观锁 + 审计落盘 + 空闲超时 + busy_timeout + stdio写回调 + list分页)

2026-08-08 18:43:38 +08:00
parent 6aa334fc9b
commit 02c8d8e5ea
5 changed files with 559 additions and 55 deletions
@@ -8,9 +8,11 @@
 //!
 //! 高风险(High)工具:tools/list 不暴露(从清单剔除),tools/call 即便绕过也由 handler 兜底拒绝。
 //! read-only:tools/list 仅留 Low,tools/call Medium/High 一律拒绝。
+//! 生命周期:空闲超时(默认 60s 无请求)自动退出,防客户端强杀后进程残留;写操作(risk != Low)可选回调(stdio 预留接线点)。

 use std::path::Path;
 use std::sync::Arc;
+use std::time::Duration;

 use df_storage::db::Database;
 use serde_json::{json, Value};
@@ -27,31 +29,62 @@ pub(crate) const PROTOCOL_VERSION: &str = "2025-06-18";
 pub(crate) const SERVER_NAME: &str = "devflow-mcp";
 pub(crate) const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION");

-/// 启动 MCP server。
+/// MCP server 运行配置。
+///
+/// 聚合 stdio 形态的运行参数,避免 `run_server` 参数膨胀。
+/// [`Default`] 即常用生产配置(可写 + 空闲超时 60s + 无写回调)。
+pub struct ServerConfig {
+    /// 只读模式:true 则仅暴露 list/get 工具,Medium/High 写操作一律拒绝。
+    pub read_only: bool,
+    /// 空闲超时:连续 N 时长无 stdin 输入则自动退出进程(防客户端强杀后孤儿残留)。
+    /// `None` = 永不因空闲退出(常驻)。默认 60s。
+    pub idle_timeout: Option<Duration>,
+    /// 写操作回调:成功执行(未被 read-only/High 拒绝)的写工具(risk != Low)触发。
+    /// 桌面内嵌形态经 server_http → on_tool_call → emit df-data-changed 刷新 GUI;
+    /// stdio 独立进程暂无 AppHandle,传 None 预留接线点,后续可注入。
+    pub on_write_call: Option<Arc<dyn Fn(&str) + Send + Sync>>,
+}
+
+impl Default for ServerConfig {
+    fn default() -> Self {
+        Self {
+            read_only: false,
+            idle_timeout: Some(Duration::from_secs(60)),
+            on_write_call: None,
+        }
+    }
+}
+
+/// 启动 MCP server(stdin/stdout 接实际进程句柄)。
 ///
 /// - `db_path`:SQLite 数据库文件路径(应用同库,开 WAL 模式)
-/// - `read_only`:true 则仅暴露 list/get 工具
+/// - `config`:运行配置(只读/空闲超时/写操作回调),见 [`ServerConfig`]
 ///
 /// 复用 [`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<()> {
+pub async fn run_server(db_path: &Path, config: ServerConfig) -> 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
+    main_loop(stdin, stdout, &ctx, &config).await
 }

 /// 可单测的主循环(参数化 stdin/stdout)。
 ///
 /// 协议正确性:每行一个 JSON-RPC 消息,Response 单行写回(末尾 \n)。
 /// Notification(id=None)不回响应。
+///
+/// 生命周期:
+/// - 空闲超时:仅在**等待下一个请求**(`read_line`)时计时;正在处理的请求不受影响,
+///   超时即 break 优雅退出(防客户端强杀后进程残留)。
+/// - 写操作回调:成功执行的写工具(risk != Low)在响应写回后触发,供外部(GUI)感知数据变更。
 pub async fn main_loop<R, W>(
    stdin: R,
    stdout: W,
    ctx: &Ctx,
-    read_only: bool,
+    config: &ServerConfig,
 ) -> anyhow::Result<()>
 where
    R: tokio::io::AsyncRead + Unpin,
@@ -63,7 +96,17 @@ where

    loop {
        line.clear();
-        let n = reader.read_line(&mut line).await?;
+        // 空闲超时:只包裹「等待下一请求」,不包裹 dispatch/写响应,处理期间绝不误杀。
+        let n = match config.idle_timeout {
+            Some(dur) => match tokio::time::timeout(dur, reader.read_line(&mut line)).await {
+                Ok(n) => n,
+                Err(_elapsed) => {
+                    tracing::info!(target: "df_mcp", idle_secs = dur.as_secs(), "空闲超时无新请求,自动退出");
+                    break;
+                }
+            }?,
+            None => reader.read_line(&mut line).await?,
+        };
        if n == 0 {
            // EOF(stdin 关闭),优雅退出
            break;
@@ -101,8 +144,16 @@ where
            continue;
        }

-        let resp = dispatch(ctx, read_only, req.id.clone(), method).await;
+        // 写回调需在 dispatch 后判定,先预取工具名(仅 tools/call 需要,避免 clone 整包请求)
+        let tool_name = match &method {
+            McpMethod::ToolsCall { name, .. } => Some(name.clone()),
+            _ => None,
+        };
+        let resp = dispatch(ctx, config.read_only, req.id.clone(), method).await;
        write_response(&mut writer, &resp).await?;
+        if let Some(name) = tool_name {
+            fire_write_hook(config, &name);
+        }
    }

    Ok(())
@@ -153,19 +204,15 @@ pub(crate) async fn dispatch(ctx: &Ctx, read_only: bool, id: Option<Value>, meth
                        .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 应用内执行。"
-                ));
+            // 执行前防御:read-only 拒 Medium/High;High 兜底拒(handler 内也会拒,双保险)。
+            // 判定收口到 should_execute,与 main_loop 写回调共用同一事实源,避免两份逻辑漂移。
+            if !should_execute(read_only, spec.risk) {
+                let msg = if read_only && spec.risk != RiskLevel::Low {
+                    format!("只读模式拒绝执行 {name}(风险等级 {:?})", spec.risk)
+                } else {
+                    format!("High 风险操作 {name} 默认拒绝,请在 DevFlow 应用内执行。")
+                };
+                let r = CallToolResult::error(msg);
                return Response::ok(id, serde_json::to_value(r).unwrap_or(Value::Null));
            }
            // Low / Medium:执行
@@ -181,6 +228,23 @@ pub(crate) async fn dispatch(ctx: &Ctx, read_only: bool, id: Option<Value>, meth
    }
 }

+/// 工具是否会被执行(dispatch 与 main_loop 写回调共用的判定)。
+///
+/// read-only 下仅 Low 可执行;非 read-only 下 High 仍兜底拒绝。两条件都过 → 可执行。
+fn should_execute(read_only: bool, risk: RiskLevel) -> bool {
+    !(read_only && risk != RiskLevel::Low) && risk != RiskLevel::High
+}
+
+/// 触发写操作回调(若有)。仅当工具为写操作(risk != Low)且未被 read-only/High 拒绝时触发,
+/// 与 dispatch 的执行判定一致。回调仅作通知(如 GUI 刷新),不承载返回结果。
+fn fire_write_hook(config: &ServerConfig, name: &str) {
+    let Some(cb) = &config.on_write_call else { return };
+    let Some(spec) = tools::find(name) else { return };
+    if spec.risk != RiskLevel::Low && should_execute(config.read_only, spec.risk) {
+        cb(name);
+    }
+}
+
 /// 工具可见性:read-only 仅 Low,否则 Low + Medium(High 永不可见)
 pub(crate) fn visible(read_only: bool, risk: RiskLevel) -> bool {
    if read_only {
@@ -213,6 +277,8 @@ async fn write_response<W: tokio::io::AsyncWrite + Unpin>(
 mod tests {
    use super::*;
    use crate::tools::RiskLevel;
+    use std::sync::Mutex;
+    use tokio::io::AsyncWriteExt;

    /// 构造内存 DB + Ctx
    async fn test_ctx() -> Ctx {
@@ -358,4 +424,122 @@ mod tests {
        assert!(!visible(true, RiskLevel::Medium));
        assert!(!visible(true, RiskLevel::High));
    }
+
+    // ── should_execute 判定(dispatch 与写回调共用)────────────────────
+
+    #[test]
+    fn should_execute_predicate() {
+        assert!(should_execute(false, RiskLevel::Low));
+        assert!(should_execute(false, RiskLevel::Medium));
+        assert!(!should_execute(false, RiskLevel::High));
+        assert!(should_execute(true, RiskLevel::Low));
+        assert!(!should_execute(true, RiskLevel::Medium));
+        assert!(!should_execute(true, RiskLevel::High));
+    }
+
+    // ── 空闲超时 / 写操作回调(main_loop 集成)────────────────────────
+
+    /// 构造测试配置:只读开关 + 默认空闲超时 + 默认无回调(字段可覆盖)
+    async fn test_config(read_only: bool) -> ServerConfig {
+        ServerConfig {
+            read_only,
+            idle_timeout: Some(Duration::from_secs(60)),
+            on_write_call: None,
+        }
+    }
+
+    #[tokio::test]
+    async fn main_loop_exits_on_idle_timeout() {
+        let ctx = test_ctx().await;
+        // 用 duplex 造一个「开着但永不写数据」的 stdin:read_line 会一直挂起,
+        // 空闲超时(100ms)触发后应正常 break 退出,而非阻塞或 panic。
+        let (_tx, rx) = tokio::io::duplex(1024);
+        let config = ServerConfig {
+            idle_timeout: Some(Duration::from_millis(100)),
+            ..test_config(false).await
+        };
+        let outer = tokio::time::timeout(
+            Duration::from_secs(2),
+            main_loop(rx, tokio::io::sink(), &ctx, &config),
+        )
+        .await;
+        let inner = outer.expect("main_loop 应在空闲超时后返回,而非一直阻塞");
+        assert!(inner.is_ok(), "空闲超时退出应为 Ok,实际: {inner:?}");
+    }
+
+    #[tokio::test]
+    async fn main_loop_fires_write_callback_for_write_tool() {
+        let ctx = test_ctx().await;
+        let calls: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let calls_cb = calls.clone();
+        let config = ServerConfig {
+            on_write_call: Some(Arc::new(move |name| {
+                calls_cb.lock().unwrap().push(name.to_string());
+            })),
+            ..test_config(false).await
+        };
+        let (mut tx, rx) = tokio::io::duplex(1024);
+        let line = r#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"create_project","arguments":{"name":"McpProj","description":"via mcp"}}}"#;
+        tx.write_all(format!("{line}\n").as_bytes()).await.unwrap();
+        drop(tx); // 关 stdin → 处理后 EOF,正常退出
+        main_loop(rx, tokio::io::sink(), &ctx, &config)
+            .await
+            .unwrap();
+        assert!(
+            calls.lock().unwrap().contains(&"create_project".to_string()),
+            "写工具 create_project 应触发写回调,实际: {:?}",
+            *calls.lock().unwrap()
+        );
+    }
+
+    #[tokio::test]
+    async fn main_loop_no_write_callback_for_read_tool() {
+        let ctx = test_ctx().await;
+        let calls: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let calls_cb = calls.clone();
+        let config = ServerConfig {
+            on_write_call: Some(Arc::new(move |name| {
+                calls_cb.lock().unwrap().push(name.to_string());
+            })),
+            ..test_config(false).await
+        };
+        let (mut tx, rx) = tokio::io::duplex(1024);
+        let line = r#"{"jsonrpc":"2.0","id":8,"method":"tools/call","params":{"name":"list_projects","arguments":{}}}"#;
+        tx.write_all(format!("{line}\n").as_bytes()).await.unwrap();
+        drop(tx);
+        main_loop(rx, tokio::io::sink(), &ctx, &config)
+            .await
+            .unwrap();
+        assert!(
+            calls.lock().unwrap().is_empty(),
+            "只读工具 list_projects 不应触发写回调,实际: {:?}",
+            *calls.lock().unwrap()
+        );
+    }
+
+    #[tokio::test]
+    async fn main_loop_no_write_callback_when_read_only_denies_write() {
+        let ctx = test_ctx().await;
+        let calls: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let calls_cb = calls.clone();
+        let config = ServerConfig {
+            read_only: true,
+            on_write_call: Some(Arc::new(move |name| {
+                calls_cb.lock().unwrap().push(name.to_string());
+            })),
+            ..test_config(true).await
+        };
+        let (mut tx, rx) = tokio::io::duplex(1024);
+        let line = r#"{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"create_project","arguments":{"name":"X","description":"d"}}}"#;
+        tx.write_all(format!("{line}\n").as_bytes()).await.unwrap();
+        drop(tx);
+        main_loop(rx, tokio::io::sink(), &ctx, &config)
+            .await
+            .unwrap();
+        assert!(
+            calls.lock().unwrap().is_empty(),
+            "read-only 下 Medium 写被拒,不应触发写回调,实际: {:?}",
+            *calls.lock().unwrap()
+        );
+    }
 }