Skip to main content

yggdrasil/mcp/tools/
media.rs

1//! MCP 媒体工具:write 上传图片,admin 管理素材。
2//!
3//! Option B 的第三通道:LLM 工具只收 `url: String`(JSON-RPC 纯文本),服务端
4//! 按 SSRF 防护抓取二进制,再走 [`crate::api::upload::process_image_upload`] 共享
5//! 入库流水线。**二进制从不进 JSON-RPC**——彻底绕开 rmcp 4MiB 请求体上限与
6//! base64 的 33% 膨胀 + 上下文窗口烧灼。
7//!
8//! 另有第二通道 `POST /api/mcp/upload`(bearer multipart)供 host/shell 直接 POST
9//! 二进制(如 Claude Code 的 Bash+curl);两条通道共用同一入库流水线。
10//!
11//! SSRF 防护(多层纵深)见 [`crate::api::url_fetch`]:强制 https、解析即锁 IP
12//! 杜绝 DNS rebinding、禁重定向、流式体积上限、超时。
13//!
14//! 本模块仅 `feature = "server"` 编译。
15
16#![cfg(feature = "server")]
17
18use rmcp::handler::server::tool::Extension;
19use rmcp::handler::server::wrapper::Parameters;
20use rmcp::model::CallToolResult;
21use rmcp::{schemars, tool, tool_router, ErrorData as McpError};
22use serde::Deserialize;
23
24use super::common::{internal, ok_json, require_admin, require_scope};
25use crate::models::asset::{AssetFilter, AssetSort};
26use crate::models::mcp_token::TokenScope;
27
28#[tool_router(router = media_router, vis = "pub")]
29impl crate::mcp::server::YggMcpServer {
30    /// 从一个图片 URL 抓取并入库(服务端转 WebP 若更小),返回可直接嵌入
31    /// Markdown 正文的 `/uploads/...` URL。要求 write 作用域。
32    ///
33    /// 仅接受 `https://` URL;服务端做 SSRF 防护(私网/回环/保留段拒绝、
34    /// DNS 锁定防 rebinding、禁重定向、体积上限)。二进制不经 JSON-RPC。
35    #[tool(
36        description = "从图片 URL 抓取并入库(服务端转 WebP 若更小),返回 asset_id、alt 和 /uploads/... URL(可直接用于 Markdown 正文 img),提供 alt 时保存到素材。仅接受 https:// URL,支持 JPEG/PNG/GIF/WebP。二进制不经 JSON-RPC。"
37    )]
38    async fn upload_media(
39        &self,
40        Parameters(p): Parameters<UploadMediaParams>,
41        Extension(parts): Extension<http::request::Parts>,
42    ) -> Result<CallToolResult, McpError> {
43        let _principal = require_scope(&parts, "upload_media", TokenScope::Write)?;
44
45        // SSRF 防护抓取 + 共享入库流水线。
46        let outcome = crate::api::url_fetch::fetch_and_ingest(&p.url, p.alt)
47            .await
48            .map_err(|e| match e {
49                crate::api::url_fetch::FetchError::Invalid(msg)
50                | crate::api::url_fetch::FetchError::BadStatus(msg) => {
51                    McpError::invalid_request(msg, None)
52                }
53                crate::api::url_fetch::FetchError::TooLarge => McpError::invalid_request(
54                    format!(
55                        "文件超过大小限制({} bytes)",
56                        crate::utils::server::MAX_FILE_SIZE
57                    ),
58                    None,
59                ),
60                crate::api::url_fetch::FetchError::Fetch(ctx) => internal(ctx, "url fetch"),
61            })?;
62
63        tracing::info!(
64            "MCP media uploaded via URL: {} ({}x{}, reused={})",
65            outcome.url,
66            outcome.width,
67            outcome.height,
68            outcome.reused
69        );
70
71        ok_json(UploadResult {
72            success: true,
73            url: outcome.url,
74            asset_id: outcome.asset_id,
75            alt: outcome.alt,
76            reused: outcome.reused,
77            width: outcome.width,
78            height: outcome.height,
79            mime: outcome.mime,
80        })
81    }
82    #[tool(
83        description = "分页查询素材,每页 60 张,返回元数据和文章/评论/头像引用明细。query 搜索文件名或 alt;filter 为 All/Used/Orphan,sort 为 CreatedDesc/SizeDesc,page 从 1 开始。图片 URL 为 /uploads/ 加 path。需要 admin 作用域。"
84    )]
85    async fn list_assets(
86        &self,
87        Parameters(p): Parameters<ListAssetsParams>,
88        Extension(parts): Extension<http::request::Parts>,
89    ) -> Result<CallToolResult, McpError> {
90        require_admin(&parts, "list_assets")?;
91        let result = crate::api::assets::list::list_assets_impl(p.filter, p.query, p.sort, p.page)
92            .await
93            .map_err(|e| internal(e, "list_assets"))?;
94        ok_json(result)
95    }
96
97    #[tool(description = "修改素材 alt,空白清除;不回写已有文章 HTML。需要 admin 作用域。")]
98    async fn update_asset_alt(
99        &self,
100        Parameters(p): Parameters<UpdateAssetAltParams>,
101        Extension(parts): Extension<http::request::Parts>,
102    ) -> Result<CallToolResult, McpError> {
103        require_admin(&parts, "update_asset_alt")?;
104        let result = crate::api::assets::delete::update_asset_alt_impl(p.id, p.alt)
105            .await
106            .map_err(|e| internal(e, "update_asset_alt"))?;
107        ok_json(result)
108    }
109
110    #[tool(
111        description = "永久删除一张无引用素材(文件、记录及缓存)。文章(含草稿与回收站)、评论(含回收站)或头像引用中的素材拒绝删除。需要 admin 作用域。"
112    )]
113    async fn delete_asset(
114        &self,
115        Parameters(p): Parameters<AssetIdParams>,
116        Extension(parts): Extension<http::request::Parts>,
117    ) -> Result<CallToolResult, McpError> {
118        require_admin(&parts, "delete_asset")?;
119        let result = crate::api::assets::delete::delete_asset_impl(p.id)
120            .await
121            .map_err(|e| internal(e, "delete_asset"))?;
122        ok_json(result)
123    }
124
125    #[tool(
126        description = "批量永久删除素材,跳过被引用项,返回删除、跳过、失败统计;每次 1..=100 个 id。需要 admin 作用域。"
127    )]
128    async fn batch_delete_assets(
129        &self,
130        Parameters(p): Parameters<AssetIdsParams>,
131        Extension(parts): Extension<http::request::Parts>,
132    ) -> Result<CallToolResult, McpError> {
133        require_admin(&parts, "batch_delete_assets")?;
134        if p.ids.is_empty() || p.ids.len() > 100 {
135            return Err(McpError::invalid_params(
136                "ids must contain 1..=100 entries",
137                None,
138            ));
139        }
140
141        let result = crate::api::assets::delete::batch_delete_assets_impl(p.ids)
142            .await
143            .map_err(|e| internal(e, "batch_delete_assets"))?;
144        ok_json(result)
145    }
146
147    #[tool(
148        description = "永久清理无引用且上传超过 7 天的素材。可先用 list_assets 查看 purgeable_count 和 purgeable_bytes。需要 admin 作用域。"
149    )]
150    async fn purge_orphan_assets(
151        &self,
152        Parameters(_p): Parameters<EmptyMediaParams>,
153        Extension(parts): Extension<http::request::Parts>,
154    ) -> Result<CallToolResult, McpError> {
155        require_admin(&parts, "purge_orphan_assets")?;
156        let result = crate::api::assets::delete::purge_orphan_assets_impl(
157            crate::api::assets::list::PURGE_GRACE_DAYS,
158        )
159        .await
160        .map_err(|e| internal(e, "purge_orphan_assets"))?;
161        ok_json(result)
162    }
163
164    #[tool(
165        description = "扫描 uploads 目录重建素材索引及文章引用,保留已有 alt 和文件名,移除文件已消失的记录。需要 admin 作用域。"
166    )]
167    async fn rebuild_assets_index(
168        &self,
169        Parameters(_p): Parameters<EmptyMediaParams>,
170        Extension(parts): Extension<http::request::Parts>,
171    ) -> Result<CallToolResult, McpError> {
172        require_admin(&parts, "rebuild_assets_index")?;
173        let result = crate::api::assets::rebuild::rebuild_assets_index_impl()
174            .await
175            .map_err(|e| internal(e, "rebuild_assets_index"))?;
176        ok_json(result)
177    }
178}
179
180// ---------------------------------------------------------------------------
181// 参数与输出结构
182// ---------------------------------------------------------------------------
183
184#[derive(Debug, Deserialize, schemars::JsonSchema)]
185pub struct UploadMediaParams {
186    /// 图片的 https URL(服务端抓取,二进制不经 JSON-RPC)。
187    pub url: String,
188    /// 保存素材 alt;重复上传时省略保留原值,空白清除,不回写已有文章。
189    #[serde(default)]
190    pub alt: Option<String>,
191}
192
193#[derive(Debug, serde::Serialize)]
194struct UploadResult {
195    success: bool,
196    asset_id: String,
197    alt: Option<String>,
198    url: String,
199    reused: bool,
200    width: u32,
201    height: u32,
202    mime: String,
203}
204
205#[derive(Debug, Deserialize, schemars::JsonSchema)]
206pub struct ListAssetsParams {
207    #[serde(default)]
208    pub filter: AssetFilter,
209    #[serde(default)]
210    pub query: String,
211    #[serde(default)]
212    pub sort: AssetSort,
213    #[serde(default = "first_page")]
214    pub page: i32,
215}
216
217fn first_page() -> i32 {
218    1
219}
220
221#[derive(Debug, Deserialize, schemars::JsonSchema)]
222pub struct UpdateAssetAltParams {
223    pub id: String,
224    pub alt: String,
225}
226
227#[derive(Debug, Deserialize, schemars::JsonSchema)]
228pub struct AssetIdParams {
229    pub id: String,
230}
231
232#[derive(Debug, Deserialize, schemars::JsonSchema)]
233pub struct AssetIdsParams {
234    pub ids: Vec<String>,
235}
236
237/// `purge_orphan_assets` / `rebuild_assets_index` 的无参入参。
238///
239/// 显式补上 `properties: {}`:schemars 对零字段结构体只生成 `{"type":"object"}`,
240/// 而部分 OpenAI 兼容上游(如 MCA)要求对象 schema 必须带 `properties` 键,
241/// 否则整个请求会被判为非法工具 schema。
242#[derive(Debug, Deserialize, schemars::JsonSchema)]
243#[schemars(extend("properties" = {}))]
244pub struct EmptyMediaParams {}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use crate::mcp::{auth::McpPrincipal, server::YggMcpServer};
250
251    fn parts(scope: Option<TokenScope>) -> http::request::Parts {
252        let (mut parts, _) = http::Request::new(()).into_parts();
253        if let Some(scope) = scope {
254            parts.extensions.insert(McpPrincipal {
255                user_id: 1,
256                scope,
257                token_id: "test".into(),
258            });
259        }
260        parts
261    }
262
263    #[tokio::test]
264    async fn management_rejects_missing_read_and_write_principals_before_io() {
265        let server = YggMcpServer;
266        for scope in [None, Some(TokenScope::Read), Some(TokenScope::Write)] {
267            let list = serde_json::from_str::<ListAssetsParams>("{}").unwrap();
268            assert!(server
269                .list_assets(Parameters(list), Extension(parts(scope)))
270                .await
271                .is_err());
272            assert!(server
273                .update_asset_alt(
274                    Parameters(UpdateAssetAltParams {
275                        id: "invalid".into(),
276                        alt: "alt".into()
277                    }),
278                    Extension(parts(scope))
279                )
280                .await
281                .is_err());
282            assert!(server
283                .delete_asset(
284                    Parameters(AssetIdParams {
285                        id: "invalid".into()
286                    }),
287                    Extension(parts(scope))
288                )
289                .await
290                .is_err());
291            assert!(server
292                .batch_delete_assets(
293                    Parameters(AssetIdsParams { ids: vec![] }),
294                    Extension(parts(scope))
295                )
296                .await
297                .is_err());
298            assert!(server
299                .purge_orphan_assets(Parameters(EmptyMediaParams {}), Extension(parts(scope)))
300                .await
301                .is_err());
302            assert!(server
303                .rebuild_assets_index(Parameters(EmptyMediaParams {}), Extension(parts(scope)))
304                .await
305                .is_err());
306        }
307    }
308
309    #[tokio::test]
310    async fn batch_delete_rejects_empty_and_oversized_batches_before_io() {
311        for ids in [vec![], vec!["id".into(); 101]] {
312            let err = YggMcpServer
313                .batch_delete_assets(
314                    Parameters(AssetIdsParams { ids }),
315                    Extension(parts(Some(TokenScope::Admin))),
316                )
317                .await
318                .unwrap_err();
319            assert!(err.message.contains("1..=100"));
320        }
321    }
322
323    #[test]
324    fn media_router_exposes_upload_and_management_tools() {
325        let tools = YggMcpServer::media_router().list_all();
326        for name in [
327            "upload_media",
328            "list_assets",
329            "update_asset_alt",
330            "delete_asset",
331            "batch_delete_assets",
332            "purge_orphan_assets",
333            "rebuild_assets_index",
334        ] {
335            assert!(tools.iter().any(|tool| tool.name == name), "missing {name}");
336        }
337        let list = tools
338            .iter()
339            .find(|tool| tool.name == "list_assets")
340            .unwrap();
341        let schema = serde_json::to_string(&list.input_schema).unwrap();
342        assert!(schema.contains("Orphan"));
343        assert!(schema.contains("SizeDesc"));
344    }
345
346    #[test]
347    fn list_defaults_and_filters_match_the_schema() {
348        let params: ListAssetsParams = serde_json::from_str("{}").unwrap();
349        assert_eq!(params.page, 1);
350        assert_eq!(params.filter, AssetFilter::All);
351        assert_eq!(params.sort, AssetSort::CreatedDesc);
352        assert!(params.query.is_empty());
353        assert!(serde_json::from_str::<ListAssetsParams>(r#"{"filter":"invalid"}"#).is_err());
354    }
355}