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