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            .await
158            .map_err(|e| internal(e, "purge_orphan_assets"))?;
159        ok_json(result)
160    }
161
162    #[tool(
163        description = "扫描 uploads 目录重建素材索引及文章引用,保留已有 alt 和文件名,移除文件已消失的记录。需要 admin 作用域。"
164    )]
165    async fn rebuild_assets_index(
166        &self,
167        Parameters(_p): Parameters<EmptyMediaParams>,
168        Extension(parts): Extension<http::request::Parts>,
169    ) -> Result<CallToolResult, McpError> {
170        require_admin(&parts, "rebuild_assets_index")?;
171        let result = crate::api::assets::rebuild::rebuild_assets_index_impl()
172            .await
173            .map_err(|e| internal(e, "rebuild_assets_index"))?;
174        ok_json(result)
175    }
176}
177
178// ---------------------------------------------------------------------------
179// 参数与输出结构
180// ---------------------------------------------------------------------------
181
182#[derive(Debug, Deserialize, schemars::JsonSchema)]
183pub struct UploadMediaParams {
184    /// 图片的 https URL(服务端抓取,二进制不经 JSON-RPC)。
185    pub url: String,
186    /// 保存素材 alt;重复上传时省略保留原值,空白清除,不回写已有文章。
187    #[serde(default)]
188    pub alt: Option<String>,
189}
190
191#[derive(Debug, serde::Serialize)]
192struct UploadResult {
193    success: bool,
194    asset_id: String,
195    alt: Option<String>,
196    url: String,
197    reused: bool,
198    width: u32,
199    height: u32,
200    mime: String,
201}
202
203#[derive(Debug, Deserialize, schemars::JsonSchema)]
204pub struct ListAssetsParams {
205    #[serde(default)]
206    pub filter: AssetFilter,
207    #[serde(default)]
208    pub query: String,
209    #[serde(default)]
210    pub sort: AssetSort,
211    #[serde(default = "first_page")]
212    pub page: i32,
213}
214
215fn first_page() -> i32 {
216    1
217}
218
219#[derive(Debug, Deserialize, schemars::JsonSchema)]
220pub struct UpdateAssetAltParams {
221    pub id: String,
222    pub alt: String,
223}
224
225#[derive(Debug, Deserialize, schemars::JsonSchema)]
226pub struct AssetIdParams {
227    pub id: String,
228}
229
230#[derive(Debug, Deserialize, schemars::JsonSchema)]
231pub struct AssetIdsParams {
232    pub ids: Vec<String>,
233}
234
235#[derive(Debug, Deserialize, schemars::JsonSchema)]
236pub struct EmptyMediaParams {}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::mcp::{auth::McpPrincipal, server::YggMcpServer};
242
243    fn parts(scope: Option<TokenScope>) -> http::request::Parts {
244        let (mut parts, _) = http::Request::new(()).into_parts();
245        if let Some(scope) = scope {
246            parts.extensions.insert(McpPrincipal {
247                user_id: 1,
248                scope,
249                token_id: "test".into(),
250            });
251        }
252        parts
253    }
254
255    #[tokio::test]
256    async fn management_rejects_missing_read_and_write_principals_before_io() {
257        let server = YggMcpServer;
258        for scope in [None, Some(TokenScope::Read), Some(TokenScope::Write)] {
259            let list = serde_json::from_str::<ListAssetsParams>("{}").unwrap();
260            assert!(server
261                .list_assets(Parameters(list), Extension(parts(scope)))
262                .await
263                .is_err());
264            assert!(server
265                .update_asset_alt(
266                    Parameters(UpdateAssetAltParams {
267                        id: "invalid".into(),
268                        alt: "alt".into()
269                    }),
270                    Extension(parts(scope))
271                )
272                .await
273                .is_err());
274            assert!(server
275                .delete_asset(
276                    Parameters(AssetIdParams {
277                        id: "invalid".into()
278                    }),
279                    Extension(parts(scope))
280                )
281                .await
282                .is_err());
283            assert!(server
284                .batch_delete_assets(
285                    Parameters(AssetIdsParams { ids: vec![] }),
286                    Extension(parts(scope))
287                )
288                .await
289                .is_err());
290            assert!(server
291                .purge_orphan_assets(Parameters(EmptyMediaParams {}), Extension(parts(scope)))
292                .await
293                .is_err());
294            assert!(server
295                .rebuild_assets_index(Parameters(EmptyMediaParams {}), Extension(parts(scope)))
296                .await
297                .is_err());
298        }
299    }
300
301    #[tokio::test]
302    async fn batch_delete_rejects_empty_and_oversized_batches_before_io() {
303        for ids in [vec![], vec!["id".into(); 101]] {
304            let err = YggMcpServer
305                .batch_delete_assets(
306                    Parameters(AssetIdsParams { ids }),
307                    Extension(parts(Some(TokenScope::Admin))),
308                )
309                .await
310                .unwrap_err();
311            assert!(err.message.contains("1..=100"));
312        }
313    }
314
315    #[test]
316    fn media_router_exposes_upload_and_management_tools() {
317        let tools = YggMcpServer::media_router().list_all();
318        for name in [
319            "upload_media",
320            "list_assets",
321            "update_asset_alt",
322            "delete_asset",
323            "batch_delete_assets",
324            "purge_orphan_assets",
325            "rebuild_assets_index",
326        ] {
327            assert!(tools.iter().any(|tool| tool.name == name), "missing {name}");
328        }
329        let list = tools
330            .iter()
331            .find(|tool| tool.name == "list_assets")
332            .unwrap();
333        let schema = serde_json::to_string(&list.input_schema).unwrap();
334        assert!(schema.contains("Orphan"));
335        assert!(schema.contains("SizeDesc"));
336    }
337
338    #[test]
339    fn list_defaults_and_filters_match_the_schema() {
340        let params: ListAssetsParams = serde_json::from_str("{}").unwrap();
341        assert_eq!(params.page, 1);
342        assert_eq!(params.filter, AssetFilter::All);
343        assert_eq!(params.sort, AssetSort::CreatedDesc);
344        assert!(params.query.is_empty());
345        assert!(serde_json::from_str::<ListAssetsParams>(r#"{"filter":"invalid"}"#).is_err());
346    }
347}