fix: scalability + OOM fixes for diffctx on large repos

Author: nikolay-e

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 test-repos
+TODO.md
 
 # C extensions
 *.so

src/treemapper/diffctx/__init__.py

@@ -102,21 +102,101 @@ def _build_preferred_revs(base_rev: str | None, head_rev: str | None) -> list[st
     return revs
 
 
+_MAX_GENERATED_FRAGMENTS = LIMITS.max_generated_fragments
+
+
+_GENERATED_FILENAME_PATTERNS = frozenset(
+    {
+        ".pb.go",
+        "_pb2.py",
+        "_pb2_grpc.py",
+        ".pb.h",
+        ".pb.cc",
+        ".pb.swift",
+        ".min.js",
+        ".min.css",
+        ".designer.cs",
+    }
+)
+
+_GENERATED_FILENAME_SUFFIXES = ("_generated.", "OuterClass.java")
+
+_GENERATED_PATH_SEGMENTS = frozenset(
+    {
+        "generated",
+        "gen-java",
+        "gen-go",
+        "gen-py",
+        "gen-cpp",
+        "gen-swift",
+        "__generated__",
+        "autogen",
+    }
+)
+
+_GENERATED_CONTENT_MARKERS = (
+    "@generated",
+    "do not edit",
+    "code generated",
+    "auto-generated",
+    "this file is generated",
+    "generated by",
+    "automatically generated",
+    "auto generated",
+)
+
+
+def _is_generated_file(path: Path, content: str) -> bool:
+    name = path.name
+    for pattern in _GENERATED_FILENAME_PATTERNS:
+        if name.endswith(pattern):
+            return True
+    for suffix in _GENERATED_FILENAME_SUFFIXES:
+        if name.endswith(suffix):
+            return True
+
+    for part in path.parts:
+        if part.lower() in _GENERATED_PATH_SEGMENTS:
+            return True
+
+    header_lower = "\n".join(content[:2000].splitlines()[:5]).lower()
+    for marker in _GENERATED_CONTENT_MARKERS:
+        if marker in header_lower:
+            return True
+
+    return False
+
+
 def _process_files_for_fragments(
     files: list[Path],
     root_dir: Path,
     preferred_revs: list[str],
     seen_frag_ids: set[FragmentId],
 ) -> list[Fragment]:
+    max_frags = LIMITS.max_fragments
     fragments: list[Fragment] = []
     for file_path in files:
         content = _read_file_content(file_path, root_dir, preferred_revs)
         if content is None:
             continue
-        for frag in fragment_file(file_path, content):
-            if frag.id not in seen_frag_ids:
-                fragments.append(frag)
-                seen_frag_ids.add(frag.id)
+        file_frags = [f for f in fragment_file(file_path, content) if f.id not in seen_frag_ids]
+
+        is_generated = _is_generated_file(file_path, content)
+        cap = _MAX_GENERATED_FRAGMENTS if is_generated else max_frags
+
+        if len(file_frags) > cap:
+            file_frags.sort(key=lambda f: f.line_count, reverse=True)
+            file_frags = file_frags[:cap]
+            logging.debug(
+                "diffctx: capped %s to %d fragments%s",
+                file_path.name,
+                cap,
+                " (generated)" if is_generated else "",
+            )
+
+        for frag in file_frags:
+            fragments.append(frag)
+            seen_frag_ids.add(frag.id)
     return fragments
 
 
@@ -301,11 +381,7 @@ def _filter_low_relevance_fragments(
     rel: dict[FragmentId, float],
 ) -> list[Fragment]:
     changed_paths = {fid.path for fid in core_ids}
-    kept = [
-        f
-        for f in fragments
-        if f.path in changed_paths or rel.get(f.id, 0.0) >= _LOW_RELEVANCE_THRESHOLD
-    ]
+    kept = [f for f in fragments if f.path in changed_paths or rel.get(f.id, 0.0) >= _LOW_RELEVANCE_THRESHOLD]
     removed = len(fragments) - len(kept)
     if removed:
         logging.debug("diffctx: filtered %d low-relevance fragments (threshold=%.4f)", removed, _LOW_RELEVANCE_THRESHOLD)
@@ -317,6 +393,8 @@ def _ensure_changed_files_represented(
     all_fragments: list[Fragment],
     changed_files: list[Path],
     remaining_budget: int,
+    root_dir: Path,
+    preferred_revs: list[str],
 ) -> list[Fragment]:
     selected_paths = {f.path for f in selected}
     changed_paths = set(changed_files)
@@ -336,6 +414,20 @@ def _ensure_changed_files_represented(
 
     for path in sorted(missing_paths):
         candidates = frags_by_path.get(path, [])
+
+        if not candidates:
+            content = _read_file_content(path, root_dir, preferred_revs)
+            if content and content.strip():
+                lines = content.splitlines()
+                frag = Fragment(
+                    id=FragmentId(path=path, start_line=1, end_line=len(lines)),
+                    kind="chunk",
+                    content=content,
+                    identifiers=extract_identifiers(content),
+                )
+                frag.token_count = count_tokens(content).count + _OVERHEAD_PER_FRAGMENT
+                candidates = [frag]
+
         if not candidates:
             continue
         best = max(candidates, key=lambda f: f.token_count if f.token_count > 0 else 0)
@@ -433,18 +525,21 @@ def build_diff_context(
     seen_frag_ids: set[FragmentId] = set()
     all_fragments = _process_files_for_fragments(changed_files, root_dir, preferred_revs, seen_frag_ids)
 
-    all_candidate_files = _collect_candidate_files(root_dir, set(changed_files), combined_spec)
+    all_candidate_files, is_large_repo = _collect_candidate_files(root_dir, set(changed_files), combined_spec)
     all_candidate_files = _filter_whitelist(all_candidate_files, root_dir, wl_spec)
 
     edge_discovered = discover_all_related_files(changed_files, all_candidate_files, root_dir)
     edge_discovered = [_normalize_path(p, root_dir) for p in edge_discovered]
     all_fragments.extend(_process_files_for_fragments(edge_discovered, root_dir, preferred_revs, seen_frag_ids))
 
-    expanded_files = _expand_universe_by_rare_identifiers(
-        root_dir, expansion_concepts, changed_files + edge_discovered, combined_spec
-    )
-    expanded_files = [_normalize_path(p, root_dir) for p in expanded_files]
-    all_fragments.extend(_process_files_for_fragments(expanded_files, root_dir, preferred_revs, seen_frag_ids))
+    if not is_large_repo:
+        expanded_files = _expand_universe_by_rare_identifiers(
+            root_dir, expansion_concepts, changed_files + edge_discovered, combined_spec
+        )
+        expanded_files = [_normalize_path(p, root_dir) for p in expanded_files]
+        all_fragments.extend(_process_files_for_fragments(expanded_files, root_dir, preferred_revs, seen_frag_ids))
+    else:
+        logging.debug("diffctx: skipping rare-identifier expansion for large repo")
 
     for frag in all_fragments:
         frag.token_count = count_tokens(frag.content).count + _OVERHEAD_PER_FRAGMENT
@@ -473,7 +568,7 @@ def build_diff_context(
         )
         effective_budget = budget_tokens if budget_tokens is not None else _UNLIMITED_BUDGET
         remaining = effective_budget - result.used_tokens
-        selected = _ensure_changed_files_represented(selected, all_fragments, changed_files, remaining)
+        selected = _ensure_changed_files_represented(selected, all_fragments, changed_files, remaining, root_dir, preferred_revs)
         _log_ppr_mode(selected, core_ids, budget_tokens, result, alpha, tau)
 
     if no_content:
@@ -717,7 +812,39 @@ def _is_candidate_file(file_path: Path, root_dir: Path, included_set: set[Path],
     return True
 
 
-def _collect_candidate_files(root_dir: Path, included_set: set[Path], combined_spec: pathspec.PathSpec) -> list[Path]:
+_MAX_CANDIDATE_FILES = LIMITS.max_candidate_files
+
+
+def _prioritize_candidates(
+    candidates: list[Path],
+    changed_files: set[Path],
+) -> list[Path]:
+    changed_dirs: set[Path] = set()
+    changed_extensions: set[str] = set()
+    for f in changed_files:
+        changed_dirs.add(f.parent)
+        if f.parent.parent != f.parent:
+            changed_dirs.add(f.parent.parent)
+        if f.suffix:
+            changed_extensions.add(f.suffix.lower())
+
+    priority: list[Path] = []
+    rest: list[Path] = []
+    for c in candidates:
+        if c.parent in changed_dirs or c.suffix.lower() in changed_extensions:
+            priority.append(c)
+        else:
+            rest.append(c)
+
+    budget = _MAX_CANDIDATE_FILES - len(priority)
+    if budget > 0:
+        priority.extend(rest[:budget])
+    return priority[:_MAX_CANDIDATE_FILES]
+
+
+def _collect_candidate_files(
+    root_dir: Path, included_set: set[Path], combined_spec: pathspec.PathSpec
+) -> tuple[list[Path], bool]:
     try:
         result = subprocess.run(
             ["git", "ls-files", "-z"],
@@ -729,18 +856,27 @@ def _collect_candidate_files(root_dir: Path, included_set: set[Path], combined_s
         if result.returncode == 0 and result.stdout:
             out = result.stdout.decode("utf-8", errors="surrogateescape")
             files = [root_dir / f for f in out.split("\0") if f]
-            return [f for f in files if _is_candidate_file(f, root_dir, included_set, combined_spec)]
+            candidates = [f for f in files if _is_candidate_file(f, root_dir, included_set, combined_spec)]
+            is_large_repo = len(candidates) > _MAX_CANDIDATE_FILES
+            if is_large_repo:
+                logging.debug(
+                    "diffctx: %d candidates exceed cap %d, prioritizing by proximity",
+                    len(candidates),
+                    _MAX_CANDIDATE_FILES,
+                )
+                candidates = _prioritize_candidates(candidates, included_set)
+            return candidates, is_large_repo
     except (subprocess.SubprocessError, OSError):
         pass
     logging.warning("diffctx: git ls-files failed, falling back to rglob (limit %d files)", _FALLBACK_MAX_FILES)
-    candidates: list[Path] = []
+    fallback: list[Path] = []
     for f in root_dir.rglob("*"):
-        if len(candidates) >= _FALLBACK_MAX_FILES:
+        if len(fallback) >= _FALLBACK_MAX_FILES:
             logging.warning("diffctx: fallback scan hit limit, results may be incomplete")
             break
         if _is_candidate_file(f, root_dir, included_set, combined_spec):
-            candidates.append(f)
-    return candidates
+            fallback.append(f)
+    return fallback, False
 
 
 def _build_ident_index(files: list[Path], concepts: frozenset[str]) -> dict[str, list[Path]]:
@@ -822,7 +958,7 @@ def _expand_universe_by_rare_identifiers(
         return []
 
     included_set = set(already_included)
-    files = _collect_candidate_files(root_dir, included_set, combined_spec)
+    files, _ = _collect_candidate_files(root_dir, included_set, combined_spec)
     inverted_index = _build_ident_index(files, concepts)
     return _collect_expansion_files(inverted_index, concepts, included_set)
 

src/treemapper/diffctx/config/limits.py

@@ -7,6 +7,9 @@
 class AlgorithmLimits:
     max_file_size: int = 100_000
     max_fragments: int = 200
+    max_generated_fragments: int = 5
+    max_candidate_files: int = 5000
+    skip_expensive_threshold: int = 2000
     rare_identifier_threshold: int = 3
     max_expansion_files: int = 20
     overhead_per_fragment: int = 18

src/treemapper/diffctx/edges/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -40,10 +41,20 @@ def get_all_builders() -> list[EdgeBuilder]:
     return [cls() for cls in all_builder_classes]
 
 
-def collect_all_edges(fragments: list[Fragment], repo_root: Path | None = None) -> tuple[EdgeDict, EdgeCategories]:
+_EXPENSIVE_CATEGORIES = frozenset({"similarity", "history"})
+
+
+def collect_all_edges(
+    fragments: list[Fragment],
+    repo_root: Path | None = None,
+    skip_expensive: bool = False,
+) -> tuple[EdgeDict, EdgeCategories]:
     all_edges: EdgeDict = {}
     edge_categories: EdgeCategories = {}
     for category, get_builders in _BUILDER_CATEGORIES:
+        if skip_expensive and category in _EXPENSIVE_CATEGORIES:
+            logging.debug("diffctx: skipping %s edge builders (skip_expensive=True)", category)
+            continue
         for cls in get_builders():
             builder = cls()
             cat = builder.category or category

src/treemapper/diffctx/graph.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 
+from .config import LIMITS
 from .edges import collect_all_edges
 from .edges.similarity.lexical import clamp_lexical_weight
 from .embeddings import _build_embedding_edges
@@ -14,6 +15,7 @@
 @dataclass
 class Graph:
     adjacency: dict[FragmentId, dict[FragmentId, float]] = field(default_factory=dict)
+    reverse_adjacency: dict[FragmentId, dict[FragmentId, float]] = field(default_factory=dict)
     nodes: set[FragmentId] = field(default_factory=set)
     edge_categories: dict[tuple[FragmentId, FragmentId], str] = field(default_factory=dict)
 
@@ -28,6 +30,12 @@ def add_edge(self, src: FragmentId, dst: FragmentId, weight: float) -> None:
             self.adjacency[src] = {}
         existing = self.adjacency[src].get(dst, 0.0)
         self.adjacency[src][dst] = max(existing, weight)
+
+        if dst not in self.reverse_adjacency:
+            self.reverse_adjacency[dst] = {}
+        existing_rev = self.reverse_adjacency[dst].get(src, 0.0)
+        self.reverse_adjacency[dst][src] = max(existing_rev, weight)
+
         self.nodes.add(src)
         self.nodes.add(dst)
 
@@ -41,20 +49,25 @@ def build_graph(fragments: list[Fragment], repo_root: Path | None = None) -> Gra
     for frag in fragments:
         graph.nodes.add(frag.id)
 
+    skip_expensive = len(fragments) > LIMITS.skip_expensive_threshold
+    if skip_expensive:
+        logging.debug("diffctx: %d fragments exceed threshold, skipping expensive edge builders", len(fragments))
+
     all_edges: dict[tuple[FragmentId, FragmentId], float] = {}
     edge_categories: dict[tuple[FragmentId, FragmentId], str] = {}
 
-    plugin_edges, plugin_categories = collect_all_edges(fragments, repo_root)
+    plugin_edges, plugin_categories = collect_all_edges(fragments, repo_root, skip_expensive=skip_expensive)
     for (src, dst), weight in plugin_edges.items():
         if weight > all_edges.get((src, dst), 0.0):
             all_edges[(src, dst)] = weight
             edge_categories[(src, dst)] = plugin_categories.get((src, dst), "generic")
 
-    embedding_edges = _build_embedding_edges(fragments, clamp_lexical_weight)
-    for (src, dst), weight in embedding_edges.items():
-        if weight > all_edges.get((src, dst), 0.0):
-            all_edges[(src, dst)] = weight
-            edge_categories[(src, dst)] = "similarity"
+    if not skip_expensive:
+        embedding_edges = _build_embedding_edges(fragments, clamp_lexical_weight)
+        for (src, dst), weight in embedding_edges.items():
+            if weight > all_edges.get((src, dst), 0.0):
+                all_edges[(src, dst)] = weight
+                edge_categories[(src, dst)] = "similarity"
 
     all_edges = _apply_hub_suppression(all_edges, edge_categories)