owlen/crates/owlen-mcp-code-server/src/lib.rs

//! MCP server exposing code execution tools with Docker sandboxing.
//!
//! This server provides:
//! - compile_project: Build projects (Rust, Node.js, Python)
//! - run_tests: Execute test suites
//! - format_code: Run code formatters
//! - lint_code: Run linters

pub mod sandbox;
pub mod tools;

use owlen_core::mcp::protocol::{
    methods, ErrorCode, InitializeParams, InitializeResult, RequestId, RpcError, RpcErrorResponse,
    RpcRequest, RpcResponse, ServerCapabilities, ServerInfo, PROTOCOL_VERSION,
};
use owlen_core::tools::{Tool, ToolResult};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::io::{self, AsyncBufReadExt, AsyncWriteExt};

use tools::{CompileProjectTool, FormatCodeTool, LintCodeTool, RunTestsTool};

/// Tool registry for the code server
#[allow(dead_code)]
struct ToolRegistry {
    tools: HashMap<String, Box<dyn Tool + Send + Sync>>,
}

#[allow(dead_code)]
impl ToolRegistry {
    fn new() -> Self {
        let mut tools: HashMap<String, Box<dyn Tool + Send + Sync>> = HashMap::new();
        tools.insert(
            "compile_project".to_string(),
            Box::new(CompileProjectTool::new()),
        );
        tools.insert("run_tests".to_string(), Box::new(RunTestsTool::new()));
        tools.insert("format_code".to_string(), Box::new(FormatCodeTool::new()));
        tools.insert("lint_code".to_string(), Box::new(LintCodeTool::new()));
        Self { tools }
    }

    fn list_tools(&self) -> Vec<owlen_core::mcp::McpToolDescriptor> {
        self.tools
            .values()
            .map(|tool| owlen_core::mcp::McpToolDescriptor {
                name: tool.name().to_string(),
                description: tool.description().to_string(),
                input_schema: tool.schema(),
                requires_network: tool.requires_network(),
                requires_filesystem: tool.requires_filesystem(),
            })
            .collect()
    }

    async fn execute(&self, name: &str, args: Value) -> Result<ToolResult, String> {
        self.tools
            .get(name)
            .ok_or_else(|| format!("Tool not found: {}", name))?
            .execute(args)
            .await
            .map_err(|e| e.to_string())
    }
}

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

    let registry = Arc::new(ToolRegistry::new());

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

                let resp = handle_request(req.clone(), registry.clone()).await;
                match resp {
                    Ok(r) => {
                        let s = serde_json::to_string(&r)?;
                        stdout.write_all(s.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                    }
                    Err(e) => {
                        let err = RpcErrorResponse::new(req.id.clone(), e);
                        let s = serde_json::to_string(&err)?;
                        stdout.write_all(s.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                    }
                }
            }
            Err(e) => {
                eprintln!("Error reading stdin: {}", e);
                break;
            }
        }
    }
    Ok(())
}

#[allow(dead_code)]
async fn handle_request(
    req: RpcRequest,
    registry: Arc<ToolRegistry>,
) -> Result<RpcResponse, RpcError> {
    match req.method.as_str() {
        methods::INITIALIZE => {
            let params: InitializeParams =
                serde_json::from_value(req.params.unwrap_or_else(|| json!({})))
                    .map_err(|e| RpcError::invalid_params(format!("Invalid init params: {}", e)))?;
            if !params.protocol_version.eq(PROTOCOL_VERSION) {
                return Err(RpcError::new(
                    ErrorCode::INVALID_REQUEST,
                    format!(
                        "Incompatible protocol version. Client: {}, Server: {}",
                        params.protocol_version, PROTOCOL_VERSION
                    ),
                ));
            }
            let result = InitializeResult {
                protocol_version: PROTOCOL_VERSION.to_string(),
                server_info: ServerInfo {
                    name: "owlen-mcp-code-server".to_string(),
                    version: env!("CARGO_PKG_VERSION").to_string(),
                },
                capabilities: ServerCapabilities {
                    supports_tools: Some(true),
                    supports_resources: Some(false),
                    supports_streaming: Some(false),
                },
            };
            Ok(RpcResponse::new(
                req.id,
                serde_json::to_value(result).unwrap(),
            ))
        }
        methods::TOOLS_LIST => {
            let tools = registry.list_tools();
            Ok(RpcResponse::new(req.id, json!(tools)))
        }
        methods::TOOLS_CALL => {
            let call = serde_json::from_value::<owlen_core::mcp::McpToolCall>(
                req.params.unwrap_or_else(|| json!({})),
            )
            .map_err(|e| RpcError::invalid_params(format!("Invalid tool call: {}", e)))?;

            let result: ToolResult = registry
                .execute(&call.name, call.arguments)
                .await
                .map_err(|e| RpcError::internal_error(format!("Tool execution failed: {}", e)))?;

            let resp = owlen_core::mcp::McpToolResponse {
                name: call.name,
                success: result.success,
                output: result.output,
                metadata: result.metadata,
                duration_ms: result.duration.as_millis() as u128,
            };
            Ok(RpcResponse::new(
                req.id,
                serde_json::to_value(resp).unwrap(),
            ))
        }
        _ => Err(RpcError::method_not_found(&req.method)),
    }
}