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

winx_code_agent/tools/
references.rs
