=== src/graph/traversal.rs ===
use crate::db::models::CodeElement;
use crate::graph::GraphEngine;
use std::collections::{HashSet, VecDeque};

pub struct ImpactAnalyzer<'a> {
    graph: &'a GraphEngine,
}

impl<'a> ImpactAnalyzer<'a> {
    pub fn new(graph: &'a GraphEngine) -> Self {
        Self { graph }
    }

    pub fn calculate_impact_radius(
        &self,
        start_file: &str,
        depth: u32,
    ) -> Result<ImpactResult, Box<dyn std::error::Error>> {
        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        let mut affected_elements = Vec::new();

        queue.push_back((start_file.to_string(), 0));
        visited.insert(start_file.to_string());

        while let Some((current, current_depth)) = queue.pop_front() {
            if current_depth >= depth {
                continue;
            }

            let relationships = self.graph.get_relationships(&current)?;

            for rel in relationships {
                if !visited.contains(&rel.target_qualified) {
                    visited.insert(rel.target_qualified.clone());
                    queue.push_back((rel.target_qualified.clone(), current_depth + 1));

                    if let Ok(Some(element)) = self.graph.find_element(&rel.target_qualified) {
                        affected_elements.push(element);
                    }
                }
            }

            let dependents = self.graph.get_dependents(&current)?;
            for rel in dependents {
                if !visited.contains(&rel.source_qualified) {
                    visited.insert(rel.source_qualified.clone());
                    queue.push_back((rel.source_qualified.clone(), current_depth + 1));

                    if let Ok(Some(element)) = self.graph.find_element(&rel.source_qualified) {

=== src/graph/query.rs ===
use crate::db::models::{BusinessLogic, CodeElement, Relationship, DocLink, TraceabilityEntry, TraceabilityReport};
use crate::db::schema::CozoDb;
use crate::graph::cache::QueryCache;
use std::sync::Arc;
use tokio::sync::RwLock;

fn escape_datalog(s: &str) -> String {
    s.replace('\\', "\\\\").replace('"', "\\\"")
}

#[derive(Clone)]
pub struct GraphEngine {
    db: CozoDb,
    cache: Arc<RwLock<QueryCache>>,
}

impl GraphEngine {
    pub fn new(db: CozoDb) -> Self {
        Self {
            db,
            cache: Arc::new(RwLock::new(QueryCache::new(300, 1000))),
        }
    }

    pub fn with_cache(db: CozoDb, cache: QueryCache) -> Self {
        Self {
            db,
            cache: Arc::new(RwLock::new(cache)),
        }
    }

    pub fn db(&self) -> &CozoDb {
        &self.db
    }

    pub fn find_element(
        &self,
        qualified_name: &str,
    ) -> Result<Option<CodeElement>, Box<dyn std::error::Error>> {
        let query = format!(
            r#"?[qualified_name, element_type, name, file_path, line_start, line_end, language, parent_qualified, metadata] := *code_elements[qualified_name, element_type, name, file_path, line_start, line_end, language, parent_qualified, metadata], qualified_name = "{}""#,
            qualified_name
        );

        let result = self.db.run_script(&query, std::collections::BTreeMap::new())?;
        let rows = result.rows;

        if rows.is_empty() {
            return Ok(None);
        }

=== src/graph/context.rs ===
use crate::db::models::CodeElement;
use crate::graph::GraphEngine;
use serde::{Deserialize, Serialize};

const DEFAULT_MAX_TOKENS: usize = 4000;
const CHARS_PER_TOKEN: usize = 4;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum ContextPriority {
    Contained = 1,
    Imported = 2,
    RecentlyChanged = 3,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextElement {
    pub element: CodeElement,
    pub priority: ContextPriority,
    pub token_count: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextResult {
    pub elements: Vec<ContextElement>,
    pub total_tokens: usize,
    pub max_tokens: usize,
    pub truncated: bool,
}

impl ContextResult {
    pub fn to_prompt(&self) -> String {
        let mut prompt = String::new();
        prompt.push_str("# Code Context\n\n");

        for ctx_elem in &self.elements {
            let elem = &ctx_elem.element;
            prompt.push_str(&format!(
                "## {}: {}\nFile: {}:{}:{}\n",
                elem.element_type,
                elem.qualified_name,
                elem.file_path,
                elem.line_start,
                elem.line_end
            ));

            if let Some(parent) = &elem.parent_qualified {
                prompt.push_str(&format!("Parent: {}\n", parent));
            }

            prompt.push('\n');

=== src/benchmark/runner.rs ===
use crate::benchmark::data::BenchmarkResult;
use std::error::Error;
use std::path::PathBuf;
use std::process::{Command, Output, Stdio};
use std::time::Duration;

fn kilo_config_path() -> PathBuf {
    dirs::config_dir()
        .unwrap_or_else(|| PathBuf::from(".config"))
        .join("kilo")
}

const KILO_MCP_WITH_LEANKG: &str = "mcp_settings_with_leankg.json";
const KILO_MCP_WITHOUT_LEANKG: &str = "mcp_settings_without_leankg.json";
const KILO_MCP_SETTINGS: &str = "kilo.json";

trait WaitWithOutputTimeout {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()>;
}

impl WaitWithOutputTimeout for std::process::Child {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()> {
        use std::thread;
        use std::time::Instant;

        let start = Instant::now();
        let handle =
            thread::spawn(move || self.wait_with_output().expect("wait_with_output failed"));

        loop {
            if handle.is_finished() {
                return handle.join().map_err(|_| ());
            }
            if start.elapsed() >= duration {
                return Err(());
            }
            thread::sleep(Duration::from_millis(100));
        }
    }
}

pub struct BenchmarkRunner {
    output_dir: PathBuf,
    cli: CliTool,
}

#[derive(Clone)]
pub enum CliTool {
    OpenCode,
    Gemini,

=== src/benchmark/data.rs ===
use serde::{Deserialize, Serialize};
use std::error::Error;
use std::path::Path;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResult {
    pub total_tokens: u32,
    pub input_tokens: u32,
    pub cached_tokens: u32,
    pub token_percent: f32,
    pub build_time_seconds: f32,
    pub success: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptTask {
    pub id: String,
    pub prompt: String,
    pub expected: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptCategory {
    pub name: String,
    pub description: String,
    pub tasks: Vec<PromptTask>,
}

impl PromptCategory {
    pub fn from_yaml(path: &Path) -> Result<Self, Box<dyn Error>> {
        let content = std::fs::read_to_string(path)?;
        let category: PromptCategory = serde_yaml::from_str(&content)?;
        Ok(category)
    }

    pub fn load_all(prompts_dir: &Path) -> Result<Vec<Self>, Box<dyn Error>> {
        let mut categories = Vec::new();
        for entry in std::fs::read_dir(prompts_dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|s| s.to_str()) == Some("yaml") {
                categories.push(Self::from_yaml(&path)?);
            }
        }
        Ok(categories)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OverheadResult {

=== src/watcher/notify_handler.rs ===
use notify::{Config, Event, RecommendedWatcher, RecursiveMode, Watcher};
use std::path::{Path, PathBuf};
use std::sync::mpsc::{channel, Receiver};
use std::time::Duration;
use tokio::sync::mpsc as tokio_mpsc;

#[derive(Debug, Clone)]
pub struct FileChange {
    pub path: PathBuf,
    pub kind: FileChangeKind,
}

#[derive(Debug, Clone)]
pub enum FileChangeKind {
    Created,
    Modified,
    Deleted,
}

impl From<&notify::EventKind> for FileChangeKind {
    fn from(kind: &notify::EventKind) -> Self {
        match kind {
            notify::EventKind::Create(_) => FileChangeKind::Created,
            notify::EventKind::Modify(_) => FileChangeKind::Modified,
            notify::EventKind::Remove(_) => FileChangeKind::Deleted,
            _ => FileChangeKind::Modified,
        }
    }
}

pub struct FileWatcher {
    watcher: RecommendedWatcher,
    watch_path: PathBuf,
    rx: Receiver<Result<Event, notify::Error>>,
}

impl FileWatcher {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, notify::Error> {
        let (tx, rx) = channel();

        let watcher = RecommendedWatcher::new(
            move |res: Result<Event, notify::Error>| {
                if let Ok(event) = res {
                    let _ = tx.send(Ok(event));
                }
            },
            Config::default().with_poll_interval(Duration::from_secs(2)),
        )?;

        let mut watcher = watcher;

=== src/benchmark/runner.rs ===
use crate::benchmark::data::BenchmarkResult;
use std::error::Error;
use std::path::PathBuf;
use std::process::{Command, Output, Stdio};
use std::time::Duration;

fn kilo_config_path() -> PathBuf {
    dirs::config_dir()
        .unwrap_or_else(|| PathBuf::from(".config"))
        .join("kilo")
}

const KILO_MCP_WITH_LEANKG: &str = "mcp_settings_with_leankg.json";
const KILO_MCP_WITHOUT_LEANKG: &str = "mcp_settings_without_leankg.json";
const KILO_MCP_SETTINGS: &str = "kilo.json";

trait WaitWithOutputTimeout {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()>;
}

impl WaitWithOutputTimeout for std::process::Child {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()> {
        use std::thread;
        use std::time::Instant;

        let start = Instant::now();
        let handle =
            thread::spawn(move || self.wait_with_output().expect("wait_with_output failed"));

        loop {
            if handle.is_finished() {
                return handle.join().map_err(|_| ());
            }
            if start.elapsed() >= duration {
                return Err(());
            }
            thread::sleep(Duration::from_millis(100));
        }
    }
}

pub struct BenchmarkRunner {
    output_dir: PathBuf,
    cli: CliTool,
}

#[derive(Clone)]
pub enum CliTool {
    OpenCode,
    Gemini,

=== src/config/project.rs ===
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProjectConfig {
    pub project: ProjectSettings,
    pub indexer: IndexerConfig,
    pub mcp: McpConfig,
    pub documentation: DocConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProjectSettings {
    pub name: String,
    pub root: PathBuf,
    pub languages: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexerConfig {
    pub exclude: Vec<String>,
    pub include: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpConfig {
    pub enabled: bool,
    pub port: u16,
    pub auth_token: String,
    pub auto_index_on_start: bool,
    pub auto_index_threshold_minutes: u64,
    pub index_on_first_call: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DocConfig {
    pub output: PathBuf,
    pub templates: Vec<String>,
}

impl Default for ProjectConfig {
    fn default() -> Self {
        Self {
            project: ProjectSettings {
                name: "my-project".to_string(),
                root: PathBuf::from("./src"),
                languages: vec![
                    "go".to_string(),
                    "typescript".to_string(),
                    "python".to_string(),

=== src/lib.rs ===
// LeanKG Library
// This library contains the core modules for the knowledge graph system

pub mod cli;
pub mod config;
pub mod db;
pub mod doc;
pub mod doc_indexer;
pub mod graph;
pub mod indexer;
pub mod mcp;
pub mod watcher;
pub mod benchmark;

=== src/watcher/mod.rs ===
pub mod notify_handler;

#[allow(unused_imports)]
pub use notify_handler::*;

=== src/web/mod.rs ===
pub mod handlers;

use axum::{
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::{get, post, put},
    Json, Router,
};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::db::schema::{init_db, CozoDb};
use crate::graph::GraphEngine;

#[derive(Clone)]
pub struct AppState {
    pub db_path: std::path::PathBuf,
    db: Arc<RwLock<Option<CozoDb>>>,
}

impl AppState {
    pub async fn new(db_path: std::path::PathBuf) -> Result<Self, Box<dyn std::error::Error>> {
        Ok(Self {
            db_path,
            db: Arc::new(RwLock::new(None)),
        })
    }

    pub async fn init_db(&self) -> Result<(), Box<dyn std::error::Error>> {
        let db = init_db(&self.db_path)?;
        let mut lock = self.db.write().await;
        *lock = Some(db);
        Ok(())
    }

    pub fn get_db(&self) -> Result<CozoDb, Box<dyn std::error::Error + Send + Sync>> {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let lock = self.db.read().await;
            lock.clone()
                .ok_or_else(|| "Database not initialized".into())
        })
    }

    pub async fn get_graph_engine(&self) -> Result<GraphEngine, Box<dyn std::error::Error + Send + Sync>> {
        let lock = self.db.read().await;
        let db = lock.clone()
            .ok_or_else(|| -> Box<dyn std::error::Error + Send + Sync> { "Database not initialized".into() })?;
        Ok(GraphEngine::new(db))

=== src/mcp/handler.rs ===
use crate::db::models::{CodeElement, Relationship};
use crate::graph::{GraphEngine, ImpactAnalyzer};
use serde_json::{json, Value};

const INSTRUCTIONS_CONTENT: &str = r#"# LeanKG Tools - Usage Instructions

## For AI Coding Agents (Cursor, OpenCode, etc.)

Use LeanKG tools **first** before performing any codebase search, navigation, or impact analysis.

---

## When to Use Each Tool

### Code Discovery & Search

| Task | Use This Tool |
|------|--------------|
| Find a file by name | `query_file` |
| Find a function definition | `find_function` |
| Search code by name/type | `search_code` |
| Get full codebase structure | `get_code_tree` |

### Dependency Analysis

| Task | Use This Tool |
|------|--------------|
| Get direct imports of a file | `get_dependencies` |
| Get files that import/use a file | `get_dependents` |
| Get function call chain (full depth) | `get_call_graph` |
| Calculate what breaks if file changes | `get_impact_radius` |

### Review & Context

| Task | Use This Tool |
|------|--------------|
| Generate focused review context | `get_review_context` |
| Get minimal AI context (token-optimized) | `get_context` |
| Find oversized functions | `find_large_functions` |

### Testing & Documentation

| Task | Use This Tool |
|------|--------------|
| Get test coverage for a function | `get_tested_by` |
| Get docs that reference a file | `get_doc_for_file` |
| Get code elements in a doc | `get_files_for_doc` |
| Get doc directory structure | `get_doc_structure` |
| Find docs related to a change | `find_related_docs` |


=== src/mcp/server.rs ===
use crate::db::schema::init_db;
use crate::graph::GraphEngine;
use crate::mcp::auth::AuthConfig;
use crate::mcp::handler::ToolHandler;
use crate::mcp::tools::ToolRegistry;
use crate::mcp::watcher::start_watcher;
use rmcp::handler::server::ServerHandler;
use rmcp::model::{
    CallToolRequestParams, CallToolResult, Content, ListToolsResult, ServerCapabilities,
    ServerInfo, Tool,
};
use rmcp::service::{serve_server, RoleServer};
use rmcp::transport::stdio;
use std::path::PathBuf;
use std::sync::Arc;
use parking_lot::RwLock;
use tokio::sync::RwLock as TokioRwLock;

pub struct MCPServer {
    auth_config: Arc<TokioRwLock<AuthConfig>>,
    db_path: Arc<RwLock<PathBuf>>,
    graph_engine: Arc<parking_lot::Mutex<Option<GraphEngine>>>,
    watch_path: Option<PathBuf>,
}

impl std::fmt::Debug for MCPServer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MCPServer")
            .field("db_path", &self.db_path)
            .finish()
    }
}

impl Clone for MCPServer {
    fn clone(&self) -> Self {
        Self {
            auth_config: self.auth_config.clone(),
            db_path: self.db_path.clone(),
            graph_engine: self.graph_engine.clone(),
            watch_path: self.watch_path.clone(),
        }
    }
}

impl MCPServer {
    pub fn new(db_path: std::path::PathBuf) -> Self {
        Self {
            auth_config: Arc::new(TokioRwLock::new(AuthConfig::default())),
            db_path: Arc::new(RwLock::new(db_path)),
            graph_engine: Arc::new(parking_lot::Mutex::new(None)),

=== src/db/models.rs ===
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum RelationshipType {
    Imports,
    Calls,
    References,
    DocumentedBy,
    TestedBy,
    Tests,
    Contains,
    Defines,
    Implements,
    Implementations,
}

impl RelationshipType {
    pub fn as_str(&self) -> &'static str {
        match self {
            RelationshipType::Imports => "imports",
            RelationshipType::Calls => "calls",
            RelationshipType::References => "references",
            RelationshipType::DocumentedBy => "documented_by",
            RelationshipType::TestedBy => "tested_by",
            RelationshipType::Tests => "tests",
            RelationshipType::Contains => "contains",
            RelationshipType::Defines => "defines",
            RelationshipType::Implements => "implements",
            RelationshipType::Implementations => "implementations",
        }
    }

    pub fn from_str(s: &str) -> Option<Self> {
        match s {
            "imports" => Some(RelationshipType::Imports),
            "calls" => Some(RelationshipType::Calls),
            "references" => Some(RelationshipType::References),
            "documented_by" => Some(RelationshipType::DocumentedBy),
            "tested_by" => Some(RelationshipType::TestedBy),
            "tests" => Some(RelationshipType::Tests),
            "contains" => Some(RelationshipType::Contains),
            "defines" => Some(RelationshipType::Defines),
            "implements" => Some(RelationshipType::Implements),
            "implementations" => Some(RelationshipType::Implementations),
            _ => None,
        }
    }
}

impl std::fmt::Display for RelationshipType {

=== src/main.rs ===
mod benchmark;
mod cli;
mod config;
mod db;
mod doc;
mod doc_indexer;
mod graph;
mod indexer;
mod mcp;
mod watcher;
mod web;

use clap::Parser;

#[derive(Parser, Debug)]
#[command(name = "leankg")]
#[command(about = "Lightweight knowledge graph for AI-assisted development")]
pub struct Args {
    #[command(subcommand)]
    pub command: cli::CLICommand,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    if !matches!(args.command, cli::CLICommand::McpStdio { watch: _ }) {
        tracing_subscriber::fmt::init();
    }

    match args.command {
        cli::CLICommand::Init { path } => {
            init_project(&path)?;
        }
        cli::CLICommand::Index {
            path,
            incremental,
            lang,
            exclude,
            verbose,
        } => {
            let project_path = find_project_root()?;
            let db_path = project_path.join(".leankg");
            tokio::fs::create_dir_all(&db_path).await?;
            let exclude_patterns: Vec<String> = exclude
                .as_ref()
                .map(|e| e.split(',').map(|s| s.trim().to_string()).collect())
                .unwrap_or_default();
            if incremental {
                incremental_index_codebase(

=== src/web/handlers.rs ===
use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::IntoResponse,
    Json,
};
use serde::{Deserialize, Serialize};

use crate::db;

use super::{ApiResponse, AppState};

#[derive(Deserialize, Serialize)]
pub struct QueryRequest {
    pub query: String,
}

#[derive(Serialize)]
pub struct QueryResponse {
    pub result: Vec<serde_json::Value>,
}

#[derive(Deserialize)]
pub struct SearchParams {
    pub q: Option<String>,
    pub element_type: Option<String>,
    pub file_path: Option<String>,
}

#[derive(Deserialize, Serialize)]
pub struct AnnotationRequest {
    pub element_qualified: String,
    pub description: String,
    pub user_story_id: Option<String>,
    pub feature_id: Option<String>,
}

#[derive(Serialize, Clone)]
pub struct GraphData {
    pub nodes: Vec<GraphNode>,
    pub edges: Vec<GraphEdge>,
    pub filtered: Option<GraphFilterInfo>,
}

#[derive(Serialize, Clone)]
pub struct GraphFilterInfo {
    pub tests_filtered: usize,
    pub message: String,
}


=== src/mcp/tools.rs ===
use serde_json::json;
use serde_json::Value;

pub struct ToolRegistry;

impl ToolRegistry {
    pub fn list_tools() -> Vec<ToolDefinition> {
        vec![
            ToolDefinition {
                name: "mcp_init".to_string(),
                description: "Initialize LeanKG project (creates .leankg/ and leankg.yaml)"
                    .to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path for LeanKG project (default: .leankg)"}
                    },
                    "required": []
                }),
            },
            ToolDefinition {
                name: "mcp_index".to_string(),
                description: "Index codebase (mirrors CLI: leankg index)".to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path to index (default: current directory)"},
                        "incremental": {"type": "boolean", "description": "Only index changed files (git-based)"},
                        "lang": {"type": "string", "description": "Filter by language (e.g., go,ts,py,rs)"},
                        "exclude": {"type": "string", "description": "Exclude patterns (comma-separated)"}
                    },
                    "required": []
                }),
            },
            ToolDefinition {
                name: "mcp_index_docs".to_string(),
                description: "Index documentation directory to create code-doc traceability edges. \
                              Run after mcp_index to populate documented_by and references relationships."
                    .to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path to docs directory (default: ./docs)"}
                    },
                    "required": []
                }),
            },
            ToolDefinition {
                name: "mcp_install".to_string(),
                description: "Create .mcp.json for MCP client configuration".to_string(),

=== src/mcp/handler.rs ===
use crate::db::models::{CodeElement, Relationship};
use crate::graph::{GraphEngine, ImpactAnalyzer};
use serde_json::{json, Value};

const INSTRUCTIONS_CONTENT: &str = r#"# LeanKG Tools - Usage Instructions

## For AI Coding Agents (Cursor, OpenCode, etc.)

Use LeanKG tools **first** before performing any codebase search, navigation, or impact analysis.

---

## When to Use Each Tool

### Code Discovery & Search

| Task | Use This Tool |
|------|--------------|
| Find a file by name | `query_file` |
| Find a function definition | `find_function` |
| Search code by name/type | `search_code` |
| Get full codebase structure | `get_code_tree` |

### Dependency Analysis

| Task | Use This Tool |
|------|--------------|
| Get direct imports of a file | `get_dependencies` |
| Get files that import/use a file | `get_dependents` |
| Get function call chain (full depth) | `get_call_graph` |
| Calculate what breaks if file changes | `get_impact_radius` |

### Review & Context

| Task | Use This Tool |
|------|--------------|
| Generate focused review context | `get_review_context` |
| Get minimal AI context (token-optimized) | `get_context` |
| Find oversized functions | `find_large_functions` |

### Testing & Documentation

| Task | Use This Tool |
|------|--------------|
| Get test coverage for a function | `get_tested_by` |
| Get docs that reference a file | `get_doc_for_file` |
| Get code elements in a doc | `get_files_for_doc` |
| Get doc directory structure | `get_doc_structure` |
| Find docs related to a change | `find_related_docs` |


=== src/benchmark/runner.rs ===
use crate::benchmark::data::BenchmarkResult;
use std::error::Error;
use std::path::PathBuf;
use std::process::{Command, Output, Stdio};
use std::time::Duration;

fn kilo_config_path() -> PathBuf {
    dirs::config_dir()
        .unwrap_or_else(|| PathBuf::from(".config"))
        .join("kilo")
}

const KILO_MCP_WITH_LEANKG: &str = "mcp_settings_with_leankg.json";
const KILO_MCP_WITHOUT_LEANKG: &str = "mcp_settings_without_leankg.json";
const KILO_MCP_SETTINGS: &str = "kilo.json";

trait WaitWithOutputTimeout {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()>;
}

impl WaitWithOutputTimeout for std::process::Child {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()> {
        use std::thread;
        use std::time::Instant;

        let start = Instant::now();
        let handle =
            thread::spawn(move || self.wait_with_output().expect("wait_with_output failed"));

        loop {
            if handle.is_finished() {
                return handle.join().map_err(|_| ());
            }
            if start.elapsed() >= duration {
                return Err(());
            }
            thread::sleep(Duration::from_millis(100));
        }
    }
}

pub struct BenchmarkRunner {
    output_dir: PathBuf,
    cli: CliTool,
}

#[derive(Clone)]
pub enum CliTool {
    OpenCode,
    Gemini,

=== src/benchmark/runner.rs ===
use crate::benchmark::data::BenchmarkResult;
use std::error::Error;
use std::path::PathBuf;
use std::process::{Command, Output, Stdio};
use std::time::Duration;

fn kilo_config_path() -> PathBuf {
    dirs::config_dir()
        .unwrap_or_else(|| PathBuf::from(".config"))
        .join("kilo")
}

const KILO_MCP_WITH_LEANKG: &str = "mcp_settings_with_leankg.json";
const KILO_MCP_WITHOUT_LEANKG: &str = "mcp_settings_without_leankg.json";
const KILO_MCP_SETTINGS: &str = "kilo.json";

trait WaitWithOutputTimeout {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()>;
}

impl WaitWithOutputTimeout for std::process::Child {
    fn wait_with_output_timeout(self, duration: Duration) -> Result<Output, ()> {
        use std::thread;
        use std::time::Instant;

        let start = Instant::now();
        let handle =
            thread::spawn(move || self.wait_with_output().expect("wait_with_output failed"));

        loop {
            if handle.is_finished() {
                return handle.join().map_err(|_| ());
            }
            if start.elapsed() >= duration {
                return Err(());
            }
            thread::sleep(Duration::from_millis(100));
        }
    }
}

pub struct BenchmarkRunner {
    output_dir: PathBuf,
    cli: CliTool,
}

#[derive(Clone)]
pub enum CliTool {
    OpenCode,
    Gemini,

=== src/benchmark/mod.rs ===
pub mod data;
pub mod runner;

use std::path::PathBuf;

pub use runner::{BenchmarkRunner, CliTool};

pub fn run(category: Option<String>, cli: CliTool) -> Result<(), Box<dyn std::error::Error>> {
    let prompts_dir = PathBuf::from("benchmark/prompts");
    let output_dir = PathBuf::from("benchmark/results");

    let categories = if let Some(cat) = category {
        vec![data::PromptCategory::from_yaml(
            &prompts_dir.join(format!("{}.yaml", cat)),
        )?]
    } else {
        data::PromptCategory::load_all(&prompts_dir)?
    };

    let runner = BenchmarkRunner::new(output_dir, cli);

    for cat in &categories {
        println!("\n=== Category: {} ===\n", cat.name);
        for task in &cat.tasks {
            println!("Running: {}", task.id);

            let with_leankg = runner.run_with_leankg(&task.prompt);
            let without_leankg = runner.run_without_leankg(&task.prompt);

            let overhead = with_leankg.overhead(&without_leankg);

            println!(
                "  With LeanKG: {} tokens (input: {}, cached: {})",
                with_leankg.total_tokens, with_leankg.input_tokens, with_leankg.cached_tokens
            );
            println!(
                "  Without: {} tokens (input: {}, cached: {})",
                without_leankg.total_tokens,
                without_leankg.input_tokens,
                without_leankg.cached_tokens
            );
            println!("  Overhead: {} tokens\n", overhead.token_delta);

            let _ = runner.save_comparison(&with_leankg, &without_leankg, &task.id);
        }
    }

    Ok(())
}

=== src/mcp/tools.rs ===
use serde_json::json;
use serde_json::Value;

pub struct ToolRegistry;

impl ToolRegistry {
    pub fn list_tools() -> Vec<ToolDefinition> {
        vec![
            ToolDefinition {
                name: "mcp_init".to_string(),
                description: "Initialize LeanKG project (creates .leankg/ and leankg.yaml)"
                    .to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path for LeanKG project (default: .leankg)"}
                    },
                    "required": []
                }),
            },
            ToolDefinition {
                name: "mcp_index".to_string(),
                description: "Index codebase (mirrors CLI: leankg index)".to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path to index (default: current directory)"},
                        "incremental": {"type": "boolean", "description": "Only index changed files (git-based)"},
                        "lang": {"type": "string", "description": "Filter by language (e.g., go,ts,py,rs)"},
                        "exclude": {"type": "string", "description": "Exclude patterns (comma-separated)"}
                    },
                    "required": []
                }),
            },
            ToolDefinition {
                name: "mcp_index_docs".to_string(),
                description: "Index documentation directory to create code-doc traceability edges. \
                              Run after mcp_index to populate documented_by and references relationships."
                    .to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path to docs directory (default: ./docs)"}
                    },
                    "required": []
                }),
            },
            ToolDefinition {
                name: "mcp_install".to_string(),
                description: "Create .mcp.json for MCP client configuration".to_string(),
