Skip to main content

headless_lms_mcp/handlers/
uploads.rs

1//! `upload_file`: stores an image or file for the course the way the CMS media upload does.
2
3use std::path::Path;
4
5use actix_web::web::Bytes;
6use async_trait::async_trait;
7use base64::Engine;
8use futures::stream;
9use headless_lms_base::error::backend_error::BackendError;
10use headless_lms_models::{
11    ModelErrorType,
12    file_uploads::{self, NewUpload},
13};
14use headless_lms_utils::file_store::{
15    GenericPayload,
16    media_uploads::{self, MediaKind, MediaOwner},
17};
18use serde_json::json;
19
20use crate::access::{authorize_course, connection};
21use crate::state::McpState;
22use crate::tools::args::UploadFileArgs;
23use crate::tools::{RequestContext, ToolError, ToolHandler, ToolOutcome, ToolResult};
24
25/// Handler of the `upload_file` tool.
26pub struct UploadFile;
27
28/// Decodes on a blocking thread: the content can be tens of megabytes.
29async fn decode(content_base64: String) -> Result<Vec<u8>, ToolError> {
30    tokio::task::spawn_blocking(move || {
31        let compact = if content_base64.contains(char::is_whitespace) {
32            content_base64.split_whitespace().collect()
33        } else {
34            content_base64
35        };
36        base64::engine::general_purpose::STANDARD.decode(compact)
37    })
38    .await
39    .map_err(|e| ToolError::Internal(format!("the decoding task failed: {e}")))?
40    .map_err(|e| ToolError::InvalidArguments {
41        message: format!("content_base64 isn't valid base64: {e}"),
42    })
43}
44
45#[async_trait(?Send)]
46impl ToolHandler<McpState> for UploadFile {
47    type Args = UploadFileArgs;
48
49    async fn call(
50        &self,
51        state: &McpState,
52        ctx: &RequestContext,
53        args: UploadFileArgs,
54    ) -> Result<ToolOutcome, ToolError> {
55        let course_id = args.course.id();
56        let file_name = Path::new(args.file_name.trim())
57            .file_name()
58            .and_then(|n| n.to_str())
59            .filter(|n| !n.is_empty())
60            .ok_or_else(|| ToolError::InvalidArguments {
61                message: "file_name must be a file name with an extension, e.g. diagram.png".into(),
62            })?
63            .to_string();
64        let contents = decode(args.content_base64).await?;
65        if contents.is_empty() {
66            return Err(ToolError::InvalidArguments {
67                message: "the file is empty".into(),
68            });
69        }
70        let mime_type = mime_guess::from_path(&file_name).first_or_octet_stream();
71        let size_limit = media_uploads::size_limit_bytes(Some(mime_type.type_().as_str()));
72        if contents.len() > size_limit as usize {
73            return Err(ToolError::InvalidArguments {
74                message: format!(
75                    "the file is {} bytes; {} files may be at most {size_limit} bytes",
76                    contents.len(),
77                    mime_type.type_()
78                ),
79            });
80        }
81
82        let user_id = ctx.principal.actor.user_id;
83        let mut conn = connection(state).await?;
84        authorize_course(&mut conn, user_id, course_id).await?;
85        drop(conn);
86        let path = media_uploads::media_path(
87            MediaOwner::Course(course_id),
88            MediaKind::from_mime_top_level(mime_type.type_().as_str()),
89            &media_uploads::stored_file_name(&file_name),
90        );
91        let size_bytes = contents.len();
92        let payload: GenericPayload =
93            Box::pin(stream::once(async move { Ok(Bytes::from(contents)) }));
94        file_uploads::store_and_insert(
95            &state.pool,
96            state.file_store.as_ref(),
97            NewUpload {
98                path: &path,
99                name: &file_name,
100                mime: mime_type.essence_str(),
101                uploader: Some(user_id),
102                size_bytes: Some(size_bytes as i64),
103            },
104            payload,
105        )
106        .await
107        .map_err(|e| match e.error_type() {
108            ModelErrorType::Util => ToolError::retryable(
109                format!("storing the file failed ({e}); the same call can be retried"),
110                2,
111            ),
112            _ => e.into(),
113        })?;
114        let url = state.file_store.get_download_url(&path, &state.app_conf);
115        Ok(ToolResult::success(json!({
116            "url": url,
117            "file_name": file_name,
118            "mime_type": mime_type.essence_str(),
119            "size_bytes": size_bytes,
120        }))
121        .into())
122    }
123}