chopratejas--headroom/headroom/proxy/interceptors/base.py

"""Protocol + registry + Transform adapter for tool_result interceptors."""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import dataclass
from typing import Any, Protocol, runtime_checkable

from headroom.cache.compression_cache import (
    _extract_tool_result_content,
    _is_tool_result_message,
    _swap_tool_result_content,
)
from headroom.config import TransformResult
from headroom.tokenizer import Tokenizer
from headroom.transforms.base import Transform

logger = logging.getLogger(__name__)

# Interceptor failure counters exposed via `interceptor_failure_counts()`.
# Incremented whenever `matches()`, `transform()`, or `progressive_disclosure_key()`
# raises an exception. Dashboards / stats endpoints can surface these to
# distinguish "nothing eligible to intercept" from "everything is crashing."
_FAILURES: dict[str, int] = {}
_FAILURES_LOCK = threading.Lock()


def _record_failure(interceptor_name: str) -> None:
    with _FAILURES_LOCK:
        _FAILURES[interceptor_name] = _FAILURES.get(interceptor_name, 0) + 1


def interceptor_failure_counts() -> dict[str, int]:
    """Return a snapshot of interceptor failure counters."""
    with _FAILURES_LOCK:
        return dict(_FAILURES)


def reset_interceptor_failure_counts() -> None:
    """Reset failure counters (used by tests)."""
    with _FAILURES_LOCK:
        _FAILURES.clear()


@runtime_checkable
class ToolResultInterceptor(Protocol):
    """A stateless rewriter for a single tool_result's text content.

    Implementations MUST be idempotent and MUST return either a strictly
    smaller string (measured in tokens) or None to pass through. Never raise
    — errors should be caught internally and logged; the pipeline always
    tolerates a no-op interceptor.

    Interceptors MAY implement `progressive_disclosure_key()` to opt into
    one-shot behavior: the framework tracks which keys have already been
    rewritten in the current conversation, and skips subsequent matches on
    the same key so that the model gets full content if it asks again.
    """

    name: str  # e.g. "ast-grep", "difft", "scc"

    def matches(
        self,
        tool_name: str | None,
        tool_input: dict[str, Any],
        tool_output: str,
    ) -> bool: ...

    def transform(
        self,
        tool_name: str | None,
        tool_input: dict[str, Any],
        tool_output: str,
    ) -> str | None: ...

    def progressive_disclosure_key(
        self,
        tool_name: str | None,
        tool_input: dict[str, Any],
    ) -> str | None:
        """Optional: return a stable content key (e.g. file path).

        If a key is returned and the same (interceptor.name, key) pair was
        already successfully rewritten earlier in the messages, subsequent
        occurrences pass through unchanged. Return None to opt out.
        """
        ...


@dataclass(frozen=True)
class TransformSpan:
    """Per-interceptor measurement emitted for dashboard/metrics."""

    tool: str
    tokens_before: int
    tokens_after: int

    @property
    def tokens_saved(self) -> int:
        return max(self.tokens_before - self.tokens_after, 0)


@dataclass(frozen=True)
class InterceptionResult:
    messages: list[dict[str, Any]]
    spans: list[TransformSpan]


INTERCEPTORS: list[ToolResultInterceptor] = []


def register(interceptor: ToolResultInterceptor) -> None:
    """Add an interceptor to the registry. Idempotent on name."""
    for existing in INTERCEPTORS:
        if existing.name == interceptor.name:
            return
    INTERCEPTORS.append(interceptor)


def _build_tool_use_index(
    messages: list[dict[str, Any]],
) -> dict[str, tuple[str | None, dict[str, Any]]]:
    """Scan once and build a dict of {tool_use_id: (tool_name, tool_input)}.

    O(total_blocks) to build, O(1) to look up — used instead of a per-message
    linear scan so `apply_to_messages()` stays linear in message count.
    """
    index: dict[str, tuple[str | None, dict[str, Any]]] = {}
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for block in content:
                if not isinstance(block, dict):
                    continue
                # Anthropic: {"type": "tool_use", "id": ..., "name": ..., "input": {...}}
                if block.get("type") == "tool_use":
                    bid = block.get("id")
                    if isinstance(bid, str):
                        index[bid] = (block.get("name"), block.get("input") or {})
        # OpenAI: assistant message with `tool_calls` list
        tool_calls = msg.get("tool_calls")
        if isinstance(tool_calls, list):
            for call in tool_calls:
                if not isinstance(call, dict):
                    continue
                cid = call.get("id")
                if not isinstance(cid, str):
                    continue
                fn = call.get("function") or {}
                args: dict[str, Any] = {}
                raw_args = fn.get("arguments")
                if isinstance(raw_args, str):
                    try:
                        args = json.loads(raw_args)
                    except json.JSONDecodeError as e:
                        # Empty dict means range-key checks fail for this call
                        # (interceptor may outline despite an explicit line range).
                        # Log so the miss is diagnosable.
                        logger.debug(
                            "tool_call %s arguments failed to JSON-decode: %s; "
                            "proceeding with empty args (range-key checks disabled)",
                            cid,
                            e,
                        )
                        args = {}
                elif isinstance(raw_args, dict):
                    args = raw_args
                index[cid] = (fn.get("name"), args)
    return index


def _tool_use_id_for_message(msg: dict[str, Any]) -> str | None:
    """Return the tool_use_id linked to a tool_result message."""
    # Anthropic format
    content = msg.get("content")
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict) and block.get("type") == "tool_result":
                tuid = block.get("tool_use_id")
                if isinstance(tuid, str):
                    return tuid
    # OpenAI format
    if msg.get("role") == "tool":
        tcid = msg.get("tool_call_id")
        if isinstance(tcid, str):
            return tcid
    return None


def apply_to_messages(
    messages: list[dict[str, Any]],
    tokenizer: Tokenizer,
    *,
    frozen_count: int = 0,
) -> InterceptionResult:
    """Run every registered interceptor against every tool_result in `messages`.

    `frozen_count`: leading messages in the provider's prefix cache that
    MUST be passed through verbatim. Their tool_uses are still scanned so
    that progressive disclosure works across the prefix / tail boundary —
    e.g., a file first Read in the frozen prefix won't be re-outlined when
    the model Reads it again in the mutable tail.

    Returns the (possibly) rewritten message list and a list of spans that
    actually saved tokens.
    """
    if not INTERCEPTORS:
        return InterceptionResult(messages=messages, spans=[])

    spans: list[TransformSpan] = []
    # Progressive disclosure: per-interceptor set of keys already rewritten
    # earlier in this message list. Prevents the second Read of the same
    # file from being outlined again — the model evidently came back for
    # more, so give it the raw content.
    fired: dict[str, set[str]] = {}
    # Build O(1) tool_use lookup index once per request (over ALL messages,
    # including the frozen prefix, so keys resolve correctly even when the
    # tool_use lives in the frozen part).
    tool_use_index = _build_tool_use_index(messages)

    # Pre-seed `fired` from the frozen prefix so that any file already Read
    # in the cached prefix counts as "already disclosed" for subsequent reads.
    for msg in messages[:frozen_count]:
        if not _is_tool_result_message(msg):
            continue
        frozen_tuid = _tool_use_id_for_message(msg)
        if not frozen_tuid:
            continue
        f_tool_name, f_tool_input = tool_use_index.get(frozen_tuid, (None, {}))
        for interceptor in INTERCEPTORS:
            key_fn = getattr(interceptor, "progressive_disclosure_key", None)
            if not callable(key_fn):
                continue
            try:
                k = key_fn(f_tool_name, f_tool_input)
            except Exception as e:  # noqa: BLE001
                logger.debug(
                    "interceptor %s key() failed on frozen prefix: %s",
                    interceptor.name,
                    e,
                )
                continue
            if k:
                fired.setdefault(interceptor.name, set()).add(k)

    # Pass frozen messages through verbatim; only the mutable tail is
    # considered for rewriting.
    new_messages: list[dict[str, Any]] = list(messages[:frozen_count])
    for msg in messages[frozen_count:]:
        if not _is_tool_result_message(msg):
            new_messages.append(msg)
            continue

        original = _extract_tool_result_content(msg)
        if not isinstance(original, str) or not original:
            new_messages.append(msg)
            continue

        tuid = _tool_use_id_for_message(msg)
        tool_name: str | None = None
        tool_input: dict[str, Any] = {}
        if tuid:
            tool_name, tool_input = tool_use_index.get(tuid, (None, {}))
            if tuid not in tool_use_index:
                # Orphaned tool_result — interceptors run without tool context.
                logger.debug("tool_result %s has no matching tool_use", tuid)

        current = original
        for interceptor in INTERCEPTORS:
            # Progressive disclosure: skip if already fired for this key.
            key: str | None = None
            key_fn = getattr(interceptor, "progressive_disclosure_key", None)
            if callable(key_fn):
                try:
                    key = key_fn(tool_name, tool_input)
                except Exception as e:  # noqa: BLE001
                    logger.warning(
                        "interceptor %s key() failed: %s",
                        interceptor.name,
                        e,
                        exc_info=True,
                    )
                    _record_failure(interceptor.name)
                    # Skip this interceptor entirely rather than firing
                    # without progressive-disclosure protection — a broken
                    # key would otherwise fire on every Read of the same file.
                    continue
            if key and key in fired.get(interceptor.name, set()):
                continue

            try:
                if not interceptor.matches(tool_name, tool_input, current):
                    continue
                rewritten = interceptor.transform(tool_name, tool_input, current)
            except Exception as e:  # noqa: BLE001 — never crash a request
                logger.warning(
                    "interceptor %s failed: %s",
                    interceptor.name,
                    e,
                    exc_info=True,
                )
                _record_failure(interceptor.name)
                continue
            if not rewritten or rewritten == current:
                continue
            before = tokenizer.count_text(current)
            after = tokenizer.count_text(rewritten)
            if after >= before:
                continue  # refuse to enlarge
            spans.append(
                TransformSpan(
                    tool=interceptor.name,
                    tokens_before=before,
                    tokens_after=after,
                )
            )
            current = rewritten
            if key:
                fired.setdefault(interceptor.name, set()).add(key)

        new_messages.append(
            _swap_tool_result_content(msg, current) if current is not original else msg
        )

    return InterceptionResult(messages=new_messages, spans=spans)


class ToolResultInterceptorTransform(Transform):
    """Pipeline-level adapter: runs interceptors as the first compression stage.

    Placed at transforms[0] so downstream compressors operate on the already-
    shrunk content. Transform names of firing interceptors are added to
    `transforms_applied` so they appear in existing dashboards/metrics.

    Honors the standard `frozen_message_count` contract: leading messages in
    the provider's prefix cache are not modified, preserving cache hits.
    """

    name = "tool_result_interceptors"

    def apply(
        self,
        messages: list[dict[str, Any]],
        tokenizer: Tokenizer,
        **kwargs: Any,
    ) -> TransformResult:
        # Measure the true baseline on the original messages — back-calculating
        # from `tokens_after + sum(saved)` would double-count per-message
        # overhead that spans don't track.
        tokens_before = tokenizer.count_messages(messages)

        # `apply_to_messages` handles the frozen/mutable split internally
        # and pre-seeds progressive-disclosure state from the frozen prefix
        # so a file already Read there isn't re-outlined in the tail.
        frozen = int(kwargs.get("frozen_message_count") or 0)
        result = apply_to_messages(messages, tokenizer, frozen_count=frozen)
        tokens_after = tokenizer.count_messages(result.messages)
        transforms_applied = [f"interceptor:{s.tool}" for s in result.spans] if result.spans else []
        return TransformResult(
            messages=result.messages,
            tokens_before=tokens_before,
            tokens_after=tokens_after,
            transforms_applied=transforms_applied,
        )