From 5e07e909e95cc52f7c859ddca6274be61c524625 Mon Sep 17 00:00:00 2001
From: ProtocolWarden <32967198+ProtocolWarden@users.noreply.github.com>
Date: Tue, 14 Jul 2026 15:31:39 -0400
Subject: [PATCH] fix(Custodian): collapse extra-root call-graph walks for
 sweep timeout

---
 src/custodian/audit_kit/passes/call_graph.py | 17 ++++-----
 tests/test_call_graph_detectors.py           | 36 +++++++++++++++++++-
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/src/custodian/audit_kit/passes/call_graph.py b/src/custodian/audit_kit/passes/call_graph.py
index 5b3742b..c25ad9c 100644
--- a/src/custodian/audit_kit/passes/call_graph.py
+++ b/src/custodian/audit_kit/passes/call_graph.py
@@ -96,7 +96,6 @@ def _collect_usages_only(tree: ast.Module, cg: CallGraph) -> None:
                     cg.constructed_names.add(base.id)
                 elif isinstance(base, ast.Attribute):
                     cg.constructed_names.add(base.attr)
-    for node in ast.walk(tree):
         if isinstance(node, ast.Call):
             func = node.func
             if isinstance(func, ast.Name):
@@ -121,6 +120,12 @@ def _collect_usages_only(tree: ast.Module, cg: CallGraph) -> None:
                     cg.constructed_names.add(kw.value.id)
                 if kw.arg:  # keyword arg name — Model(field=value) records "field"
                     cg.kw_arg_names.add(kw.arg)
+            if (
+                isinstance(func, ast.Attribute)
+                and func.attr.startswith("model_validate")
+                and isinstance(func.value, ast.Name)
+            ):
+                cg.model_validate_classes.add(func.value.id)
         if isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load):
             cg.accessed_attrs.add(node.attr)
             # ClassName.method(...) or EnumClass.MEMBER — treat as "class is in active use"
@@ -128,16 +133,6 @@ def _collect_usages_only(tree: ast.Module, cg: CallGraph) -> None:
                 cg.constructed_names.add(node.value.id)
         if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load):
             cg.called_names.add(node.id)
-    # ClassName.model_validate*(...) — deserialized from external data; all fields are schema fields
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Call):
-            func = node.func
-            if (
-                isinstance(func, ast.Attribute)
-                and func.attr.startswith("model_validate")
-                and isinstance(func.value, ast.Name)
-            ):
-                cg.model_validate_classes.add(func.value.id)
 
 
 def _collect_from_module(tree: ast.Module, cg: CallGraph) -> None:
diff --git a/tests/test_call_graph_detectors.py b/tests/test_call_graph_detectors.py
index cf3a33a..16d365f 100644
--- a/tests/test_call_graph_detectors.py
+++ b/tests/test_call_graph_detectors.py
@@ -5,10 +5,11 @@
 
 import textwrap
 from pathlib import Path
+from unittest.mock import patch
 
 from custodian.audit_kit.detector import AnalysisGraph, AuditContext
 from custodian.audit_kit.detectors.dead_code import detect_d1, detect_f1
-from custodian.audit_kit.passes.call_graph import build_call_graph
+from custodian.audit_kit.passes.call_graph import _collect_usages_only, build_call_graph
 
 
 # ── helpers ───────────────────────────────────────────────────────────────────
@@ -86,6 +87,39 @@ def foo(): pass
         cg = build_call_graph(tmp_path / "src")
         assert "my_decorator" in cg.decorated_names
 
+    def test_collect_usages_only_walks_tree_once(self):
+        tree = __import__("ast").parse(
+            textwrap.dedent(
+                """
+                class Child(Base):
+                    pass
+
+                Foo.model_validate(payload)
+                obj.run()
+                value = obj.field
+                """
+            )
+        )
+        walk_calls = 0
+        import ast
+
+        real_walk = ast.walk
+
+        def _counting_walk(node):
+            nonlocal walk_calls
+            walk_calls += 1
+            return real_walk(node)
+
+        with patch("custodian.audit_kit.passes.call_graph.ast.walk", side_effect=_counting_walk):
+            cg = build_call_graph(Path("/tmp/does-not-matter"))
+            _collect_usages_only(tree, cg)
+
+        assert walk_calls == 1
+        assert "Base" in cg.constructed_names
+        assert "run" in cg.called_attrs
+        assert "field" in cg.accessed_attrs
+        assert "Foo" in cg.model_validate_classes
+
 
 # ── D1 tests ──────────────────────────────────────────────────────────────────