1#![forbid(unsafe_code)]
20
21use async_trait::async_trait;
22
23use serde_json::{Value, json};
24use std::collections::{HashMap, VecDeque};
25use std::sync::Arc;
26use std::sync::Mutex;
27use wm_core::{Context, EffectRow, Gana, Resource, Tool, ToolStats};
28
29const DEFAULT_MAX_FILES: usize = 50_000;
31const MAX_FILE_BYTES: u64 = 1_000_000;
33const SKIP_DIRS: &[&str] = &[
35 "target",
36 "node_modules",
37 ".git",
38 "build",
39 "dist",
40 "vendor",
41 ".venv",
42 "__pycache__",
43 "coverage",
44];
45
46#[derive(Debug, Clone)]
48pub struct CodeNode {
49 pub id: String,
51 pub name: String,
53 pub node_type: String,
55 pub file: String,
57 pub line: usize,
59 pub language: String,
61}
62
63#[derive(Debug, Clone)]
65pub struct CodeEdge {
66 pub source_id: String,
67 pub target_id: String,
68 pub edge_type: String,
70}
71
72#[derive(Debug, Default)]
74pub struct CodeGraph {
75 nodes: Vec<CodeNode>,
76 edges: Vec<CodeEdge>,
77 project_root: String,
78 built: bool,
79}
80
81impl CodeGraph {
82 #[must_use]
84 pub fn new() -> Self {
85 Self::default()
86 }
87
88 #[must_use]
90 pub const fn is_built(&self) -> bool {
91 self.built
92 }
93
94 #[must_use]
96 pub fn root(&self) -> &str {
97 &self.project_root
98 }
99
100 pub fn build(&mut self, project_root: &str, max_files: usize) -> Result<usize, String> {
102 let root = std::path::Path::new(project_root);
103 if !root.is_dir() {
104 return Err(format!("project_root is not a directory: {project_root}"));
105 }
106 let mut files: Vec<std::path::PathBuf> = Vec::new();
107 let mut walked = 0usize;
108 collect_files(root, root, max_files, &mut files, &mut walked)?;
109
110 let mut nodes: Vec<CodeNode> = Vec::new();
111 let mut edges: Vec<CodeEdge> = Vec::new();
112 struct FileInfo {
116 rel: String,
117 language: String,
118 contents: String,
119 }
120 let mut files_info: Vec<FileInfo> = Vec::new();
121
122 for file in &files {
123 let rel = file
124 .strip_prefix(root)
125 .unwrap_or(file)
126 .to_string_lossy()
127 .replace('\\', "/");
128 let language = language_for(&rel);
129 if language.is_empty() {
130 continue;
131 }
132 let contents = read_bounded(file);
133 let Some(contents) = contents else { continue };
134
135 let (patterns, inline) = symbol_patterns(&language);
137 for (line_no, line) in contents.lines().enumerate() {
138 let line_no = line_no + 1;
139 for (pat, ty) in &patterns {
140 for cap in regex_captures(pat, line) {
141 let name = cap.clone();
142 let id = format!("{rel}:{name}:{ty}");
143 nodes.push(CodeNode {
144 id: id.clone(),
145 name: name.clone(),
146 node_type: (*ty).to_string(),
147 file: rel.clone(),
148 line: line_no,
149 language: language.clone(),
150 });
151 if *ty == "class" {
153 for parent in inherit_targets(line) {
154 if !parent.is_empty() {
155 edges.push(CodeEdge {
156 source_id: id.clone(),
157 target_id: format!("{rel}:{parent}:class"),
158 edge_type: "inherits".into(),
159 });
160 }
161 }
162 }
163 }
164 }
165 for (pat, ty) in &inline {
166 for cap in regex_captures(pat, line) {
167 let name = cap.clone();
168 nodes.push(CodeNode {
169 id: format!("{rel}:{name}:{ty}"),
170 name,
171 node_type: (*ty).to_string(),
172 file: rel.clone(),
173 line: line_no,
174 language: language.clone(),
175 });
176 }
177 }
178 }
179 files_info.push(FileInfo {
180 rel: rel.clone(),
181 language,
182 contents,
183 });
184 }
185
186 let name_index: HashMap<&str, Vec<usize>> = {
190 let mut index: HashMap<&str, Vec<usize>> = HashMap::new();
191 for (i, node) in nodes.iter().enumerate() {
192 index.entry(node.name.as_str()).or_default().push(i);
193 }
194 index
195 };
196 for info in &files_info {
197 for (line_no, line) in info.contents.lines().enumerate() {
198 let Some(caller) = enclosing_function(&info.contents, line_no) else {
199 continue;
200 };
201 let source_id = format!("{}:{caller}:function", info.rel);
202 for name in call_sites(line) {
203 let Some(mut candidates) = name_index.get(name.as_str()).cloned() else {
204 continue;
205 };
206 candidates.sort_by_key(|&i| usize::from(nodes[i].file != info.rel));
208 for i in candidates.into_iter().take(3) {
209 edges.push(CodeEdge {
210 source_id: source_id.clone(),
211 target_id: nodes[i].id.clone(),
212 edge_type: "calls".into(),
213 });
214 }
215 }
216 }
217 for line in info.contents.lines() {
219 for target in import_targets(&info.language, line) {
220 edges.push(CodeEdge {
221 source_id: format!("{}:*:import", info.rel),
222 target_id: target,
223 edge_type: "imports".into(),
224 });
225 }
226 }
227 }
228
229 let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
231 nodes.retain(|n| seen.insert(n.id.clone()));
232 let mut edge_seen: std::collections::HashSet<String> = std::collections::HashSet::new();
233 edges.retain(|e| {
234 edge_seen.insert(format!("{}|{}|{}", e.source_id, e.target_id, e.edge_type))
235 });
236
237 self.nodes = nodes;
238 self.edges = edges;
239 self.project_root = project_root.to_string();
240 self.built = true;
241 Ok(files.len())
242 }
243
244 #[must_use]
246 pub fn search(&self, query: &str, limit: usize) -> Vec<Value> {
247 let q = query.to_ascii_lowercase();
248 self.nodes
249 .iter()
250 .filter(|n| n.name.to_ascii_lowercase().contains(&q))
251 .take(limit)
252 .map(node_json)
253 .collect()
254 }
255
256 #[must_use]
258 pub fn callers(&self, symbol: &str, limit: usize) -> Vec<Value> {
259 let targets: Vec<&CodeNode> = self.nodes.iter().filter(|n| n.name == symbol).collect();
260 let mut out = Vec::new();
261 for target in targets {
262 for edge in &self.edges {
263 if edge.edge_type == "calls" && edge.target_id == target.id {
264 if let Some(src) = self.node(&edge.source_id) {
265 out.push(node_json(src));
266 if out.len() >= limit {
267 return out;
268 }
269 }
270 }
271 }
272 }
273 out
274 }
275
276 #[must_use]
278 pub fn callees(&self, symbol: &str, limit: usize) -> Vec<Value> {
279 let sources: Vec<&CodeNode> = self.nodes.iter().filter(|n| n.name == symbol).collect();
280 let mut out = Vec::new();
281 for source in sources {
282 for edge in &self.edges {
283 if edge.edge_type == "calls" && edge.source_id == source.id {
284 if let Some(tgt) = self.node(&edge.target_id) {
285 out.push(node_json(tgt));
286 if out.len() >= limit {
287 return out;
288 }
289 }
290 }
291 }
292 }
293 out
294 }
295
296 #[must_use]
298 pub fn explain(&self, symbol: &str) -> Option<Value> {
299 let node = self.nodes.iter().find(|n| n.name == symbol).or_else(|| {
300 self.nodes
301 .iter()
302 .find(|n| n.name == symbol && n.node_type == "class")
303 })?;
304 let in_deg = self.edges.iter().filter(|e| e.target_id == node.id).count();
305 let out_deg = self.edges.iter().filter(|e| e.source_id == node.id).count();
306 let incoming: Vec<Value> = self
307 .edges
308 .iter()
309 .filter(|e| e.target_id == node.id && e.edge_type == "calls")
310 .take(20)
311 .filter_map(|e| self.node(&e.source_id))
312 .map(node_json)
313 .collect();
314 let outgoing: Vec<Value> = self
315 .edges
316 .iter()
317 .filter(|e| e.source_id == node.id && e.edge_type == "calls")
318 .take(20)
319 .filter_map(|e| self.node(&e.target_id))
320 .map(node_json)
321 .collect();
322 Some(json!({
323 "symbol": node.name,
324 "node_type": node.node_type,
325 "file": node.file,
326 "line": node.line,
327 "language": node.language,
328 "degree": in_deg + out_deg,
329 "in_degree": in_deg,
330 "out_degree": out_deg,
331 "incoming": incoming,
332 "outgoing": outgoing,
333 }))
334 }
335
336 #[must_use]
338 pub fn path(&self, symbol_a: &str, symbol_b: &str, max_hops: usize) -> Value {
339 let node_a = self.nodes.iter().find(|n| n.name == symbol_a);
340 let node_b = self.nodes.iter().find(|n| n.name == symbol_b);
341 let (Some(node_a), Some(node_b)) = (node_a, node_b) else {
342 let missing = if node_a.is_none() { symbol_a } else { symbol_b };
343 return json!({
344 "status": "error",
345 "error": format!("symbol not found: {missing}"),
346 });
347 };
348 if node_a.id == node_b.id {
349 return json!({
350 "status": "success",
351 "path": [symbol_a],
352 "hops": 0,
353 });
354 }
355 let mut adj: HashMap<String, Vec<String>> = HashMap::new();
356 for e in &self.edges {
357 if e.edge_type == "calls" {
358 adj.entry(e.source_id.clone())
359 .or_default()
360 .push(e.target_id.clone());
361 }
362 }
363 let mut visited = std::collections::HashSet::new();
364 let mut queue: VecDeque<(String, Vec<String>)> = VecDeque::new();
365 visited.insert(node_a.id.clone());
366 queue.push_back((node_a.id.clone(), vec![node_a.id.clone()]));
367 while let Some((current, path)) = queue.pop_front() {
368 if current == node_b.id {
369 let names: Vec<String> = path
370 .iter()
371 .map(|id| self.node(id).map_or_else(|| id.clone(), |n| n.name.clone()))
372 .collect();
373 return json!({
374 "status": "success",
375 "path": names,
376 "hops": path.len() - 1,
377 });
378 }
379 if path.len() > max_hops {
380 continue;
381 }
382 if let Some(neighbors) = adj.get(¤t) {
383 for neighbor in neighbors {
384 if !visited.contains(neighbor) {
385 visited.insert(neighbor.clone());
386 let mut next = path.clone();
387 next.push(neighbor.clone());
388 queue.push_back((neighbor.clone(), next));
389 }
390 }
391 }
392 }
393 json!({
394 "status": "no_path",
395 "message": format!("no call path between {symbol_a} and {symbol_b} in {max_hops} hops"),
396 })
397 }
398
399 #[must_use]
402 pub fn affected_by(&self, symbol: &str, max_depth: usize) -> Value {
403 let roots: Vec<&CodeNode> = self.nodes.iter().filter(|n| n.name == symbol).collect();
404 if roots.is_empty() {
405 return json!({
406 "status": "error",
407 "error": format!("symbol not found: {symbol}"),
408 });
409 }
410 let mut rev: HashMap<String, Vec<String>> = HashMap::new();
412 for e in &self.edges {
413 if e.edge_type == "calls" {
414 rev.entry(e.target_id.clone())
415 .or_default()
416 .push(e.source_id.clone());
417 }
418 }
419 let mut affected: Vec<Value> = Vec::new();
420 let mut visited = std::collections::HashSet::new();
421 let mut queue: VecDeque<(String, usize)> = VecDeque::new();
422 for root in roots {
423 queue.push_back((root.id.clone(), 0));
424 visited.insert(root.id.clone());
425 }
426 while let Some((current, depth)) = queue.pop_front() {
427 if depth > 0 {
428 if let Some(node) = self.node(¤t) {
429 affected.push(json!({
430 "symbol": node.name,
431 "node_type": node.node_type,
432 "file": node.file,
433 "line": node.line,
434 "depth": depth,
435 }));
436 }
437 }
438 if depth >= max_depth {
439 continue;
440 }
441 if let Some(callers) = rev.get(¤t) {
442 for caller in callers {
443 if !visited.contains(caller) {
444 visited.insert(caller.clone());
445 queue.push_back((caller.clone(), depth + 1));
446 }
447 }
448 }
449 }
450 json!({
451 "status": "success",
452 "symbol": symbol,
453 "affected_count": affected.len(),
454 "max_depth": max_depth,
455 "affected": affected,
456 })
457 }
458
459 #[must_use]
461 pub fn god_nodes(&self, limit: usize) -> Vec<Value> {
462 let mut degrees: HashMap<String, usize> = HashMap::new();
463 for e in &self.edges {
464 if e.edge_type == "calls" {
465 *degrees.entry(e.source_id.clone()).or_default() += 1;
466 *degrees.entry(e.target_id.clone()).or_default() += 1;
467 }
468 }
469 let mut ranked: Vec<(&CodeNode, usize)> = degrees
470 .iter()
471 .filter_map(|(id, d)| self.node(id).map(|n| (n, *d)))
472 .collect();
473 ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.name.cmp(&b.0.name)));
474 ranked
475 .into_iter()
476 .take(limit)
477 .map(|(n, d)| {
478 let mut v = node_json(n);
479 v["degree"] = json!(d);
480 v
481 })
482 .collect()
483 }
484
485 #[must_use]
488 pub fn fragment_search(&self, query: &str, max_results: usize) -> Vec<Value> {
489 let mut out = Vec::new();
490 let root = std::path::Path::new(&self.project_root);
491 let mut files: Vec<std::path::PathBuf> = Vec::new();
492 let mut walked = 0usize;
493 if collect_files(root, root, 20_000, &mut files, &mut walked).is_err() {
494 return out;
495 }
496 for file in files {
497 let rel = file
498 .strip_prefix(root)
499 .unwrap_or(&file)
500 .to_string_lossy()
501 .replace('\\', "/");
502 if language_for(&rel).is_empty() {
503 continue;
504 }
505 let Some(contents) = read_bounded(&file) else {
506 continue;
507 };
508 for (line_no, line) in contents.lines().enumerate() {
509 if line
510 .to_ascii_lowercase()
511 .contains(&query.to_ascii_lowercase())
512 {
513 out.push(json!({
514 "file": rel,
515 "line": line_no + 1,
516 "content": line.trim().to_string(),
517 }));
518 if out.len() >= max_results {
519 return out;
520 }
521 }
522 }
523 }
524 out
525 }
526
527 #[must_use]
529 pub fn stats(&self) -> Value {
530 let mut languages: HashMap<String, (usize, usize)> = HashMap::new();
531 for node in &self.nodes {
532 let e = languages.entry(node.language.clone()).or_default();
533 e.0 += 1;
534 }
535 for edge in &self.edges {
536 if edge.edge_type == "calls" {
537 if let Some(n) = self.node(&edge.source_id) {
538 let e = languages.entry(n.language.clone()).or_default();
539 e.1 += 1;
540 }
541 }
542 }
543 let languages: Vec<Value> = languages
544 .into_iter()
545 .map(|(lang, (nodes, calls))| json!({"language": lang, "nodes": nodes, "call_edges": calls}))
546 .collect();
547 json!({
548 "built": self.built,
549 "project_root": self.project_root,
550 "nodes": self.nodes.len(),
551 "edges": self.edges.len(),
552 "languages": languages,
553 })
554 }
555
556 fn node(&self, id: &str) -> Option<&CodeNode> {
557 self.nodes.iter().find(|n| n.id == id)
558 }
559}
560
561fn node_json(node: &CodeNode) -> Value {
562 json!({
563 "name": node.name,
564 "node_type": node.node_type,
565 "file": node.file,
566 "line": node.line,
567 "language": node.language,
568 })
569}
570
571#[allow(clippy::only_used_in_recursion)]
573fn collect_files(
574 root: &std::path::Path,
575 dir: &std::path::Path,
576 max_files: usize,
577 out: &mut Vec<std::path::PathBuf>,
578 walked: &mut usize,
579) -> Result<(), String> {
580 if out.len() >= max_files {
581 return Ok(());
582 }
583 let entries =
584 std::fs::read_dir(dir).map_err(|e| format!("cannot read dir {}: {e}", dir.display()))?;
585 for entry in entries.flatten() {
586 let path = entry.path();
587 if path.is_dir() {
588 let name = entry.file_name().to_string_lossy().to_string();
589 if SKIP_DIRS.contains(&name.as_str()) {
590 continue;
591 }
592 collect_files(root, &path, max_files, out, walked)?;
593 if out.len() >= max_files {
594 return Ok(());
595 }
596 } else if path.is_file() {
597 *walked += 1;
598 if language_for(&path.to_string_lossy()).is_empty() {
599 continue;
600 }
601 if entry.metadata().map_or(true, |m| m.len() > MAX_FILE_BYTES) {
602 continue;
603 }
604 out.push(path);
605 }
606 }
607 Ok(())
608}
609
610fn read_bounded(path: &std::path::Path) -> Option<String> {
611 use std::io::Read;
612 let file = std::fs::File::open(path).ok()?;
613 let mut buf = Vec::new();
614 std::io::Read::take(file, MAX_FILE_BYTES)
615 .read_to_end(&mut buf)
616 .ok()?;
617 Some(String::from_utf8_lossy(&buf).into_owned())
618}
619
620fn language_for(path: &str) -> String {
621 let ext = path.rsplit('.').next().unwrap_or("").to_ascii_lowercase();
622 match ext.as_str() {
623 "rs" => "rust".to_string(),
624 "py" => "python".to_string(),
625 "js" | "jsx" | "mjs" | "cjs" => "javascript".to_string(),
626 "ts" | "tsx" => "typescript".to_string(),
627 "go" => "go".to_string(),
628 "java" => "java".to_string(),
629 "c" | "h" => "c".to_string(),
630 "cpp" | "cc" | "hpp" => "cpp".to_string(),
631 "rb" => "ruby".to_string(),
632 "sh" => "shell".to_string(),
633 "zig" => "zig".to_string(),
634 "jl" => "julia".to_string(),
635 _ => String::new(),
636 }
637}
638
639type PatternList = Vec<(&'static str, &'static str)>;
642fn symbol_patterns(language: &str) -> (PatternList, PatternList) {
643 match language {
644 "rust" => (
645 vec![
646 (r"fn\s+([a-zA-Z_][a-zA-Z0-9_]*)", "function"),
647 (r"struct\s+([a-zA-Z_][a-zA-Z0-9_]*)", "struct"),
648 (r"enum\s+([a-zA-Z_][a-zA-Z0-9_]*)", "enum"),
649 (r"impl\s+([a-zA-Z_][a-zA-Z0-9_]*)", "impl"),
650 (r"trait\s+([a-zA-Z_][a-zA-Z0-9_]*)", "trait"),
651 (r"mod\s+([a-zA-Z_][a-zA-Z0-9_]*)", "module"),
652 ],
653 vec![(r"pub\s+fn\s+([a-zA-Z_][a-zA-Z0-9_]*)", "function")],
654 ),
655 "python" => (
656 vec![
657 (r"def\s+([a-zA-Z_][a-zA-Z0-9_]*)", "function"),
658 (r"class\s+([a-zA-Z_][a-zA-Z0-9_]*)", "class"),
659 ],
660 vec![],
661 ),
662 "javascript" | "typescript" => (
663 vec![
664 (r"function\s+([a-zA-Z_$][a-zA-Z0-9_$]*)", "function"),
665 (r"class\s+([a-zA-Z_$][a-zA-Z0-9_$]*)", "class"),
666 ],
667 vec![(
668 r"(?:const|let|var)\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\s*=",
669 "const",
670 )],
671 ),
672 "go" => (
673 vec![
674 (r"func\s+([a-zA-Z_][a-zA-Z0-9_]*)", "function"),
675 (r"type\s+([a-zA-Z_][a-zA-Z0-9_]*)\s+struct", "struct"),
676 ],
677 vec![],
678 ),
679 "java" => (
680 vec![
681 (
682 r"(?:public|private|protected)\s+(?:static\s+)?[a-zA-Z0-9_<>\[\]]+\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(",
683 "method",
684 ),
685 (r"class\s+([a-zA-Z_][a-zA-Z0-9_]*)", "class"),
686 ],
687 vec![],
688 ),
689 _ => (
690 vec![(
691 r"(?:fn|def|function)\s+([a-zA-Z_][a-zA-Z0-9_]*)",
692 "function",
693 )],
694 vec![],
695 ),
696 }
697}
698
699fn regex_captures(pattern: &str, line: &str) -> Vec<String> {
703 let mut captures = Vec::new();
704 let line_lower = line;
705 let open = pattern.find('(').unwrap_or(pattern.len());
708 let literal = &pattern[..open];
709 let trimmed = literal.trim_end_matches("\\s+");
710 let mut rest = line_lower;
711 while let Some(idx) = rest.find(trimmed) {
712 let after = &rest[idx + trimmed.len()..];
713 let after = after.trim_start_matches([' ', '\t']);
715 if let Some(name) = leading_identifier(after) {
717 captures.push(name);
718 }
719 rest = after;
720 }
721 captures
722}
723
724fn leading_identifier(s: &str) -> Option<String> {
726 let mut name = String::new();
727 for c in s.chars() {
728 if name.is_empty() {
729 if c.is_ascii_alphabetic() || c == '_' || c == '$' {
730 name.push(c);
731 } else {
732 return None;
733 }
734 } else if c.is_ascii_alphanumeric() || c == '_' || c == '$' {
735 name.push(c);
736 } else {
737 break;
738 }
739 }
740 if name.is_empty() { None } else { Some(name) }
741}
742
743fn inherit_targets(line: &str) -> Vec<String> {
745 let Some(open) = line.find('(') else {
746 return Vec::new();
747 };
748 let Some(close) = line[open..].find(')') else {
749 return Vec::new();
750 };
751 line[open + 1..open + close]
752 .split(',')
753 .map(|s| s.trim().to_string())
754 .filter(|s| !s.is_empty())
755 .collect()
756}
757
758fn call_sites(line: &str) -> Vec<String> {
763 let mut out = Vec::new();
764 let mut rest = line;
765 while !rest.is_empty() {
766 let skipped = rest
768 .chars()
769 .take_while(|c| !(c.is_ascii_alphabetic() || *c == '_' || *c == '$'))
770 .count();
771 rest = &rest[skipped..];
772 let Some(name) = leading_identifier(rest) else {
773 break;
774 };
775 let consumed = line.len() - rest.len();
776 let before = line[..consumed].chars().last();
777 let boundary_ok = !before.is_some_and(|c| c.is_ascii_alphanumeric() || c == '_');
779 if boundary_ok {
780 let after = &rest[name.len()..];
781 if after.starts_with('(') {
782 let def_marker = line[..consumed].rsplit([' ', '\t']).next().unwrap_or("");
784 if !(def_marker == "fn"
785 || def_marker == "def"
786 || def_marker == "function"
787 || out.contains(&name))
788 {
789 out.push(name.clone());
790 }
791 }
792 }
793 let skip = name.len().max(1);
795 rest = &rest[skip..];
796 }
797 out
798}
799
800fn enclosing_function(contents: &str, line_idx: usize) -> Option<String> {
804 let mut found: Option<String> = None;
805 for (i, line) in contents.lines().enumerate() {
806 if i > line_idx {
807 break;
808 }
809 for pat in ["fn ", "def ", "function "] {
810 if let Some(idx) = line.find(pat) {
811 if let Some(name) = leading_identifier(&line[idx + pat.len()..]) {
812 found = Some(name);
813 }
814 }
815 }
816 }
817 found
818}
819
820fn import_targets(language: &str, line: &str) -> Vec<String> {
822 let mut out = Vec::new();
823 match language {
824 "rust" => {
825 if let Some(rest) = line.trim().strip_prefix("use ") {
826 let target = rest.trim_end_matches(';').trim();
827 if !target.starts_with("crate::") {
828 out.push(format!("import:{target}"));
829 }
830 }
831 }
832 "python" => {
833 if let Some(rest) = line.trim().strip_prefix("import ") {
834 for part in rest.split(',') {
835 let module = part.trim().split('.').next().unwrap_or("").to_string();
836 if !module.is_empty() {
837 out.push(format!("import:{module}"));
838 }
839 }
840 } else if let Some(rest) = line.trim().strip_prefix("from ") {
841 let module = rest.split(" import ").next().unwrap_or("").trim();
842 if !module.is_empty() {
843 out.push(format!("import:{module}"));
844 }
845 }
846 }
847 "javascript" | "typescript" => {
848 if let Some(rest) = line.trim().strip_prefix("import ") {
849 if let Some(from) = rest.split(" from ").nth(1) {
850 let module = from.trim().trim_matches(['\'', '"']).to_string();
851 out.push(format!("import:{module}"));
852 }
853 }
854 }
855 "go" => {
856 if let Some(rest) = line.trim().strip_prefix("import \"") {
857 let module = rest.trim_end_matches('"').to_string();
858 out.push(format!("import:{module}"));
859 }
860 }
861 _ => {}
862 }
863 out
864}
865
866pub struct CodeGraphTool {
870 graph: Arc<Mutex<CodeGraph>>,
871 stats: ToolStats,
872 effects: EffectRow,
873}
874
875impl CodeGraphTool {
876 #[must_use]
877 pub fn new(graph: Arc<Mutex<CodeGraph>>) -> Self {
878 Self {
879 graph,
880 stats: ToolStats::default(),
881 effects: EffectRow::read_only(vec![Resource::Filesystem]),
882 }
883 }
884}
885
886#[async_trait]
887impl Tool for CodeGraphTool {
888 fn name(&self) -> &str {
889 "code.graph"
890 }
891 fn gana(&self) -> Gana {
892 Gana::Chariot
893 }
894 fn effects(&self) -> &EffectRow {
895 &self.effects
896 }
897 fn description(&self) -> &str {
898 "Build (or refresh) the code structure graph for a project. Args: project_root (required), max_files (default 50000)."
899 }
900 async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
901 let project_root = args
902 .get("project_root")
903 .and_then(Value::as_str)
904 .ok_or_else(|| wm_core::CoreError::InvalidArgs("project_root is required".into()))?;
905 let max_files = args
906 .get("max_files")
907 .and_then(Value::as_u64)
908 .unwrap_or(DEFAULT_MAX_FILES as u64) as usize;
909 let project_root = project_root.to_string();
910 let mut graph = self
911 .graph
912 .lock()
913 .map_err(|e| wm_core::CoreError::Tool(format!("code graph lock: {e}")))?;
914 let files = graph
915 .build(&project_root, max_files)
916 .map_err(wm_core::CoreError::Tool)?;
917 let stats = graph.stats();
918 let mut result = json!({
919 "status": "success",
920 "files_scanned": files,
921 });
922 for (k, v) in stats.as_object().unwrap() {
923 result[k.clone()] = v.clone();
924 }
925 Ok(result)
926 }
927 fn stats(&self) -> &ToolStats {
928 &self.stats
929 }
930}
931
932pub struct CodeQueryTool {
936 graph: Arc<Mutex<CodeGraph>>,
937 stats: ToolStats,
938 effects: EffectRow,
939}
940
941impl CodeQueryTool {
942 #[must_use]
943 pub fn new(graph: Arc<Mutex<CodeGraph>>) -> Self {
944 Self {
945 graph,
946 stats: ToolStats::default(),
947 effects: EffectRow::read_only(vec![Resource::Filesystem]),
948 }
949 }
950}
951
952#[async_trait]
953impl Tool for CodeQueryTool {
954 fn name(&self) -> &str {
955 "code.query"
956 }
957 fn gana(&self) -> Gana {
958 Gana::Chariot
959 }
960 fn effects(&self) -> &EffectRow {
961 &self.effects
962 }
963 fn description(&self) -> &str {
964 "Query the code graph with natural language: 'what calls X', 'what does X call', 'path from A to B', 'explain X', 'god nodes', or a symbol search. Args: query (required), limit (default 20). Build the graph first with code.graph."
965 }
966 async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
967 let query = args
968 .get("query")
969 .and_then(Value::as_str)
970 .ok_or_else(|| wm_core::CoreError::InvalidArgs("query is required".into()))?;
971 let limit = args.get("limit").and_then(Value::as_u64).unwrap_or(20) as usize;
972 let graph = self
973 .graph
974 .lock()
975 .map_err(|e| wm_core::CoreError::Tool(format!("code graph lock: {e}")))?;
976 if !graph.is_built() {
977 return Ok(json!({
978 "status": "error",
979 "error": "code graph not built — run code.graph first",
980 }));
981 }
982 let q = query.to_ascii_lowercase();
983 let result = if let Some(rest) = q.strip_prefix("what calls ") {
984 json!({"status": "success", "query": query, "callers": graph.callers(rest.trim(), limit)})
985 } else if q.contains("what does") && q.contains("call") {
986 let symbol = q
987 .split("what does")
988 .nth(1)
989 .and_then(|s| s.split("call").next())
990 .unwrap_or("")
991 .trim();
992 json!({"status": "success", "query": query, "callees": graph.callees(symbol, limit)})
993 } else if q.contains("path from") && q.contains(" to ") {
994 let parts: Vec<&str> = q.split(" to ").collect();
995 let a = parts[0].strip_prefix("path from").unwrap_or("").trim();
996 let b = parts[1].trim();
997 graph.path(a, b, 5)
998 } else if let Some(rest) = q.strip_prefix("explain ") {
999 match graph.explain(rest.trim()) {
1000 Some(expl) => json!({"status": "success", "query": query, "explanation": expl}),
1001 None => {
1002 json!({"status": "error", "error": format!("symbol not found: {}", rest.trim())})
1003 }
1004 }
1005 } else if q.contains("god") || q.contains("most connected") {
1006 json!({"status": "success", "query": query, "god_nodes": graph.god_nodes(limit)})
1007 } else if q == "stats" || q.contains("stats") {
1008 graph.stats()
1009 } else {
1010 json!({"status": "success", "query": query, "matches": graph.search(query, limit)})
1011 };
1012 Ok(result)
1013 }
1014 fn stats(&self) -> &ToolStats {
1015 &self.stats
1016 }
1017}
1018
1019pub struct CodeAffectedByTool {
1023 graph: Arc<Mutex<CodeGraph>>,
1024 stats: ToolStats,
1025 effects: EffectRow,
1026}
1027
1028impl CodeAffectedByTool {
1029 #[must_use]
1030 pub fn new(graph: Arc<Mutex<CodeGraph>>) -> Self {
1031 Self {
1032 graph,
1033 stats: ToolStats::default(),
1034 effects: EffectRow::read_only(vec![Resource::Filesystem]),
1035 }
1036 }
1037}
1038
1039#[async_trait]
1040impl Tool for CodeAffectedByTool {
1041 fn name(&self) -> &str {
1042 "code.affected_by"
1043 }
1044 fn gana(&self) -> Gana {
1045 Gana::Chariot
1046 }
1047 fn effects(&self) -> &EffectRow {
1048 &self.effects
1049 }
1050 fn description(&self) -> &str {
1051 "Find all symbols transitively affected by a change to the given symbol (reverse call-graph BFS). Args: symbol (required), max_depth (default 3)."
1052 }
1053 async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
1054 let symbol = args
1055 .get("symbol")
1056 .and_then(Value::as_str)
1057 .ok_or_else(|| wm_core::CoreError::InvalidArgs("symbol is required".into()))?;
1058 let max_depth = args
1059 .get("max_depth")
1060 .and_then(Value::as_u64)
1061 .unwrap_or(3)
1062 .clamp(1, 10) as usize;
1063 let graph = self
1064 .graph
1065 .lock()
1066 .map_err(|e| wm_core::CoreError::Tool(format!("code graph lock: {e}")))?;
1067 if !graph.is_built() {
1068 return Ok(json!({
1069 "status": "error",
1070 "error": "code graph not built — run code.graph first",
1071 }));
1072 }
1073 Ok(graph.affected_by(symbol, max_depth))
1074 }
1075 fn stats(&self) -> &ToolStats {
1076 &self.stats
1077 }
1078}
1079
1080pub struct FragmentSearchTool {
1084 graph: Arc<Mutex<CodeGraph>>,
1085 stats: ToolStats,
1086 effects: EffectRow,
1087}
1088
1089impl FragmentSearchTool {
1090 #[must_use]
1091 pub fn new(graph: Arc<Mutex<CodeGraph>>) -> Self {
1092 Self {
1093 graph,
1094 stats: ToolStats::default(),
1095 effects: EffectRow::read_only(vec![Resource::Filesystem]),
1096 }
1097 }
1098}
1099
1100#[async_trait]
1101impl Tool for FragmentSearchTool {
1102 fn name(&self) -> &str {
1103 "fragment.search"
1104 }
1105 fn gana(&self) -> Gana {
1106 Gana::WinnowingBasket
1107 }
1108 fn effects(&self) -> &EffectRow {
1109 &self.effects
1110 }
1111 fn description(&self) -> &str {
1112 "Locate file/line fragments mentioning a query in the built code graph's project. Args: query (required), max_results (default 20)."
1113 }
1114 async fn call(&self, _ctx: &mut Context, args: Value) -> wm_core::Result<Value> {
1115 let query = args
1116 .get("query")
1117 .and_then(Value::as_str)
1118 .ok_or_else(|| wm_core::CoreError::InvalidArgs("query is required".into()))?;
1119 let max_results = args
1120 .get("max_results")
1121 .and_then(Value::as_u64)
1122 .unwrap_or(20) as usize;
1123 let graph = self
1124 .graph
1125 .lock()
1126 .map_err(|e| wm_core::CoreError::Tool(format!("code graph lock: {e}")))?;
1127 if !graph.is_built() {
1128 return Ok(json!({
1129 "status": "error",
1130 "error": "code graph not built — run code.graph first",
1131 }));
1132 }
1133 let symbols = graph.search(query, max_results);
1134 let fragments = graph.fragment_search(query, max_results);
1135 Ok(json!({
1136 "status": "success",
1137 "query": query,
1138 "symbol_matches": symbols.len(),
1139 "fragment_matches": fragments.len(),
1140 "symbols": symbols,
1141 "fragments": fragments,
1142 }))
1143 }
1144 fn stats(&self) -> &ToolStats {
1145 &self.stats
1146 }
1147}
1148
1149#[must_use]
1151pub fn register_code(
1152 registry: &wm_dispatch::ToolRegistry,
1153 graph: Arc<Mutex<CodeGraph>>,
1154) -> wm_dispatch::ToolRegistry {
1155 registry
1156 .register(Arc::new(CodeGraphTool::new(graph.clone())))
1157 .register(Arc::new(CodeQueryTool::new(graph.clone())))
1158 .register(Arc::new(CodeAffectedByTool::new(graph.clone())))
1159 .register(Arc::new(FragmentSearchTool::new(graph)))
1160}
1161
1162#[cfg(test)]
1163mod tests {
1164 use super::*;
1165
1166 fn write_project(dir: &std::path::Path) {
1167 std::fs::create_dir_all(dir.join("src")).unwrap();
1168 std::fs::write(
1169 dir.join("src/main.rs"),
1170 "mod utils;\nfn main() {\n let x = utils::double(21);\n println!(\"{x}\");\n}\n",
1171 )
1172 .unwrap();
1173 std::fs::write(
1174 dir.join("src/utils.rs"),
1175 "pub fn double(x: i32) -> i32 {\n x * 2\n}\n",
1176 )
1177 .unwrap();
1178 }
1179
1180 #[test]
1181 fn build_extracts_symbols_and_calls() {
1182 let dir = tempfile::tempdir().unwrap();
1183 write_project(dir.path());
1184 let mut graph = CodeGraph::new();
1185 let files = graph.build(dir.path().to_str().unwrap(), 100).unwrap();
1186 assert!(files >= 2);
1187 assert!(
1188 graph
1189 .nodes
1190 .iter()
1191 .any(|n| n.name == "main" && n.node_type == "function")
1192 );
1193 assert!(
1194 graph
1195 .nodes
1196 .iter()
1197 .any(|n| n.name == "double" && n.node_type == "function")
1198 );
1199 assert!(!graph.callers("double", 10).is_empty());
1201 assert!(!graph.callees("main", 10).is_empty());
1202 }
1203
1204 #[test]
1205 fn affected_by_traces_reverse_bfs() {
1206 let dir = tempfile::tempdir().unwrap();
1207 std::fs::create_dir_all(dir.path().join("src")).unwrap();
1208 std::fs::write(
1209 dir.path().join("src/a.rs"),
1210 "fn top() { mid(); }\nfn mid() { leaf(); }\nfn leaf() {}\n",
1211 )
1212 .unwrap();
1213 let mut graph = CodeGraph::new();
1214 graph.build(dir.path().to_str().unwrap(), 100).unwrap();
1215 let result = graph.affected_by("leaf", 3);
1216 assert_eq!(result["status"], "success");
1217 let affected = result["affected"].as_array().unwrap();
1218 let names: Vec<&str> = affected
1219 .iter()
1220 .filter_map(|a| a.get("symbol").and_then(Value::as_str))
1221 .collect();
1222 assert!(names.contains(&"mid"));
1223 assert!(names.contains(&"top"));
1224 }
1225
1226 #[test]
1227 fn path_finds_connection() {
1228 let dir = tempfile::tempdir().unwrap();
1229 std::fs::create_dir_all(dir.path().join("src")).unwrap();
1230 std::fs::write(dir.path().join("src/a.rs"), "fn a() { b(); }\nfn b() {}\n").unwrap();
1231 let mut graph = CodeGraph::new();
1232 graph.build(dir.path().to_str().unwrap(), 100).unwrap();
1233 let result = graph.path("a", "b", 3);
1234 assert_eq!(result["status"], "success");
1235 assert_eq!(result["hops"], 1);
1236 }
1237
1238 #[test]
1239 fn fragment_search_finds_lines() {
1240 let dir = tempfile::tempdir().unwrap();
1241 std::fs::create_dir_all(dir.path().join("src")).unwrap();
1242 std::fs::write(
1243 dir.path().join("src/a.rs"),
1244 "fn a() {}\n// the unique marker word\nfn b() {}\n",
1245 )
1246 .unwrap();
1247 let mut graph = CodeGraph::new();
1248 graph.build(dir.path().to_str().unwrap(), 100).unwrap();
1249 let fragments = graph.fragment_search("unique marker", 10);
1250 assert_eq!(fragments.len(), 1);
1251 assert_eq!(fragments[0]["line"], 2);
1252 }
1253
1254 #[test]
1255 fn skip_dirs_are_not_scanned() {
1256 let dir = tempfile::tempdir().unwrap();
1257 std::fs::create_dir_all(dir.path().join("target")).unwrap();
1258 std::fs::write(dir.path().join("target/bad.rs"), "fn bad() {}\n").unwrap();
1259 std::fs::write(dir.path().join("good.rs"), "fn good() {}\n").unwrap();
1260 let mut graph = CodeGraph::new();
1261 graph.build(dir.path().to_str().unwrap(), 100).unwrap();
1262 assert!(!graph.nodes.iter().any(|n| n.name == "bad"));
1263 assert!(graph.nodes.iter().any(|n| n.name == "good"));
1264 }
1265
1266 #[tokio::test]
1267 async fn tools_require_built_graph() {
1268 let graph = Arc::new(Mutex::new(CodeGraph::new()));
1269 let tool = CodeQueryTool::new(graph);
1270 let mut ctx = Context::default();
1271 let result = tool
1272 .call(&mut ctx, json!({"query": "what calls main"}))
1273 .await
1274 .unwrap();
1275 assert_eq!(result["status"], "error");
1276 assert_eq!(
1277 result["error"],
1278 "code graph not built — run code.graph first"
1279 );
1280 }
1281}