use std::path::PathBuf;
use std::sync::Mutex;
use rmcp::model::{
CallToolRequestParam, CallToolResult, ErrorCode, Implementation, InitializeRequestParam,
InitializeResult, ListToolsResult, PaginatedRequestParam, ProtocolVersion, ServerCapabilities,
ServerInfo,
};
use rmcp::service::{RequestContext, RoleServer};
use rmcp::transport::io::stdio;
use rmcp::{ServerHandler, ServiceExt};
use crate::analytics::Analytics;
use crate::db::Database;
use crate::index;
use super::tools;
pub struct CtxServer {
root: PathBuf,
pub(crate) db: Mutex<Database>,
pub(crate) analytics: Option<Mutex<Analytics>>,
}
impl CtxServer {
pub fn new(root: PathBuf) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
let db = index::open_database(&root)
.map_err(|e| format!("Failed to open database: {}. Run 'ctx index' first.", e))?;
let analytics = Analytics::open(&root).ok().map(Mutex::new);
Ok(Self {
root,
db: Mutex::new(db),
analytics,
})
}
pub fn root(&self) -> &PathBuf {
&self.root
}
pub fn with_db<F, R>(&self, f: F) -> R
where
F: FnOnce(&Database) -> R,
{
let db = self.db.lock().unwrap();
f(&db)
}
#[allow(dead_code)]
pub fn with_analytics<F, R>(&self, f: F) -> Option<R>
where
F: FnOnce(&Analytics) -> R,
{
self.analytics.as_ref().map(|a| {
let analytics = a.lock().unwrap();
f(&analytics)
})
}
}
impl ServerHandler for CtxServer {
fn get_info(&self) -> ServerInfo {
ServerInfo {
protocol_version: ProtocolVersion::LATEST,
capabilities: ServerCapabilities::builder().enable_tools().build(),
server_info: Implementation {
name: "ctx".into(),
version: env!("CARGO_PKG_VERSION").into(),
title: None,
icons: None,
website_url: None,
},
instructions: None,
}
}
async fn initialize(
&self,
_request: InitializeRequestParam,
_context: RequestContext<RoleServer>,
) -> Result<InitializeResult, rmcp::ErrorData> {
Ok(self.get_info())
}
async fn list_tools(
&self,
_request: Option<PaginatedRequestParam>,
_context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, rmcp::ErrorData> {
Ok(ListToolsResult {
tools: tools::get_all_tools(),
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
request: CallToolRequestParam,
_context: RequestContext<RoleServer>,
) -> Result<CallToolResult, rmcp::ErrorData> {
let name = request.name.as_ref();
let args = request.arguments.as_ref();
match name {
"search_symbols" => tools::search::search_symbols(self, args).await,
"get_definition" => tools::search::get_definition(self, args).await,
"find_references" => tools::search::find_references(self, args).await,
"get_file" => tools::files::get_file(self, args).await,
"get_file_tree" => tools::files::get_file_tree(self, args).await,
"get_callers" => tools::analysis::get_callers(self, args).await,
"get_callees" => tools::analysis::get_callees(self, args).await,
"smart_context" => tools::analysis::smart_context(self, args).await,
_ => Err(rmcp::ErrorData::new(
ErrorCode::METHOD_NOT_FOUND,
format!("Unknown tool: {}", name),
None,
)),
}
}
}
pub async fn run_mcp_server(root: PathBuf) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let server = CtxServer::new(root)?;
let transport = stdio();
let service = server.serve(transport).await?;
service.waiting().await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::TempDir;
use crate::index::Indexer;
fn setup_test_project() -> (TempDir, PathBuf) {
let temp_dir = TempDir::new().unwrap();
let root = temp_dir.path().to_path_buf();
let src_dir = root.join("src");
fs::create_dir_all(&src_dir).unwrap();
fs::write(
src_dir.join("lib.rs"),
r#"
/// A test function.
pub fn hello_world() -> String {
"Hello, World!".to_string()
}
/// Another function that calls hello_world.
pub fn greet() -> String {
hello_world()
}
"#,
)
.unwrap();
let mut indexer =
Indexer::with_config(&root, false, crate::walker::WalkerConfig::default()).unwrap();
indexer.index().unwrap();
(temp_dir, root)
}
#[test]
fn test_mcp_server_startup() {
let (_temp_dir, root) = setup_test_project();
let server =
CtxServer::new(root.clone()).expect("Server should start with indexed database");
assert_eq!(server.root(), &root);
let symbol_count = server.with_db(|db| db.find_symbols("hello", 10).unwrap().len());
assert!(symbol_count > 0, "Should find indexed symbols");
}
#[test]
fn test_mcp_server_without_index_fails() {
let temp_dir = TempDir::new().unwrap();
let root = temp_dir.path().to_path_buf();
fs::create_dir_all(root.join("src")).unwrap();
fs::write(root.join("src/lib.rs"), "fn test() {}").unwrap();
let result = CtxServer::new(root);
assert!(
result.is_err(),
"Server should fail without indexed database"
);
}
#[test]
fn test_server_get_info() {
let (_temp_dir, root) = setup_test_project();
let server = CtxServer::new(root).unwrap();
let info = server.get_info();
assert_eq!(info.server_info.name.as_str(), "ctx");
assert!(!info.server_info.version.is_empty());
}
}