src/main.rs

use std::time::Duration;

use rmcp::{serde_json, transport::StreamableHttpServer};

use rmcp::{Error as McpError, ServerHandler, model::*, schemars, tool};
use tracing_subscriber::{
    layer::SubscriberExt,
    util::SubscriberInitExt,
    {self},
};

const BIND_ADDRESS: &str = "0.0.0.0:9090";

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "hf_file_mcp=debug".to_string().into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let ct = StreamableHttpServer::serve(BIND_ADDRESS.parse()?)
        .await?
        .with_service(HfFile::new);

    tokio::signal::ctrl_c().await?;
    ct.cancel();
    Ok(())
}
#[derive(Clone)]
pub struct HfFile {
    http_client: reqwest::Client,
}

#[tool(tool_box)]
impl HfFile {
    pub fn new() -> Self {
        let http_client = reqwest::Client::builder()
            .timeout(Duration::from_secs(10))
            .build()
            .expect("Failed to create HTTP client");
        Self {
            http_client: http_client,
        }
    }
    #[tool(description = "List files in HF Hub")]
    async fn list_repo_file(
        &self,
        #[tool(param)]
        #[schemars(description = "repo")]
        repository: String,
        #[tool(param)]
        #[schemars(description = "repo type (models, datasets, spaces.)")]
        repo_type: String,
    ) -> Result<CallToolResult, McpError> {
        let response = self
            .http_client
            .get(format!(
                "https://huggingface.co/api/{repo_type}/{repository}"
            ))
            .send()
            .await
            .unwrap();

        if !response.status().is_success() {
            return Err(McpError::internal_error("Failed to list files", None));
        }
        let body = response.json::<serde_json::Value>().await.unwrap();
        let siblings = body.get("siblings").unwrap().as_array().unwrap();
        let mut contents = Vec::new();
        for sibling in siblings {
            contents.push(sibling.get("rfilename").unwrap().as_str().unwrap());
        }
        Ok(CallToolResult::success(vec![Content::json(contents)?]))
    }

    #[tool(description = "List files in HF Hub")]
    async fn get_file_content(
        &self,
        #[tool(param)]
        #[schemars(description = "repo")]
        repository: String,
        #[tool(param)]
        #[schemars(description = "repo type (models, datasets, spaces.)")]
        repo_type: String,
        #[tool(param)]
        #[schemars(description = "file path")]
        file_path: String,
    ) -> Result<CallToolResult, McpError> {
        let response = self
            .http_client
            .get(format!(
                "https://huggingface.co/{repo_type}/{repository}/resolve/main/{file_path}"
            ))
            .send()
            .await
            .unwrap();

        if !response.status().is_success() {
            return Err(McpError::internal_error("Failed to get file content", None));
        }
        Ok(CallToolResult::success(vec![Content::text(
            response.text().await.unwrap(),
        )]))
    }
}

#[tool(tool_box)]
impl ServerHandler for HfFile {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2025_03_26,
            capabilities: ServerCapabilities::builder()
                .enable_tools()
                .build(),
            server_info: Implementation::from_build_env(),
            instructions: Some("This server provides a Hugging Face file tool that can list and get files from Hugging Face Hub.".to_string()),
        }
    }
}