feat: add new tool for reading media files (Image / Audio) (#43)

* roots support

* feat: handle roots from client

* chore: better messaging when allowed_directory list is empty

* chore: better messaging

* update tests and cleanup

* update docs

* typo

* clippy

* add read_media tool

* add read multiple media files

* typo

* cleanup

Author: hashemix

Cargo.lock

@@ -37,6 +37,8 @@ futures = "0.3"
 tokio-util = "0.7"
 async_zip = { version = "0.0", features = ["full"] }
 grep = "0.3"
+base64 =  "0.22"
+infer = "0.19.0"
 
 [dev-dependencies]
 tempfile = "3.2"

Cargo.toml

@@ -34,4 +34,10 @@ pub enum ServiceError {
     ZipError(#[from] ZipError),
     #[error("{0}")]
     GlobPatternError(#[from] PatternError),
+    #[error("File size exceeds the maximum allowed limit of {0} bytes")]
+    FileTooLarge(usize),
+    #[error("File size is below the minimum required limit of {0} bytes")]
+    FileTooSmall(usize),
+    #[error("The file is either not an image/audio type or is unsupported (mime:{0}).")]
+    InvalidMediaFile(String),
 }

src/error.rs

@@ -5,7 +5,9 @@ use crate::{
     tools::EditOperation,
 };
 use async_zip::tokio::{read::seek::ZipFileReader, write::ZipFileWriter};
+use base64::{engine::general_purpose, write::EncoderWriter};
 use file_info::FileInfo;
+use futures::{StreamExt, stream};
 use glob::Pattern;
 use grep::{
     matcher::{Match, Matcher},
@@ -19,12 +21,13 @@ use std::{
     collections::HashSet,
     env,
     fs::{self},
+    io::Write,
     path::{Path, PathBuf},
     sync::Arc,
 };
 use tokio::{
-    fs::File,
-    io::{AsyncWriteExt, BufReader},
+    fs::{File, metadata},
+    io::{AsyncReadExt, AsyncWriteExt, BufReader},
     sync::RwLock,
 };
 use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
@@ -36,6 +39,7 @@ use walkdir::WalkDir;
 
 const SNIPPET_MAX_LENGTH: usize = 200;
 const SNIPPET_BACKWARD_CHARS: usize = 30;
+const MAX_CONCURRENT_FILE_READ: usize = 5;
 
 type PathResultList = Vec<Result<PathBuf, ServiceError>>;
 
@@ -432,7 +436,108 @@ impl FileSystemService {
         Ok(result_message)
     }
 
-    pub async fn read_file(&self, file_path: &Path) -> ServiceResult<String> {
+    pub fn mime_from_path(&self, path: &Path) -> ServiceResult<infer::Type> {
+        let is_svg = path
+            .extension()
+            .is_some_and(|e| e.to_str().is_some_and(|s| s == "svg"));
+        // consider it is a svg file as we cannot detect svg from bytes pattern
+        if is_svg {
+            return Ok(infer::Type::new(
+                infer::MatcherType::Image,
+                "image/svg+xml",
+                "svg",
+                |_: &[u8]| true,
+            ));
+
+            // infer::Type::new(infer::MatcherType::Image, "", "svg",);
+        }
+        let kind = infer::get_from_path(path)?.ok_or(ServiceError::FromString(
+            "File tyle is unknown!".to_string(),
+        ))?;
+        Ok(kind)
+    }
+
+    pub async fn validate_file_size<P: AsRef<Path>>(
+        &self,
+        path: P,
+        min_bytes: Option<usize>,
+        max_bytes: Option<usize>,
+    ) -> ServiceResult<()> {
+        if min_bytes.is_none() && max_bytes.is_none() {
+            return Ok(());
+        }
+
+        let file_size = metadata(&path).await?.len() as usize;
+
+        match (min_bytes, max_bytes) {
+            (_, Some(max)) if file_size > max => Err(ServiceError::FileTooLarge(max)),
+            (Some(min), _) if file_size < min => Err(ServiceError::FileTooSmall(min)),
+            _ => Ok(()),
+        }
+    }
+
+    pub async fn read_media_files(
+        &self,
+        paths: Vec<String>,
+        max_bytes: Option<usize>,
+    ) -> ServiceResult<Vec<(infer::Type, String)>> {
+        let results = stream::iter(paths)
+            .map(|path| async {
+                self.read_media_file(Path::new(&path), max_bytes)
+                    .await
+                    .map_err(|e| (path, e))
+            })
+            .buffer_unordered(MAX_CONCURRENT_FILE_READ) // Process up to MAX_CONCURRENT_FILE_READ files concurrently
+            .filter_map(|result| async move { result.ok() })
+            .collect::<Vec<_>>()
+            .await;
+        Ok(results)
+    }
+
+    pub async fn read_media_file(
+        &self,
+        file_path: &Path,
+        max_bytes: Option<usize>,
+    ) -> ServiceResult<(infer::Type, String)> {
+        let allowed_directories = self.allowed_directories().await;
+        let valid_path = self.validate_path(file_path, allowed_directories)?;
+        self.validate_file_size(&valid_path, None, max_bytes)
+            .await?;
+        let kind = self.mime_from_path(&valid_path)?;
+        let content = self.read_file_as_base64(&valid_path).await?;
+        Ok((kind, content))
+    }
+
+    // reads file as base64 efficiently in a streaming manner
+    async fn read_file_as_base64(&self, file_path: &Path) -> ServiceResult<String> {
+        let file = File::open(file_path).await?;
+        let mut reader = BufReader::new(file);
+
+        let mut output = Vec::new();
+        {
+            // Wrap output Vec<u8> in a Base64 encoder writer
+            let mut encoder = EncoderWriter::new(&mut output, &general_purpose::STANDARD);
+
+            let mut buffer = [0u8; 8192];
+            loop {
+                let n = reader.read(&mut buffer).await?;
+                if n == 0 {
+                    break;
+                }
+                // Write raw bytes to the Base64 encoder
+                encoder.write_all(&buffer[..n])?;
+            }
+            // Make sure to flush any remaining bytes
+            encoder.flush()?;
+        } // drop encoder before consuming output
+
+        // Convert the Base64 bytes to String (safe UTF-8)
+        let base64_string =
+            String::from_utf8(output).map_err(|err| ServiceError::FromString(format!("{err}")))?;
+        Ok(base64_string)
+    }
+
+    pub async fn read_text_file(&self, file_path: &Path) -> ServiceResult<String> {
         let allowed_directories = self.allowed_directories().await;
         let valid_path = self.validate_path(file_path, allowed_directories)?;
         let content = tokio::fs::read_to_string(valid_path).await?;

src/fs_service.rs

@@ -88,47 +88,45 @@ impl FileSystemHandler {
             let fs_service = self.fs_service.clone();
             let mcp_roots_support = self.mcp_roots_support;
             // retrieve roots from the client and update the allowed directories accordingly
-            tokio::spawn(async move {
-                let roots = match runtime.clone().list_roots(None).await {
-                    Ok(roots_result) => roots_result.roots,
-                    Err(_err) => {
-                        vec![]
-                    }
-                };
-
-                let valid_roots = if roots.is_empty() {
+            let roots = match runtime.clone().list_roots(None).await {
+                Ok(roots_result) => roots_result.roots,
+                Err(_err) => {
                     vec![]
-                } else {
-                    let roots: Vec<_> = roots.iter().map(|v| v.uri.as_str()).collect();
-
-                    match fs_service.valid_roots(roots) {
-                        Ok((roots, skipped)) => {
-                            if let Some(message) = skipped {
-                                let _ = runtime.stderr_message(message.to_string()).await;
-                            }
-                            roots
+                }
+            };
+
+            let valid_roots = if roots.is_empty() {
+                vec![]
+            } else {
+                let roots: Vec<_> = roots.iter().map(|v| v.uri.as_str()).collect();
+
+                match fs_service.valid_roots(roots) {
+                    Ok((roots, skipped)) => {
+                        if let Some(message) = skipped {
+                            let _ = runtime.stderr_message(message.to_string()).await;
                         }
-                        Err(_err) => vec![],
+                        roots
                     }
-                };
+                    Err(_err) => vec![],
+                }
+            };
 
-                if valid_roots.is_empty() && !mcp_roots_support {
-                    let message = if allowed_directories.is_empty() {
-                        "Server cannot operate: No allowed directories available. Server was started without command-line directories and client provided empty roots. Please either: 1) Start server with directory arguments, or 2) Use a client that supports MCP roots protocol and provides valid root directories."
-                    } else {
-                        "Client provided empty roots. Allowed directories passed from command-line will be used."
-                    };
-                    let _ = runtime.stderr_message(message.to_string()).await;
+            if valid_roots.is_empty() && !mcp_roots_support {
+                let message = if allowed_directories.is_empty() {
+                    "Server cannot operate: No allowed directories available. Server was started without command-line directories and client provided empty roots. Please either: 1) Start server with directory arguments, or 2) Use a client that supports MCP roots protocol and provides valid root directories."
                 } else {
-                    let num_valid_roots = valid_roots.len();
+                    "Client provided empty roots. Allowed directories passed from command-line will be used."
+                };
+                let _ = runtime.stderr_message(message.to_string()).await;
+            } else {
+                let num_valid_roots = valid_roots.len();
 
-                    fs_service.update_allowed_paths(valid_roots).await;
-                    let message = format!(
-                        "Updated allowed directories from MCP roots: {num_valid_roots} valid directories",
-                    );
-                    let _ = runtime.stderr_message(message.to_string()).await;
-                }
-            });
+                fs_service.update_allowed_paths(valid_roots).await;
+                let message = format!(
+                    "Updated allowed directories from MCP roots: {num_valid_roots} valid directories",
+                );
+                let _ = runtime.stderr_message(message.to_string()).await;
+            }
         }
     }
 }
@@ -196,11 +194,17 @@ impl ServerHandler for FileSystemHandler {
         }
 
         match tool_params {
-            FileSystemTools::ReadFileTool(params) => {
-                ReadFileTool::run_tool(params, &self.fs_service).await
+            FileSystemTools::ReadMediaFileTool(params) => {
+                ReadMediaFileTool::run_tool(params, &self.fs_service).await
+            }
+            FileSystemTools::ReadMultipleMediaFilesTool(params) => {
+                ReadMultipleMediaFilesTool::run_tool(params, &self.fs_service).await
+            }
+            FileSystemTools::ReadTextFileTool(params) => {
+                ReadTextFileTool::run_tool(params, &self.fs_service).await
             }
-            FileSystemTools::ReadMultipleFilesTool(params) => {
-                ReadMultipleFilesTool::run_tool(params, &self.fs_service).await
+            FileSystemTools::ReadMultipleTextFilesTool(params) => {
+                ReadMultipleTextFilesTool::run_tool(params, &self.fs_service).await
             }
             FileSystemTools::WriteFileTool(params) => {
                 WriteFileTool::run_tool(params, &self.fs_service).await