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

winx_code_agent/tools/
references.rs
