winx_code_agent/tools/
references.rs

1//! Implementation of the `FindReferences` tool — tree-sitter symbol occurrences.
2//!
3//! Finds where an identifier is defined and referenced (called / used) across
4//! the workspace, by name. It is a *name-based* semantic lookup (it counts only
5//! real identifier occurrences, never matches inside strings or comments like a
6//! raw grep would), built on the same `tree-sitter-tags` engine as `Outline`.
7//! Read-only, gitignore-aware, workspace-confined; works in every mode.
8
9use std::collections::HashMap;
10use std::fmt::Write as FmtWrite;
11use std::path::{Path, PathBuf};
12use std::sync::Arc;
13
14use tokio::sync::Mutex;
15use tree_sitter_tags::{TagsConfiguration, TagsContext};
16
17use crate::errors::{Result, WinxError};
18use crate::state::bash_state::BashState;
19use crate::types::{FindReferences, ReferenceHit, ReferencesOutput};
20use crate::utils::mmap::read_file_to_string;
21use crate::utils::path::resolve_in_workspace;
22use crate::utils::repo::walk_workspace_files;
23use crate::utils::symbols;
24
25/// Cap on occurrences returned when the caller passes 0.
26const DEFAULT_MAX_HITS: usize = 200;
27/// Skip files larger than this when scanning.
28const MAX_FILE_SIZE: u64 = 2_000_000;
29/// Stop after reading+parsing this many files (a scan-budget cap).
30const MAX_FILES_SCANNED: usize = 20_000;
31
32type Configs = HashMap<String, Option<TagsConfiguration>>;
33
34pub async fn handle_tool_call(
35    bash_state_arc: &Arc<Mutex<Option<BashState>>>,
36    args: FindReferences,
37) -> Result<(String, serde_json::Value)> {
38    let (cwd, workspace_root) = {
39        let guard = bash_state_arc.lock().await;
40        let bash_state = guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;
41        (bash_state.cwd.clone(), bash_state.workspace_root.clone())
42    };
43    let workspace_root = workspace_root.canonicalize().unwrap_or(workspace_root);
44
45    if args.name.trim().is_empty() {
46        return Err(WinxError::ArgumentParseError("Symbol name must not be empty.".to_string()));
47    }
48
49    let root = resolve_in_workspace(&args.path, &cwd, &workspace_root).map_err(|e| {
50        WinxError::PathSecurityError { path: PathBuf::from(&args.path), message: e.to_string() }
51    })?;
52    let max = if args.max_results == 0 { DEFAULT_MAX_HITS } else { args.max_results };
53
54    // Don't silently degrade a typo'd file path into an empty whole-workspace
55    // scan (mirrors Outline).
56    let candidates: Vec<PathBuf> = if root.is_file() {
57        vec![root.clone()]
58    } else if root.is_dir() {
59        walk_workspace_files(&root)
60    } else {
61        return Err(WinxError::FileAccessError {
62            path: root.clone(),
63            message: "path not found (or not a regular file/directory)".to_string(),
64        });
65    };
66
67    let mut context = TagsContext::new();
68    let mut configs: Configs = HashMap::new();
69    // Two buckets: definitions are collected UNCAPPED (there are few, and losing
70    // one is the worst failure), only references are capped at `max`. The walk is
71    // stopped only by the file budget — never by the hit count — so a definition
72    // in a late-walked file is never missed (which the naive top-of-loop cap did).
73    let mut defs: Vec<ReferenceHit> = Vec::new();
74    let mut refs: Vec<ReferenceHit> = Vec::new();
75    let mut scanned = 0usize;
76    let mut truncated = false;
77
78    for abs in candidates {
79        if scanned >= MAX_FILES_SCANNED {
80            truncated = true;
81            break;
82        }
83        let ext = ext_of(&abs);
84        if !symbols::supports(&ext) {
85            continue;
86        }
87        let Ok(text) = read_file_to_string(&abs, MAX_FILE_SIZE) else { continue };
88        scanned += 1;
89        let rel = rel_of(&abs, &workspace_root);
90        let syms = match config_for(&mut configs, &ext) {
91            Some(config) => symbols::extract_all(&mut context, config, &text),
92            None => continue,
93        };
94        for s in syms {
95            if s.name != args.name {
96                continue;
97            }
98            let hit = ReferenceHit {
99                file: rel.clone(),
100                line: s.line,
101                kind: s.kind,
102                is_definition: s.is_definition,
103            };
104            if s.is_definition {
105                defs.push(hit);
106            } else if refs.len() < max {
107                refs.push(hit);
108            } else {
109                truncated = true; // more references than the cap allows
110            }
111        }
112    }
113
114    let by_loc =
115        |a: &ReferenceHit, b: &ReferenceHit| a.file.cmp(&b.file).then_with(|| a.line.cmp(&b.line));
116    defs.sort_by(by_loc);
117    refs.sort_by(by_loc);
118    let definitions = defs.len();
119    let references = refs.len();
120    let mut hits = defs;
121    hits.extend(refs); // definitions first, then references
122
123    let out = render(&args.name, &hits, definitions, references, truncated, &root);
124    let structured = ReferencesOutput { name: args.name, definitions, references, truncated, hits };
125    Ok((out, crate::tools::structured_json(&structured)?))
126}
127
128fn render(
129    name: &str,
130    hits: &[ReferenceHit],
131    definitions: usize,
132    references: usize,
133    truncated: bool,
134    root: &Path,
135) -> String {
136    if hits.is_empty() {
137        return format!("No occurrences of `{name}` found under {}.", root.display());
138    }
139    let mut out = format!("`{name}` — {definitions} definition(s), {references} reference(s):\n");
140    for h in hits {
141        let tag = if h.is_definition { "def" } else { "ref" };
142        let loc = format!("{}:{}", h.file, h.line);
143        let _ = writeln!(out, "  {tag}  {loc:<40} {:<9} {name}", h.kind);
144    }
145    if truncated {
146        let _ = write!(out, "(...capped; narrow `path` or raise max_results)");
147    }
148    out
149}
150
151/// Lowercase extension of `path` ("" if none).
152fn ext_of(path: &Path) -> String {
153    path.extension().and_then(|e| e.to_str()).unwrap_or_default().to_lowercase()
154}
155
156/// Workspace-relative display path.
157fn rel_of(path: &Path, workspace_root: &Path) -> String {
158    path.strip_prefix(workspace_root).unwrap_or(path).to_string_lossy().to_string()
159}
160
161/// Compile-once-per-language cached config lookup.
162fn config_for<'a>(configs: &'a mut Configs, ext: &str) -> Option<&'a TagsConfiguration> {
163    configs.entry(ext.to_string()).or_insert_with(|| symbols::config_for(ext)).as_ref()
164}
165
166#[cfg(test)]
167mod tests {
168    #![allow(clippy::unwrap_used)]
169    use super::*;
170    use tempfile::TempDir;
171
172    fn state_in(dir: &TempDir) -> Arc<Mutex<Option<BashState>>> {
173        let mut bs = BashState::new();
174        let root = dir.path().canonicalize().unwrap();
175        bs.cwd = root.clone();
176        bs.workspace_root = root;
177        Arc::new(Mutex::new(Some(bs)))
178    }
179
180    fn args(name: &str) -> FindReferences {
181        FindReferences {
182            name: name.to_string(),
183            path: String::new(),
184            max_results: 0,
185            thread_id: String::new(),
186        }
187    }
188
189    #[tokio::test]
190    async fn finds_definition_and_references() {
191        let dir = TempDir::new().unwrap();
192        std::fs::write(
193            dir.path().join("a.rs"),
194            "fn target() {}\nfn caller() { target(); target(); }\n",
195        )
196        .unwrap();
197        let st = state_in(&dir);
198        let (out, structured) = handle_tool_call(&st, args("target")).await.unwrap();
199        assert!(out.contains("target"));
200        assert_eq!(structured["definitions"], 1);
201        assert!(structured["references"].as_u64().unwrap() >= 1);
202        // definition is listed first
203        assert_eq!(structured["hits"][0]["is_definition"], true);
204    }
205
206    #[tokio::test]
207    async fn empty_name_errors() {
208        let dir = TempDir::new().unwrap();
209        let st = state_in(&dir);
210        assert!(handle_tool_call(&st, args("  ")).await.is_err());
211    }
212
213    #[tokio::test]
214    async fn reports_no_occurrences() {
215        let dir = TempDir::new().unwrap();
216        std::fs::write(dir.path().join("a.rs"), "fn something() {}\n").unwrap();
217        let st = state_in(&dir);
218        let (out, _) = handle_tool_call(&st, args("nonexistent_symbol")).await.unwrap();
219        assert!(out.to_lowercase().contains("no occurrences"));
220    }
221
222    #[tokio::test]
223    async fn definition_survives_when_references_exceed_cap() {
224        // Regression (B1): definitions are uncapped, so the def is found and kept
225        // even when far more reference hits than `max` are collected, regardless
226        // of filesystem walk order.
227        let dir = TempDir::new().unwrap();
228        for i in 0..20 {
229            std::fs::write(
230                dir.path().join(format!("ref_{i:02}.rs")),
231                "fn caller() { target(); }\n",
232            )
233            .unwrap();
234        }
235        std::fs::write(dir.path().join("zzz_def.rs"), "fn target() {}\n").unwrap();
236        let st = state_in(&dir);
237        let mut a = args("target");
238        a.max_results = 3; // far fewer than the 20 reference hits
239        let (_, structured) = handle_tool_call(&st, a).await.unwrap();
240        assert_eq!(
241            structured["definitions"], 1,
242            "definition must never be dropped by the reference cap"
243        );
244        assert!(structured["references"].as_u64().unwrap() <= 3);
245        assert_eq!(structured["truncated"], true);
246        assert_eq!(structured["hits"][0]["is_definition"], true); // definitions listed first
247    }
248
249    #[tokio::test]
250    async fn bare_name_matches_same_name_across_types() {
251        // Documented limitation: a name-based lookup can't distinguish A::run from
252        // B::run — both count. Useful as a superset for rename/impact analysis.
253        let dir = TempDir::new().unwrap();
254        std::fs::write(
255            dir.path().join("a.rs"),
256            "struct A;\nstruct B;\nimpl A { fn run(&self) {} }\nimpl B { fn run(&self) {} }\n",
257        )
258        .unwrap();
259        let st = state_in(&dir);
260        let (_, structured) = handle_tool_call(&st, args("run")).await.unwrap();
261        assert_eq!(structured["definitions"], 2);
262    }
263
264    #[tokio::test]
265    async fn nonexistent_path_errors() {
266        // B3: a typo'd path must error, not silently scan the whole workspace.
267        let dir = TempDir::new().unwrap();
268        std::fs::write(dir.path().join("real.rs"), "fn x() {}\n").unwrap();
269        let st = state_in(&dir);
270        let mut a = args("x");
271        a.path = "nope_typo.rs".to_string();
272        assert!(handle_tool_call(&st, a).await.is_err());
273    }
274}
winx_code_agent/tools/references.rs

winx_code_agent/tools/
references.rs
