telos_agent/knowledge/code_index/
mod.rs

1//! Lightweight local code index for fast path/line search and context lookup.
2
3use serde::{Deserialize, Serialize};
4use std::path::{Path, PathBuf};
5
6const INDEX_REL_PATH: &[&str] = &[".telos", "index", "code_index.json"];
7const MAX_FILE_BYTES: u64 = 512 * 1024;
8
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
10pub struct CodeIndex {
11    pub root: PathBuf,
12    pub files: Vec<IndexedFile>,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
16pub struct IndexedFile {
17    pub path: String,
18    pub language: String,
19    pub modified_secs: u64,
20    pub size: u64,
21    pub lines: Vec<String>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
25pub struct CodeSearchMatch {
26    pub path: String,
27    pub line: usize,
28    pub text: String,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32pub struct CodeContextLine {
33    pub line: usize,
34    pub text: String,
35}
36
37impl CodeIndex {
38    pub fn load_or_refresh(root: impl Into<PathBuf>) -> std::io::Result<Self> {
39        let root = root.into();
40        match Self::load(&root) {
41            Ok(index) => Ok(index),
42            Err(_) => Self::refresh(root),
43        }
44    }
45
46    pub fn load(root: impl AsRef<Path>) -> std::io::Result<Self> {
47        let root = root.as_ref();
48        let content = std::fs::read_to_string(index_path(root))?;
49        let mut index: Self = serde_json::from_str(&content)
50            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
51        index.root = root.to_path_buf();
52        Ok(index)
53    }
54
55    pub fn refresh(root: impl Into<PathBuf>) -> std::io::Result<Self> {
56        let root = root.into();
57        let mut index = Self { root: root.clone(), files: Vec::new() };
58        visit_dir(&root, &root, &mut index.files)?;
59        index.files.sort_by(|a, b| a.path.cmp(&b.path));
60        let path = index_path(&root);
61        if let Some(parent) = path.parent() {
62            std::fs::create_dir_all(parent)?;
63        }
64        let content = serde_json::to_string_pretty(&index)?;
65        std::fs::write(path, content)?;
66        Ok(index)
67    }
68
69    pub fn search(
70        &self,
71        query: &str,
72        path_prefix: Option<&str>,
73        max_results: usize,
74        case_sensitive: bool,
75    ) -> Vec<CodeSearchMatch> {
76        let needle = if case_sensitive { query.to_string() } else { query.to_lowercase() };
77        let mut matches = Vec::new();
78        for file in &self.files {
79            if let Some(prefix) = path_prefix
80                && !file.path.contains(prefix)
81            {
82                continue;
83            }
84            for (idx, line) in file.lines.iter().enumerate() {
85                let haystack = if case_sensitive { line.to_string() } else { line.to_lowercase() };
86                if haystack.contains(&needle) {
87                    matches.push(CodeSearchMatch {
88                        path: file.path.clone(),
89                        line: idx + 1,
90                        text: line.clone(),
91                    });
92                    if matches.len() >= max_results {
93                        return matches;
94                    }
95                }
96            }
97        }
98        matches
99    }
100
101    pub fn context(
102        &self,
103        path: &str,
104        line: usize,
105        before: usize,
106        after: usize,
107    ) -> Option<Vec<CodeContextLine>> {
108        let file = self.files.iter().find(|file| file.path == path)?;
109        if file.lines.is_empty() || line == 0 {
110            return Some(Vec::new());
111        }
112        let start = line.saturating_sub(before).max(1);
113        let end = (line + after).min(file.lines.len());
114        Some(
115            (start..=end)
116                .map(|line_no| CodeContextLine {
117                    line: line_no,
118                    text: file.lines[line_no - 1].clone(),
119                })
120                .collect(),
121        )
122    }
123
124    pub fn index_path(root: impl AsRef<Path>) -> PathBuf {
125        index_path(root.as_ref())
126    }
127}
128
129fn index_path(root: &Path) -> PathBuf {
130    INDEX_REL_PATH.iter().fold(root.to_path_buf(), |path, component| path.join(component))
131}
132
133fn visit_dir(root: &Path, dir: &Path, files: &mut Vec<IndexedFile>) -> std::io::Result<()> {
134    for entry in std::fs::read_dir(dir)? {
135        let entry = entry?;
136        let path = entry.path();
137        let name = entry.file_name();
138        let name = name.to_string_lossy();
139        if should_skip_name(&name) {
140            continue;
141        }
142        let metadata = entry.metadata()?;
143        if metadata.is_dir() {
144            visit_dir(root, &path, files)?;
145        } else if metadata.is_file()
146            && metadata.len() <= MAX_FILE_BYTES
147            && let Some(file) = index_file(root, &path, metadata.len())?
148        {
149            files.push(file);
150        }
151    }
152    Ok(())
153}
154
155fn index_file(root: &Path, path: &Path, size: u64) -> std::io::Result<Option<IndexedFile>> {
156    let bytes = std::fs::read(path)?;
157    if bytes.contains(&0) {
158        return Ok(None);
159    }
160    let Ok(content) = String::from_utf8(bytes) else {
161        return Ok(None);
162    };
163    let relative = path.strip_prefix(root).unwrap_or(path).to_string_lossy().replace('\\', "/");
164    let modified_secs = std::fs::metadata(path)?
165        .modified()
166        .ok()
167        .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
168        .map(|duration| duration.as_secs())
169        .unwrap_or(0);
170    Ok(Some(IndexedFile {
171        language: language_for(path),
172        path: relative,
173        modified_secs,
174        size,
175        lines: content.lines().map(str::to_string).collect(),
176    }))
177}
178
179fn should_skip_name(name: &str) -> bool {
180    matches!(name, ".git" | ".telos" | "target" | "node_modules" | ".next" | "dist" | "build")
181}
182
183fn language_for(path: &Path) -> String {
184    match path.extension().and_then(|ext| ext.to_str()).unwrap_or_default() {
185        "rs" => "rust",
186        "toml" => "toml",
187        "md" => "markdown",
188        "json" => "json",
189        "yaml" | "yml" => "yaml",
190        "js" | "jsx" => "javascript",
191        "ts" | "tsx" => "typescript",
192        "py" => "python",
193        "go" => "go",
194        "java" => "java",
195        "c" | "h" | "cc" | "cpp" | "hpp" => "cpp",
196        "" => "text",
197        other => other,
198    }
199    .to_string()
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn refresh_indexes_text_files_and_skips_telos_dir() {
208        let dir = tempfile::tempdir().unwrap();
209        std::fs::create_dir_all(dir.path().join("src")).unwrap();
210        std::fs::create_dir_all(dir.path().join(".telos")).unwrap();
211        std::fs::write(dir.path().join("src/lib.rs"), "fn target() {}\nfn other() {}\n").unwrap();
212        std::fs::write(dir.path().join(".telos/ignored.rs"), "fn target() {}\n").unwrap();
213
214        let index = CodeIndex::refresh(dir.path().to_path_buf()).unwrap();
215
216        assert_eq!(index.files.len(), 1);
217        assert_eq!(index.files[0].path, "src/lib.rs");
218        assert!(CodeIndex::index_path(dir.path()).exists());
219    }
220
221    #[test]
222    fn search_returns_path_line_and_context() {
223        let dir = tempfile::tempdir().unwrap();
224        std::fs::create_dir_all(dir.path().join("src")).unwrap();
225        std::fs::write(dir.path().join("src/lib.rs"), "first\nfn target() {}\nlast\n").unwrap();
226        let index = CodeIndex::refresh(dir.path().to_path_buf()).unwrap();
227
228        let matches = index.search("TARGET", None, 10, false);
229        assert_eq!(matches.len(), 1);
230        assert_eq!(matches[0].path, "src/lib.rs");
231        assert_eq!(matches[0].line, 2);
232
233        let context = index.context("src/lib.rs", 2, 1, 1).unwrap();
234        assert_eq!(context.len(), 3);
235        assert_eq!(context[1].text, "fn target() {}");
236    }
237}
telos_agent/knowledge/code_index/mod.rs

telos_agent/knowledge/code_index/
mod.rs
