Skip to main content

aptu_coder/
lib.rs

1// SPDX-FileCopyrightText: 2026 aptu-coder contributors
2// SPDX-License-Identifier: Apache-2.0
3//! Rust MCP server for code structure analysis using tree-sitter.
4//!
5//! This crate exposes seven MCP tools for multiple programming languages:
6//!
7//! **Analyze family:**
8//! - **`analyze_directory`**: Directory tree with file counts and structure
9//! - **`analyze_file`**: Semantic extraction (functions, classes, imports)
10//! - **`analyze_symbol`**: Call graph analysis (callers and callees)
11//! - **`analyze_module`**: Lightweight function and import index
12//!
13//! **Edit family:**
14//! - **`edit_overwrite`**: Create or overwrite files
15//! - **`edit_replace`**: Replace text blocks in files
16//!
17//! **Exec family:**
18//! - **`exec_command`**: Run shell commands with progress notifications
19//!
20//! Key entry points:
21//! - [`analyze::analyze_directory`]: Analyze entire directory tree
22//! - [`analyze::analyze_file`]: Analyze single file
23//!
24//! Languages supported: Astro, C/C++, C#, CSS, Fortran, Go, HTML, Java, JavaScript, JSON, Kotlin, Markdown, Python, Rust, TOML, TSX, TypeScript, YAML.
25
26#![cfg_attr(test, allow(clippy::unwrap_used))]
27#![cfg_attr(test, allow(clippy::expect_used))]
28
29mod filters;
30pub(crate) mod heredoc_validation;
31pub(crate) mod metrics;
32pub(crate) mod metrics_export;
33pub(crate) mod otel;
34pub(crate) mod shell;
35pub(crate) mod shell_scan;
36/// Heredoc and shell file-write pattern detection (pre-spawn guard for exec_command).
37pub(crate) mod shell_write;
38pub(crate) mod tools;
39pub(crate) mod validation;
40
41pub use metrics::{MetricEvent, MetricsSender, MetricsWriter, migrate_legacy_metrics_dir};
42pub use otel::{
43    ClientMetadata, extract_and_set_trace_context, init_log_appender, init_meter, init_otel,
44};
45
46use aptu_coder_core::analyze;
47use aptu_coder_core::{cache, completion, types};
48use validation::validate_path;
49
50use crate::tools::common::{err_to_tool_result, no_cache_meta};
51
52pub const STDIN_MAX_BYTES: usize = 1_048_576;
53
54#[non_exhaustive]
55#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
56pub struct ExecCommandParams {
57    /// Shell command to execute via sh -c (or $SHELL if set).
58    pub command: String,
59    /// Working directory for the command. Set this instead of prepending cd to the command string. Validated against path traversal; does not sandbox the process.
60    pub working_dir: Option<String>,
61    /// UTF-8 content to pipe into the process stdin (max `STDIN_MAX_BYTES` = 1 MB). When None, stdin is closed (null).
62    pub stdin: Option<String>,
63    /// Maximum execution time in seconds. When the command exceeds this limit, the
64    /// child process is killed and the response indicates `timed_out: true`.
65    /// A value of 0 or None means no timeout (unlimited execution).
66    #[serde(default)]
67    pub timeout_secs: Option<i64>,
68    /// Drain timeout in milliseconds after the child process exits. When the child
69    /// exits but a background subprocess holds pipes open, the drain collects
70    /// buffered output for this many milliseconds before returning
71    /// `output_truncated: true`. Default: 500ms when omitted or 0.
72    /// Positive values override the default. Negative values are rejected with
73    /// INVALID_PARAMS.
74    #[serde(default)]
75    pub drain_timeout_secs: Option<i64>,
76}
77
78impl ExecCommandParams {
79    /// Creates a new ExecCommandParams with the given command.
80    #[must_use]
81    pub fn new(command: String, working_dir: Option<String>) -> Self {
82        Self {
83            command,
84            working_dir,
85            ..Default::default()
86        }
87    }
88}
89
90#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
91pub struct ShellOutput {
92    /// Standard output from the command.
93    pub stdout: String,
94    /// Standard error from the command.
95    pub stderr: String,
96    /// Stdout and stderr interleaved in arrival order.
97    pub interleaved: String,
98    /// Exit code; null if the process could not be waited on (e.g. drain timeout from a background process holding pipes).
99    pub exit_code: Option<i32>,
100    /// True if the post-exit drain timed out (backgrounded process kept pipes open).
101    /// When true, any available output is still included; use the overflow file path
102    /// from the truncation notice Content block to recover the full output.
103    pub output_truncated: bool,
104    /// Set when the post-exit drain timed out because a background process held the
105    /// pipes open. Distinct from `output_truncated` (size cap) -- this indicates a
106    /// drain timeout rather than a size overflow.
107    pub output_collection_error: Option<String>,
108    /// Path to the slot file containing full stdout (if output was persisted).
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub stdout_path: Option<String>,
111    /// Path to the slot file containing full stderr (if output was persisted).
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub stderr_path: Option<String>,
114    /// Path to the slot file containing full interleaved output (if output was persisted).
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub interleaved_path: Option<String>,
117    /// Description of the filter applied to stdout (if any).
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub filter_applied: Option<String>,
120    /// True when the command was killed due to exceeding `timeout_secs`.
121    /// When true, exit_code is None and no partial output is available.
122    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
123    pub timed_out: bool,
124}
125
126impl ShellOutput {
127    /// Creates a new ShellOutput with the given parameters.
128    #[must_use]
129    pub fn new(
130        stdout: String,
131        stderr: String,
132        interleaved: String,
133        exit_code: Option<i32>,
134        output_truncated: bool,
135    ) -> Self {
136        Self {
137            stdout,
138            stderr,
139            interleaved,
140            exit_code,
141            output_truncated,
142            output_collection_error: None,
143            stdout_path: None,
144            stderr_path: None,
145            interleaved_path: None,
146            filter_applied: None,
147            timed_out: false,
148        }
149    }
150}
151
152#[cfg(test)]
153use aptu_coder_core::cache::CacheTier;
154use aptu_coder_core::cache::{AnalysisCache, CallGraphCache};
155use aptu_coder_core::types::{
156    AnalyzeDirectoryParams, AnalyzeFileParams, AnalyzeModuleParams, AnalyzeSymbolParams,
157    EditOverwriteOutput, EditOverwriteParams, EditReplaceOutput, EditReplaceParams,
158};
159use filters::CompiledRule;
160
161use rmcp::handler::server::tool::{ToolRouter, schema_for_type};
162use rmcp::handler::server::wrapper::Parameters;
163use rmcp::model::{
164    CallToolResponse, CallToolResult, CancelledNotificationParam, CompleteRequestParams,
165    CompleteResult, CompletionInfo, ContentBlock, DiscoverResult, ErrorData, Implementation,
166    InitializeRequestParams, InitializeResult, ListResourceTemplatesResult, ListResourcesResult,
167    ReadResourceRequestParams, ReadResourceResponse, ServerCapabilities,
168};
169use rmcp::service::{NotificationContext, RequestContext};
170use rmcp::{Peer, RoleServer, ServerHandler, tool, tool_handler, tool_router};
171
172use std::collections::HashMap;
173use std::path::Path;
174use std::sync::{Arc, Mutex};
175use tokio::sync::{Mutex as TokioMutex, RwLock};
176use tracing::instrument;
177
178static GLOBAL_SESSION_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
179
180// 5_000 chars fires at ~150-180 files at depth=2 (~28-33 chars/file).
181// Empirical data (684 calls, Jun 2026): max observed output was 4,882 chars; the old
182// 50_000 threshold never triggered once. At 5_000, auto-summary engages for repos that
183// would otherwise produce an overwhelming flat response.
184pub(crate) const SIZE_LIMIT: usize = 5_000;
185
186pub(crate) fn err_to_tool_result_from_pagination(
187    e: aptu_coder_core::pagination::PaginationError,
188) -> CallToolResult {
189    let msg = format!("Pagination error: {}", e);
190    CallToolResult::error(vec![ContentBlock::text(msg)]).with_meta(Some(no_cache_meta()))
191}
192
193/// MCP server handler that wires the four analysis tools to the rmcp transport.
194///
195/// Holds shared state: tool router, analysis cache, peer connection, log-level filter,
196/// log event channel, metrics sender, and per-session sequence tracking.
197#[derive(Clone)]
198pub struct CodeAnalyzer {
199    // Read lock acquired by list_tools/call_tool; write lock acquired during on_initialized
200    // for bind_peer_notifier.
201    // IMPORTANT: Do not perform long-running I/O while holding the write lock.
202    pub(crate) tool_router: Arc<RwLock<ToolRouter<Self>>>,
203    cache: AnalysisCache,
204    disk_cache: std::sync::Arc<cache::DiskCache>,
205    peer: Arc<TokioMutex<Option<Peer<RoleServer>>>>,
206    metrics_tx: crate::metrics::MetricsSender,
207    session_call_seq: Arc<std::sync::atomic::AtomicU32>,
208    session_id: Arc<TokioMutex<Option<String>>>,
209    client_name: Arc<TokioMutex<Option<String>>>,
210    client_version: Arc<TokioMutex<Option<String>>>,
211    // Resolved login shell PATH, captured once at startup via login shell invocation.
212    // Arc<Option<String>> is immutable after init; no lock needed.
213    resolved_path: Arc<Option<String>>,
214    // Compiled filter rules table (built-in + project-local from .aptu/filters.toml).
215    // Immutable after init; no lock needed.
216    filter_table: Arc<Vec<CompiledRule>>,
217    // L1 in-memory LRU cache for call graph results (analyze_symbol).
218    // Capacity controlled by APTU_CODER_SYMBOL_CACHE_CAPACITY env var (default 32).
219    call_graph_cache: CallGraphCache,
220    // Per-(session_id, canonical_path) consecutive edit_replace failure counter.
221    // Used to detect stale LLM context and return a directive error instead of
222    // repeatedly trying an old_text that no longer matches the file content.
223    edit_failure_counts: Arc<Mutex<HashMap<(String, String), u8>>>,
224    // On-disk graph store for structural knowledge graph resources.
225    graph_store: std::sync::Arc<aptu_coder_core::graph::GraphDiskStore>,
226}
227
228#[tool_router]
229impl CodeAnalyzer {
230    #[must_use]
231    pub fn list_tools() -> Vec<rmcp::model::Tool> {
232        Self::tool_router().list_all()
233    }
234
235    pub fn new(
236        peer: Arc<TokioMutex<Option<Peer<RoleServer>>>>,
237        metrics_tx: crate::metrics::MetricsSender,
238    ) -> Self {
239        crate::tools::server::build_analyzer(peer, metrics_tx)
240    }
241
242    /// Emit a "received" metric event for the given tool name.
243    /// Increments the session call sequence, locks the session ID, and sends
244    /// the metric event via the channel. Returns the (seq, sid) pair for use
245    /// by the caller in exit metrics, preserving per-call seq uniqueness.
246    async fn emit_received_metric(&self, tool: &'static str) -> (u32, Option<String>) {
247        crate::tools::server::emit_received_metric(
248            &self.metrics_tx,
249            &self.session_id,
250            &self.session_call_seq,
251            tool,
252        )
253        .await
254    }
255
256    /// Delegates to [`tools::server::handle_overview_mode`].
257    /// Kept for test access; production path goes through `analyze_directory` shim.
258    #[cfg(test)]
259    pub(crate) async fn handle_overview_mode(
260        &self,
261        params: &AnalyzeDirectoryParams,
262        ct: tokio_util::sync::CancellationToken,
263    ) -> Result<(std::sync::Arc<analyze::AnalysisOutput>, CacheTier), ErrorData> {
264        let ctx = crate::tools::AnalyzeDirectoryContext {
265            cache: self.cache.clone(),
266            disk_cache: self.disk_cache.clone(),
267            metrics_tx: self.metrics_tx.clone(),
268            peer: self.peer.clone(),
269            sid: self.session_id.lock().await.clone(),
270        };
271        crate::tools::server::handle_overview_mode(&ctx, params, ct).await
272    }
273
274    /// Delegates to [`tools::server::handle_file_details_mode`].
275    /// Kept for test access; production path goes through `analyze_file` shim.
276    #[cfg(test)]
277    pub(crate) async fn handle_file_details_mode(
278        &self,
279        params: &aptu_coder_core::types::AnalyzeFileParams,
280    ) -> Result<(std::sync::Arc<analyze::FileAnalysisOutput>, CacheTier), ErrorData> {
281        crate::tools::server::handle_file_details_mode(
282            self.cache.clone(),
283            self.disk_cache.clone(),
284            self.metrics_tx.clone(),
285            self.session_id.lock().await.clone(),
286            params,
287        )
288        .await
289    }
290
291    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, path = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty, cache_tier = tracing::field::Empty))]
292    #[tool(
293        name = "analyze_directory",
294        title = "Analyze Directory",
295        description = "Tree-view of directory with LOC, function/class counts, test markers. Respects .gitignore. Paginates with next_cursor. Default max_depth=3; pass 0 for unlimited. Large dirs (1000+ files) auto-compact to summary; pass summary=false for per-file list (summary and cursor are mutually exclusive). git_ref restricts to files changed since a branch/tag/commit. Empty directories return zero counts.",
296        output_schema = schema_for_type::<analyze::AnalysisOutput>(),
297        annotations(
298            title = "Analyze Directory",
299            read_only_hint = true,
300            destructive_hint = false,
301            idempotent_hint = true,
302            open_world_hint = false
303        )
304    )]
305    async fn analyze_directory(
306        &self,
307        params: Parameters<AnalyzeDirectoryParams>,
308        context: RequestContext<RoleServer>,
309    ) -> Result<CallToolResult, ErrorData> {
310        let mut params = params.0;
311        params.max_depth = params.max_depth.or(Some(3));
312        let t_start = std::time::Instant::now();
313        let (seq, sid) = self.emit_received_metric("analyze_directory").await;
314        let session_id = self.session_id.lock().await.clone();
315        let client_name = self.client_name.lock().await.clone();
316        let client_version = self.client_version.lock().await.clone();
317        extract_and_set_trace_context(
318            Some(&context.meta),
319            ClientMetadata {
320                session_id,
321                client_name,
322                client_version,
323            },
324        );
325        let span = tracing::Span::current();
326        span.record("gen_ai.system", "mcp");
327        span.record("gen_ai.operation.name", "execute_tool");
328        span.record("gen_ai.tool.name", "analyze_directory");
329        span.record("path", &params.path);
330        let _validated_path = match validate_path(&params.path, true) {
331            Ok(p) => p,
332            Err(e) => {
333                span.record("error", true);
334                span.record("error.type", "invalid_params");
335                return Ok(err_to_tool_result(e));
336            }
337        };
338        let ct = context.ct.clone();
339        let param_path = params.path.clone();
340        let max_depth_val = params.max_depth;
341        let ctx = tools::AnalyzeDirectoryContext {
342            cache: self.cache.clone(),
343            disk_cache: self.disk_cache.clone(),
344            metrics_tx: self.metrics_tx.clone(),
345            peer: self.peer.clone(),
346            sid: sid.clone(),
347        };
348        tools::analyze_directory::analyze_directory_handler(
349            &ctx,
350            params,
351            tools::DirectoryHandlerCall {
352                seq,
353                sid,
354                t_start,
355                param_path,
356                max_depth_val,
357                ct,
358            },
359            &span,
360        )
361        .await
362    }
363
364    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, path = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty, cache_tier = tracing::field::Empty))]
365    #[tool(
366        name = "analyze_file",
367        title = "Analyze File",
368        description = "Functions, types, classes, and imports from a single source file. Fails if directory path supplied; use analyze_directory instead. Paginates with cursor/page_size; use fields=[\"functions\",\"classes\",\"imports\"] to limit sections. summary=true and cursor are mutually exclusive. git_ref not supported. Use analyze_module for a lightweight function/import index (~75% smaller). Supported: Astro, C/C++, C#, CSS, Fortran, Go, HTML, Java, JavaScript, JSON, Kotlin, Markdown, Python, Rust, TOML, TSX, TypeScript, YAML.",
369        output_schema = schema_for_type::<analyze::FileAnalysisOutput>(),
370        annotations(
371            title = "Analyze File",
372            read_only_hint = true,
373            destructive_hint = false,
374            idempotent_hint = true,
375            open_world_hint = false
376        )
377    )]
378    async fn analyze_file(
379        &self,
380        params: Parameters<AnalyzeFileParams>,
381        context: RequestContext<RoleServer>,
382    ) -> Result<CallToolResult, ErrorData> {
383        let params = params.0;
384        let t_start = std::time::Instant::now();
385        let (seq, sid) = self.emit_received_metric("analyze_file").await;
386        let session_id = self.session_id.lock().await.clone();
387        let client_name = self.client_name.lock().await.clone();
388        let client_version = self.client_version.lock().await.clone();
389        extract_and_set_trace_context(
390            Some(&context.meta),
391            ClientMetadata {
392                session_id,
393                client_name,
394                client_version,
395            },
396        );
397        let span = tracing::Span::current();
398        span.record("gen_ai.system", "mcp");
399        span.record("gen_ai.operation.name", "execute_tool");
400        span.record("gen_ai.tool.name", "analyze_file");
401        span.record("path", &params.path);
402        let _validated_path = match validate_path(&params.path, true) {
403            Ok(p) => p,
404            Err(e) => {
405                span.record("error", true);
406                span.record("error.type", "invalid_params");
407                return Ok(err_to_tool_result(e));
408            }
409        };
410        let param_path = params.path.clone();
411        let ctx = tools::AnalyzeFileContext {
412            cache: self.cache.clone(),
413            disk_cache: self.disk_cache.clone(),
414            metrics_tx: self.metrics_tx.clone(),
415            sid: sid.clone(),
416        };
417        tools::analyze_file::analyze_file_handler(
418            &ctx, params, seq, sid, t_start, param_path, &span,
419        )
420        .await
421    }
422
423    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, symbol = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty, cache_tier = tracing::field::Empty))]
424    #[tool(
425        name = "analyze_symbol",
426        title = "Analyze Symbol",
427        description = "Call graph for a named symbol across all files in a directory. Use for \"who calls X\", transitive chains, or files importing a module path. Prefer over analyze_file when the question is \"who calls X\" or \"what does X call\" rather than \"what is in this file\". Modes: call graph (default), import_lookup (files importing a module path), def_use (write/read sites). Fails if file path supplied; fails if impl_only=true on non-Rust directory; fails if import_lookup=true with empty symbol; fails if summary=true and cursor. match_mode controls name matching (exact/insensitive/prefix/contains). git_ref restricts to changed files.",
428        output_schema = schema_for_type::<analyze::FocusedAnalysisOutput>(),
429        annotations(
430            title = "Analyze Symbol",
431            read_only_hint = true,
432            destructive_hint = false,
433            idempotent_hint = true,
434            open_world_hint = false
435        )
436    )]
437    async fn analyze_symbol(
438        &self,
439        params: Parameters<AnalyzeSymbolParams>,
440        context: RequestContext<RoleServer>,
441    ) -> Result<CallToolResult, ErrorData> {
442        let params = params.0;
443        let t_start = std::time::Instant::now();
444        let (seq, sid) = self.emit_received_metric("analyze_symbol").await;
445        let session_id = self.session_id.lock().await.clone();
446        let client_name = self.client_name.lock().await.clone();
447        let client_version = self.client_version.lock().await.clone();
448        extract_and_set_trace_context(
449            Some(&context.meta),
450            ClientMetadata {
451                session_id,
452                client_name,
453                client_version,
454            },
455        );
456        let span = tracing::Span::current();
457        span.record("gen_ai.system", "mcp");
458        span.record("gen_ai.operation.name", "execute_tool");
459        span.record("gen_ai.tool.name", "analyze_symbol");
460        span.record("symbol", &params.symbol);
461        let _validated_path = match validate_path(&params.path, true) {
462            Ok(p) => p,
463            Err(e) => {
464                span.record("error", true);
465                span.record("error.type", "invalid_params");
466                return Ok(err_to_tool_result(e));
467            }
468        };
469        let ct = context.ct.clone();
470        let param_path = params.path.clone();
471        let max_depth_val = params.follow_depth;
472        let ctx = tools::AnalyzeSymbolContext {
473            metrics_tx: self.metrics_tx.clone(),
474            call_graph_cache: self.call_graph_cache.clone(),
475            disk_cache: self.disk_cache.clone(),
476            sid: sid.clone(),
477            seq,
478        };
479        let call = tools::AnalyzeSymbolCall {
480            ct,
481            param_path,
482            max_depth_val,
483            span,
484            t_start,
485        };
486        tools::analyze_symbol::analyze_symbol_handler(ctx, params, call).await
487    }
488
489    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, path = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty, cache_tier = tracing::field::Empty))]
490    #[tool(
491        name = "analyze_module",
492        title = "Analyze Module",
493        description = "Lightweight function and import index for a single source file with minimal token cost: name, line_count, language, function names with line numbers, import list only (~75% smaller than analyze_file). Fails if directory path supplied. Pagination and git_ref not supported. Use analyze_file for signatures, types, or class details. Supported: Astro, C/C++, C#, CSS, Fortran, Go, HTML, Java, JavaScript, JSON, Kotlin, Markdown, Python, Rust, TOML, TSX, TypeScript, YAML.",
494        output_schema = schema_for_type::<types::ModuleInfo>(),
495        annotations(
496            title = "Analyze Module",
497            read_only_hint = true,
498            destructive_hint = false,
499            idempotent_hint = true,
500            open_world_hint = false
501        )
502    )]
503    async fn analyze_module(
504        &self,
505        params: Parameters<AnalyzeModuleParams>,
506        context: RequestContext<RoleServer>,
507    ) -> Result<CallToolResult, ErrorData> {
508        let params = params.0;
509        let t_start = std::time::Instant::now();
510        let (seq, sid) = self.emit_received_metric("analyze_module").await;
511        let session_id = self.session_id.lock().await.clone();
512        let client_name = self.client_name.lock().await.clone();
513        let client_version = self.client_version.lock().await.clone();
514        extract_and_set_trace_context(
515            Some(&context.meta),
516            ClientMetadata {
517                session_id,
518                client_name,
519                client_version,
520            },
521        );
522        let span = tracing::Span::current();
523        span.record("gen_ai.system", "mcp");
524        span.record("gen_ai.operation.name", "execute_tool");
525        span.record("gen_ai.tool.name", "analyze_module");
526        span.record("path", &params.path);
527        let _validated_path = match validate_path(&params.path, true) {
528            Ok(p) => p,
529            Err(e) => {
530                span.record("error", true);
531                span.record("error.type", "invalid_params");
532                return Ok(err_to_tool_result(e));
533            }
534        };
535        let param_path = params.path.clone();
536        let ctx = tools::AnalyzeModuleContext {
537            disk_cache: self.disk_cache.clone(),
538            metrics_tx: self.metrics_tx.clone(),
539            sid: sid.clone(),
540            seq,
541        };
542        tools::analyze_module::analyze_module_handler(ctx, params, param_path, &span, t_start).await
543    }
544
545    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, path = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty))]
546    #[tool(
547        name = "edit_overwrite",
548        title = "Edit Overwrite",
549        description = "Creates or overwrites a file with UTF-8 content; creates parent directories if needed. Works on any file type. Use edit_replace for targeted single-block edits. working_dir sets the base directory for path resolution (default: server CWD).",
550        output_schema = schema_for_type::<EditOverwriteOutput>(),
551        annotations(
552            title = "Edit Overwrite",
553            read_only_hint = false,
554            destructive_hint = true,
555            idempotent_hint = false,
556            open_world_hint = false
557        )
558    )]
559    async fn edit_overwrite(
560        &self,
561        params: Parameters<EditOverwriteParams>,
562        context: RequestContext<RoleServer>,
563    ) -> Result<CallToolResult, ErrorData> {
564        let params = params.0;
565        let t_start = std::time::Instant::now();
566        let (seq, sid) = self.emit_received_metric("edit_overwrite").await;
567        // Extract W3C Trace Context from request _meta if present
568        let session_id = self.session_id.lock().await.clone();
569        let client_name = self.client_name.lock().await.clone();
570        let client_version = self.client_version.lock().await.clone();
571        extract_and_set_trace_context(
572            Some(&context.meta),
573            ClientMetadata {
574                session_id,
575                client_name,
576                client_version,
577            },
578        );
579        let span = tracing::Span::current();
580        span.record("gen_ai.system", "mcp");
581        span.record("gen_ai.operation.name", "execute_tool");
582        tools::edit_overwrite::edit_overwrite(
583            params,
584            tools::EditHandlerContext {
585                sid,
586                seq,
587                cache: &self.cache,
588                metrics_tx: &self.metrics_tx,
589                edit_failure_counts: &self.edit_failure_counts,
590            },
591            &span,
592            t_start,
593        )
594        .await
595    }
596
597    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, path = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty))]
598    #[tool(
599        name = "edit_replace",
600        title = "Edit Replace",
601        description = "Replaces an exact text block; old_text must appear exactly once. Fails if zero or multiple matches (extend old_text to disambiguate). Set replace_all=true to replace every non-overlapping occurrence; old_text must be non-empty. Pass empty new_text to delete. CRLF in old_text normalized to LF; all other whitespace matched exactly. If invalid_params, re-read the file with analyze_file or analyze_module before retrying. Use edit_overwrite to replace the whole file. working_dir sets the base directory for path resolution (default: server CWD).",
602        output_schema = schema_for_type::<EditReplaceOutput>(),
603        annotations(
604            title = "Edit Replace",
605            read_only_hint = false,
606            destructive_hint = true,
607            idempotent_hint = false,
608            open_world_hint = false
609        )
610    )]
611    async fn edit_replace(
612        &self,
613        params: Parameters<EditReplaceParams>,
614        context: RequestContext<RoleServer>,
615    ) -> Result<CallToolResult, ErrorData> {
616        let params = params.0;
617        let t_start = std::time::Instant::now();
618        let (seq, sid) = self.emit_received_metric("edit_replace").await;
619        // Extract W3C Trace Context from request _meta if present
620        let session_id = self.session_id.lock().await.clone();
621        let client_name = self.client_name.lock().await.clone();
622        let client_version = self.client_version.lock().await.clone();
623        extract_and_set_trace_context(
624            Some(&context.meta),
625            ClientMetadata {
626                session_id,
627                client_name,
628                client_version,
629            },
630        );
631        let span = tracing::Span::current();
632        span.record("gen_ai.system", "mcp");
633        span.record("gen_ai.operation.name", "execute_tool");
634        tools::edit_replace::edit_replace(
635            params,
636            tools::EditHandlerContext {
637                sid,
638                seq,
639                cache: &self.cache,
640                metrics_tx: &self.metrics_tx,
641                edit_failure_counts: &self.edit_failure_counts,
642            },
643            &span,
644            t_start,
645        )
646        .await
647    }
648
649    #[tool(
650        name = "exec_command",
651        title = "Exec Command",
652        description = "Execute shell command via sh -c (or $SHELL if set). Output capped at 30 KB stdout / 10 KB stderr / 2000 lines. Set working_dir to the target directory; write commands with relative paths only. Pass stdin to pipe UTF-8 content (max 1 MB); heredoc syntax is rejected. For file writes use edit_overwrite or edit_replace. Prefer machine-readable output flags (e.g. --json) to reduce tokens.",
653        output_schema = schema_for_type::<ShellOutput>(),
654        annotations(
655            title = "Exec Command",
656            read_only_hint = false,
657            destructive_hint = true,
658            idempotent_hint = false,
659            open_world_hint = true
660        )
661    )]
662    #[instrument(skip(self, context), fields(gen_ai.system = tracing::field::Empty, gen_ai.operation.name = tracing::field::Empty, gen_ai.tool.name = tracing::field::Empty, error = tracing::field::Empty, error.type = tracing::field::Empty, command = tracing::field::Empty, exit_code = tracing::field::Empty, output_truncated = tracing::field::Empty, mcp.session.id = tracing::field::Empty, client.name = tracing::field::Empty, client.version = tracing::field::Empty, mcp.client.session.id = tracing::field::Empty))]
663    pub async fn exec_command(
664        &self,
665        params: Parameters<ExecCommandParams>,
666        context: RequestContext<RoleServer>,
667    ) -> Result<CallToolResult, ErrorData> {
668        let t_start = std::time::Instant::now();
669        let (seq, sid) = self.emit_received_metric("exec_command").await;
670        let params = params.0;
671        let session_id = self.session_id.lock().await.clone();
672        let client_name = self.client_name.lock().await.clone();
673        let client_version = self.client_version.lock().await.clone();
674        let ctx = crate::tools::exec_command::ExecContext {
675            seq,
676            sid,
677            session_id,
678            client_name,
679            client_version,
680            resolved_path: self.resolved_path.as_ref().as_deref().map(str::to_owned),
681            filter_table: self.filter_table.clone(),
682            metrics_tx: self.metrics_tx.clone(),
683            t_start,
684        };
685        crate::tools::exec_command::exec_command_impl(params, context, ctx).await
686    }
687}
688
689#[tool_handler]
690impl ServerHandler for CodeAnalyzer {
691    #[instrument(skip(self, _context), fields(service.name = tracing::field::Empty, service.version = tracing::field::Empty))]
692    async fn initialize(
693        &self,
694        request: InitializeRequestParams,
695        _context: RequestContext<RoleServer>,
696    ) -> Result<InitializeResult, ErrorData> {
697        let span = tracing::Span::current();
698        span.record("service.name", "aptu-coder");
699        span.record("service.version", env!("CARGO_PKG_VERSION"));
700
701        // Store client_info from the initialize request
702        {
703            let mut client_name_lock = self.client_name.lock().await;
704            *client_name_lock = Some(request.client_info.name.clone());
705        }
706        {
707            let mut client_version_lock = self.client_version.lock().await;
708            *client_version_lock = Some(request.client_info.version.clone());
709        }
710        Ok(self.get_info())
711    }
712
713    /// Returns server discovery information including supported protocol versions
714    /// and capabilities.
715    async fn discover(
716        &self,
717        _context: RequestContext<RoleServer>,
718    ) -> Result<DiscoverResult, ErrorData> {
719        Ok(DiscoverResult::from_server_info(
720            self.supported_protocol_versions().into_owned(),
721            self.get_info(),
722        ))
723    }
724
725    fn get_info(&self) -> InitializeResult {
726        let excluded = aptu_coder_core::EXCLUDED_DIRS.join(", ");
727        let instructions = format!(
728            "Recommended workflow:\n\
729            1. Start with analyze_directory(path=<repo_root>, max_depth=2, summary=true) to identify source package (largest by file count; exclude {excluded}).\n\
730            2. Re-run analyze_directory(path=<source_package>, max_depth=2, summary=true) for module map. Include test directories (tests/, *_test.go, test_*.py, test_*.rs, *.spec.ts, *.spec.js).\n\
731            3. For key files, prefer analyze_module for function/import index; use analyze_file for signatures and types.\n\
732            4. Use analyze_symbol to trace call graphs.\n\
733            Prefer summary=true on 1000+ files. Set max_depth=2; increase if packages too large. Paginate with cursor/page_size. For subagents: DISABLE_PROMPT_CACHING=1.\n\
734            JSONL metrics at $HOME/.local/share/aptu-coder/ (or $XDG_DATA_HOME/aptu-coder/). Always cd there before jq glob queries."
735        );
736        let capabilities = ServerCapabilities::builder()
737            .enable_tools()
738            .enable_tool_list_changed()
739            .enable_completions()
740            .enable_resources()
741            .build();
742        let server_info = Implementation::new("aptu-coder", env!("CARGO_PKG_VERSION"))
743            .with_title("Aptu Coder")
744            .with_description("MCP server for code structure analysis using tree-sitter");
745        InitializeResult::new(capabilities)
746            .with_server_info(server_info)
747            .with_instructions(&instructions)
748    }
749
750    async fn list_tools(
751        &self,
752        _request: Option<rmcp::model::PaginatedRequestParams>,
753        _context: RequestContext<RoleServer>,
754    ) -> Result<rmcp::model::ListToolsResult, ErrorData> {
755        let router = self.tool_router.read().await;
756        Ok(rmcp::model::ListToolsResult::with_all_items(
757            router.list_all(),
758        ))
759    }
760
761    async fn call_tool(
762        &self,
763        request: rmcp::model::CallToolRequestParams,
764        context: RequestContext<RoleServer>,
765    ) -> Result<CallToolResponse, ErrorData> {
766        let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context);
767        let router = self.tool_router.read().await;
768        router.call(tcc).await
769    }
770
771    async fn list_resources(
772        &self,
773        _request: Option<rmcp::model::PaginatedRequestParams>,
774        _context: RequestContext<RoleServer>,
775    ) -> Result<ListResourcesResult, ErrorData> {
776        tools::resources::list_resources_impl(_request, &_context)
777    }
778
779    async fn list_resource_templates(
780        &self,
781        _request: Option<rmcp::model::PaginatedRequestParams>,
782        _context: RequestContext<RoleServer>,
783    ) -> Result<ListResourceTemplatesResult, ErrorData> {
784        tools::resources::list_resource_templates_impl(_request, &_context)
785    }
786
787    async fn read_resource(
788        &self,
789        request: ReadResourceRequestParams,
790        _context: RequestContext<RoleServer>,
791    ) -> Result<ReadResourceResponse, ErrorData> {
792        tools::resources::read_resource_impl(request, &self.graph_store)
793    }
794
795    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
796        crate::tools::server::on_initialized_impl(
797            self.peer.clone(),
798            self.session_id.clone(),
799            self.session_call_seq.clone(),
800            self.tool_router.clone(),
801            &context.peer,
802        )
803        .await;
804    }
805
806    #[instrument(skip(self, _context))]
807    async fn on_cancelled(
808        &self,
809        notification: CancelledNotificationParam,
810        _context: NotificationContext<RoleServer>,
811    ) {
812        tracing::info!(
813            request_id = ?notification.request_id,
814            reason = ?notification.reason,
815            "Received cancellation notification"
816        );
817    }
818
819    #[instrument(skip(self, _context))]
820    async fn complete(
821        &self,
822        request: CompleteRequestParams,
823        _context: RequestContext<RoleServer>,
824    ) -> Result<CompleteResult, ErrorData> {
825        // Dispatch on argument name: "path" or "symbol"
826        let argument_name = &request.argument.name;
827        let argument_value = &request.argument.value;
828
829        let completions = match argument_name.as_str() {
830            "path" => {
831                // Path completions: use current directory as root
832                let root = Path::new(".");
833                completion::path_completions(root, argument_value)
834            }
835            "symbol" => {
836                // Symbol completions: need the path argument from context
837                let path_arg = request
838                    .context
839                    .as_ref()
840                    .and_then(|ctx| ctx.get_argument("path"));
841
842                match path_arg {
843                    Some(path_str) => {
844                        let path = Path::new(path_str);
845                        completion::symbol_completions(&self.cache, path, argument_value)
846                    }
847                    None => Vec::new(),
848                }
849            }
850            _ => Vec::new(),
851        };
852
853        // Create CompletionInfo with has_more flag if >100 results
854        let total_count = u32::try_from(completions.len()).unwrap_or(u32::MAX);
855        let (values, has_more) = if completions.len() > 100 {
856            (completions.into_iter().take(100).collect(), true)
857        } else {
858            (completions, false)
859        };
860
861        let completion_info =
862            match CompletionInfo::with_pagination(values, Some(total_count), has_more) {
863                Ok(info) => info,
864                Err(_) => {
865                    // Graceful degradation: return empty on error
866                    CompletionInfo::with_all_values(Vec::new())
867                        .unwrap_or_else(|_| CompletionInfo::new(Vec::new()).unwrap())
868                }
869            };
870
871        Ok(CompleteResult::new(completion_info))
872    }
873}
874
875#[cfg(test)]
876mod tests;