=== 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/cli/mod.rs ===
use clap::Subcommand;

#[derive(Subcommand, Debug)]
pub enum CLICommand {
    /// Initialize a new LeanKG project
    Init {
        #[arg(long, default_value = ".leankg")]
        path: String,
    },
    /// Index the codebase
    Index {
        /// Path to index
        path: Option<String>,
        #[arg(long, short)]
        incremental: bool,
        /// Filter by language (e.g., go,ts,py)
        #[arg(long, short)]
        lang: Option<String>,
        /// Exclude patterns (comma-separated)
        #[arg(long)]
        exclude: Option<String>,
        /// Verbose output
        #[arg(long, short)]
        verbose: bool,
    },
    /// Query the knowledge graph
    Query {
        /// Query string
        query: String,
        /// Query type: name, type, rel, pattern
        #[arg(long, default_value = "name")]
        kind: String,
    },
    /// Generate documentation
    Generate {
        #[arg(long, short)]
        template: Option<String>,
    },
    /// Start web UI server (deprecated - use 'web' command instead)
    Serve,
    /// Start the embedded web UI server
    Web,
    /// Start MCP server with stdio transport (for opencode integration)
    McpStdio {
        /// Enable auto-indexing with file watcher
        #[arg(long)]
        watch: bool,
    },
    /// Calculate impact radius
    Impact {

=== 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/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/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/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/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/mod.rs ===
pub mod cache;
pub mod context;
pub mod query;
pub mod traversal;

#[allow(unused_imports)]
pub use cache::*;
#[allow(unused_imports)]
pub use context::{ContextElement, ContextPriority, ContextProvider, ContextResult};
#[allow(unused_imports)]
pub use query::*;
#[allow(unused_imports)]
pub use traversal::*;

=== 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/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/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/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/cli/mod.rs ===
use clap::Subcommand;

#[derive(Subcommand, Debug)]
pub enum CLICommand {
    /// Initialize a new LeanKG project
    Init {
        #[arg(long, default_value = ".leankg")]
        path: String,
    },
    /// Index the codebase
    Index {
        /// Path to index
        path: Option<String>,
        #[arg(long, short)]
        incremental: bool,
        /// Filter by language (e.g., go,ts,py)
        #[arg(long, short)]
        lang: Option<String>,
        /// Exclude patterns (comma-separated)
        #[arg(long)]
        exclude: Option<String>,
        /// Verbose output
        #[arg(long, short)]
        verbose: bool,
    },
    /// Query the knowledge graph
    Query {
        /// Query string
        query: String,
        /// Query type: name, type, rel, pattern
        #[arg(long, default_value = "name")]
        kind: String,
    },
    /// Generate documentation
    Generate {
        #[arg(long, short)]
        template: Option<String>,
    },
    /// Start web UI server (deprecated - use 'web' command instead)
    Serve,
    /// Start the embedded web UI server
    Web,
    /// Start MCP server with stdio transport (for opencode integration)
    McpStdio {
        /// Enable auto-indexing with file watcher
        #[arg(long)]
        watch: bool,
    },
    /// Calculate impact radius
    Impact {

=== 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/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/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/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/mod.rs ===
pub mod cache;
pub mod context;
pub mod query;
pub mod traversal;

#[allow(unused_imports)]
pub use cache::*;
#[allow(unused_imports)]
pub use context::{ContextElement, ContextPriority, ContextProvider, ContextResult};
#[allow(unused_imports)]
pub use query::*;
#[allow(unused_imports)]
pub use traversal::*;

=== 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/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/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;
