context_creator/mcp_server/
mod.rs

1//! MCP (Model Context Protocol) server implementation for context-creator
2//!
3//! This module provides a JSON-RPC server that allows AI agents to
4//! analyze codebases programmatically.
5
6use anyhow::Result;
7use jsonrpsee::{
8    core::RpcResult,
9    proc_macros::rpc,
10    server::{Server, ServerHandle as JsonRpcServerHandle},
11};
12use schemars::JsonSchema;
13use serde::{Deserialize, Serialize};
14use std::net::SocketAddr;
15
16pub mod cache;
17pub mod handlers;
18pub mod rmcp_handlers;
19pub mod rmcp_server;
20
21/// Health check response structure
22#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
23pub struct HealthResponse {
24    pub status: String,
25    pub timestamp: u64,
26    pub version: String,
27}
28
29/// RPC trait for health check
30#[rpc(server)]
31pub trait HealthRpc {
32    /// Returns the current health status of the server
33    #[method(name = "health_check")]
34    async fn health_check(&self) -> RpcResult<HealthResponse>;
35}
36
37/// Request structure for process_local_codebase
38#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
39pub struct ProcessLocalRequest {
40    /// The question/prompt to ask about the codebase (primary feature)
41    pub prompt: String,
42    /// Path to the codebase to analyze
43    pub path: std::path::PathBuf,
44    /// Optional: specific file patterns to include
45    pub include_patterns: Vec<String>,
46    /// Optional: patterns to ignore
47    pub ignore_patterns: Vec<String>,
48    /// Optional: whether to trace imports
49    pub include_imports: bool,
50    /// Optional: max tokens for context (auto-calculated based on LLM if not specified)
51    pub max_tokens: Option<u32>,
52    /// Optional: LLM tool to use (default: gemini)
53    pub llm_tool: Option<String>,
54    /// Optional: return markdown context along with answer
55    pub include_context: Option<bool>,
56}
57
58/// Response structure for process_local_codebase
59#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
60pub struct ProcessLocalResponse {
61    /// The LLM's answer to the prompt
62    pub answer: String,
63    /// Optional: the markdown context used (if include_context is true)
64    pub context: Option<String>,
65    /// Number of files analyzed
66    pub file_count: usize,
67    /// Token count of the context
68    pub token_count: usize,
69    /// Processing time in milliseconds
70    pub processing_time_ms: u64,
71    /// LLM tool used
72    pub llm_tool: String,
73}
74
75/// Request structure for process_remote_repo
76#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
77pub struct ProcessRemoteRequest {
78    /// The question/prompt to ask about the repository (primary feature)
79    pub prompt: String,
80    /// Git repository URL to analyze
81    pub repo_url: String,
82    /// Optional: specific file patterns to include
83    pub include_patterns: Vec<String>,
84    /// Optional: patterns to ignore
85    pub ignore_patterns: Vec<String>,
86    /// Optional: whether to trace imports
87    pub include_imports: bool,
88    /// Optional: max tokens for context (auto-calculated based on LLM if not specified)
89    pub max_tokens: Option<u32>,
90    /// Optional: LLM tool to use (default: gemini)
91    pub llm_tool: Option<String>,
92    /// Optional: return markdown context along with answer
93    pub include_context: Option<bool>,
94}
95
96/// Response structure for process_remote_repo
97#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
98pub struct ProcessRemoteResponse {
99    /// The LLM's answer to the prompt
100    pub answer: String,
101    /// Optional: the markdown context used (if include_context is true)
102    pub context: Option<String>,
103    /// Number of files analyzed
104    pub file_count: usize,
105    /// Token count of the context
106    pub token_count: usize,
107    /// Processing time in milliseconds
108    pub processing_time_ms: u64,
109    /// Repository name
110    pub repo_name: String,
111    /// LLM tool used
112    pub llm_tool: String,
113}
114
115/// Request structure for get_file_metadata
116#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
117pub struct GetFileMetadataRequest {
118    pub file_path: std::path::PathBuf,
119}
120
121/// Response structure for get_file_metadata
122#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
123pub struct GetFileMetadataResponse {
124    pub path: std::path::PathBuf,
125    pub size: u64,
126    pub modified: u64,
127    pub is_symlink: bool,
128    pub language: Option<String>,
129}
130
131/// Request structure for search_codebase
132#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
133pub struct SearchCodebaseRequest {
134    pub path: std::path::PathBuf,
135    pub query: String,
136    pub max_results: Option<u32>,
137    pub file_pattern: Option<String>,
138}
139
140/// Search result entry
141#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
142pub struct SearchResult {
143    pub file_path: std::path::PathBuf,
144    pub line_number: usize,
145    pub line_content: String,
146    pub match_context: String,
147}
148
149/// Response structure for search_codebase
150#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
151pub struct SearchCodebaseResponse {
152    pub results: Vec<SearchResult>,
153    pub total_matches: usize,
154    pub files_searched: usize,
155    pub search_time_ms: u64,
156}
157
158/// Request structure for diff_files
159#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
160pub struct DiffFilesRequest {
161    pub file1_path: std::path::PathBuf,
162    pub file2_path: std::path::PathBuf,
163    pub context_lines: Option<u32>,
164}
165
166/// Diff hunk structure
167#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
168pub struct DiffHunk {
169    pub old_start: usize,
170    pub old_lines: usize,
171    pub new_start: usize,
172    pub new_lines: usize,
173    pub content: String,
174}
175
176/// Response structure for diff_files
177#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
178pub struct DiffFilesResponse {
179    pub file1_path: std::path::PathBuf,
180    pub file2_path: std::path::PathBuf,
181    pub hunks: Vec<DiffHunk>,
182    pub added_lines: usize,
183    pub removed_lines: usize,
184    pub is_binary: bool,
185}
186
187/// Request structure for semantic_search
188#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
189pub struct SemanticSearchRequest {
190    pub path: std::path::PathBuf,
191    pub query: String,
192    pub search_type: SemanticSearchType,
193    pub max_results: Option<u32>,
194}
195
196/// Type of semantic search
197#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
198#[serde(rename_all = "snake_case")]
199pub enum SemanticSearchType {
200    /// Find functions/methods by name
201    Functions,
202    /// Find classes/structs/interfaces by name
203    Types,
204    /// Find imports/dependencies
205    Imports,
206    /// Find all references to a symbol
207    References,
208}
209
210/// Semantic search result
211#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
212pub struct SemanticSearchResult {
213    pub file_path: std::path::PathBuf,
214    pub symbol_name: String,
215    pub symbol_type: String,
216    pub line_number: usize,
217    pub context: String,
218}
219
220/// Response structure for semantic_search
221#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
222pub struct SemanticSearchResponse {
223    pub results: Vec<SemanticSearchResult>,
224    pub total_matches: usize,
225    pub files_analyzed: usize,
226    pub search_time_ms: u64,
227}
228
229/// RPC trait for codebase processing
230#[rpc(server)]
231pub trait CodebaseRpc {
232    /// Process a local codebase directory
233    #[method(name = "process_local_codebase")]
234    async fn process_local_codebase(
235        &self,
236        request: ProcessLocalRequest,
237    ) -> RpcResult<ProcessLocalResponse>;
238
239    /// Process a remote repository
240    #[method(name = "process_remote_repo")]
241    async fn process_remote_repo(
242        &self,
243        request: ProcessRemoteRequest,
244    ) -> RpcResult<ProcessRemoteResponse>;
245
246    /// Get metadata for a specific file
247    #[method(name = "get_file_metadata")]
248    async fn get_file_metadata(
249        &self,
250        request: GetFileMetadataRequest,
251    ) -> RpcResult<GetFileMetadataResponse>;
252
253    /// Search codebase for a query string
254    #[method(name = "search_codebase")]
255    async fn search_codebase(
256        &self,
257        request: SearchCodebaseRequest,
258    ) -> RpcResult<SearchCodebaseResponse>;
259
260    /// Get diff between two files
261    #[method(name = "diff_files")]
262    async fn diff_files(&self, request: DiffFilesRequest) -> RpcResult<DiffFilesResponse>;
263
264    /// Perform semantic search across codebase
265    #[method(name = "semantic_search")]
266    async fn semantic_search(
267        &self,
268        request: SemanticSearchRequest,
269    ) -> RpcResult<SemanticSearchResponse>;
270}
271
272/// Server handle wrapper for managing the MCP server lifecycle
273pub struct ServerHandle {
274    inner: JsonRpcServerHandle,
275    local_addr: SocketAddr,
276}
277
278impl ServerHandle {
279    /// Get the local address the server is listening on
280    pub fn local_addr(&self) -> Result<SocketAddr> {
281        Ok(self.local_addr)
282    }
283
284    /// Stop the server gracefully
285    pub fn stop(self) -> Result<()> {
286        self.inner.stop()?;
287        Ok(())
288    }
289}
290
291/// Start the MCP server on the specified address
292pub async fn start_server(addr: &str) -> Result<ServerHandle> {
293    let addr: SocketAddr = addr.parse()?;
294
295    // Build the server
296    let server = Server::builder().build(addr).await?;
297
298    // Get the actual address (in case port 0 was used)
299    let local_addr = server.local_addr()?;
300
301    // Create shared cache
302    let cache = std::sync::Arc::new(cache::McpCache::new());
303
304    // Create and register the handlers
305    let health_impl = handlers::HealthRpcImpl;
306    let codebase_impl = handlers::CodebaseRpcImpl::new(cache);
307
308    // Merge RPC modules
309    let mut rpc_module = health_impl.into_rpc();
310    rpc_module.merge(codebase_impl.into_rpc())?;
311
312    // Start the server in the background
313    let handle = server.start(rpc_module);
314
315    Ok(ServerHandle {
316        inner: handle,
317        local_addr,
318    })
319}