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/// Leave room for the shared `CodeMap` envelope and source-scope metadata.
32const PAYLOAD_FIXED_RESERVE: usize = 1_024;
33
34type Configs = HashMap<String, Option<TagsConfiguration>>;
35
36pub async fn handle_tool_call(
37    bash_state_arc: &Arc<Mutex<Option<BashState>>>,
38    args: FindReferences,
39) -> Result<(String, serde_json::Value)> {
40    handle_tool_call_with_budget(
41        bash_state_arc,
42        args,
43        crate::tools::outline::CODE_MAP_PAYLOAD_MAX_BYTES,
44    )
45    .await
46}
47
48pub(crate) async fn handle_tool_call_with_budget(
49    bash_state_arc: &Arc<Mutex<Option<BashState>>>,
50    args: FindReferences,
51    payload_max_bytes: usize,
52) -> Result<(String, serde_json::Value)> {
53    let (cwd, workspace_root) = {
54        let guard = bash_state_arc.lock().await;
55        let bash_state = guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;
56        (bash_state.cwd.clone(), bash_state.workspace_root.clone())
57    };
58
59    tokio::task::spawn_blocking(move || {
60        references_with_paths(&args, &cwd, workspace_root, payload_max_bytes)
61    })
62    .await
63    .map_err(|error| {
64        WinxError::CommandExecutionError(format!("CodeMap references worker failed: {error}"))
65    })?
66}
67
68fn references_with_paths(
69    args: &FindReferences,
70    cwd: &Path,
71    workspace_root: PathBuf,
72    payload_max_bytes: usize,
73) -> Result<(String, serde_json::Value)> {
74    let workspace_root = workspace_root.canonicalize().unwrap_or(workspace_root);
75
76    if args.name.trim().is_empty() {
77        return Err(WinxError::ArgumentParseError("Symbol name must not be empty.".to_string()));
78    }
79
80    let root = resolve_in_workspace(&args.path, cwd, &workspace_root).map_err(|e| {
81        WinxError::PathSecurityError { path: PathBuf::from(&args.path), message: e.to_string() }
82    })?;
83    let max = if args.max_results == 0 { DEFAULT_MAX_HITS } else { args.max_results };
84
85    // Don't silently degrade a typo'd file path into an empty whole-workspace
86    // scan (mirrors Outline).
87    let candidates: Vec<PathBuf> = if root.is_file() {
88        vec![root.clone()]
89    } else if root.is_dir() {
90        walk_workspace_files(&root)
91    } else {
92        return Err(WinxError::FileAccessError {
93            path: root.clone(),
94            message: "path not found (or not a regular file/directory)".to_string(),
95        });
96    };
97
98    let mut context = TagsContext::new();
99    let mut configs: Configs = HashMap::new();
100    // Two buckets: definitions are collected UNCAPPED (there are few, and losing
101    // one is the worst failure), only references are capped at `max`. The walk is
102    // stopped only by the file budget — never by the hit count — so a definition
103    // in a late-walked file is never missed (which the naive top-of-loop cap did).
104    let mut defs: Vec<ReferenceHit> = Vec::new();
105    let mut refs: Vec<ReferenceHit> = Vec::new();
106    let mut scanned = 0usize;
107    let mut truncated = false;
108
109    for abs in candidates {
110        if scanned >= MAX_FILES_SCANNED {
111            truncated = true;
112            break;
113        }
114        let ext = ext_of(&abs);
115        if !symbols::supports(&ext) {
116            continue;
117        }
118        let Ok(text) = read_file_to_string(&abs, MAX_FILE_SIZE) else { continue };
119        scanned += 1;
120        let rel = rel_of(&abs, &workspace_root);
121        let syms = match config_for(&mut configs, &ext) {
122            Some(config) => symbols::extract_all(&mut context, config, &text),
123            None => continue,
124        };
125        for s in syms {
126            if s.name != args.name {
127                continue;
128            }
129            let hit = ReferenceHit {
130                file: rel.clone(),
131                line: s.line,
132                kind: s.kind,
133                is_definition: s.is_definition,
134            };
135            if s.is_definition {
136                defs.push(hit);
137            } else if refs.len() < max {
138                refs.push(hit);
139            } else {
140                truncated = true; // more references than the cap allows
141            }
142        }
143    }
144
145    let by_loc =
146        |a: &ReferenceHit, b: &ReferenceHit| a.file.cmp(&b.file).then_with(|| a.line.cmp(&b.line));
147    defs.sort_by(by_loc);
148    refs.sort_by(by_loc);
149    let mut hits = defs;
150    hits.extend(refs); // definitions first, then references
151
152    bounded_output(&args.name, &hits, truncated, &root, payload_max_bytes)
153}
154
155fn bounded_output(
156    name: &str,
157    hits: &[ReferenceHit],
158    already_truncated: bool,
159    root: &Path,
160    payload_max_bytes: usize,
161) -> Result<(String, serde_json::Value)> {
162    let budget = payload_max_bytes.saturating_sub(PAYLOAD_FIXED_RESERVE);
163    let mut low = 0usize;
164    let mut high = hits.len();
165    while low < high {
166        let middle = low + (high - low).div_ceil(2);
167        let selected = &hits[..middle];
168        let candidate = reference_output(name, selected, already_truncated || middle < hits.len());
169        let text = render(
170            name,
171            selected,
172            candidate.definitions,
173            candidate.references,
174            candidate.truncated,
175            root,
176            !hits.is_empty(),
177        );
178        let bytes = text.len().saturating_add(
179            serde_json::to_vec(&candidate).map_or(usize::MAX, |serialized| serialized.len()),
180        );
181        if bytes <= budget {
182            low = middle;
183        } else {
184            high = middle - 1;
185        }
186    }
187
188    let selected = &hits[..low];
189    let mut output =
190        reference_output(name, selected, already_truncated || selected.len() < hits.len());
191    let text = render(
192        name,
193        selected,
194        output.definitions,
195        output.references,
196        output.truncated,
197        root,
198        !hits.is_empty(),
199    );
200    for _ in 0..3 {
201        let serialized = serde_json::to_vec(&output)
202            .map_err(|error| WinxError::SerializationError(error.to_string()))?;
203        let measured = text.len().saturating_add(serialized.len());
204        if measured == output.payload_bytes {
205            break;
206        }
207        output.payload_bytes = measured;
208    }
209    Ok((text, crate::tools::structured_json(&output)?))
210}
211
212fn reference_output(name: &str, hits: &[ReferenceHit], truncated: bool) -> ReferencesOutput {
213    ReferencesOutput {
214        name: name.to_string(),
215        definitions: hits.iter().filter(|hit| hit.is_definition).count(),
216        references: hits.iter().filter(|hit| !hit.is_definition).count(),
217        truncated,
218        payload_bytes: 0,
219        hits: hits.to_vec(),
220    }
221}
222
223fn render(
224    name: &str,
225    hits: &[ReferenceHit],
226    definitions: usize,
227    references: usize,
228    truncated: bool,
229    root: &Path,
230    had_hits: bool,
231) -> String {
232    if hits.is_empty() {
233        if had_hits {
234            return format!(
235                "Occurrences of `{name}` were found under {}, but the first result exceeds the \
236                 CodeMap response budget. Narrow `path` or use ReadFiles for exact source.",
237                root.display()
238            );
239        }
240        return format!("No occurrences of `{name}` found under {}.", root.display());
241    }
242    let mut out = format!("`{name}` — {definitions} definition(s), {references} reference(s):\n");
243    for h in hits {
244        let tag = if h.is_definition { "def" } else { "ref" };
245        let loc = format!("{}:{}", h.file, h.line);
246        let _ = writeln!(out, "  {tag}  {loc:<40} {:<9} {name}", h.kind);
247    }
248    if truncated {
249        let _ = write!(out, "(...capped; narrow `path` or raise max_results)");
250    }
251    out
252}
253
254/// Lowercase extension of `path` ("" if none).
255fn ext_of(path: &Path) -> String {
256    path.extension().and_then(|e| e.to_str()).unwrap_or_default().to_lowercase()
257}
258
259/// Workspace-relative display path.
260fn rel_of(path: &Path, workspace_root: &Path) -> String {
261    path.strip_prefix(workspace_root).unwrap_or(path).to_string_lossy().to_string()
262}
263
264/// Compile-once-per-language cached config lookup.
265fn config_for<'a>(configs: &'a mut Configs, ext: &str) -> Option<&'a TagsConfiguration> {
266    configs.entry(ext.to_string()).or_insert_with(|| symbols::config_for(ext)).as_ref()
267}
268
269#[cfg(test)]
270mod tests {
271    #![allow(clippy::unwrap_used)]
272    use super::*;
273    use tempfile::TempDir;
274
275    fn state_in(dir: &TempDir) -> Arc<Mutex<Option<BashState>>> {
276        let mut bs = BashState::new();
277        let root = dir.path().canonicalize().unwrap();
278        bs.cwd = root.clone();
279        bs.workspace_root = root;
280        Arc::new(Mutex::new(Some(bs)))
281    }
282
283    fn args(name: &str) -> FindReferences {
284        FindReferences {
285            name: name.to_string(),
286            path: String::new(),
287            max_results: 0,
288            thread_id: String::new(),
289        }
290    }
291
292    #[tokio::test]
293    async fn finds_definition_and_references() {
294        let dir = TempDir::new().unwrap();
295        std::fs::write(
296            dir.path().join("a.rs"),
297            "fn target() {}\nfn caller() { target(); target(); }\n",
298        )
299        .unwrap();
300        let st = state_in(&dir);
301        let (out, structured) = handle_tool_call(&st, args("target")).await.unwrap();
302        assert!(out.contains("target"));
303        assert_eq!(structured["definitions"], 1);
304        assert!(structured["references"].as_u64().unwrap() >= 1);
305        // definition is listed first
306        assert_eq!(structured["hits"][0]["is_definition"], true);
307    }
308
309    #[tokio::test]
310    async fn finds_python_and_elixir_definitions_and_calls() {
311        let dir = TempDir::new().unwrap();
312        std::fs::write(
313            dir.path().join("service.py"),
314            "def python_target():\n    return 1\n\npython_target()\n",
315        )
316        .unwrap();
317        std::fs::write(
318            dir.path().join("worker.ex"),
319            "defmodule Worker do\n  def elixir_target(), do: :ok\n  def run(), do: elixir_target()\nend\n",
320        )
321        .unwrap();
322        let st = state_in(&dir);
323
324        let (_, python) = handle_tool_call(&st, args("python_target")).await.unwrap();
325        assert_eq!(python["definitions"], 1);
326        assert!(python["references"].as_u64().unwrap() >= 1, "{python}");
327
328        let (_, elixir) = handle_tool_call(&st, args("elixir_target")).await.unwrap();
329        assert_eq!(elixir["definitions"], 1);
330        assert!(elixir["references"].as_u64().unwrap() >= 1, "{elixir}");
331    }
332
333    #[tokio::test]
334    async fn empty_name_errors() {
335        let dir = TempDir::new().unwrap();
336        let st = state_in(&dir);
337        assert!(handle_tool_call(&st, args("  ")).await.is_err());
338    }
339
340    #[tokio::test]
341    async fn reports_no_occurrences() {
342        let dir = TempDir::new().unwrap();
343        std::fs::write(dir.path().join("a.rs"), "fn something() {}\n").unwrap();
344        let st = state_in(&dir);
345        let (out, _) = handle_tool_call(&st, args("nonexistent_symbol")).await.unwrap();
346        assert!(out.to_lowercase().contains("no occurrences"));
347    }
348
349    #[tokio::test]
350    async fn definition_survives_when_references_exceed_cap() {
351        // Regression (B1): definitions are uncapped, so the def is found and kept
352        // even when far more reference hits than `max` are collected, regardless
353        // of filesystem walk order.
354        let dir = TempDir::new().unwrap();
355        for i in 0..20 {
356            std::fs::write(
357                dir.path().join(format!("ref_{i:02}.rs")),
358                "fn caller() { target(); }\n",
359            )
360            .unwrap();
361        }
362        std::fs::write(dir.path().join("zzz_def.rs"), "fn target() {}\n").unwrap();
363        let st = state_in(&dir);
364        let mut a = args("target");
365        a.max_results = 3; // far fewer than the 20 reference hits
366        let (_, structured) = handle_tool_call(&st, a).await.unwrap();
367        assert_eq!(
368            structured["definitions"], 1,
369            "definition must never be dropped by the reference cap"
370        );
371        assert!(structured["references"].as_u64().unwrap() <= 3);
372        assert_eq!(structured["truncated"], true);
373        assert_eq!(structured["hits"][0]["is_definition"], true); // definitions listed first
374    }
375
376    #[tokio::test]
377    async fn bare_name_matches_same_name_across_types() {
378        // Documented limitation: a name-based lookup can't distinguish A::run from
379        // B::run — both count. Useful as a superset for rename/impact analysis.
380        let dir = TempDir::new().unwrap();
381        std::fs::write(
382            dir.path().join("a.rs"),
383            "struct A;\nstruct B;\nimpl A { fn run(&self) {} }\nimpl B { fn run(&self) {} }\n",
384        )
385        .unwrap();
386        let st = state_in(&dir);
387        let (_, structured) = handle_tool_call(&st, args("run")).await.unwrap();
388        assert_eq!(structured["definitions"], 2);
389    }
390
391    #[tokio::test]
392    async fn dense_references_respect_the_combined_payload_budget() {
393        let dir = TempDir::new().unwrap();
394        let mut source = "fn target() {}\nfn caller() {\n".to_string();
395        for _ in 0..800 {
396            source.push_str("    target();\n");
397        }
398        source.push_str("}\n");
399        std::fs::write(dir.path().join("dense.rs"), source).unwrap();
400        let st = state_in(&dir);
401        let mut request = args("target");
402        request.max_results = 1_000;
403
404        let (text, structured) =
405            handle_tool_call_with_budget(&st, request, 8 * 1024).await.unwrap();
406        let bytes = text.len() + serde_json::to_vec(&structured).unwrap().len();
407        assert!(bytes <= 8 * 1024, "payload used {bytes} bytes");
408        assert_eq!(structured["truncated"], true);
409        assert_eq!(structured["hits"][0]["is_definition"], true);
410    }
411
412    #[tokio::test]
413    async fn nonexistent_path_errors() {
414        // B3: a typo'd path must error, not silently scan the whole workspace.
415        let dir = TempDir::new().unwrap();
416        std::fs::write(dir.path().join("real.rs"), "fn x() {}\n").unwrap();
417        let st = state_in(&dir);
418        let mut a = args("x");
419        a.path = "nope_typo.rs".to_string();
420        assert!(handle_tool_call(&st, a).await.is_err());
421    }
422}
winx_code_agent/tools/references.rs

winx_code_agent/tools/
references.rs
