owlen/crates/owlen-mcp-server/src/main.rs

use owlen_core::mcp::protocol::{
    ErrorCode, InitializeParams, InitializeResult, PROTOCOL_VERSION, RequestId, RpcError,
    RpcErrorResponse, RpcRequest, RpcResponse, ServerCapabilities, ServerInfo, is_compatible,
};
use path_clean::PathClean;
use serde::Deserialize;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use tokio::io::{self, AsyncBufReadExt, AsyncWriteExt};

#[derive(Deserialize)]
struct FileArgs {
    path: String,
}

#[derive(Deserialize)]
struct WriteArgs {
    path: String,
    content: String,
}

async fn handle_request(req: &RpcRequest, root: &Path) -> Result<serde_json::Value, RpcError> {
    match req.method.as_str() {
        "initialize" => {
            let params = req
                .params
                .as_ref()
                .ok_or_else(|| RpcError::invalid_params("Missing params for initialize"))?;

            let init_params: InitializeParams =
                serde_json::from_value(params.clone()).map_err(|e| {
                    RpcError::invalid_params(format!("Invalid initialize params: {}", e))
                })?;

            // Check protocol version compatibility
            if !is_compatible(&init_params.protocol_version, PROTOCOL_VERSION) {
                return Err(RpcError::new(
                    ErrorCode::INVALID_REQUEST,
                    format!(
                        "Incompatible protocol version. Client: {}, Server: {}",
                        init_params.protocol_version, PROTOCOL_VERSION
                    ),
                ));
            }

            // Build initialization result
            let result = InitializeResult {
                protocol_version: PROTOCOL_VERSION.to_string(),
                server_info: ServerInfo {
                    name: "owlen-mcp-server".to_string(),
                    version: env!("CARGO_PKG_VERSION").to_string(),
                },
                capabilities: ServerCapabilities {
                    supports_tools: Some(false),
                    supports_resources: Some(true), // Supports read, write, delete
                    supports_streaming: Some(false),
                },
            };

            Ok(serde_json::to_value(result).map_err(|e| {
                RpcError::internal_error(format!("Failed to serialize result: {}", e))
            })?)
        }
        "resources/list" => {
            let params = req
                .params
                .as_ref()
                .ok_or_else(|| RpcError::invalid_params("Missing params"))?;
            let args: FileArgs = serde_json::from_value(params.clone())
                .map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
            resources_list(&args.path, root).await
        }
        "resources/get" => {
            let params = req
                .params
                .as_ref()
                .ok_or_else(|| RpcError::invalid_params("Missing params"))?;
            let args: FileArgs = serde_json::from_value(params.clone())
                .map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
            resources_get(&args.path, root).await
        }
        "resources/write" => {
            let params = req
                .params
                .as_ref()
                .ok_or_else(|| RpcError::invalid_params("Missing params"))?;
            let args: WriteArgs = serde_json::from_value(params.clone())
                .map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
            resources_write(&args.path, &args.content, root).await
        }
        "resources/delete" => {
            let params = req
                .params
                .as_ref()
                .ok_or_else(|| RpcError::invalid_params("Missing params"))?;
            let args: FileArgs = serde_json::from_value(params.clone())
                .map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
            resources_delete(&args.path, root).await
        }
        _ => Err(RpcError::method_not_found(&req.method)),
    }
}

fn sanitize_path(path: &str, root: &Path) -> Result<PathBuf, RpcError> {
    let path = Path::new(path);
    let path = if path.is_absolute() {
        path.strip_prefix("/")
            .map_err(|_| RpcError::invalid_params("Invalid path"))?
            .to_path_buf()
    } else {
        path.to_path_buf()
    };

    let full_path = root.join(path).clean();

    if !full_path.starts_with(root) {
        return Err(RpcError::path_traversal());
    }

    Ok(full_path)
}

async fn resources_list(path: &str, root: &Path) -> Result<serde_json::Value, RpcError> {
    let full_path = sanitize_path(path, root)?;

    let entries = fs::read_dir(full_path).map_err(|e| {
        RpcError::new(
            ErrorCode::RESOURCE_NOT_FOUND,
            format!("Failed to read directory: {}", e),
        )
    })?;

    let mut result = Vec::new();
    for entry in entries {
        let entry = entry.map_err(|e| {
            RpcError::internal_error(format!("Failed to read directory entry: {}", e))
        })?;
        result.push(entry.file_name().to_string_lossy().to_string());
    }

    Ok(serde_json::json!(result))
}

async fn resources_get(path: &str, root: &Path) -> Result<serde_json::Value, RpcError> {
    let full_path = sanitize_path(path, root)?;

    let content = fs::read_to_string(full_path).map_err(|e| {
        RpcError::new(
            ErrorCode::RESOURCE_NOT_FOUND,
            format!("Failed to read file: {}", e),
        )
    })?;

    Ok(serde_json::json!(content))
}

async fn resources_write(
    path: &str,
    content: &str,
    root: &Path,
) -> Result<serde_json::Value, RpcError> {
    let full_path = sanitize_path(path, root)?;
    // Ensure parent directory exists
    if let Some(parent) = full_path.parent() {
        std::fs::create_dir_all(parent).map_err(|e| {
            RpcError::internal_error(format!("Failed to create parent directories: {}", e))
        })?;
    }
    std::fs::write(full_path, content)
        .map_err(|e| RpcError::internal_error(format!("Failed to write file: {}", e)))?;
    Ok(serde_json::json!(null))
}

async fn resources_delete(path: &str, root: &Path) -> Result<serde_json::Value, RpcError> {
    let full_path = sanitize_path(path, root)?;
    if full_path.is_file() {
        std::fs::remove_file(full_path)
            .map_err(|e| RpcError::internal_error(format!("Failed to delete file: {}", e)))?;
        Ok(serde_json::json!(null))
    } else {
        Err(RpcError::new(
            ErrorCode::RESOURCE_NOT_FOUND,
            "Path does not refer to a file",
        ))
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let root = env::current_dir()?;
    let mut stdin = io::BufReader::new(io::stdin());
    let mut stdout = io::stdout();

    loop {
        let mut line = String::new();
        match stdin.read_line(&mut line).await {
            Ok(0) => {
                // EOF
                break;
            }
            Ok(_) => {
                let req: RpcRequest = match serde_json::from_str(&line) {
                    Ok(req) => req,
                    Err(e) => {
                        let err_resp = RpcErrorResponse::new(
                            RequestId::Number(0),
                            RpcError::parse_error(format!("Parse error: {}", e)),
                        );
                        let resp_str = serde_json::to_string(&err_resp)?;
                        stdout.write_all(resp_str.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                        continue;
                    }
                };

                let request_id = req.id.clone();

                match handle_request(&req, &root).await {
                    Ok(result) => {
                        let resp = RpcResponse::new(request_id, result);
                        let resp_str = serde_json::to_string(&resp)?;
                        stdout.write_all(resp_str.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                    }
                    Err(error) => {
                        let err_resp = RpcErrorResponse::new(request_id, error);
                        let resp_str = serde_json::to_string(&err_resp)?;
                        stdout.write_all(resp_str.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                    }
                }
            }
            Err(e) => {
                // Handle read error
                eprintln!("Error reading from stdin: {}", e);
                break;
            }
        }
    }

    Ok(())
}