From 5e0bb697254b7ec87e3191fa73031bde9321a79c Mon Sep 17 00:00:00 2001 From: Rocker Zhang Date: Tue, 23 Jun 2026 12:00:57 +0800 Subject: [PATCH] fix(code): verify a real parse in tree-sitter availability check (#1231) (#1299) ## Description `is_tree_sitter_available()` / `_check_tree_sitter_available()` in `headroom/transforms/code_compressor.py` return `True` based on importing `tree_sitter_language_pack` alone, without ever constructing a parser or attempting a parse. When the installed pack/parser combination is ABI-incompatible, `get_parser`/`parse` raises at runtime; the caller catches it and silently falls back to the lossy text compressor, while the availability flag and startup banner still report code-aware as on. This is the defensive half that the `<1.0` pin in #1234 does not cover: if that cap is ever lifted, the availability signal silently lies again. Follow-up to #1231. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Make `_check_tree_sitter_available()` construct a parser and parse a tiny snippet, returning `True` only if it yields a real `module` AST instead of trusting an import. - Add `_tree_sitter_importable()` for the cheap import-only probe, and use it to guard parser construction so the real-parse check cannot recurse. - Add tests asserting the check is `False` when parsing raises and `True` on a real parse, plus that AST compression runs for python/rust without falling back. ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text # pytest tests/test_transforms/test_code_compressor.py -> passed locally (tree-sitter-language-pack 0.13.0) # ruff check . and ruff format --check . pass locally on the rebased branch. # Full pytest suite / mypy not run locally; left to CI. ``` ## Real Behavior Proof - Environment: local repo on tree-sitter-language-pack 0.13.0, tree-sitter 0.25.2, Python 3.12, Linux - Exact command / steps: call `is_tree_sitter_available()`, then run `pytest tests/test_transforms/test_code_compressor.py` - Observed result: with a working pack the probe parses and returns `True` (code-aware runs, strategy `CODE_AWARE` rather than the kompress fallback); the new `test_check_tree_sitter_available_false_when_parse_broken` confirms that when parsing raises the check now returns `False` instead of the old import-only `True`, so the lossy fallback is no longer entered silently. - Not tested: reproducing the specific ABI-incompatible 1.x pack combo against a live install (covered instead by a mocked broken parse in the test) ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable --- headroom/transforms/code_compressor.py | 44 ++++++- tests/test_transforms/test_code_compressor.py | 116 ++++++++++++++++++ 2 files changed, 154 insertions(+), 6 deletions(-) diff --git a/headroom/transforms/code_compressor.py b/headroom/transforms/code_compressor.py index 984661d34..9cd338219 100644 --- a/headroom/transforms/code_compressor.py +++ b/headroom/transforms/code_compressor.py @@ -61,18 +61,47 @@ _tree_sitter_local = threading.local() def _check_tree_sitter_available() -> bool: - """Check if tree-sitter packages are available.""" + """Check if tree-sitter is available *and actually parses*. + + The mere presence of ``tree_sitter_language_pack`` is not enough: prior + versions of this code green-lit a code path that raised ``TypeError`` at + parse time and silently fell back to a lossy stripper. To stop misleading + callers, we now verify an end-to-end parse of a tiny snippet and only + return ``True`` if it yields a real AST. + """ global _tree_sitter_available if _tree_sitter_available is None: try: - import tree_sitter_language_pack # noqa: F401 - - _tree_sitter_available = True - except ImportError: + parser = _get_parser("python") + tree = parser.parse(b"def _probe():\n return 1\n") + root = tree.root_node + # A real parse yields a non-error root with children. + _tree_sitter_available = ( + root is not None + and root.type == "module" + and root.child_count > 0 + and not _has_syntax_issues(root) + ) + except Exception: _tree_sitter_available = False return _tree_sitter_available +def _tree_sitter_importable() -> bool: + """Return True if the tree-sitter grammar pack can be imported. + + This only checks importability (cheap, no parse). Use + :func:`_check_tree_sitter_available` for the stronger "parsing actually + works" guarantee. + """ + try: + import tree_sitter_language_pack # noqa: F401 + + return True + except ImportError: + return False + + def _get_parser(language: str) -> Any: """Get a tree-sitter parser for the given language. @@ -100,7 +129,10 @@ def _get_parser(language: str) -> Any: ImportError: If tree-sitter is not installed. ValueError: If language is not supported. """ - if not _check_tree_sitter_available(): + # NOTE: guard on importability (not _check_tree_sitter_available), because + # _check_tree_sitter_available now performs a real end-to-end parse via + # _get_parser; guarding on it here would recurse. + if not _tree_sitter_importable(): raise ImportError( "tree-sitter is not installed. Install with: pip install headroom-ai[code]\n" "This adds ~50MB for tree-sitter grammars." diff --git a/tests/test_transforms/test_code_compressor.py b/tests/test_transforms/test_code_compressor.py index 95da801af..d7bdaa111 100644 --- a/tests/test_transforms/test_code_compressor.py +++ b/tests/test_transforms/test_code_compressor.py @@ -1280,3 +1280,119 @@ function _internalDebug(msg) { """semantic_analysis is True by default in config.""" config = CodeCompressorConfig() assert config.semantic_analysis is True + + +# ============================================================================= +# Regression: tree-sitter ABI mismatch (real AST must run, no silent fallback) +# ============================================================================= + + +@pytest.mark.skipif(not TREE_SITTER_INSTALLED, reason="tree-sitter grammar pack not installed") +class TestRealASTRuns: + """Guards against the regression where the code-aware compressor silently + fell back to a lossy stripper because ``_get_parser`` built a stock + ``tree_sitter.Parser`` and assigned it a foreign grammar-pack ``Language`` + (raising ``TypeError`` that was swallowed into a fallback). + """ + + def _compressor(self): + return CodeAwareCompressor( + CodeCompressorConfig( + min_tokens_for_compression=10, + enable_ccr=False, + ) + ) + + def test_get_parser_returns_stock_node_api(self): + """The parser must yield nodes with the stock tree_sitter property API + that the tree-walking code relies on (``.type``/``.children``/...).""" + from headroom.transforms.code_compressor import _get_parser + + parser = _get_parser("python") + tree = parser.parse(b"def foo(x):\n return x + 1\n") + root = tree.root_node + + # Property access (NOT method calls) — the old pack binding exposed + # methods like ``.kind()`` which would break every call site. + assert root.type == "module" + assert root.child_count >= 1 + assert isinstance(root.children, list) + + func = root.children[0] + assert func.type == "function_definition" + assert isinstance(func.start_byte, int) + assert isinstance(func.end_byte, int) + # start_point must be index-able like a (row, col) tuple. + assert func.start_point[0] == 0 + assert b"def foo" in func.text + + def test_check_tree_sitter_available_verifies_real_parse(self): + """``_check_tree_sitter_available`` must only return True when an actual + parse succeeds — not merely when the package imports.""" + import headroom.transforms.code_compressor as cc + + cc._tree_sitter_available = None # reset memoized result + assert cc._check_tree_sitter_available() is True + + def test_check_tree_sitter_available_false_when_parse_broken(self): + """If parsing raises (e.g. the old foreign-Language bug), availability + must report False instead of green-lighting the broken path.""" + import headroom.transforms.code_compressor as cc + + cc._tree_sitter_available = None + with patch.object(cc, "_get_parser", side_effect=TypeError("boom")): + assert cc._check_tree_sitter_available() is False + cc._tree_sitter_available = None # reset for other tests + + def test_ast_runs_for_python_no_fallback(self): + """A supported language must be compressed via real AST, not the + UNKNOWN-language Kompress fallback.""" + result = self._compressor().compress(_payment_processing_code(), language="python") + + # The fallback path forces language=UNKNOWN and syntax_valid=False. + # Real AST keeps the detected language and guarantees valid syntax. + assert result.language == CodeLanguage.PYTHON + assert result.syntax_valid is True + compile(result.compressed, "", "exec") + + def test_ast_preserves_structure_for_python(self): + """AST output retains signatures/scopes/imports (unlike the old + whitespace garble).""" + code = ( + "import math\n" + "\n" + "def compute(values):\n" + " total = 0\n" + " for v in values:\n" + " total += v * v\n" + " total -= 1\n" + " total *= 2\n" + " total //= 3\n" + " return math.sqrt(total)\n" + ) + result = self._compressor().compress(code, language="python") + + assert result.language == CodeLanguage.PYTHON + assert result.syntax_valid is True + # Structure markers survive compression. + assert "import math" in result.compressed + assert "def compute(values):" in result.compressed + # Output is still valid Python. + compile(result.compressed, "", "exec") + + def test_ast_runs_for_rust_no_fallback(self): + """A second supported language (Rust) also runs through real AST.""" + code = ( + "pub fn add(a: i64, b: i64) -> i64 {\n" + " let mut acc = a;\n" + " acc += b;\n" + " acc -= 0;\n" + " acc\n" + "}\n" + ) + result = self._compressor().compress(code, language="rust") + + assert result.language == CodeLanguage.RUST + assert result.syntax_valid is True + # Signature is preserved verbatim. + assert "pub fn add(a: i64, b: i64) -> i64" in result.compressed