yggdrasil/mcp/tools/
runner.rs

1//! 代码运行器 MCP 工具：在容器沙箱内执行代码，同步返回完整输出。
2//!
3//! 与 `src/api/code_runner/execute.rs::start_exec` 的执行链一致（语言白名单 →
4//! 源码大小 → 信号量限并发 → clamp_limits → run_in_container），但 **同步返回**
5//! 完整 stdout/stderr，不引入 task_id / 轮询 / SSE 机制——MCP 工具返回单一结果。
6//!
7//! 鉴权与限流：MCP 走 bearer token → admin 作用域。admin 跳过 IP 速率限制
8//! （与 web 的 `check_rate_limit_for_user` admin 放行一致），但仍受并发槽、
9//! 资源钳制与源码大小校验约束。Docker daemon 不可用时返回明确错误（与
10//! `get_docker()` 的 NotFound 脱敏路径一致）。
11
12#![cfg(feature = "server")]
13
14use rmcp::handler::server::tool::Extension;
15use rmcp::handler::server::wrapper::Parameters;
16use rmcp::model::{CallToolResult, ContentBlock, TextContent};
17use rmcp::{schemars, tool, tool_router, ErrorData as McpError};
18
19use super::common::require_admin;
20use serde::{Deserialize, Serialize};
21use std::time::Duration;
22
23use crate::api::code_runner::execute::RUNNER_SEMAPHORE;
24use crate::api::code_runner::languages::{is_supported_lang, normalize_lang, LANGUAGES};
25use crate::infra::docker::run_in_container;
26use crate::infra::runner_config::{clamp_limits, RUNNER_CONFIG};
27
28/// `run_code` 入参。
29#[derive(Debug, Deserialize, schemars::JsonSchema)]
30pub struct RunCodeParams {
31    /// 语言标识（支持别名归一化：js/javascript→node、ts/typescript→bun、rs→rust；
32    /// canonical: python/node/go/rust/bun）。
33    pub language: String,
34    /// 源代码（受 `CODE_RUNNER_MAX_SOURCE_BYTES` 限制，默认 64KB）。
35    pub source: String,
36}
37
38/// `run_code` 返回的执行结果。
39#[derive(Debug, Serialize)]
40struct RunResult {
41    /// 执行终态：success / error / oom / timeout / unavailable。
42    status: &'static str,
43    /// 进程退出码（容器未跑完 / 超时 / daemon 不可用时为 null）。
44    exit_code: Option<i64>,
45    /// 标准输出（已按 output_bytes 上限截断）。
46    stdout: String,
47    /// 标准错误 / 失败原因描述（已按 output_bytes 上限截断）。
48    stderr: String,
49    /// 执行耗时（毫秒）。
50    duration_ms: u64,
51    /// 归一化后的 canonical 语言 key（如 python/node/go/rust/bun）。
52    language: String,
53}
54
55#[tool_router(router = runner_router, vis = "pub")]
56impl crate::mcp::server::YggMcpServer {
57    /// 在容器沙箱内执行代码并返回输出。要求 admin 作用域。
58    #[tool(
59        description = "在 Docker 沙箱内执行代码（支持 python/node/go/rust/bun），返回 stdout/stderr。需要 admin 作用域。"
60    )]
61    async fn run_code(
62        &self,
63        Parameters(RunCodeParams { language, source }): Parameters<RunCodeParams>,
64        Extension(parts): Extension<http::request::Parts>,
65    ) -> Result<CallToolResult, McpError> {
66        require_admin(&parts, "run_code")?;
67
68        // 1. 语言白名单（与 validate_exec_request 一致；is_supported_lang 内含归一化）。
69        if !is_supported_lang(&language) {
70            return Err(McpError::invalid_request(
71                "unsupported language: use one of python/node/go/rust/bun (js/rs/ts aliases accepted)",
72                None,
73            ));
74        }
75
76        // 2. 源码大小限制。
77        if source.len() > RUNNER_CONFIG.max_source_bytes as usize {
78            return Err(McpError::invalid_request(
79                format!(
80                    "source too large: {} bytes > limit {}",
81                    source.len(),
82                    RUNNER_CONFIG.max_source_bytes
83                ),
84                None,
85            ));
86        }
87
88        let result = execute_in_container(&language, &source)
89            .await
90            .map_err(|e| McpError::internal_error(format!("code execution failed: {e}"), None))?;
91
92        let text = serde_json::to_string_pretty(&result)
93            .map_err(|e| McpError::internal_error(format!("encode failed: {e}"), None))?;
94        Ok(CallToolResult::success(vec![ContentBlock::Text(
95            TextContent::new(text),
96        )]))
97    }
98}
99
100/// 执行一次容器内代码运行（同步收集输出）。
101///
102/// 镜像 start_exec 的后台 spawn 体：排队信号量 → 归一化语言 → clamp_limits →
103/// run_in_container，但把结果直接返回而非写入 EXEC_TASKS。
104async fn execute_in_container(language: &str, source: &str) -> Result<RunResult, String> {
105    let lang_key = normalize_lang(language);
106
107    let lang_def = LANGUAGES
108        .get(&lang_key)
109        .ok_or_else(|| format!("language not registered: {lang_key}"))?;
110
111    // 排队等待可用容器槽（与 start_exec 一致的 queue_timeout_secs）。
112    let ticket = tokio::time::timeout(
113        Duration::from_secs(RUNNER_CONFIG.queue_timeout_secs),
114        RUNNER_SEMAPHORE.acquire(),
115    )
116    .await
117    .map_err(|_| "container queue timeout: too many concurrent executions".to_string())?
118    .map_err(|e| format!("semaphore acquire: {e}"))?;
119
120    let final_limits = clamp_limits(lang_def.default_limits.clone(), lang_def.allow_network);
121
122    let start_time = chrono::Utc::now();
123    let res = run_in_container(
124        &lang_def.image,
125        &lang_def.run_cmd,
126        source,
127        &lang_def.extension,
128        final_limits,
129    )
130    .await;
131    let duration_ms = (chrono::Utc::now() - start_time).num_milliseconds().max(0) as u64;
132
133    drop(ticket); // 显式释放信号量
134
135    match res {
136        Ok((exit_code, stdout, stderr, oom_killed)) => {
137            let status = if oom_killed {
138                "oom"
139            } else if exit_code == Some(0) {
140                "success"
141            } else {
142                "error"
143            };
144            Ok(RunResult {
145                status,
146                exit_code,
147                stdout,
148                stderr,
149                duration_ms,
150                language: lang_key,
151            })
152        }
153        Err(e) => {
154            // 脱敏与 start_exec 的失败分支一致：日志记详情，对外给通用消息。
155            let s = e.to_string();
156            let is_timeout = s.contains("TimedOut");
157            tracing::error!(error = ?e, "MCP container execution failed");
158            // bollard IOError{NotFound} = Docker daemon 不可用 → 明确的 unavailable。
159            let is_daemon_down = s.contains("Docker daemon") || s.contains("NotFound");
160            let status = if is_daemon_down {
161                "unavailable"
162            } else if is_timeout {
163                "timeout"
164            } else {
165                "error"
166            };
167            let stderr = if is_daemon_down {
168                "code runner unavailable: Docker daemon not running".to_string()
169            } else if is_timeout {
170                "execution timed out".to_string()
171            } else {
172                "code runner temporarily unavailable".to_string()
173            };
174            Ok(RunResult {
175                status,
176                exit_code: None,
177                stdout: String::new(),
178                stderr,
179                duration_ms,
180                language: lang_key,
181            })
182        }
183    }
184}
yggdrasil/mcp/tools/runner.rs

yggdrasil/mcp/tools/
runner.rs
