use anyhow::{Context, Result};
use async_trait::async_trait;
use serde_json::json;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use crate::config::TurboPropConfig;
use crate::index::PersistentChunkIndex;
use crate::types::{ConnectionLimit, Port, TimeoutSeconds};
use super::index_manager::IndexManager;
use super::protocol::{
constants, InitializeParams, InitializeResult, JsonRpcError, JsonRpcRequest, JsonRpcResponse,
ServerCapabilities, ServerInfo, ToolsCapability,
};
use super::tools::Tools;
use super::transport::StdioTransport;
#[derive(Debug, Clone)]
pub struct McpServerConfig {
pub address: String,
pub port: Port,
pub max_connections: ConnectionLimit,
pub request_timeout: TimeoutSeconds,
}
impl Default for McpServerConfig {
fn default() -> Self {
Self {
address: "127.0.0.1".to_string(),
port: Port::dynamic(), max_connections: ConnectionLimit::default(),
request_timeout: TimeoutSeconds::default(),
}
}
}
#[derive(Debug, Clone, PartialEq, Default)]
pub enum InitializationState {
#[default]
NotStarted,
InProgress,
Ready,
Failed {
error: String,
retry_count: u32,
last_attempt: std::time::SystemTime,
},
}
impl InitializationState {
pub fn is_ready(&self) -> bool {
matches!(self, InitializationState::Ready)
}
pub fn is_in_progress(&self) -> bool {
matches!(self, InitializationState::InProgress)
}
pub fn is_failed(&self) -> bool {
matches!(self, InitializationState::Failed { .. })
}
pub fn error_message(&self) -> Option<&str> {
match self {
InitializationState::Failed { error, .. } => Some(error),
_ => None,
}
}
pub fn retry_count(&self) -> u32 {
match self {
InitializationState::Failed { retry_count, .. } => *retry_count,
_ => 0,
}
}
pub fn can_retry(&self, retry_delay_seconds: u64) -> bool {
match self {
InitializationState::Failed { last_attempt, .. } => {
std::time::SystemTime::now()
.duration_since(*last_attempt)
.map(|d| d.as_secs() >= retry_delay_seconds)
.unwrap_or(true)
}
_ => false,
}
}
}
#[async_trait]
pub trait McpServerTrait {
async fn initialize(&mut self, params: InitializeParams) -> Result<InitializeResult>;
async fn handle_request(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse>;
fn is_running(&self) -> bool;
}
pub struct McpServer {
repo_path: PathBuf,
config: TurboPropConfig,
server_config: McpServerConfig,
tools: Tools,
index: Arc<RwLock<Option<PersistentChunkIndex>>>,
initialization_state: Arc<RwLock<InitializationState>>,
running: Arc<RwLock<bool>>,
index_manager: Arc<RwLock<Option<IndexManager>>>,
}
impl McpServer {
pub async fn new(repo_path: &Path, config: &TurboPropConfig) -> Result<Self> {
info!("Initializing MCP server for {}", repo_path.display());
let tools = Tools::with_search_tool(
repo_path.to_path_buf(),
repo_path.to_path_buf(),
config.clone(),
);
let server = Self {
repo_path: repo_path.to_path_buf(),
config: config.clone(),
server_config: McpServerConfig::default(),
tools,
index: Arc::new(RwLock::new(None)),
initialization_state: Arc::new(RwLock::new(InitializationState::default())),
running: Arc::new(RwLock::new(false)),
index_manager: Arc::new(RwLock::new(None)),
};
Ok(server)
}
pub fn with_config_and_tools(server_config: McpServerConfig, tools: Tools) -> Self {
Self {
repo_path: PathBuf::new(), config: TurboPropConfig::default(),
server_config,
tools,
index: Arc::new(RwLock::new(None)),
initialization_state: Arc::new(RwLock::new(InitializationState::default())),
running: Arc::new(RwLock::new(false)),
index_manager: Arc::new(RwLock::new(None)),
}
}
pub async fn run(self) -> Result<()> {
let server = Arc::new(self);
let mut transport = StdioTransport::new();
let mut index_manager = match IndexManager::new(
&server.repo_path,
&server.config,
None, )
.await
{
Ok(manager) => manager,
Err(e) => {
error!(
"Failed to create IndexManager for repository '{}': {}",
server.repo_path.display(),
e
);
if !server.repo_path.exists() {
return Err(anyhow::anyhow!(
"Repository path '{}' does not exist",
server.repo_path.display()
));
}
if !server.repo_path.is_dir() {
return Err(anyhow::anyhow!(
"Repository path '{}' is not a directory",
server.repo_path.display()
));
}
if let Err(perm_err) = std::fs::read_dir(&server.repo_path) {
return Err(anyhow::anyhow!(
"Cannot access repository directory '{}': {}",
server.repo_path.display(), perm_err
));
}
return Err(e.context(format!(
"Failed to create IndexManager for repository '{}'",
server.repo_path.display()
)));
}
};
if let Err(e) = index_manager.start().await {
error!("Failed to start IndexManager background tasks: {}", e);
return Err(e.context(
"IndexManager initialization succeeded but background tasks failed to start"
));
}
{
let mut manager_guard = server.index_manager.write().await;
*manager_guard = Some(index_manager);
}
if let Err(e) = server.initialize_index_with_manager().await {
warn!("Initial index build failed: {}. The server will continue to run, but search functionality may be limited until the index is built.", e);
} else {
info!("Index initialization completed successfully");
}
{
let mut running_guard = server.running.write().await;
*running_guard = true;
}
info!(
"MCP server ready and listening on stdio (timeout: {}s, max_connections: {})",
server.server_config.request_timeout, server.server_config.max_connections
);
loop {
match transport.receive_request().await {
Some(Ok(request)) => {
let response = server.handle_request_internal(request).await;
if let Err(e) = transport.send_response(response).await {
error!("Failed to send response: {}", e);
break;
}
}
Some(Err(e)) => {
error!("Error receiving request: {}", e);
let error_response = StdioTransport::create_error_response(
None,
JsonRpcError::parse_error(e.to_string()),
);
let _ = transport.send_response(error_response).await;
}
None => {
info!("STDIN closed, shutting down MCP server");
break;
}
}
}
{
let mut running_guard = server.running.write().await;
*running_guard = false;
}
if let Some(manager) = server.index_manager.read().await.as_ref() {
let _ = manager.stop().await;
}
info!("MCP server shutdown complete");
Ok(())
}
async fn handle_request_internal(&self, request: JsonRpcRequest) -> JsonRpcResponse {
debug!(
"Handling request: method={}, id={:?}",
request.method, request.id
);
if let Err(error) = request.validate() {
return request.create_error_response(error);
}
match request.method.as_str() {
constants::methods::INITIALIZE => self.handle_initialize(request).await,
constants::methods::TOOLS_LIST => self.handle_tools_list(request).await,
constants::methods::TOOLS_CALL => self.handle_tools_call(request).await,
_ => {
let error = JsonRpcError::method_not_found(request.method.clone());
request.create_error_response(error)
}
}
}
async fn handle_initialize(&self, request: JsonRpcRequest) -> JsonRpcResponse {
debug!("Handling initialize request");
let params = match &request.params {
Some(params) => match serde_json::from_value::<InitializeParams>(params.clone()) {
Ok(params) => params,
Err(e) => {
let error =
JsonRpcError::invalid_params(format!("Invalid initialize params: {}", e));
return request.create_error_response(error);
}
},
None => {
let error =
JsonRpcError::invalid_params("Missing initialize parameters".to_string());
return request.create_error_response(error);
}
};
info!(
"Initializing MCP server for client: {} v{}",
params.client_info.name, params.client_info.version
);
if params.protocol_version != constants::PROTOCOL_VERSION {
warn!(
"Client protocol version {} differs from server version {}",
params.protocol_version,
constants::PROTOCOL_VERSION
);
}
let index_clone = Arc::clone(&self.index);
let repo_path = self.repo_path.clone();
let config = self.config.clone();
let init_state_clone = Arc::clone(&self.initialization_state);
{
let mut state_guard = init_state_clone.write().await;
*state_guard = InitializationState::InProgress;
}
tokio::spawn(async move {
const MAX_RETRIES: u32 = 3;
const RETRY_DELAY_SECONDS: u64 = 10;
let mut retry_count = 0;
loop {
match Self::initialize_index(&repo_path, &config).await {
Ok(index) => {
{
let mut index_guard = index_clone.write().await;
*index_guard = Some(index);
}
{
let mut state_guard = init_state_clone.write().await;
*state_guard = InitializationState::Ready;
}
info!("Index initialization completed successfully");
break;
}
Err(e) => {
retry_count += 1;
let error_msg = format!("Failed to initialize index: {}", e);
if retry_count >= MAX_RETRIES {
error!("{} (attempt {}/{}). Giving up.", error_msg, retry_count, MAX_RETRIES);
let mut state_guard = init_state_clone.write().await;
*state_guard = InitializationState::Failed {
error: error_msg,
retry_count,
last_attempt: std::time::SystemTime::now(),
};
break;
} else {
warn!("{} (attempt {}/{}). Retrying in {} seconds...",
error_msg, retry_count, MAX_RETRIES, RETRY_DELAY_SECONDS);
tokio::time::sleep(tokio::time::Duration::from_secs(RETRY_DELAY_SECONDS)).await;
}
}
}
}
});
let result = InitializeResult {
protocol_version: constants::PROTOCOL_VERSION.to_string(),
server_info: ServerInfo {
name: constants::SERVER_NAME.to_string(),
version: constants::SERVER_VERSION.to_string(),
},
capabilities: ServerCapabilities {
tools: Some(ToolsCapability {
list_changed: false, }),
experimental: std::collections::HashMap::new(),
},
};
match serde_json::to_value(result) {
Ok(result_value) => {
info!("MCP server initialized successfully");
request.create_success_response(result_value)
}
Err(e) => {
let error =
JsonRpcError::internal_error(format!("Failed to serialize result: {}", e));
request.create_error_response(error)
}
}
}
async fn handle_tools_list(&self, request: JsonRpcRequest) -> JsonRpcResponse {
debug!("Handling tools/list request");
let init_state = {
let state_guard = self.initialization_state.read().await;
state_guard.clone()
};
if !init_state.is_ready() {
let error = match init_state {
InitializationState::InProgress => {
JsonRpcError::index_not_ready()
}
InitializationState::Failed { error, retry_count, .. } => {
JsonRpcError::internal_error(format!(
"Server initialization failed after {} attempts: {}",
retry_count, error
))
}
_ => JsonRpcError::internal_error("Server not initialized")
};
return request.create_error_response(error);
}
let tools = self.tools.list_tools();
let result = json!({
"tools": tools
});
request.create_success_response(result)
}
async fn handle_tools_call(&self, request: JsonRpcRequest) -> JsonRpcResponse {
debug!("Handling tools/call request");
let init_state = {
let state_guard = self.initialization_state.read().await;
state_guard.clone()
};
if !init_state.is_ready() {
let error = match init_state {
InitializationState::InProgress => {
JsonRpcError::index_not_ready()
}
InitializationState::Failed { error, retry_count, .. } => {
JsonRpcError::application_error(
-32003, format!(
"Index initialization failed after {} attempts: {}. Please restart the server or check logs for details.",
retry_count, error
)
)
}
_ => JsonRpcError::index_not_ready()
};
return request.create_error_response(error);
}
let params = match &request.params {
Some(params) => params.clone(),
None => {
let error =
JsonRpcError::invalid_params("Missing tool call parameters".to_string());
return request.create_error_response(error);
}
};
let tool_name = match params.get("name").and_then(|v| v.as_str()) {
Some(name) => name,
None => {
let error = JsonRpcError::invalid_params("Missing tool name".to_string());
return request.create_error_response(error);
}
};
let arguments = params.get("arguments").cloned().unwrap_or(json!({}));
let tool_call_request = crate::mcp::tools::ToolCallRequest {
name: tool_name.to_string(),
arguments: serde_json::from_value(arguments).unwrap_or_default(),
};
match self.tools.execute_tool(tool_call_request).await {
Ok(tool_response) => {
if tool_response.success {
debug!("Tool executed successfully: {}", tool_name);
let result = tool_response.content.unwrap_or(json!({}));
request.create_success_response(result)
} else {
error!(
"Tool execution failed: {}",
tool_response
.error
.as_ref()
.unwrap_or(&"Unknown error".to_string())
);
let error = JsonRpcError::tool_execution_error(
tool_response
.error
.unwrap_or_else(|| "Unknown error".to_string()),
);
request.create_error_response(error)
}
}
Err(e) => {
error!("Tool execution failed: {}", e);
let error = JsonRpcError::tool_execution_error(e.to_string());
request.create_error_response(error)
}
}
}
async fn initialize_index(
repo_path: &Path,
config: &TurboPropConfig,
) -> Result<PersistentChunkIndex> {
info!("Initializing search index for {}", repo_path.display());
let index = if let Ok(existing_index) = PersistentChunkIndex::load(repo_path) {
info!("Loading existing index from {}", repo_path.display());
existing_index
} else {
info!("Creating new index");
let index = crate::commands::index::build_index(repo_path, config)
.await
.context("Failed to build initial index")?;
info!("Index created successfully with {} chunks", index.len());
index
};
Ok(index)
}
async fn initialize_index_with_manager(&self) -> Result<()> {
info!("Starting index initialization");
let index = Self::initialize_index(&self.repo_path, &self.config).await?;
if let Some(manager) = self.index_manager.read().await.as_ref() {
manager.set_index(index.clone()).await;
}
{
let mut index_guard = self.index.write().await;
*index_guard = Some(index);
}
{
let mut state_guard = self.initialization_state.write().await;
*state_guard = InitializationState::Ready;
}
info!("Index initialization completed with file watching enabled");
Ok(())
}
pub async fn initialize(&mut self, params: InitializeParams) -> Result<InitializeResult> {
info!(
"Initializing MCP server for client: {} v{}",
params.client_info.name, params.client_info.version
);
if params.protocol_version != constants::PROTOCOL_VERSION {
anyhow::bail!(
"Unsupported protocol version: {} (expected: {})",
params.protocol_version,
constants::PROTOCOL_VERSION
);
}
params
.validate()
.map_err(|e| anyhow::anyhow!("Invalid initialization parameters: {:?}", e))?;
{
let mut state_guard = self.initialization_state.write().await;
*state_guard = InitializationState::Ready;
}
let result = InitializeResult {
protocol_version: constants::PROTOCOL_VERSION.to_string(),
server_info: ServerInfo {
name: constants::SERVER_NAME.to_string(),
version: constants::SERVER_VERSION.to_string(),
},
capabilities: ServerCapabilities {
tools: Some(ToolsCapability {
list_changed: false, }),
experimental: std::collections::HashMap::new(),
},
};
info!("MCP server initialized successfully");
Ok(result)
}
pub async fn is_running(&self) -> bool {
*self.running.read().await
}
}
pub struct McpServerBuilder {
repo_path: Option<PathBuf>,
config: Option<TurboPropConfig>,
}
impl McpServerBuilder {
pub fn new() -> Self {
Self {
repo_path: None,
config: None,
}
}
pub fn repo_path<P: AsRef<Path>>(mut self, path: P) -> Self {
self.repo_path = Some(path.as_ref().to_path_buf());
self
}
pub fn config(mut self, config: TurboPropConfig) -> Self {
self.config = Some(config);
self
}
pub async fn build(self) -> Result<McpServer> {
let repo_path = self
.repo_path
.ok_or_else(|| anyhow::anyhow!("Repository path is required"))?;
let config = self
.config
.ok_or_else(|| anyhow::anyhow!("Configuration is required"))?;
McpServer::new(&repo_path, &config).await
}
}
impl Default for McpServerBuilder {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl McpServerTrait for McpServer {
async fn initialize(&mut self, params: InitializeParams) -> Result<InitializeResult> {
self.initialize(params).await
}
async fn handle_request(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse> {
let response = self.handle_request_internal(request).await;
if response.error.is_some() {
let error_msg = response
.error
.as_ref()
.map(|e| e.message.clone())
.unwrap_or_else(|| "Unknown error".to_string());
Err(anyhow::anyhow!("Request failed: {}", error_msg))
} else {
Ok(response)
}
}
fn is_running(&self) -> bool {
self.running.try_read().map(|guard| *guard).unwrap_or(false)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use tempfile::TempDir;
#[tokio::test]
async fn test_mcp_server_creation() {
let temp_dir = TempDir::new().unwrap();
let config = TurboPropConfig::default();
let server = McpServer::new(temp_dir.path(), &config).await;
assert!(server.is_ok());
}
#[tokio::test]
async fn test_initialize_request() {
let temp_dir = TempDir::new().unwrap();
let config = TurboPropConfig::default();
let server = McpServer::new(temp_dir.path(), &config).await.unwrap();
let request = JsonRpcRequest::new(
constants::methods::INITIALIZE.to_string(),
Some(json!({
"protocol_version": constants::PROTOCOL_VERSION,
"client_info": {
"name": "test-client",
"version": "1.0.0"
},
"capabilities": {}
})),
);
let response = server.handle_initialize(request).await;
if let Some(error) = &response.error {
println!("Error: {:?}", error);
}
assert!(response.error.is_none());
assert!(response.result.is_some());
let result = response.result.unwrap();
println!("Result: {:?}", result);
assert_eq!(result["protocol_version"], constants::PROTOCOL_VERSION);
assert_eq!(result["server_info"]["name"], constants::SERVER_NAME);
}
#[tokio::test]
async fn test_tools_list_request() {
let temp_dir = TempDir::new().unwrap();
let config = TurboPropConfig::default();
let server = McpServer::new(temp_dir.path(), &config).await.unwrap();
{
let mut state_guard = server.initialization_state.write().await;
*state_guard = InitializationState::Ready;
}
let request = JsonRpcRequest::new(constants::methods::TOOLS_LIST.to_string(), None);
let response = server.handle_tools_list(request).await;
assert!(response.error.is_none());
assert!(response.result.is_some());
let result = response.result.unwrap();
let tools = result["tools"].as_array().unwrap();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0]["name"], "semantic_search");
}
#[test]
fn test_server_builder() {
let temp_dir = TempDir::new().unwrap();
let config = TurboPropConfig::default();
let builder = McpServerBuilder::new()
.repo_path(temp_dir.path())
.config(config);
drop(builder);
}
#[tokio::test]
async fn test_invalid_method_request() {
let temp_dir = TempDir::new().unwrap();
let config = TurboPropConfig::default();
let server = McpServer::new(temp_dir.path(), &config).await.unwrap();
let request = JsonRpcRequest::new("invalid_method".to_string(), None);
let response = server.handle_request_internal(request).await;
assert!(response.result.is_none());
assert!(response.error.is_some());
assert_eq!(response.error.unwrap().code, -32601); }
#[tokio::test]
async fn test_tools_call_before_initialization() {
let temp_dir = TempDir::new().unwrap();
let config = TurboPropConfig::default();
let server = McpServer::new(temp_dir.path(), &config).await.unwrap();
let request = JsonRpcRequest::new(
constants::methods::TOOLS_CALL.to_string(),
Some(json!({
"name": "semantic_search",
"arguments": {
"query": "test"
}
})),
);
let response = server.handle_tools_call(request).await;
assert!(response.result.is_none());
assert!(response.error.is_some());
}
}