Skip to main content

kdo_context/
lib.rs

1//! Context generation for kdo projects.
2//!
3//! Uses tree-sitter to extract public API signatures (no bodies) and generates
4//! structured CONTEXT.md files with token budget enforcement.
5
6mod extract;
7mod generate;
8
9pub use extract::{extract_signatures, Signature, SignatureKind};
10pub use generate::{generate_context, render_context_md, ContextBundle};
11
12use kdo_core::KdoError;
13use kdo_graph::WorkspaceGraph;
14
15/// Context generator — thin wrapper providing the MCP-facing API.
16#[derive(Debug)]
17pub struct ContextGenerator;
18
19impl ContextGenerator {
20    /// Create a new context generator.
21    pub fn new() -> Self {
22        Self
23    }
24
25    /// Generate a token-budgeted context bundle and render it as markdown.
26    pub fn generate_bundle(
27        &self,
28        project: &str,
29        token_budget: usize,
30        graph: &WorkspaceGraph,
31    ) -> Result<String, KdoError> {
32        let bundle = generate_context(graph, project, token_budget)?;
33        Ok(render_context_md(&bundle))
34    }
35
36    /// Read a specific symbol's source from a project.
37    ///
38    /// Searches all extracted signatures for a matching name and returns the text.
39    pub fn read_symbol(
40        &self,
41        project_name: &str,
42        symbol: &str,
43        graph: &WorkspaceGraph,
44    ) -> Result<String, KdoError> {
45        let project = graph.get_project(project_name)?;
46        let source_files = collect_source_files(&project.path, &project.language);
47
48        for file in &source_files {
49            let sigs = extract_signatures(file, &project.language);
50            for sig in &sigs {
51                if sig.text.contains(symbol) {
52                    return Ok(format!("// {}:{}\n{}", sig.file, sig.line, sig.text));
53                }
54            }
55        }
56
57        // If not found via signatures, try searching file contents directly
58        for file in &source_files {
59            if let Ok(content) = std::fs::read_to_string(file) {
60                if content.contains(symbol) {
61                    // Find the relevant block
62                    for (i, line) in content.lines().enumerate() {
63                        if line.contains(symbol) {
64                            let start = i.saturating_sub(2);
65                            let end = (i + 20).min(content.lines().count());
66                            let snippet: String = content
67                                .lines()
68                                .skip(start)
69                                .take(end - start)
70                                .collect::<Vec<_>>()
71                                .join("\n");
72                            return Ok(format!("// {}:{}\n{}", file.display(), start + 1, snippet));
73                        }
74                    }
75                }
76            }
77        }
78
79        Err(KdoError::ProjectNotFound(format!(
80            "symbol '{symbol}' not found in project '{project_name}'"
81        )))
82    }
83}
84
85impl Default for ContextGenerator {
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91/// Collect source files for a project based on its language.
92fn collect_source_files(
93    project_path: &std::path::Path,
94    language: &kdo_core::Language,
95) -> Vec<std::path::PathBuf> {
96    let extensions: &[&str] = match language {
97        kdo_core::Language::Rust | kdo_core::Language::Anchor => &["rs"],
98        kdo_core::Language::TypeScript => &["ts", "tsx"],
99        kdo_core::Language::JavaScript => &["js", "jsx"],
100        kdo_core::Language::Python => &["py"],
101        kdo_core::Language::Go => &["go"],
102    };
103
104    let walker = ignore::WalkBuilder::new(project_path)
105        .hidden(true)
106        .git_ignore(true)
107        .add_custom_ignore_filename(".kdoignore")
108        .build();
109
110    let mut result = Vec::new();
111    for entry in walker.flatten() {
112        if !entry.file_type().map(|ft| ft.is_file()).unwrap_or(false) {
113            continue;
114        }
115        let matches_ext = entry
116            .path()
117            .extension()
118            .and_then(|ext| ext.to_str())
119            .map(|ext| extensions.contains(&ext))
120            .unwrap_or(false);
121        if matches_ext {
122            result.push(entry.into_path());
123        }
124    }
125    result
126}