//! MCP Server implementation for CoderLib
//!
//! This module provides MCP server capabilities, allowing CoderLib to act as an MCP server
//! and expose its tools, resources, and prompts to other MCP clients.
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use std::future::Future;
use tokio::sync::RwLock;
use serde_json::{json, Value as JsonValue};
use walkdir::WalkDir;
#[cfg(feature = "mcp")]
use rmcp::{
Error as McpError, RoleServer, ServerHandler,
model::*,
service::RequestContext,
};
use crate::tools::{ToolRegistry, router::ToolRouter};
use crate::integration::HostIntegration;
/// CoderLib MCP Server that exposes tools, resources, and prompts
pub struct CoderLibServer {
/// Tool registry containing all available tools
tool_registry: Arc<RwLock<ToolRegistry>>,
/// Host integration for accessing file system and other resources
host: Arc<dyn HostIntegration>,
/// Server configuration
config: ServerConfig,
/// Resource providers
resources: Arc<RwLock<HashMap<String, Box<dyn ResourceProvider>>>>,
/// Prompt templates
prompts: Arc<RwLock<HashMap<String, PromptTemplate>>>,
}
/// Configuration for the CoderLib MCP Server
#[derive(Debug, Clone)]
pub struct ServerConfig {
/// Server name
pub name: String,
/// Server version
pub version: String,
/// Server description
pub description: String,
/// Instructions for using the server
pub instructions: Option<String>,
/// Whether to enable tools
pub enable_tools: bool,
/// Whether to enable resources
pub enable_resources: bool,
/// Whether to enable prompts
pub enable_prompts: bool,
}
impl Default for ServerConfig {
fn default() -> Self {
Self {
name: "CoderLib".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
description: "CoderLib MCP Server - AI-powered code generation and analysis tools".to_string(),
instructions: Some("This server provides comprehensive code generation, analysis, and manipulation tools. Use the available tools to read files, analyze code, generate content, and perform various development tasks.".to_string()),
enable_tools: true,
enable_resources: true,
enable_prompts: true,
}
}
}
/// Trait for providing resources to the MCP server
#[async_trait]
pub trait ResourceProvider: Send + Sync {
/// List available resources
async fn list_resources(&self, cursor: Option<String>) -> Result<(Vec<Resource>, Option<String>), McpError>;
/// Read a specific resource
async fn read_resource(&self, uri: &str) -> Result<Vec<ResourceContents>, McpError>;
/// Get the URI scheme this provider handles
fn scheme(&self) -> &str;
}
/// Template for generating prompts
#[derive(Debug, Clone)]
pub struct PromptTemplate {
/// Template name
pub name: String,
/// Template description
pub description: Option<String>,
/// Template arguments
pub arguments: Vec<PromptArgument>,
/// Template content generator
pub generator: fn(&HashMap<String, JsonValue>) -> Result<Vec<PromptMessage>, McpError>,
}
impl CoderLibServer {
/// Create a new CoderLib MCP server
pub fn new(
tool_registry: Arc<RwLock<ToolRegistry>>,
host: Arc<dyn HostIntegration>,
config: ServerConfig,
) -> Self {
Self {
tool_registry,
host,
config,
resources: Arc::new(RwLock::new(HashMap::new())),
prompts: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Register a resource provider
pub async fn register_resource_provider<P: ResourceProvider + 'static>(&self, provider: P) {
let scheme = provider.scheme().to_string();
self.resources.write().await.insert(scheme, Box::new(provider));
}
/// Register a prompt template
pub async fn register_prompt_template(&self, template: PromptTemplate) {
let name = template.name.clone();
self.prompts.write().await.insert(name, template);
}
/// Get all available tools from the registry
async fn get_available_tools(&self) -> Result<Vec<Tool>, McpError> {
let registry = self.tool_registry.read().await;
let tool_schemas = registry.get_tool_schemas();
let tools = tool_schemas
.into_iter()
.map(|(name, schema)| Tool {
name: name.into(),
description: Some(schema.description.into()),
input_schema: std::sync::Arc::new(
schema.input_schema.as_object().unwrap_or(&serde_json::Map::new()).clone()
),
annotations: None,
})
.collect();
Ok(tools)
}
/// Execute a tool by name
async fn execute_tool(&self, name: &str, arguments: Option<JsonValue>) -> Result<rmcp::model::CallToolResult, McpError> {
let registry = self.tool_registry.read().await;
let enhanced_router = registry.enhanced_router();
let params = arguments.unwrap_or(JsonValue::Null);
match enhanced_router.execute_tool(name, params, self.host.as_ref()).await {
Ok(result) => {
// Convert our CallToolResult to rmcp's CallToolResult
let content = result.content
.into_iter()
.map(|c| match c {
crate::tools::router::Content::Text { text } => {
rmcp::model::Content::text(text)
}
crate::tools::router::Content::Image { data, mime_type } => {
rmcp::model::Content::image(data, mime_type)
}
crate::tools::router::Content::Resource { resource } => {
rmcp::model::Content::text(format!("Resource: {}", resource.uri))
}
})
.collect();
Ok(rmcp::model::CallToolResult {
content,
is_error: Some(result.is_error),
})
}
Err(e) => {
tracing::error!("Tool execution failed: {}", e);
Ok(rmcp::model::CallToolResult {
content: vec![rmcp::model::Content::text(format!("Tool execution failed: {}", e))],
is_error: Some(true),
})
}
}
}
}
impl ServerHandler for CoderLibServer {
fn get_info(&self) -> ServerInfo {
let capabilities = ServerCapabilities {
tools: if self.config.enable_tools {
Some(rmcp::model::ToolsCapability { list_changed: None })
} else { None },
resources: if self.config.enable_resources {
Some(rmcp::model::ResourcesCapability {
list_changed: None,
subscribe: None,
})
} else { None },
prompts: if self.config.enable_prompts {
Some(rmcp::model::PromptsCapability { list_changed: None })
} else { None },
completions: None,
logging: None,
experimental: None,
};
ServerInfo {
protocol_version: ProtocolVersion::V_2024_11_05,
capabilities,
server_info: Implementation {
name: self.config.name.clone(),
version: self.config.version.clone(),
},
instructions: self.config.instructions.clone(),
}
}
fn list_tools(
&self,
_request: Option<PaginatedRequestParam>,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
async move {
if !self.config.enable_tools {
return Ok(ListToolsResult {
tools: vec![],
next_cursor: None,
});
}
let tools = self.get_available_tools().await?;
Ok(ListToolsResult {
tools,
next_cursor: None,
})
}
}
fn call_tool(
&self,
CallToolRequestParam { name, arguments }: CallToolRequestParam,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CallToolResult, McpError>> + Send + '_ {
async move {
if !self.config.enable_tools {
return Err(McpError::invalid_params("Tools are disabled", None));
}
let args = arguments.map(|args| serde_json::Value::Object(args));
self.execute_tool(&name, args).await
}
}
fn list_resources(
&self,
request: Option<PaginatedRequestParam>,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListResourcesResult, McpError>> + Send + '_ {
async move {
if !self.config.enable_resources {
return Ok(ListResourcesResult {
resources: vec![],
next_cursor: None,
});
}
let cursor = request.and_then(|r| r.cursor);
let mut all_resources = Vec::new();
let mut next_cursor = None;
let providers = self.resources.read().await;
for provider in providers.values() {
let (resources, provider_cursor) = provider.list_resources(cursor.clone()).await?;
all_resources.extend(resources);
if provider_cursor.is_some() {
next_cursor = provider_cursor;
}
}
Ok(ListResourcesResult {
resources: all_resources,
next_cursor,
})
}
}
fn read_resource(
&self,
ReadResourceRequestParam { uri }: ReadResourceRequestParam,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ReadResourceResult, McpError>> + Send + '_ {
async move {
if !self.config.enable_resources {
return Err(McpError::invalid_params("Resources are disabled", None));
}
// Parse URI to determine scheme
let scheme = uri.split(':').next().unwrap_or("");
let providers = self.resources.read().await;
if let Some(provider) = providers.get(scheme) {
let contents = provider.read_resource(&uri).await?;
Ok(ReadResourceResult { contents })
} else {
Err(McpError::resource_not_found(
"Resource not found",
Some(json!({ "uri": uri })),
))
}
}
}
fn list_prompts(
&self,
_request: Option<PaginatedRequestParam>,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListPromptsResult, McpError>> + Send + '_ {
async move {
if !self.config.enable_prompts {
return Ok(ListPromptsResult {
prompts: vec![],
next_cursor: None,
});
}
let prompts_map = self.prompts.read().await;
let prompts = prompts_map
.values()
.map(|template| Prompt {
name: template.name.clone(),
description: template.description.clone(),
arguments: Some(template.arguments.clone()),
})
.collect();
Ok(ListPromptsResult {
prompts,
next_cursor: None,
})
}
}
fn get_prompt(
&self,
GetPromptRequestParam { name, arguments }: GetPromptRequestParam,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetPromptResult, McpError>> + Send + '_ {
async move {
if !self.config.enable_prompts {
return Err(McpError::invalid_params("Prompts are disabled", None));
}
let prompts_map = self.prompts.read().await;
if let Some(template) = prompts_map.get(&name) {
let args = arguments.unwrap_or_default();
let args_map: HashMap<String, JsonValue> = args
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
let messages = (template.generator)(&args_map)?;
Ok(GetPromptResult {
description: template.description.clone(),
messages,
})
} else {
Err(McpError::invalid_params(
format!("Prompt '{}' not found", name),
None,
))
}
}
}
fn list_resource_templates(
&self,
_request: Option<PaginatedRequestParam>,
_: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListResourceTemplatesResult, McpError>> + Send + '_ {
async move {
// For now, return empty list - can be extended later
Ok(ListResourceTemplatesResult {
resource_templates: Vec::new(),
next_cursor: None,
})
}
}
fn initialize(
&self,
_request: InitializeRequestParam,
_context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<InitializeResult, McpError>> + Send + '_ {
async move {
tracing::info!("CoderLib MCP Server initialized");
Ok(self.get_info())
}
}
}
/// File system resource provider
pub struct FileSystemResourceProvider {
/// Base path for file system access
base_path: std::path::PathBuf,
}
impl FileSystemResourceProvider {
/// Create a new file system resource provider
pub fn new(base_path: std::path::PathBuf) -> Self {
Self { base_path }
}
}
#[async_trait]
impl ResourceProvider for FileSystemResourceProvider {
async fn list_resources(&self, _cursor: Option<String>) -> Result<(Vec<Resource>, Option<String>), McpError> {
use std::fs;
use walkdir::WalkDir;
let mut resources = Vec::new();
// Walk through the base directory and collect files
for entry in WalkDir::new(&self.base_path)
.max_depth(3) // Limit depth to avoid too many files
.into_iter()
.filter_map(|e| e.ok())
.filter(|e| e.file_type().is_file())
{
let path = entry.path();
let relative_path = path.strip_prefix(&self.base_path)
.unwrap_or(path)
.to_string_lossy();
let uri = format!("file://{}", relative_path);
let name = path.file_name()
.and_then(|n| n.to_str())
.unwrap_or("unknown")
.to_string();
let resource = rmcp::model::Resource {
raw: RawResource {
uri: uri.clone(),
name,
description: Some(format!("File: {}", relative_path)),
mime_type: Some(Self::guess_mime_type(path)),
size: None,
},
annotations: None,
};
resources.push(resource);
}
Ok((resources, None))
}
async fn read_resource(&self, uri: &str) -> Result<Vec<ResourceContents>, McpError> {
use std::fs;
// Parse the file:// URI
let path_str = uri.strip_prefix("file://").unwrap_or(uri);
let full_path = self.base_path.join(path_str);
// Security check: ensure the path is within the base directory
if !full_path.starts_with(&self.base_path) {
return Err(McpError::invalid_params(
"Path traversal not allowed",
Some(json!({ "uri": uri, "path": path_str })),
));
}
// Check if file exists
if !full_path.exists() {
return Err(McpError::resource_not_found(
"File not found",
Some(json!({ "uri": uri, "path": full_path.display().to_string() })),
));
}
// Read file content
match fs::read_to_string(&full_path) {
Ok(content) => {
let mime_type = Self::guess_mime_type(&full_path);
Ok(vec![ResourceContents::TextResourceContents {
text: content,
uri: uri.to_string(),
mime_type: Some(mime_type),
}])
}
Err(e) => {
Err(McpError::internal_error(
format!("Failed to read file: {}", e),
Some(json!({ "uri": uri, "error": e.to_string() })),
))
}
}
}
fn scheme(&self) -> &str {
"file"
}
}
impl FileSystemResourceProvider {
/// Guess MIME type based on file extension
fn guess_mime_type(path: &std::path::Path) -> String {
match path.extension().and_then(|ext| ext.to_str()) {
Some("rs") => "text/rust".to_string(),
Some("py") => "text/python".to_string(),
Some("js") => "text/javascript".to_string(),
Some("ts") => "text/typescript".to_string(),
Some("html") => "text/html".to_string(),
Some("css") => "text/css".to_string(),
Some("json") => "application/json".to_string(),
Some("xml") => "application/xml".to_string(),
Some("md") => "text/markdown".to_string(),
Some("txt") => "text/plain".to_string(),
Some("toml") => "text/toml".to_string(),
Some("yaml") | Some("yml") => "text/yaml".to_string(),
Some("png") => "image/png".to_string(),
Some("jpg") | Some("jpeg") => "image/jpeg".to_string(),
Some("gif") => "image/gif".to_string(),
Some("svg") => "image/svg+xml".to_string(),
_ => "text/plain".to_string(),
}
}
}
/// Memory-based resource provider for storing dynamic content
pub struct MemoryResourceProvider {
/// In-memory storage for resources
resources: Arc<RwLock<HashMap<String, MemoryResource>>>,
}
/// A resource stored in memory
#[derive(Debug, Clone)]
pub struct MemoryResource {
/// Resource name
pub name: String,
/// Resource content
pub content: String,
/// MIME type
pub mime_type: String,
/// Description
pub description: Option<String>,
}
impl MemoryResourceProvider {
/// Create a new memory resource provider
pub fn new() -> Self {
Self {
resources: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Add a resource to memory
pub async fn add_resource(&self, uri: String, resource: MemoryResource) {
self.resources.write().await.insert(uri, resource);
}
/// Remove a resource from memory
pub async fn remove_resource(&self, uri: &str) -> Option<MemoryResource> {
self.resources.write().await.remove(uri)
}
/// Update a resource in memory
pub async fn update_resource(&self, uri: &str, content: String) -> Result<(), McpError> {
let mut resources = self.resources.write().await;
if let Some(resource) = resources.get_mut(uri) {
resource.content = content;
Ok(())
} else {
Err(McpError::resource_not_found(
"Memory resource not found",
Some(json!({ "uri": uri })),
))
}
}
}
impl Default for MemoryResourceProvider {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl ResourceProvider for MemoryResourceProvider {
async fn list_resources(&self, _cursor: Option<String>) -> Result<(Vec<Resource>, Option<String>), McpError> {
let resources_map = self.resources.read().await;
let resources = resources_map
.iter()
.map(|(uri, resource)| {
rmcp::model::Resource {
raw: RawResource {
uri: uri.clone(),
name: resource.name.clone(),
description: resource.description.clone(),
mime_type: Some(resource.mime_type.clone()),
size: None,
},
annotations: None,
}
})
.collect();
Ok((resources, None))
}
async fn read_resource(&self, uri: &str) -> Result<Vec<ResourceContents>, McpError> {
let resources_map = self.resources.read().await;
if let Some(resource) = resources_map.get(uri) {
Ok(vec![ResourceContents::TextResourceContents {
text: resource.content.clone(),
uri: uri.to_string(),
mime_type: Some(resource.mime_type.clone()),
}])
} else {
Err(McpError::resource_not_found(
"Memory resource not found",
Some(json!({ "uri": uri })),
))
}
}
fn scheme(&self) -> &str {
"memory"
}
}
/// Default prompt templates for CoderLib
impl CoderLibServer {
/// Register default prompt templates
pub async fn register_default_prompts(&self) {
// Code analysis prompt
self.register_prompt_template(PromptTemplate {
name: "analyze_code".to_string(),
description: Some("Analyze code for quality, patterns, and potential improvements".to_string()),
arguments: vec![
PromptArgument {
name: "code".to_string(),
description: Some("The code to analyze".to_string()),
required: Some(true),
},
PromptArgument {
name: "language".to_string(),
description: Some("Programming language (optional)".to_string()),
required: Some(false),
},
PromptArgument {
name: "focus".to_string(),
description: Some("Specific aspect to focus on (performance, security, readability, etc.)".to_string()),
required: Some(false),
},
],
generator: |args| {
let code = args.get("code")
.and_then(|v| v.as_str())
.ok_or_else(|| McpError::invalid_params("Missing 'code' parameter", None))?;
let language = args.get("language")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let focus = args.get("focus")
.and_then(|v| v.as_str())
.unwrap_or("general quality");
let prompt = format!(
"Please analyze the following {} code with a focus on {}:\n\n```{}\n{}\n```\n\nProvide insights on:\n1. Code quality and structure\n2. Potential improvements\n3. Best practices adherence\n4. Any issues or concerns",
language, focus, language, code
);
Ok(vec![PromptMessage {
role: PromptMessageRole::User,
content: PromptMessageContent::text(prompt),
}])
},
}).await;
// Code generation prompt
self.register_prompt_template(PromptTemplate {
name: "generate_code".to_string(),
description: Some("Generate code based on requirements and specifications".to_string()),
arguments: vec![
PromptArgument {
name: "requirements".to_string(),
description: Some("Description of what the code should do".to_string()),
required: Some(true),
},
PromptArgument {
name: "language".to_string(),
description: Some("Target programming language".to_string()),
required: Some(true),
},
PromptArgument {
name: "style".to_string(),
description: Some("Coding style or framework preferences".to_string()),
required: Some(false),
},
],
generator: |args| {
let requirements = args.get("requirements")
.and_then(|v| v.as_str())
.ok_or_else(|| McpError::invalid_params("Missing 'requirements' parameter", None))?;
let language = args.get("language")
.and_then(|v| v.as_str())
.ok_or_else(|| McpError::invalid_params("Missing 'language' parameter", None))?;
let style = args.get("style")
.and_then(|v| v.as_str())
.unwrap_or("standard");
let prompt = format!(
"Generate {} code that meets the following requirements:\n\n{}\n\nPlease follow {} style conventions and include:\n1. Clear, readable code\n2. Appropriate comments\n3. Error handling where needed\n4. Best practices for {}",
language, requirements, style, language
);
Ok(vec![PromptMessage {
role: PromptMessageRole::User,
content: PromptMessageContent::text(prompt),
}])
},
}).await;
// Code review prompt
self.register_prompt_template(PromptTemplate {
name: "review_code".to_string(),
description: Some("Perform a comprehensive code review".to_string()),
arguments: vec![
PromptArgument {
name: "code".to_string(),
description: Some("The code to review".to_string()),
required: Some(true),
},
PromptArgument {
name: "context".to_string(),
description: Some("Additional context about the code's purpose".to_string()),
required: Some(false),
},
],
generator: |args| {
let code = args.get("code")
.and_then(|v| v.as_str())
.ok_or_else(|| McpError::invalid_params("Missing 'code' parameter", None))?;
let context = args.get("context")
.and_then(|v| v.as_str())
.unwrap_or("No additional context provided");
let prompt = format!(
"Please perform a thorough code review of the following code:\n\nContext: {}\n\n```\n{}\n```\n\nPlease evaluate:\n1. **Correctness**: Does the code work as intended?\n2. **Performance**: Are there any performance concerns?\n3. **Security**: Any security vulnerabilities?\n4. **Maintainability**: Is the code easy to understand and modify?\n5. **Style**: Does it follow good coding practices?\n6. **Testing**: Are there adequate tests or test considerations?\n\nProvide specific suggestions for improvement.",
context, code
);
Ok(vec![PromptMessage {
role: PromptMessageRole::User,
content: PromptMessageContent::text(prompt),
}])
},
}).await;
// Documentation prompt
self.register_prompt_template(PromptTemplate {
name: "document_code".to_string(),
description: Some("Generate documentation for code".to_string()),
arguments: vec![
PromptArgument {
name: "code".to_string(),
description: Some("The code to document".to_string()),
required: Some(true),
},
PromptArgument {
name: "format".to_string(),
description: Some("Documentation format (markdown, rst, etc.)".to_string()),
required: Some(false),
},
],
generator: |args| {
let code = args.get("code")
.and_then(|v| v.as_str())
.ok_or_else(|| McpError::invalid_params("Missing 'code' parameter", None))?;
let format = args.get("format")
.and_then(|v| v.as_str())
.unwrap_or("markdown");
let prompt = format!(
"Generate comprehensive {} documentation for the following code:\n\n```\n{}\n```\n\nInclude:\n1. Overview and purpose\n2. Function/method descriptions\n3. Parameter explanations\n4. Return value descriptions\n5. Usage examples\n6. Any important notes or warnings",
format, code
);
Ok(vec![PromptMessage {
role: PromptMessageRole::User,
content: PromptMessageContent::text(prompt),
}])
},
}).await;
// Test generation prompt
self.register_prompt_template(PromptTemplate {
name: "generate_tests".to_string(),
description: Some("Generate unit tests for code".to_string()),
arguments: vec![
PromptArgument {
name: "code".to_string(),
description: Some("The code to test".to_string()),
required: Some(true),
},
PromptArgument {
name: "framework".to_string(),
description: Some("Testing framework to use".to_string()),
required: Some(false),
},
],
generator: |args| {
let code = args.get("code")
.and_then(|v| v.as_str())
.ok_or_else(|| McpError::invalid_params("Missing 'code' parameter", None))?;
let framework = args.get("framework")
.and_then(|v| v.as_str())
.unwrap_or("standard");
let prompt = format!(
"Generate comprehensive unit tests for the following code using {} testing framework:\n\n```\n{}\n```\n\nInclude tests for:\n1. Normal operation cases\n2. Edge cases\n3. Error conditions\n4. Boundary values\n5. Integration scenarios where applicable\n\nEnsure good test coverage and clear test names.",
framework, code
);
Ok(vec![PromptMessage {
role: PromptMessageRole::User,
content: PromptMessageContent::text(prompt),
}])
},
}).await;
}
}