fix: address multiple diffctx bugs and tighten similarity thresholds

- Fix PPR iterator consumption bug (dict.values() exhausted after comprehension)
- Differentiate TS vs TSX parsing (use correct tree-sitter grammar)
- Fix Language double-wrapping in tree-sitter parser
- Add backtick support in bracket balancing for JS template strings
- Fix Go edge builder case mismatch in package name matching
- Replace rglob with git ls-files for faster candidate file collection
- Tighten lexical similarity: min_similarity 0.1->0.30, top_k 10->5
- Restrict config-to-code edges: require exact match, min key length 6
- Fix JSON fragmentation: handle files with no top-level keys

Author: nikolay-e

src/treemapper/diffctx/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import subprocess
 from collections import defaultdict
 from pathlib import Path
 from typing import Any
@@ -340,6 +341,19 @@ def _is_candidate_file(file_path: Path, root_dir: Path, included_set: set[Path],
 
 
 def _collect_candidate_files(root_dir: Path, included_set: set[Path], combined_spec: pathspec.PathSpec) -> list[Path]:
+    try:
+        result = subprocess.run(
+            ["git", "ls-files", "-z"],
+            cwd=root_dir,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode == 0 and result.stdout:
+            files = [root_dir / f for f in result.stdout.split("\0") if f]
+            return [f for f in files if _is_candidate_file(f, root_dir, included_set, combined_spec)]
+    except (subprocess.SubprocessError, OSError):
+        pass
     return [f for f in root_dir.rglob("*") if _is_candidate_file(f, root_dir, included_set, combined_spec)]
 
 

src/treemapper/diffctx/config/limits.py

@@ -22,15 +22,15 @@ class PPRConfig:
 
 @dataclass(frozen=True)
 class LexicalConfig:
-    min_similarity: float = 0.1
+    min_similarity: float = 0.30
     hub_percentile: float = 0.95
-    top_k_neighbors: int = 10
-    max_df_ratio: float = 0.20
-    min_idf: float = 1.6
-    max_postings: int = 200
-    weight_min: float = 0.1
-    weight_max: float = 0.2
-    backward_factor: float = 0.7
+    top_k_neighbors: int = 5
+    max_df_ratio: float = 0.15
+    min_idf: float = 2.0
+    max_postings: int = 100
+    weight_min: float = 0.05
+    weight_max: float = 0.15
+    backward_factor: float = 0.5
 
 
 @dataclass(frozen=True)

src/treemapper/diffctx/constants.py

@@ -14,16 +14,6 @@
 
 
 def expand_config_key(key: str) -> set[str]:
-    expanded: set[str] = {key}
-    parts = key.split("_")
-    for part in parts:
-        if len(part) >= 3:
-            expanded.add(part)
-    for prefix in CONFIG_KEY_COMMON_PREFIXES:
-        if key.startswith(prefix + "_") and len(key) > len(prefix) + 1:
-            stripped = key[len(prefix) + 1 :]
-            expanded.add(stripped)
-            for sub in stripped.split("_"):
-                if len(sub) >= 3:
-                    expanded.add(sub)
-    return expanded
+    if len(key) < 6:
+        return set()
+    return {key}

src/treemapper/diffctx/edges/config/cicd.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import re
 from pathlib import Path
 
@@ -18,8 +19,10 @@
 _JENKINS_SH_RE = re.compile(r"sh\s*(?:\(['\"]|['\"])(.+?)['\"]\)?", re.MULTILINE | re.DOTALL)
 _JENKINS_SCRIPT_RE = re.compile(r"script\s*\{([^}]+)\}", re.MULTILINE | re.DOTALL)
 
-_SCRIPT_CALL_RE = re.compile(r"(?:bash|sh|python|python3|node|npm|yarn|pnpm|make|go|cargo|dotnet|mvn|gradle)\s+([^\s;&|]+)")
-_FILE_REF_RE = re.compile(r"(?:\.\/|scripts\/|bin\/|tools\/)([a-zA-Z0-9_.-]+(?:\.(?:sh|py|js|ts|rb))?)")
+_SCRIPT_CALL_RE = re.compile(
+    r"(?:bash|sh|python|python3|node|npm|yarn|pnpm|make|go|cargo|dotnet|mvn|gradle|pytest|ruff|mypy|black|isort|flake8)\s+([^\s;&|]+)"
+)
+_FILE_REF_RE = re.compile(r"(?:\.\/|scripts\/|bin\/|tools\/|src\/|tests\/)([a-zA-Z0-9_.-]+(?:\.(?:sh|py|js|ts|rb))?)")
 
 
 def _is_github_actions(path: Path) -> bool:
@@ -49,6 +52,14 @@ def _is_azure_pipelines(path: Path) -> bool:
     return name in {"azure-pipelines.yml", "azure-pipelines.yaml"} or name.startswith("azure-pipeline")
 
 
+def _is_tox(path: Path) -> bool:
+    return path.name.lower() == "tox.ini"
+
+
+def _is_nox(path: Path) -> bool:
+    return path.name.lower() == "noxfile.py"
+
+
 def _is_ci_file(path: Path) -> bool:
     return any(
         [
@@ -58,6 +69,8 @@ def _is_ci_file(path: Path) -> bool:
             _is_circleci(path),
             _is_travis(path),
             _is_azure_pipelines(path),
+            _is_tox(path),
+            _is_nox(path),
         ]
     )
 
@@ -112,6 +125,63 @@ def _extract_jenkins_refs(content: str) -> set[str]:
     return refs
 
 
+def _extract_tox_refs(content: str) -> set[str]:
+    refs: set[str] = set()
+    # Extract deps
+    for match in re.finditer(r"^\s*deps\s*=\s*(.+)$", content, re.MULTILINE):
+        deps = match.group(1).split()
+        refs.update(d.strip() for d in deps if d.strip())
+
+    # Extract commands
+    for match in re.finditer(r"^\s*commands\s*=\s*(.+)$", content, re.MULTILINE):
+        cmd = match.group(1)
+        # Split by whitespace to find potential paths
+        # Examples:
+        #   commands = pytest {posargs}
+        #   commands = ruff check src/
+        #   commands = python -m pytest tests/
+        parts = cmd.split()
+        for p in parts:
+            # Strip quoting
+            p = p.strip("'\"")
+            # Remove tox specific vars like {posargs}
+            p = re.sub(r"\{[^}]+\}", "", p)
+
+            if not p or p.startswith("-"):
+                continue
+
+            # Heuristic: if it looks like a path (contains / or .py/.ini etc)
+            # or if we are permissive and just add everything that looks like an ident
+            # Given discover_files_by_refs filters candidates, being permissive is safer.
+            refs.add(p)
+
+        refs.update(_extract_script_refs(cmd))
+
+    return refs
+
+
+def _extract_nox_refs(content: str) -> set[str]:
+    refs: set[str] = set()
+    # Extract session.run calls
+    # session.run("cmd", "arg1", ...)
+    for match in re.finditer(r"session\.run\(([^)]+)\)", content):
+        args_str = match.group(1)
+        # Simple extraction of string literals
+        args = re.findall(r'["\']([^"\']+)["\']', args_str)
+        for arg in args:
+            refs.add(arg)
+            refs.update(_extract_script_refs(arg))
+
+    # Extract session.install calls
+    for match in re.finditer(r"session\.install\(([^)]+)\)", content):
+        args_str = match.group(1)
+        args = re.findall(r'["\']([^"\']+)["\']', args_str)
+        for arg in args:
+            refs.add(arg)
+
+    return refs
+
+
 class CICDEdgeBuilder(EdgeBuilder):
     weight = 0.55
     script_weight = 0.60
@@ -135,19 +205,29 @@ def discover_related_files(
             except (OSError, UnicodeDecodeError):
                 continue
 
+            local_refs = set()
             if _is_github_actions(ci):
-                refs.update(_extract_gha_refs(content))
+                local_refs.update(_extract_gha_refs(content))
             elif _is_gitlab_ci(ci):
-                refs.update(_extract_gitlab_refs(content))
+                local_refs.update(_extract_gitlab_refs(content))
             elif _is_jenkinsfile(ci):
-                refs.update(_extract_jenkins_refs(content))
+                local_refs.update(_extract_jenkins_refs(content))
+            elif _is_tox(ci):
+                local_refs.update(_extract_tox_refs(content))
+            elif _is_nox(ci):
+                local_refs.update(_extract_nox_refs(content))
             else:
-                refs.update(_extract_script_refs(content))
+                local_refs.update(_extract_script_refs(content))
 
             if any(cmd in content.lower() for cmd in ["npm", "yarn", "pnpm"]):
-                refs.add("package.json")
+                local_refs.add("package.json")
+
+            logging.debug("CICD refs for %s: %s", ci.name, local_refs)
+            refs.update(local_refs)
 
-        return discover_files_by_refs(refs, changed_files, all_candidate_files)
+        discovered = discover_files_by_refs(refs, changed_files, all_candidate_files)
+        logging.debug("CICD discovered for %s: %s", [c.name for c in ci_files], [d.name for d in discovered])
+        return discovered
 
     def build(self, fragments: list[Fragment], repo_root: Path | None = None) -> EdgeDict:
         ci_frags = [f for f in fragments if _is_ci_file(f.path)]
@@ -171,6 +251,10 @@ def _extract_refs(self, ci: Fragment) -> set[str]:
             return _extract_gitlab_refs(ci.content)
         if _is_jenkinsfile(ci.path):
             return _extract_jenkins_refs(ci.content)
+        if _is_tox(ci.path):
+            return _extract_tox_refs(ci.content)
+        if _is_nox(ci.path):
+            return _extract_nox_refs(ci.content)
         return _extract_script_refs(ci.content)
 
     def _link_refs(self, ci_id: FragmentId, refs: set[str], idx: FragmentIndex, edges: EdgeDict) -> None:

src/treemapper/diffctx/edges/config/generic.py

@@ -86,9 +86,7 @@ def discover_related_files(
     def _build_key_patterns(self, keys: set[str]) -> dict[str, re.Pattern[str]]:
         patterns: dict[str, re.Pattern[str]] = {}
         for key in keys:
-            if len(key) >= 4:
-                patterns[key] = re.compile(rf"\b\w*{re.escape(key)}\w*\b", re.IGNORECASE)
-            else:
+            if len(key) >= 6:
                 patterns[key] = re.compile(rf"\b{re.escape(key)}\b", re.IGNORECASE)
         return patterns
 

src/treemapper/diffctx/edges/semantic/go.py

@@ -88,7 +88,7 @@ def _build_indices(
         func_defs: dict[str, list[FragmentId]] = defaultdict(list)
 
         for f in go_frags:
-            pkg = _get_package_name(f.path)
+            pkg = _get_package_name(f.path).lower()
             pkg_to_frags[pkg].append(f.id)
 
             if repo_root:
@@ -143,7 +143,7 @@ def _link_import_by_package(
         pkg_to_frags: dict[str, list[FragmentId]],
         edges: EdgeDict,
     ) -> None:
-        imp_pkg = imp.split("/")[-1]
+        imp_pkg = imp.split("/")[-1].lower()
         for pkg, frag_ids in pkg_to_frags.items():
             if pkg == imp_pkg:
                 self.add_edges_from_ids(gf_id, frag_ids, self.import_weight, edges)
@@ -190,7 +190,7 @@ def _link_same_package(
         pkg_to_frags: dict[str, list[FragmentId]],
         edges: EdgeDict,
     ) -> None:
-        current_pkg = _get_package_name(gf.path)
+        current_pkg = _get_package_name(gf.path).lower()
         for fid in pkg_to_frags.get(current_pkg, []):
             if fid != gf.id:
                 self.add_edge(edges, gf.id, fid, self.same_package_weight)

src/treemapper/diffctx/languages.py

@@ -164,11 +164,11 @@
     ".pyw": "python",
     ".pyi": "python",
     ".js": "javascript",
-    ".jsx": "javascript",
+    ".jsx": "jsx",
     ".mjs": "javascript",
     ".cjs": "javascript",
     ".ts": "typescript",
-    ".tsx": "typescript",
+    ".tsx": "tsx",
     ".mts": "typescript",
     ".cts": "typescript",
     ".go": "go",

src/treemapper/diffctx/parsers/base.py

@@ -35,7 +35,7 @@ def _process_char_in_string(char: str, string_char: str, escape_count: int) -> t
 
 
 def _process_char_outside_string(char: str, stack: list[str]) -> tuple[bool, str]:
-    if char in ('"', "'"):
+    if char in ('"', "'", "`"):
         return True, char
     if char in _BRACKET_PAIRS:
         stack.append(_BRACKET_PAIRS[char])

src/treemapper/diffctx/parsers/config.py

@@ -82,15 +82,16 @@ def fragment(self, path: Path, content: str) -> list[Fragment]:
         elif suffix == ".toml":
             key_re = re.compile(r"^\[([a-zA-Z_][a-zA-Z0-9_.-]*)\]")
         else:
-            key_re = re.compile(r'^\s{0,2}"([^"]+)":\s*')
+            key_re = re.compile(r'^\s{0,4}"([^"]+)":\s*')
 
         boundaries: list[int] = []
         for i, line in enumerate(lines):
             if key_re.match(line):
                 boundaries.append(i)
 
-        if len(boundaries) < 2:
-            return []
+        if not boundaries:
+            frag = create_fragment_from_lines(path, lines, 1, len(lines), "config", "data")
+            return [frag] if frag else []
 
         fragments: list[Fragment] = []
         boundaries.append(len(lines))

src/treemapper/diffctx/parsers/tree_sitter.py

@@ -13,6 +13,7 @@
 _DEFINITION_NODE_TYPES = {
     "python": {"function_definition", "class_definition", "decorated_definition"},
     "javascript": {"function_declaration", "class_declaration", "method_definition", "arrow_function", "variable_declarator"},
+    "jsx": {"function_declaration", "class_declaration", "method_definition", "arrow_function", "variable_declarator"},
     "typescript": {
         "function_declaration",
         "class_declaration",
@@ -23,6 +24,16 @@
         "enum_declaration",
         "variable_declarator",
     },
+    "tsx": {
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "arrow_function",
+        "interface_declaration",
+        "type_alias_declaration",
+        "enum_declaration",
+        "variable_declarator",
+    },
     "go": {"function_declaration", "method_declaration", "type_declaration"},
     "rust": {"function_item", "impl_item", "struct_item", "enum_item", "trait_item"},
     "java": {"method_declaration", "class_declaration", "interface_declaration", "enum_declaration"},
@@ -69,7 +80,9 @@
 _LANG_MODULES = {
     "python": "tree_sitter_python",
     "javascript": "tree_sitter_javascript",
+    "jsx": "tree_sitter_javascript",
     "typescript": "tree_sitter_typescript",
+    "tsx": "tree_sitter_typescript",
     "go": "tree_sitter_go",
     "rust": "tree_sitter_rust",
     "java": "tree_sitter_java",
@@ -98,15 +111,20 @@ def _get_parser(self, lang: str) -> Parser:
         module_name = _LANG_MODULES[lang]
         ts_lang_module = importlib.import_module(module_name)
 
-        if lang == "typescript":
+        if lang == "tsx":
             ts_lang = ts_lang_module.language_tsx()
+        elif lang == "typescript":
+            ts_lang = ts_lang_module.language_typescript()
         elif hasattr(ts_lang_module, "language"):
             ts_lang = ts_lang_module.language()
         else:
             ts_lang = ts_lang_module
 
         parser = Parser()
-        parser.language = Language(ts_lang)
+        if isinstance(ts_lang, Language):
+            parser.language = ts_lang
+        else:
+            parser.language = Language(ts_lang)
         self._parsers[lang] = parser
         return parser