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
This commit is contained in:
Rocker Zhang 2026-06-23 12:00:57 +08:00 committed by GitHub
parent cdfeeacc63
commit 5e0bb69725
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 154 additions and 6 deletions

View file

@ -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."

View file

@ -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, "<test>", "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, "<test>", "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