diff --git a/CHANGELOG.md b/CHANGELOG.md index b78bca69e..92cc8de4a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Bug Fixes +* **code:** keep Python `from __future__` imports before executable code during AST compression and validate compressed Python with `compile(..., "exec")` so compile-time syntax rules are enforced ([#1233](https://github.com/chopratejas/headroom/issues/1233)). * **proxy:** report real input tokens on the streaming `message_start` event for LiteLLM/Bedrock-backed requests. LiteLLM streaming never surfaces prompt tokens mid-stream, so `message_start.usage.input_tokens` was always `0`; Anthropic clients (e.g. Claude Code) read input-token metrics from that event, underreporting token usage by ~99% in OTel/CloudWatch dashboards. The Bedrock streamer now backfills `input_tokens` with the count Headroom actually sent upstream when the backend leaves it unset, preserving any non-zero value the backend genuinely reports ([#1132](https://github.com/chopratejas/headroom/issues/1132)). * **proxy:** force Responses API `store=true` when Headroom injects memory tools so `previous_response_id` continuations work after memory tool calls from clients that requested `store=false` ([#1103](https://github.com/chopratejas/headroom/pull/1103)). * **proxy:** build SSL contexts for custom CA bundles so enterprise/private PKI roots work with Python/OpenSSL strict verification. diff --git a/headroom/transforms/code_compressor.py b/headroom/transforms/code_compressor.py index 9cd338219..d762cd669 100644 --- a/headroom/transforms/code_compressor.py +++ b/headroom/transforms/code_compressor.py @@ -261,7 +261,9 @@ class LangConfig: _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = { CodeLanguage.PYTHON: LangConfig( - import_nodes=frozenset({"import_statement", "import_from_statement"}), + import_nodes=frozenset( + {"future_import_statement", "import_statement", "import_from_statement"} + ), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset({"class_definition"}), type_nodes=frozenset({"type_alias_statement"}), @@ -1660,7 +1662,9 @@ class CodeAwareCompressor(Transform): method_compressed = None for deco_child in child.children: if deco_child.type == "decorator": - decorator_lines.append(_get_node_text(deco_child, code)) + deco_start = deco_child.start_point[0] + deco_end = deco_child.end_point[0] + decorator_lines.append("\n".join(code_lines[deco_start : deco_end + 1])) elif deco_child.type in lang_config.function_nodes: method_compressed = self._compress_function_ast( deco_child, code, language, lang_config, body_limits, analysis @@ -1765,6 +1769,12 @@ class CodeAwareCompressor(Transform): (tokens the parser expected but didn't find). """ try: + if language == CodeLanguage.PYTHON: + import ast + + ast.parse(code) + compile(code, "", "exec") + parser = _get_parser(language.value) tree = parser.parse(bytes(code, "utf-8")) return not _has_syntax_issues(tree.root_node) diff --git a/tests/test_transforms/test_code_compressor.py b/tests/test_transforms/test_code_compressor.py index d7bdaa111..a94a4a235 100644 --- a/tests/test_transforms/test_code_compressor.py +++ b/tests/test_transforms/test_code_compressor.py @@ -9,6 +9,7 @@ Comprehensive tests covering: - Edge cases: Empty content, unavailable dependency, fallbacks """ +import textwrap from unittest.mock import patch import pytest @@ -852,6 +853,85 @@ def main(): except SyntaxError: pytest.fail("Compressed output has invalid Python syntax") + def test_python_future_import_stays_at_module_start(self): + """Compressed Python keeps future imports before executable statements.""" + config = CodeCompressorConfig( + min_tokens_for_compression=10, + target_compression_rate=0.2, + max_body_lines=3, + enable_ccr=False, + ) + compressor = CodeAwareCompressor(config) + code = textwrap.dedent( + """ + from __future__ import annotations + + from dataclasses import dataclass + from typing import Any, Callable, Iterable + + + def traced(label: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + def decorate(fn: Callable[..., Any]) -> Callable[..., Any]: + async def wrapper(*args: Any, **kwargs: Any) -> Any: + return await fn(*args, **kwargs) + + return wrapper + + return decorate + + + @dataclass(slots=True) + class Event: + kind: str + payload: dict[str, Any] + retries: int = 0 + + @property + def important(self) -> bool: + return self.kind in {"error", "retry"} or self.retries > 2 + + + class EventRouter: + def __init__(self, sinks: dict[str, Callable[[Event], Any]]) -> None: + self.sinks = sinks + self.history: list[tuple[str, bool]] = [] + + @traced("route") + async def route(self, events: Iterable[Event]) -> list[str]: + accepted: list[str] = [] + for event in events: + match event: + case Event(kind="error", payload={"code": code, "message": msg}, retries=r) if r > 1: + destination = "pager" + accepted.append(f"{destination}:{code}:{msg}") + case Event(kind=kind, payload=payload) if (route := payload.get("route")): + destination = str(route) + accepted.append(f"{destination}:{kind}") + case _: + destination = "dead_letter" + accepted.append(destination) + + self.history.append((destination, event.important)) + + return [item for item in accepted if item] + """ + ) + + result = compressor.compress(code, language="python") + + assert result.syntax_valid is True + future_import_index = result.compressed.index("from __future__ import annotations") + first_executable_index = min( + result.compressed.index("@dataclass"), + result.compressed.index("def traced"), + result.compressed.index("class EventRouter"), + ) + assert future_import_index < first_executable_index + try: + compile(result.compressed, "", "exec") + except SyntaxError as exc: + pytest.fail(f"Compressed output has invalid Python syntax: {exc}\n{result.compressed}") + def test_tree_sitter_loaded_after_compression(self): """Parser is loaded after compression.""" config = CodeCompressorConfig(