implement token reduction with headroom

.env.example

@@ -70,3 +70,23 @@ WEB_ALLOWED_USERS=octocat,hubot
 # WEB_CLONE_CACHE_TTL_SECONDS=604800
 # Shallow-fetch depth of the PR head.
 # WEB_CLONE_DEPTH=50
+
+# --- Context compression (optional; headroom-ai) ---
+# Compress token-heavy context (tool outputs, older turns) before each LLM
+# call. Opt-in: install the extra with `pip install '.[headroom]'`. No-op if
+# the package is missing or a compression call fails.
+# HEADROOM_COMPRESS=false
+# Keep-ratio for text compression (e.g. 0.5 keeps ~50%). Empty = let headroom decide.
+# HEADROOM_TARGET_RATIO=
+# "true" to also compress user messages (the annotated diff). Off keeps cited lines intact.
+# HEADROOM_COMPRESS_USER_MESSAGES=false
+# "false" to leave system messages uncompressed.
+# HEADROOM_COMPRESS_SYSTEM_MESSAGES=true
+# Number of most-recent messages never compressed.
+# HEADROOM_PROTECT_RECENT=4
+# Skip compressing messages shorter than this many tokens.
+# HEADROOM_MIN_TOKENS=250
+# Kompress model id for ML compression, or "disabled" to skip it.
+# HEADROOM_KOMPRESS_MODEL=
+# Model context window (tokens) used for sizing.
+# HEADROOM_MODEL_LIMIT=200000

README.md

@@ -132,10 +132,38 @@ All modes share the same settings — as **Action inputs** (Mode 1) or
 | `REVIEW_EVENT` | `COMMENT` | Fallback if the LLM omits one |
 | `MAX_DIFF_CHARS` | `200000` | Diffs larger than this are truncated |
 | `REVIEW_RULES_PATH` | `.ai/review-rules.md` | Repo rules, read from default branch |
+| `HEADROOM_COMPRESS` | `false` | Compress context before each LLM call (see below) |
 
 See [`action.yml`](action.yml) and [`.env.example`](.env.example) for the full
 list (billing routing, streaming, reasoning effort, browse-tool limits, etc.).
 
+### Context compression (optional)
+
+Long agentic reviews accumulate token-heavy tool outputs (file reads, grep
+dumps) and assistant turns. Set `HEADROOM_COMPRESS=true` to compress that
+context with [headroom](https://github.com/chopratejas/headroom) before each
+LLM call. It's an opt-in extra — install it with `pip install
+'reviewbot[headroom]'` (the Action installs it automatically when the input is
+on). If the package is missing or a compression call fails, messages are sent
+uncompressed, so a review never breaks on it.
+
+By default the annotated diff (a `user` message, whose line numbers the model
+must cite) and the most recent turns are left intact; only tool outputs and
+older turns shrink. Compression is model-aware — the resolved model id drives
+token counting and the context limit, so it works for both OpenAI- and
+Anthropic-family models over the OpenAI-compatible protocol.
+
+| Env var / input | Default | Notes |
+| --------------- | ------- | ----- |
+| `HEADROOM_COMPRESS` | `false` | Master switch |
+| `HEADROOM_TARGET_RATIO` | — | Keep-ratio for text compression (e.g. `0.5`) |
+| `HEADROOM_COMPRESS_USER_MESSAGES` | `false` | Also compress the diff — off by default to keep cited lines intact |
+| `HEADROOM_COMPRESS_SYSTEM_MESSAGES` | `true` | Compress system messages |
+| `HEADROOM_PROTECT_RECENT` | `4` | Never compress the last N messages |
+| `HEADROOM_MIN_TOKENS` | `250` | Skip messages shorter than this |
+| `HEADROOM_KOMPRESS_MODEL` | — | Kompress model id, or `disabled` to skip ML compression |
+| `HEADROOM_MODEL_LIMIT` | `200000` | Model context window used for sizing |
+
 **LLM compatibility:** any service exposing
 `POST {base}/chat/completions` with `response_format: json_object` works —
 OpenAI, Hugging Face Router, Anthropic's OpenAI shim, vLLM, TGI, llama.cpp,

action.yml

@@ -67,6 +67,36 @@ inputs:
     description: 'Maximum number of tool-calling rounds before forcing the model to produce a final review. Each iteration is a separate LLM call.'
     required: false
     default: '8'
+  headroom_compress:
+    description: 'Set to "true" to compress context with headroom-ai before each LLM call (cuts tokens on tool outputs and older turns). Requires the headroom-ai package; install reviewbot with the [headroom] extra. No-op if the package is missing.'
+    required: false
+    default: 'false'
+  headroom_target_ratio:
+    description: 'Optional headroom keep-ratio for text compression (e.g. 0.5 keeps ~50%). Leave empty to let headroom decide.'
+    required: false
+  headroom_compress_user_messages:
+    description: 'Set to "true" to also compress user messages (the annotated diff). Off by default so cited line numbers stay intact.'
+    required: false
+    default: 'false'
+  headroom_compress_system_messages:
+    description: 'Set to "false" to leave system messages uncompressed. On by default.'
+    required: false
+    default: 'true'
+  headroom_protect_recent:
+    description: 'Number of most-recent messages never compressed.'
+    required: false
+    default: '4'
+  headroom_min_tokens:
+    description: 'Skip compressing messages shorter than this many tokens.'
+    required: false
+    default: '250'
+  headroom_kompress_model:
+    description: 'Kompress model id for ML compression, or "disabled" to skip it. Leave empty for the headroom default.'
+    required: false
+  headroom_model_limit:
+    description: 'Model context window (tokens) used for sizing.'
+    required: false
+    default: '200000'
   github_token:
     description: 'Token with pull-requests:write. Defaults to the job token.'
     required: false
@@ -86,7 +116,15 @@ runs:
 
     - name: Install reviewbot
       shell: bash
-      run: pip install --disable-pip-version-check "${{ github.action_path }}"
+      env:
+        HEADROOM_COMPRESS: ${{ inputs.headroom_compress }}
+      run: |
+        # Pull in the (heavy) headroom-ai extra only when compression is on.
+        if [ "${HEADROOM_COMPRESS,,}" = "true" ] || [ "$HEADROOM_COMPRESS" = "1" ]; then
+          pip install --disable-pip-version-check "${{ github.action_path }}[headroom]"
+        else
+          pip install --disable-pip-version-check "${{ github.action_path }}"
+        fi
 
     - name: Run AI reviewer
       shell: bash
@@ -108,4 +146,12 @@ runs:
         CONTEXT_SCRIPT_TIMEOUT: ${{ inputs.context_script_timeout }}
         REPO_CHECKOUT_PATH: ${{ inputs.repo_checkout_path }}
         TOOL_MAX_ITERATIONS: ${{ inputs.tool_max_iterations }}
+        HEADROOM_COMPRESS: ${{ inputs.headroom_compress }}
+        HEADROOM_TARGET_RATIO: ${{ inputs.headroom_target_ratio }}
+        HEADROOM_COMPRESS_USER_MESSAGES: ${{ inputs.headroom_compress_user_messages }}
+        HEADROOM_COMPRESS_SYSTEM_MESSAGES: ${{ inputs.headroom_compress_system_messages }}
+        HEADROOM_PROTECT_RECENT: ${{ inputs.headroom_protect_recent }}
+        HEADROOM_MIN_TOKENS: ${{ inputs.headroom_min_tokens }}
+        HEADROOM_KOMPRESS_MODEL: ${{ inputs.headroom_kompress_model }}
+        HEADROOM_MODEL_LIMIT: ${{ inputs.headroom_model_limit }}
       run: reviewbot-action

pyproject.toml

@@ -30,6 +30,10 @@ web = [
   "httpx>=0.27",
   "itsdangerous>=2.1",
 ]
+# Optional context compression. Enable at runtime with HEADROOM_COMPRESS=1.
+headroom = [
+  "headroom-ai>=0.22",
+]
 
 [project.scripts]
 reviewbot-action = "reviewbot.action_runner:main"

reviewbot/compression.py

@@ -0,0 +1,156 @@
+"""Opt-in context compression for LLM calls via the ``headroom-ai`` package.
+
+Runs only when ``HEADROOM_COMPRESS`` is set and ``headroom`` is importable;
+otherwise messages pass through unchanged. Any failure falls back to the
+original messages so compression can never block a review.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+log = logging.getLogger(__name__)
+
+
+def _bool_env(name: str, default: bool = False) -> bool:
+    raw = (os.environ.get(name) or "").strip().lower()
+    if not raw:
+        return default
+    return raw in ("1", "true", "yes", "on")
+
+
+def _int_env(name: str, default: int) -> int:
+    raw = (os.environ.get(name) or "").strip()
+    try:
+        return int(raw) if raw else default
+    except ValueError:
+        return default
+
+
+def _float_env(name: str, default: float | None = None) -> float | None:
+    raw = (os.environ.get(name) or "").strip()
+    if not raw:
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        return default
+
+
+@dataclass
+class CompressionConfig:
+    # Defaults are tuned for a review agent: leave user/system messages alone
+    # and protect recent turns, so only tool outputs and older assistant turns
+    # get compressed. Fields mirror headroom.CompressConfig.
+    enabled: bool = False
+    compress_user_messages: bool = False
+    compress_system_messages: bool = True
+    protect_recent: int = 4
+    target_ratio: Optional[float] = None
+    min_tokens_to_compress: int = 250
+    # None = headroom default Kompress model; "disabled" = no ML compression.
+    kompress_model: Optional[str] = None
+    # Model context window, forwarded as compress(model_limit=...).
+    model_limit: int = 200_000
+
+    @classmethod
+    def from_env(cls) -> "CompressionConfig":
+        return cls(
+            enabled=_bool_env("HEADROOM_COMPRESS", False),
+            compress_user_messages=_bool_env("HEADROOM_COMPRESS_USER_MESSAGES", False),
+            compress_system_messages=_bool_env(
+                "HEADROOM_COMPRESS_SYSTEM_MESSAGES", True
+            ),
+            protect_recent=_int_env("HEADROOM_PROTECT_RECENT", 4),
+            target_ratio=_float_env("HEADROOM_TARGET_RATIO", None),
+            min_tokens_to_compress=_int_env("HEADROOM_MIN_TOKENS", 250),
+            kompress_model=(os.environ.get("HEADROOM_KOMPRESS_MODEL") or "").strip()
+            or None,
+            model_limit=_int_env("HEADROOM_MODEL_LIMIT", 200_000),
+        )
+
+
+class MessageCompressor:
+    """Lazily-loaded wrapper around ``headroom.compress``."""
+
+    def __init__(self, config: Optional[CompressionConfig] = None):
+        self.config = config or CompressionConfig()
+        self._headroom: Optional[tuple[Callable[..., Any], Any]] = None
+        self._unavailable = False
+        self._warned_unavailable = False
+
+    @classmethod
+    def from_env(cls) -> "MessageCompressor":
+        return cls(CompressionConfig.from_env())
+
+    @property
+    def active(self) -> bool:
+        return self.config.enabled
+
+    def _load(self) -> Optional[tuple[Callable[..., Any], Any]]:
+        # Import once; cache success or failure so we don't retry every call.
+        if self._unavailable:
+            return None
+        if self._headroom is not None:
+            return self._headroom
+        try:
+            from headroom import CompressConfig, compress  # type: ignore
+        except Exception as exc:  # noqa: BLE001 — ImportError or broken install
+            self._unavailable = True
+            if not self._warned_unavailable:
+                log.warning(
+                    "HEADROOM_COMPRESS is set but 'headroom-ai' is not importable "
+                    "(%s); sending messages uncompressed. Install with: "
+                    "pip install 'reviewbot[headroom]'",
+                    exc,
+                )
+                self._warned_unavailable = True
+            return None
+        self._headroom = (compress, CompressConfig)
+        return self._headroom
+
+    def compress(
+        self, messages: list[dict[str, Any]], *, model: Optional[str] = None
+    ) -> list[dict[str, Any]]:
+        if not self.config.enabled or not messages:
+            return messages
+        loaded = self._load()
+        if loaded is None:
+            return messages
+        compress_fn, compress_config_cls = loaded
+        cfg = compress_config_cls(
+            compress_user_messages=self.config.compress_user_messages,
+            compress_system_messages=self.config.compress_system_messages,
+            protect_recent=self.config.protect_recent,
+            target_ratio=self.config.target_ratio,
+            min_tokens_to_compress=self.config.min_tokens_to_compress,
+            kompress_model=self.config.kompress_model,
+        )
+        try:
+            # `model` is only for token counting / context limit — the request
+            # still goes to the bot's own OpenAI-compatible endpoint.
+            result = compress_fn(
+                messages,
+                model=model or "gpt-4o",
+                model_limit=self.config.model_limit,
+                config=cfg,
+            )
+        except Exception:  # noqa: BLE001 — never break a review on compression
+            log.warning(
+                "headroom compression failed; using original messages", exc_info=True
+            )
+            return messages
+
+        saved = getattr(result, "tokens_saved", 0) or 0
+        if saved > 0:
+            log.info(
+                "headroom compressed context: %d -> %d tokens (saved %d, %.0f%%)",
+                getattr(result, "tokens_before", 0),
+                getattr(result, "tokens_after", 0),
+                saved,
+                100.0 * (getattr(result, "compression_ratio", 0.0) or 0.0),
+            )
+        return getattr(result, "messages", messages) or messages

reviewbot/llm_client.py

@@ -6,6 +6,8 @@
 
 import requests
 
+from .compression import MessageCompressor
+
 log = logging.getLogger(__name__)
 
 
@@ -82,12 +84,14 @@ def __init__(
         model: Optional[str] = None,
         bill_to: Optional[str] = None,
         stream: bool = False,
+        compressor: Optional["MessageCompressor"] = None,
     ):
         self.api_base = api_base.rstrip("/")
         self.api_key = api_key
         self.model = model
         self.bill_to = bill_to or None
         self.stream = stream
+        self.compressor = compressor
 
     def _api_base_v1(self) -> str:
         if self.api_base.endswith("/v1"):
@@ -151,8 +155,11 @@ def complete(
         extra: Optional[dict[str, Any]] = None,
         chunk_callback: Optional[Callable[[str, str], None]] = None,
     ) -> ChatResult:
+        model = self._resolve_model()
+        if self.compressor is not None:
+            messages = self.compressor.compress(messages, model=model)
         payload: dict[str, Any] = {
-            "model": self._resolve_model(),
+            "model": model,
             "messages": messages,
             "temperature": temperature,
             "max_tokens": max_tokens,

reviewbot/reviewer.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Callable, Optional
 
+from .compression import MessageCompressor
 from .config import Config
 from .context_script import run_context_script
 from .github_client import GitHubClient
@@ -1117,6 +1118,7 @@ def _emit(kind: str, text: str) -> None:
         cfg.llm_model,
         bill_to=cfg.llm_bill_to,
         stream=cfg.llm_stream,
+        compressor=MessageCompressor.from_env(),
     )
     system_prompt = build_system_prompt(
         review_rules, tools_enabled=tool_env is not None
@@ -1461,6 +1463,7 @@ def run_followup(cfg: Config, gh: GitHubClient, req: ReviewRequest) -> None:
         cfg.llm_model,
         bill_to=cfg.llm_bill_to,
         stream=cfg.llm_stream,
+        compressor=MessageCompressor.from_env(),
     )
 
     system_prompt = build_followup_system_prompt(