fix(compression): keep container bodies compressible in code handler (#890)

## Description

Two bugs in `CodeStructureHandler`'s tree-sitter path. (1) Container
nodes (class/impl/trait/decorated definitions) were marked structural
over their full span and `_spans_to_mask` never un-marks, so every
method body inside a class was preserved and compression silently
no-opped at confidence 0.95. (2) Discovered while testing:
`tree-sitter-language-pack >= 1.0` switched to a Rust binding (methods,
not attributes; `parse(str)`), so the handler raised `TypeError` on
every call and silently fell back to regex.

Closes # <!-- compression-handler review -->

## 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

- `headroom/compression/handlers/code_handler.py`: containers emit a
signature-only span (start to body start); recursion gives nested
functions their own signature/body split; decorated definitions emit no
whole-node span.
- `headroom/compression/handlers/code_handler.py`: small compat shim
supporting both the classic attribute API and the new Rust-binding
method API.
- `tests/test_compression/test_code_handler.py`: new file (the handler
had zero dedicated tests) covering class/decorated/impl body
compressibility, regex fallback, and a preservation-ratio bound.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest tests/test_compression/ -q
92 passed, 8 skipped
```

## Real Behavior Proof

- Environment: local macOS, Python 3.11, tree-sitter-language-pack 1.8.1
installed, branch `fix/code-container-bodies`.
- Exact command / steps: `pytest tests/test_compression/ -q`.
- Observed result: Class method bodies are now compressible
(preservation ratio drops from ~1.0 to roughly the signature fraction);
the tree-sitter path runs instead of falling back to regex.
- Not tested: End-to-end through the live proxy pipeline.

## 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
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Screenshots (if applicable)

N/A — library change. See Test Output.

## Additional Notes

PR 3 of 7; branched fresh from main (independent of #887/#889).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: agaonker

headroom/compression/handlers/code_handler.py

@@ -63,6 +63,50 @@ def _get_parser(language: str) -> Any:
         return _tree_sitter_parsers[language]
 
 
+# tree-sitter API compatibility. tree-sitter-language-pack switched to a
+# Rust binding (>=1.0) where node accessors are METHODS (kind(),
+# start_byte(), child(i)) and parse() takes str; the classic pybind API
+# uses attributes (.type, .start_byte, .children) and parse(bytes).
+# Without this shim the tree-sitter path raises TypeError on modern
+# installs and silently falls back to regex.
+
+
+def _ts_parse(parser: Any, content: str) -> Any:
+    try:
+        return parser.parse(content.encode("utf-8"))
+    except TypeError:
+        return parser.parse(content)
+
+
+def _ts_root(tree: Any) -> Any:
+    root = tree.root_node
+    return root() if callable(root) else root
+
+
+def _ts_kind(node: Any) -> str:
+    kind = getattr(node, "type", None)
+    if isinstance(kind, str):
+        return kind
+    return str(node.kind())
+
+
+def _ts_start_byte(node: Any) -> int:
+    start = node.start_byte
+    return int(start()) if callable(start) else int(start)
+
+
+def _ts_end_byte(node: Any) -> int:
+    end = node.end_byte
+    return int(end()) if callable(end) else int(end)
+
+
+def _ts_children(node: Any) -> list[Any]:
+    children = getattr(node, "children", None)
+    if children is not None and not callable(children):
+        return list(children)
+    return [node.child(i) for i in range(node.child_count())]
+
+
 class CodeLanguage(Enum):
     """Supported programming languages."""
 
@@ -175,6 +219,23 @@ class CodeSpan:
     ],
 }
 
+# Body child node types for container definitions (classes, impls,
+# traits). A container's span up to its body is structural (the
+# signature); the body itself is NOT marked — recursion into the body
+# emits signature spans for nested functions/methods, leaving their
+# bodies compressible.
+_CONTAINER_BODY_TYPES: frozenset[str] = frozenset(
+    {
+        "block",  # python class body
+        "statement_block",  # js/ts
+        "compound_statement",  # c/cpp
+        "class_body",  # js/ts/java class body
+        "interface_body",  # java/ts interface body
+        "declaration_list",  # rust impl/trait body
+        "enum_body",  # java enum body
+    }
+)
+
 # Import patterns for fallback
 _IMPORT_PATTERNS: dict[str, re.Pattern[str]] = {
     "python": re.compile(r"^\s*(import\s+\w+|from\s+\w+\s+import)", re.MULTILINE),
@@ -301,42 +362,43 @@ def _extract_with_tree_sitter(
             HandlerResult with mask.
         """
         parser = _get_parser(language)
-        tree = parser.parse(content.encode("utf-8"))
+        tree = _ts_parse(parser, content)
 
         # Collect structural spans
         spans: list[CodeSpan] = []
 
         def visit_node(node: Any, depth: int = 0) -> None:
             """Visit AST node and collect structural spans."""
-            node_type = node.type
+            node_type = _ts_kind(node)
             structural_types = _STRUCTURAL_NODE_TYPES.get(language, set())
+            children = _ts_children(node)
 
             # Check if this is a structural node type
             if node_type in structural_types:
                 # For functions, only the signature is structural
                 if "function" in node_type or "method" in node_type:
                     # Find the body node and exclude it
                     body_node = None
-                    for child in node.children:
-                        if child.type in ("block", "statement_block", "compound_statement"):
+                    for child in children:
+                        if _ts_kind(child) in ("block", "statement_block", "compound_statement"):
                             body_node = child
                             break
 
                     if body_node:
                         # Signature is from start to body start
                         spans.append(
                             CodeSpan(
-                                start=node.start_byte,
-                                end=body_node.start_byte,
+                                start=_ts_start_byte(node),
+                                end=_ts_start_byte(body_node),
                                 role="signature",
                                 is_structural=True,
                             )
                         )
                         # Body is compressible
                         spans.append(
                             CodeSpan(
-                                start=body_node.start_byte,
-                                end=body_node.end_byte,
+                                start=_ts_start_byte(body_node),
+                                end=_ts_end_byte(body_node),
                                 role="body",
                                 is_structural=False,
                             )
@@ -345,37 +407,75 @@ def visit_node(node: Any, depth: int = 0) -> None:
                         # No body found, preserve whole thing
                         spans.append(
                             CodeSpan(
-                                start=node.start_byte,
-                                end=node.end_byte,
+                                start=_ts_start_byte(node),
+                                end=_ts_end_byte(node),
                                 role=node_type,
                                 is_structural=True,
                             )
                         )
+                elif node_type == "decorated_definition":
+                    # Wrapper around decorator(s) + definition. Emit no
+                    # span: recursion marks the decorators and gives the
+                    # inner function its signature/body split. A whole-
+                    # node span here would preserve the function body.
+                    pass
                 else:
-                    # Non-function structural nodes
-                    spans.append(
-                        CodeSpan(
-                            start=node.start_byte,
-                            end=node.end_byte,
-                            role=node_type,
-                            is_structural=True,
+                    # Container definitions (class, impl, trait): the
+                    # signature runs to the body start; the body is NOT
+                    # marked, so nested function bodies stay compressible
+                    # (recursion emits their signature spans). Leaf
+                    # declarations (imports, type aliases, structs) have
+                    # no such body child and are preserved whole.
+                    body_node = None
+                    for child in children:
+                        if _ts_kind(child) in _CONTAINER_BODY_TYPES:
+                            body_node = child
+                            break
+
+                    if body_node is not None:
+                        spans.append(
+                            CodeSpan(
+                                start=_ts_start_byte(node),
+                                end=_ts_start_byte(body_node),
+                                role="signature",
+                                is_structural=True,
+                            )
                         )
+                    else:
+                        spans.append(
+                            CodeSpan(
+                                start=_ts_start_byte(node),
+                                end=_ts_end_byte(node),
+                                role=node_type,
+                                is_structural=True,
+                            )
+                        )
+            elif node_type == "decorator":
+                # Decorators are structural (preserved) on their own so
+                # the decorated_definition wrapper doesn't need a span.
+                spans.append(
+                    CodeSpan(
+                        start=_ts_start_byte(node),
+                        end=_ts_end_byte(node),
+                        role="decorator",
+                        is_structural=True,
                     )
+                )
             elif node_type == "comment" and self.preserve_comments:
                 spans.append(
                     CodeSpan(
-                        start=node.start_byte,
-                        end=node.end_byte,
+                        start=_ts_start_byte(node),
+                        end=_ts_end_byte(node),
                         role="comment",
                         is_structural=True,
                     )
                 )
 
             # Recurse into children
-            for child in node.children:
+            for child in children:
                 visit_node(child, depth + 1)
 
-        visit_node(tree.root_node)
+        visit_node(_ts_root(tree))
 
         # Build mask from spans
         mask = self._spans_to_mask(spans, len(content))

tests/test_compression/test_code_handler.py

@@ -0,0 +1,151 @@
+"""Tests for code structure handler."""
+
+import pytest
+
+from headroom.compression.handlers.code_handler import (
+    CodeStructureHandler,
+    is_tree_sitter_available,
+)
+
+requires_tree_sitter = pytest.mark.skipif(
+    not is_tree_sitter_available(),
+    reason="tree-sitter-language-pack not installed",
+)
+
+
+class TestCanHandle:
+    @pytest.fixture
+    def handler(self):
+        return CodeStructureHandler()
+
+    def test_detects_python(self, handler):
+        assert handler.can_handle("def foo():\n    pass\n") is True
+
+    def test_detects_javascript(self, handler):
+        assert handler.can_handle("function foo() { return 1; }") is True
+
+    def test_rejects_prose(self, handler):
+        assert handler.can_handle("This is a plain sentence.") is False
+
+
+class TestRegexFallback:
+    """Regex path runs regardless of tree-sitter availability."""
+
+    @pytest.fixture
+    def handler(self):
+        return CodeStructureHandler(use_tree_sitter=False)
+
+    def test_python_signature_preserved_body_compressible(self, handler):
+        code = "def hello(name: str) -> str:\n    message = name\n    return message\n"
+        result = handler.get_mask(code, language="python")
+
+        assert result.metadata["parser"] == "regex"
+        sig = "def hello(name: str) -> str:"
+        start = code.index(sig)
+        assert all(result.mask.mask[i] for i in range(start, start + len(sig)))
+
+        body_char = code.index("message = name")
+        assert result.mask.mask[body_char] is False
+
+    def test_python_import_preserved(self, handler):
+        code = "import os\n\nx = 1\n"
+        result = handler.get_mask(code, language="python")
+        assert all(result.mask.mask[i] for i in range(len("import os")))
+
+
+@requires_tree_sitter
+class TestTreeSitterContainers:
+    """Container bodies must stay compressible (signature-only spans).
+
+    Regression: class_definition / decorated_definition / impl_item were
+    marked structural over their FULL span, so every method body inside a
+    class (i.e. most real code) was preserved and compression no-opped at
+    confidence 0.95.
+    """
+
+    @pytest.fixture
+    def handler(self):
+        return CodeStructureHandler()
+
+    def test_class_method_bodies_compressible(self, handler):
+        code = (
+            "class Foo:\n"
+            "    def method_a(self):\n"
+            "        body_line_a = 1\n"
+            "        return body_line_a\n"
+            "\n"
+            "    def method_b(self):\n"
+            "        body_line_b = 2\n"
+            "        return body_line_b\n"
+        )
+        result = handler.get_mask(code, language="python")
+        assert result.metadata["parser"] == "tree-sitter"
+
+        # Class signature and method signatures preserved
+        assert all(result.mask.mask[i] for i in range(len("class Foo:")))
+        sig = "def method_a(self):"
+        start = code.index(sig)
+        assert all(result.mask.mask[i] for i in range(start, start + len(sig)))
+
+        # Method bodies compressible
+        for body in ("body_line_a = 1", "body_line_b = 2"):
+            start = code.index(body)
+            assert not any(result.mask.mask[i] for i in range(start, start + len(body))), (
+                f"method body {body!r} must be compressible"
+            )
+
+    def test_decorated_function_body_compressible(self, handler):
+        code = "@decorator\ndef decorated():\n    body_line = 4\n    return body_line\n"
+        result = handler.get_mask(code, language="python")
+
+        # Decorator and signature preserved
+        assert all(result.mask.mask[i] for i in range(len("@decorator")))
+        sig = "def decorated():"
+        start = code.index(sig)
+        assert all(result.mask.mask[i] for i in range(start, start + len(sig)))
+
+        # Body compressible
+        start = code.index("body_line = 4")
+        assert not any(result.mask.mask[i] for i in range(start, start + len("body_line = 4"))), (
+            "decorated function body must be compressible"
+        )
+
+    def test_module_function_body_compressible(self, handler):
+        code = "def standalone():\n    body_line = 3\n    return body_line\n"
+        result = handler.get_mask(code, language="python")
+
+        start = code.index("body_line = 3")
+        assert not any(result.mask.mask[i] for i in range(start, start + len("body_line = 3")))
+
+    def test_rust_impl_method_bodies_compressible(self, handler):
+        code = (
+            "struct Foo { x: i32 }\n"
+            "impl Foo {\n"
+            "    fn method(&self) -> i32 {\n"
+            "        let body_line = 5;\n"
+            "        body_line\n"
+            "    }\n"
+            "}\n"
+        )
+        result = handler.get_mask(code, language="rust")
+
+        # impl signature preserved
+        start = code.index("impl Foo")
+        assert all(result.mask.mask[i] for i in range(start, start + len("impl Foo")))
+
+        # method body compressible
+        start = code.index("let body_line = 5;")
+        assert not any(
+            result.mask.mask[i] for i in range(start, start + len("let body_line = 5;"))
+        ), "impl method body must be compressible"
+
+    def test_preservation_ratio_sane_for_class_code(self, handler):
+        """A class with substantial method bodies should NOT preserve
+        everything — the whole point of the handler."""
+        body = "\n".join(f"        line_{i} = {i}" for i in range(20))
+        code = f"class Big:\n    def method(self):\n{body}\n        return 0\n"
+        result = handler.get_mask(code, language="python")
+        assert result.preservation_ratio < 0.5, (
+            f"class code preserved {result.preservation_ratio:.0%} — "
+            "container bodies are leaking into the structural mask"
+        )