chopratejas--headroom/headroom/proxy/helpers.py

"""Top-level helper functions and constants for the Headroom proxy.

Contains lazy loaders, file logging setup, request body decompression,
and safety-limit constants.

Extracted from server.py for maintainability.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
import random
import re
import threading
import time
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, cast

from headroom import paths as _paths
from headroom._subprocess import run
from headroom.proxy import (
    diagnostic_decode_policy,
    memory_injection_mode_policy,
    query_log_policy,
    request_limit_policy,
    sse_byte_buffer_policy,
    wire_debug_format_policy,
    wire_debug_redaction_policy,
)
from headroom.proxy.beta_header_policy import (
    BETA_HEADER_STICKY_DEFAULT,
    BETA_HEADER_STICKY_ENV,
    BETA_TRACKER_MAX_SESSIONS_DEFAULT,
    BETA_TRACKER_MAX_SESSIONS_ENV,
    BetaHeaderStickyMode,
    resolve_beta_header_sticky_mode,
    resolve_beta_tracker_max_sessions,
)
from headroom.proxy.body_forwarding import (
    BodyMutationTracker as BodyMutationTracker,  # noqa: F401 - compatibility export
)
from headroom.proxy.body_forwarding import (
    PythonForwarderMode as PythonForwarderMode,  # noqa: F401 - compatibility export
)
from headroom.proxy.body_forwarding import (
    get_python_forwarder_mode as get_python_forwarder_mode,  # noqa: F401 - compatibility export
)
from headroom.proxy.body_forwarding import (
    prepare_outbound_body_bytes as prepare_outbound_body_bytes,  # noqa: F401 - compatibility export
)
from headroom.proxy.body_forwarding import serialize_body_canonical
from headroom.proxy.ccr_golden_policy import (
    create_fresh_ccr_tool_definition,
    replay_golden_ccr_tool_definition,
)
from headroom.proxy.ccr_marker_policy import (
    has_new_ccr_markers as _has_new_ccr_markers,
)
from headroom.proxy.ccr_marker_policy import (
    should_inject_ccr_tool as _should_inject_ccr_tool,
)
from headroom.proxy.ccr_session_tracker import SessionCcrTracker as _SessionCcrTracker
from headroom.proxy.internal_header_policy import (
    INTERNAL_HEADER_PREFIX,
    STRIP_INTERNAL_HEADERS_DEFAULT,
    STRIP_INTERNAL_HEADERS_ENV,
    StripInternalHeadersMode,
    resolve_strip_internal_headers_mode,
    strip_internal_headers,
)
from headroom.proxy.memory_golden_policy import (
    replay_golden_memory_tool_definition,
    serialize_memory_tool_definition_canonical,
)
from headroom.proxy.tool_injection_config import (
    ToolInjectionStickyMode,
)
from headroom.proxy.tool_injection_config import (
    get_tool_injection_sticky_mode as _get_tool_injection_sticky_mode,
)
from headroom.proxy.tool_injection_config import (
    get_tool_tracker_max_sessions as _get_tool_tracker_max_sessions,
)
from headroom.proxy.tool_injection_logging import (
    ToolInjectionDecision,
)
from headroom.proxy.tool_injection_logging import (
    log_tool_injection_decision as _log_tool_injection_decision,
)
from headroom.proxy.tool_injection_tracker import SessionToolTracker as _SessionToolTracker
from headroom.proxy.tool_name_policy import extract_tool_name

if TYPE_CHECKING:
    import httpx
    from fastapi import Request

logger = logging.getLogger("headroom.proxy")

_CODEX_WIRE_DEBUG_ENV = "HEADROOM_CODEX_WIRE_DEBUG"
_CODEX_WIRE_DEBUG_DIR_ENV = "HEADROOM_CODEX_WIRE_DEBUG_DIR"
_CODEX_WIRE_REDACTED = wire_debug_redaction_policy.WIRE_DEBUG_REDACTED
_CODEX_WIRE_SECRET_KEYS = wire_debug_redaction_policy.WIRE_DEBUG_SECRET_KEYS


def codex_wire_debug_enabled() -> bool:
    """Return whether opt-in Codex wire capture is enabled."""

    return os.environ.get(_CODEX_WIRE_DEBUG_ENV, "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )


def _codex_wire_debug_dir() -> Path:
    explicit = os.environ.get(_CODEX_WIRE_DEBUG_DIR_ENV, "").strip()
    if explicit:
        return Path(explicit).expanduser()
    return _paths.codex_wire_debug_dir()


def _should_redact_key(key: str) -> bool:
    return wire_debug_redaction_policy.should_redact_key(key)


def _redact_value(value: Any) -> Any:
    return wire_debug_redaction_policy.redact_for_wire_debug(value)


def redact_for_wire_debug(value: Any) -> Any:
    """Redact obvious secrets while preserving request/response shape."""
    return wire_debug_redaction_policy.redact_for_wire_debug(value)


def _safe_event_name(event: str) -> str:
    return wire_debug_format_policy.safe_wire_debug_name(event)


def _wire_debug_preview(value: Any, *, max_chars: int | None = None) -> str:
    """Return the redacted wire payload for proxy.log.

    This is intentionally not truncated. During Codex WS debugging we need the
    proxy log itself to show the complete frame so we can decide later where a
    deliberate trim boundary belongs.
    """

    return wire_debug_format_policy.wire_debug_preview(value, max_chars=max_chars)


def capture_codex_wire_debug(
    event: str,
    *,
    request_id: str | None = None,
    session_id: str | None = None,
    transport: str,
    direction: str,
    method: str | None = None,
    url: str | None = None,
    headers: dict[str, Any] | None = None,
    body: Any = None,
    raw_text: str | None = None,
    status_code: int | None = None,
    metadata: dict[str, Any] | None = None,
) -> Path | None:
    """Write an opt-in redacted Codex wire snapshot to disk.

    This is intentionally file-based rather than log-based: real Codex
    requests can be large, and operators need the exact envelope shape without
    mixing it into normal proxy logs. Header/body secret-looking keys are
    redacted, but request content is otherwise preserved because this mode is
    explicitly for local debugging.
    """

    if not codex_wire_debug_enabled():
        return None

    try:
        out_dir = _codex_wire_debug_dir()
        out_dir.mkdir(parents=True, exist_ok=True)
        ts_ns = time.time_ns()
        req = request_id or "no_request"
        safe_req = _safe_event_name(req)
        safe_event = _safe_event_name(event)
        path = out_dir / f"{ts_ns}_{safe_req}_{safe_event}.json"
        payload = {
            "event": event,
            "timestamp_ns": ts_ns,
            "request_id": request_id,
            "session_id": session_id,
            "transport": transport,
            "direction": direction,
            "method": method,
            "url": url,
            "status_code": status_code,
            "headers": redact_for_wire_debug(headers or {}),
            "body": redact_for_wire_debug(body),
            "raw_text": raw_text,
            "metadata": redact_for_wire_debug(metadata or {}),
        }
        path.write_text(
            json.dumps(payload, indent=2, ensure_ascii=False, default=str), encoding="utf-8"
        )
        logger.info(
            "event=codex_wire_debug_capture path=%s request_id=%s wire_event=%s",
            path,
            request_id or "",
            event,
        )
        preview_source = redact_for_wire_debug(body) if body is not None else raw_text
        preview = _wire_debug_preview(preview_source)
        meta_keys = ",".join(sorted((metadata or {}).keys()))
        logger.info(
            "event=codex_wire_debug_frame request_id=%s session_id=%s wire_event=%s "
            "transport=%s direction=%s status_code=%s meta_keys=%s preview=%s",
            request_id or "",
            session_id or "",
            event,
            transport,
            direction,
            status_code if status_code is not None else "",
            meta_keys,
            preview,
        )
        return path
    except Exception as exc:  # pragma: no cover - debug path must never break traffic
        logger.warning("event=codex_wire_debug_capture_failed error=%s", exc)
        return None


# Memory injection mode (P0-1 fix in PR-A2).
#
# Values:
#   - "live_zone_tail" (default): Memory context appends to the first text block
#     of the latest non-frozen user message. Cache hot zone (system + frozen
#     prefix) is never mutated.
#   - "disabled": Memory context lookup is skipped entirely; the request
#     forwards untouched.
#
# Configurable via HEADROOM_MEMORY_INJECTION_MODE env var. There is no
# "system_prompt" option — that path is permanently retired by I2 (cache hot
# zone never modified). See REALIGNMENT/02-architecture.md §2.2.
_MEMORY_INJECTION_MODE_ENV = memory_injection_mode_policy.MEMORY_INJECTION_MODE_ENV
_MEMORY_INJECTION_MODE_DEFAULT = memory_injection_mode_policy.MEMORY_INJECTION_MODE_DEFAULT
MemoryInjectionMode = memory_injection_mode_policy.MemoryInjectionMode


def get_memory_injection_mode() -> MemoryInjectionMode:
    """Return the active memory-injection routing mode.

    Read at request time so the env var can be flipped without restart for
    smoke tests. Unknown values are rejected loudly (no silent fallback).
    """
    return memory_injection_mode_policy.resolve_memory_injection_mode(
        os.environ.get(_MEMORY_INJECTION_MODE_ENV)
    )


def hash_query_for_log(query: str) -> str:
    """Stable short hash of a memory-context query, safe to log.

    Uses BLAKE2b truncated to 16 hex chars. Never logs the raw query content.
    """
    return query_log_policy.hash_query_for_log(query)


def extract_tags(headers: Any) -> dict[str, str]:
    """Extract ``x-headroom-*`` tags from inbound headers.

    Pure function (no I/O, no state). Used by every handler at request
    entry to capture operator slicing tags into the per-request
    ``RequestOutcome.tags``. Free function rather than a mixin method so
    handler mixins instantiated in isolation (tests using
    ``object.__new__(OpenAIHandlerMixin)``) don't need a shim
    implementation.

    Header name match is case-insensitive; the returned key has the
    ``x-headroom-`` prefix stripped.
    """
    return {
        k.lower().replace("x-headroom-", ""): v
        for k, v in headers.items()
        if k.lower().startswith("x-headroom-")
    }


def _headroom_bypass_enabled(headers: Any) -> bool:
    """Return True when inbound headers request full Headroom passthrough.

    This is transport-neutral policy: HTTP and WebSocket handlers both call
    it on original inbound headers before request-body mutation.
    """

    try:
        bypass = str(headers.get("x-headroom-bypass", "")).strip().lower() == "true"
        passthrough = str(headers.get("x-headroom-mode", "")).strip().lower() == "passthrough"
    except AttributeError:
        return False
    return bypass or passthrough


def log_outbound_request(
    *,
    forwarder: str,
    method: str,
    path: str,
    body_bytes_count: int,
    body_mutated: bool,
    mutation_reasons: list[str],
    request_id: str | None,
    source: str,
) -> None:
    """Structured log line for every outbound forwarder call.

    Per realignment build constraints: every cache-affecting decision is
    logged. Never includes ``Authorization``/``x-api-key`` content or full
    body bytes.
    """
    logger.info(
        "event=outbound_request forwarder=%s method=%s path=%s body_bytes=%d "
        "body_mutated=%s mutation_reasons=%s source=%s request_id=%s",
        forwarder,
        method,
        path,
        body_bytes_count,
        "true" if body_mutated else "false",
        ",".join(mutation_reasons) if mutation_reasons else "",
        source,
        request_id or "",
    )


def log_memory_injection(
    *,
    request_id: str,
    session_id: str | None,
    decision: str,
    bytes_injected: int,
    query: str | None = None,
) -> None:
    """Emit a structured log line for every memory-context routing decision.

    Per realignment build constraints: log every cache-affecting decision.
    Never log raw query content or Authorization header — only a stable
    hash of the query.
    """
    query_hash = hash_query_for_log(query) if query else ""
    logger.info(
        "event=memory_injection request_id=%s session_id=%s decision=%s "
        "bytes_injected=%d query_hash=%s",
        request_id,
        session_id or "",
        decision,
        bytes_injected,
        query_hash,
    )


def append_text_to_latest_user_chat_message(
    messages: list[dict[str, Any]],
    context_text: str,
) -> tuple[list[dict[str, Any]], int]:
    """Append context text to the first text block of the latest user chat message.

    OpenAI Chat Completions ``body["messages"]`` shape: each message is
    ``{"role": ..., "content": str | list[{"type": "text"|"input_text", "text": ...}]}``.

    This is the OpenAI Chat Completions analog of
    ``_append_context_to_latest_non_frozen_user_turn`` (Anthropic) and
    ``append_text_to_latest_user_input_item`` (OpenAI Responses). Used by
    PR-A3 to retire the legacy system-prepend memory-injection path
    (P0-equivalent for /v1/chat/completions).

    Returns ``(new_messages, bytes_appended)``. ``bytes_appended == 0``
    when no eligible user message was found (no mutation occurred).
    """
    if not messages or not context_text:
        return messages, 0

    new_messages = list(messages)
    for idx in range(len(new_messages) - 1, -1, -1):
        msg = new_messages[idx]
        if not isinstance(msg, dict):
            continue
        if msg.get("role") != "user":
            continue

        content = msg.get("content")
        if isinstance(content, str):
            updated_msg = {**msg, "content": content + "\n\n" + context_text}
            new_messages[idx] = updated_msg
            return new_messages, len(context_text)

        if isinstance(content, list) and content:
            new_content: list[dict[str, Any]] = []
            appended = False
            for part in content:
                if (
                    not appended
                    and isinstance(part, dict)
                    and part.get("type") in ("text", "input_text")
                ):
                    existing_text = part.get("text", "")
                    new_part = {**part, "text": existing_text + "\n\n" + context_text}
                    new_content.append(new_part)
                    appended = True
                else:
                    new_content.append(part)
            if appended:
                updated_msg = {**msg, "content": new_content}
                new_messages[idx] = updated_msg
                return new_messages, len(context_text)

        # User message but no eligible text block — leave untouched and stop.
        return messages, 0

    return messages, 0


def append_text_to_latest_user_input_item(
    body_input: list[dict[str, Any]],
    context_text: str,
) -> tuple[list[dict[str, Any]], int]:
    """Append context text to the first text block of the latest user input item.

    Mirrors ``_append_context_to_latest_non_frozen_user_turn`` but for the
    OpenAI Responses API ``body["input"]`` shape, which uses a flat item list
    where each user item's content is a list like
    ``[{"type": "input_text", "text": "..."}]``.

    Returns a tuple ``(new_input, bytes_appended)`` where ``bytes_appended``
    is 0 when the item list was unchanged (no eligible user item).
    """
    if not body_input or not context_text:
        return body_input, 0

    new_input = list(body_input)

    for idx in range(len(new_input) - 1, -1, -1):
        item = new_input[idx]
        if not isinstance(item, dict):
            continue
        if item.get("role") != "user":
            continue

        content = item.get("content")
        if isinstance(content, str):
            updated_item = {**item, "content": content + "\n\n" + context_text}
            new_input[idx] = updated_item
            return new_input, len(context_text)

        if isinstance(content, list) and content:
            new_content: list[dict[str, Any]] = []
            appended = False
            for part in content:
                if (
                    not appended
                    and isinstance(part, dict)
                    and part.get("type") in ("input_text", "text")
                ):
                    existing_text = part.get("text", "")
                    new_part = {**part, "text": existing_text + "\n\n" + context_text}
                    new_content.append(new_part)
                    appended = True
                else:
                    new_content.append(part)
            if appended:
                updated_item = {**item, "content": new_content}
                new_input[idx] = updated_item
                return new_input, len(context_text)

        # User item but no eligible text block — leave untouched and stop.
        return body_input, 0

    return body_input, 0


_CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL"
_CONTEXT_TOOL_RTK = "rtk"
_CONTEXT_TOOL_LEAN_CTX = "lean-ctx"
_RTK_GAIN_SCOPE_ENV = "HEADROOM_RTK_GAIN_SCOPE"
_RTK_GAIN_SCOPE_GLOBAL = "global"
_RTK_GAIN_SCOPE_PROJECT = "project"
_RTK_GAIN_SCOPES = {_RTK_GAIN_SCOPE_GLOBAL, _RTK_GAIN_SCOPE_PROJECT}

RTK_STATS_CACHE_TTL_SECONDS = float(os.environ.get("HEADROOM_CONTEXT_TOOL_STATS_TTL_SECONDS", "60"))
CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS = RTK_STATS_CACHE_TTL_SECONDS
_context_tool_stats_cache_lock = threading.Lock()
_context_tool_stats_cache: dict[str, Any] = {
    "expires_at": 0.0,
    "has_value": False,
    "tool": None,
    "value": None,
}
_context_tool_session_baseline: dict[str, Any] = {
    "initialized": False,
    "tool": None,
    "total_commands": 0,
    "input_tokens": 0,
    "output_tokens": 0,
    "tokens_saved": 0,
    "total_time_ms": 0,
    "captured_at": 0.0,
}
_rtk_stats_cache_lock = _context_tool_stats_cache_lock
_rtk_stats_cache = _context_tool_stats_cache
_rtk_session_baseline = _context_tool_session_baseline

# Maximum request body size (100MB - increased to support image-heavy requests)
MAX_REQUEST_BODY_SIZE = 100 * 1024 * 1024

# Maximum SSE buffer size (10MB - prevents memory exhaustion from malformed streams)
MAX_SSE_BUFFER_SIZE = 10 * 1024 * 1024

# Per-event SSE size cap (PR-A8 / P1-8). Configurable via
# HEADROOM_SSE_BUFFER_MAX_BYTES. Guards against pathological huge events
# (a single event > 1 MB by default is treated as an upstream protocol bug
# and surfaces loudly rather than silently growing the buffer).
_SSE_EVENT_MAX_BYTES_ENV = request_limit_policy.SSE_EVENT_MAX_BYTES_ENV
_SSE_EVENT_MAX_BYTES_DEFAULT = request_limit_policy.SSE_EVENT_MAX_BYTES_DEFAULT


def get_sse_event_max_bytes() -> int:
    """Return the per-event SSE size cap.

    Read at request time so operators can flip the env var without a
    restart. Negative values are rejected loudly (no silent fallback).
    """
    return request_limit_policy.resolve_sse_event_max_bytes(
        os.environ.get(_SSE_EVENT_MAX_BYTES_ENV)
    )


# Body-too-large status code (PR-A8 / P5-59). Default 413 (RFC 7231 §6.5.11).
# Configurable via HEADROOM_PROXY_BODY_TOO_LARGE_STATUS for operators who need
# to override (no expected production use; documentation knob).
_BODY_TOO_LARGE_STATUS_ENV = request_limit_policy.BODY_TOO_LARGE_STATUS_ENV
_BODY_TOO_LARGE_STATUS_DEFAULT = request_limit_policy.BODY_TOO_LARGE_STATUS_DEFAULT


def get_body_too_large_status() -> int:
    """Return the HTTP status code for body-too-large rejections."""
    return request_limit_policy.resolve_body_too_large_status(
        os.environ.get(_BODY_TOO_LARGE_STATUS_ENV)
    )


_SSE_EVENT_TERMINATORS = sse_byte_buffer_policy.SSE_EVENT_TERMINATORS


def _find_sse_event_terminator(buf: bytearray) -> tuple[int, int] | None:
    """Return the earliest complete SSE event terminator in ``buf``."""
    return sse_byte_buffer_policy.find_sse_event_terminator(buf)


_SSE_EVENT_LINE_PREFIX = b"event:"
_SSE_DATA_LINE_PREFIX = b"data:"


def safe_decode_for_logging(raw: bytes, *, max_bytes: int | None = None) -> str:
    """Decode bytes to a string for **log/diagnostic display only**.

    PR-A8 / P1-8: the SSE wire path forbids ``errors="ignore"`` /
    ``errors="replace"`` because corrupting bytes silently busts cache
    safety. Diagnostic logs (e.g. error response bodies) are fine to
    show with a replacement character because the bytes are already
    discarded; this helper centralizes that single legitimate use of
    the lossy decoder so a project-wide grep stays clean.

    Use ``parse_sse_events_from_byte_buffer`` for SSE parsing instead.
    """
    return diagnostic_decode_policy.safe_decode_for_logging(raw, max_bytes=max_bytes)


def parse_sse_events_from_byte_buffer(
    buf: bytearray,
) -> list[tuple[str | None, str]]:
    """Drain complete ``event:`` + ``data:`` events from a bytes buffer.

    Returns list of ``(event_name, data_str)`` tuples for complete events.
    Mutates ``buf`` in-place to leave only partial-event tail bytes.

    Operates on bytes; only decodes complete events as UTF-8 (raises if a
    *complete* event has invalid UTF-8 — that's an upstream protocol bug
    we want loud, not silent).

    Per PR-A8 / P1-8: this is the canonical SSE event splitter. NEVER use
    ``decode("utf-8", errors="ignore")`` on a partial buffer; UTF-8
    multi-byte characters split across TCP reads will corrupt content.
    """
    return sse_byte_buffer_policy.parse_sse_events_from_byte_buffer(buf)


# Maximum message array length (prevents DoS from deeply nested payloads)
MAX_MESSAGE_ARRAY_LENGTH = 10000

# Compression pipeline timeout in seconds. Override via the
# HEADROOM_COMPRESSION_TIMEOUT_SECONDS env var for slow CPUs or long Claude Code
# conversations (GH #946). Falls back to 30 on an unparseable value.
try:
    COMPRESSION_TIMEOUT_SECONDS = float(
        os.environ.get("HEADROOM_COMPRESSION_TIMEOUT_SECONDS", "30")
    )
except ValueError:
    COMPRESSION_TIMEOUT_SECONDS = 30.0

# Eager startup preload timeout in seconds. The preload (compressor/parser models,
# cache-only, allow_download=False) runs off the event loop during startup; this
# bound only fires on a true hang or an uncatchable native stall so the proxy still
# binds its port instead of never opening (GH #790). Override via
# HEADROOM_EAGER_PRELOAD_TIMEOUT_SECONDS. Falls back to 120 on an unparseable value.
try:
    EAGER_PRELOAD_TIMEOUT_SECONDS = float(
        os.environ.get("HEADROOM_EAGER_PRELOAD_TIMEOUT_SECONDS", "120")
    )
except ValueError:
    EAGER_PRELOAD_TIMEOUT_SECONDS = 120.0

# Maximum compression cache sessions (prevents unbounded memory growth)
MAX_COMPRESSION_CACHE_SESSIONS = 500


# ---------------------------------------------------------------------------
# Compression-failure escape hatch
# ---------------------------------------------------------------------------
# When the proxy's compression stage fails (timeout, exception) on a frame
# Headroom thought was large enough to compress, the legacy behaviour was to
# fall through and forward the *original* uncompressed frame to the upstream.
# That fail-open turned a recoverable timeout into a context-window overflow
# downstream: Codex's auto-compaction reads ``total_usage_tokens`` from
# upstream (which Headroom's earlier successful compressions shrunk), then
# the un-compressed retry overflows the model context and the client
# locks up.
#
# Default behaviour is now fail-CLOSED: refuse to forward, close the client
# WS with code 1009 (or return HTTP 413) so the client knows to compact and
# retry. Operators who want the old behaviour can set
# ``HEADROOM_WS_FAIL_OPEN_ON_COMPRESSION_FAILURE=1``. The oversize threshold
# below which transient errors still fall through to passthrough is
# configurable via ``HEADROOM_WS_COMPRESSION_FAIL_THRESHOLD_BYTES``
# (default 256 KiB ≈ 64K tokens).
WS_COMPRESSION_FAIL_OPEN_ENV = "HEADROOM_WS_FAIL_OPEN_ON_COMPRESSION_FAILURE"
WS_COMPRESSION_OVERSIZE_BYTES_ENV = "HEADROOM_WS_COMPRESSION_FAIL_THRESHOLD_BYTES"
WS_COMPRESSION_OVERSIZE_BYTES_DEFAULT = 256 * 1024


@dataclass(frozen=True)
class CompressionFailureAction:
    """Decision returned by :func:`decide_compression_failure_action`."""

    refuse: bool
    """If True, the caller MUST NOT forward the original frame. Close the
    client connection with a clear error code instead."""

    reason: str
    """Short machine-readable label for telemetry. One of:
    ``timeout``, ``oversize:bytes=<n>>threshold=<m>``,
    ``small_frame_transient``, ``client_override:codex``, or
    ``env_override:fail_open``."""

    frame_bytes: int
    """Original frame size in bytes (for logging / metrics)."""


def decide_compression_failure_action(
    exception: BaseException,
    frame_bytes: int,
    *,
    client: str | None = None,
) -> CompressionFailureAction:
    """Decide whether to refuse-and-close vs forward-original after the
    proxy's compression pipeline fails on a Realtime WebSocket frame
    (or analogous HTTP body).

    Decision matrix:

    * env :data:`WS_COMPRESSION_FAIL_OPEN_ENV` truthy → forward (legacy
      behaviour, opt-in for debugging or strict compatibility).
    * Codex client compression timeout → forward. Codex currently treats
      the proxy's 1009/413 refusal path as a hard connection failure, so
      fail-open is safer for Codex sessions even when the proxy is run
      standalone rather than through ``headroom wrap codex``.
    * exception is :class:`asyncio.TimeoutError` → refuse (the compression
      stage hit its own timeout, which only fires on frames Headroom
      thought were big enough to need compression in the first place).
    * ``frame_bytes`` > :data:`WS_COMPRESSION_OVERSIZE_BYTES_ENV`
      (default 256 KiB) → refuse (large + any compression failure is a
      strong signal the upstream will reject the original).
    * otherwise → forward (a transient pipeline error on a small frame
      shouldn't break the request).
    """
    fail_open = os.environ.get(WS_COMPRESSION_FAIL_OPEN_ENV, "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )
    if fail_open:
        return CompressionFailureAction(
            refuse=False,
            reason="env_override:fail_open",
            frame_bytes=frame_bytes,
        )

    if (client or "").strip().lower() == "codex" and isinstance(exception, asyncio.TimeoutError):
        return CompressionFailureAction(
            refuse=False,
            reason="client_override:codex",
            frame_bytes=frame_bytes,
        )

    threshold = WS_COMPRESSION_OVERSIZE_BYTES_DEFAULT
    raw_threshold = os.environ.get(WS_COMPRESSION_OVERSIZE_BYTES_ENV, "").strip()
    if raw_threshold:
        try:
            parsed = int(raw_threshold)
            if parsed > 0:
                threshold = parsed
        except ValueError:
            # Operator typo'd the env value — keep the default rather than
            # raise on every WS frame. Loud warning instead.
            logger.warning(
                "Ignoring non-integer %s=%r; using default %d",
                WS_COMPRESSION_OVERSIZE_BYTES_ENV,
                raw_threshold,
                WS_COMPRESSION_OVERSIZE_BYTES_DEFAULT,
            )

    if isinstance(exception, asyncio.TimeoutError):
        return CompressionFailureAction(refuse=True, reason="timeout", frame_bytes=frame_bytes)
    if frame_bytes > threshold:
        return CompressionFailureAction(
            refuse=True,
            reason=f"oversize:bytes={frame_bytes}>threshold={threshold}",
            frame_bytes=frame_bytes,
        )
    return CompressionFailureAction(
        refuse=False, reason="small_frame_transient", frame_bytes=frame_bytes
    )


def jitter_delay_ms(base_ms: int, max_ms: int, attempt: int) -> float:
    """Exponential backoff with 50-150% jitter.

    Returns ``min(base_ms * 2**attempt, max_ms) * (0.5 + random())`` — the
    canonical formula used across proxy retry loops. Extracted so every
    retry site shares one implementation.
    """
    capped: float = min(base_ms * (2**attempt), max_ms)
    return capped * (0.5 + random.random())


def retry_after_ms(response: httpx.Response, max_ms: int) -> float | None:
    """Parse an HTTP ``Retry-After`` header into a millisecond delay, capped at ``max_ms``.

    Returns the delay in ms for a numeric ``seconds`` value or an HTTP-date, or
    ``None`` when the header is absent or unparseable so the caller falls back to
    exponential backoff. Anthropic sends integer seconds; the HTTP-date branch
    covers other upstreams. Fails open on any parse error.
    """
    value = response.headers.get("retry-after")
    if not value:
        return None
    try:
        seconds = float(value)
    except ValueError:
        try:
            from datetime import datetime
            from email.utils import parsedate_to_datetime

            retry_at = parsedate_to_datetime(value)
            seconds = (retry_at - datetime.now(retry_at.tzinfo)).total_seconds()
        except (TypeError, ValueError):
            return None
    return min(max(seconds, 0.0) * 1000.0, float(max_ms))


# Transient upstream statuses worth retrying with backoff: 429 (rate limit) and
# 529 (Anthropic ``overloaded_error``). Both mean "the server is temporarily
# limiting/overloaded — try again shortly", unlike other 4xx which signal a
# problem with the request itself. Single source of truth so the streaming and
# non-streaming forwarders agree on what is retriable.
RETRYABLE_OVERLOAD_STATUSES: frozenset[int] = frozenset({429, 529})


async def request_with_transient_retry(
    client: httpx.AsyncClient,
    *,
    request_id: str | None = None,
    max_retries: int = 1,
    **request_kwargs: Any,
) -> httpx.Response:
    """Issue a buffered httpx request, retrying once on a transient close.

    ``httpx.RemoteProtocolError`` ("peer closed connection without sending
    complete message body (incomplete chunked read)") is raised when an
    upstream closes a pooled keep-alive connection that httpx then reuses for
    the next request. A direct ``curl`` never hits this because it opens a
    fresh connection per call; Headroom reuses pooled connections, so the
    first request issued on a stale connection fails even though the upstream
    is healthy (it answers a fresh connection with 200). Retrying opens a new
    connection and succeeds, mirroring curl's behaviour. See GH #1112.

    Only ``httpx.RemoteProtocolError`` is retried — the specific stale
    keep-alive symptom; every other exception (``ConnectError``, timeouts,
    HTTP status errors) propagates immediately so existing handling is
    unchanged. Use this for buffered (non-streaming) requests only: a streamed
    response cannot be safely replayed once bytes have reached the client.
    """
    import httpx

    attempt = 0
    while True:
        try:
            return await client.request(**request_kwargs)
        except httpx.RemoteProtocolError as exc:
            if attempt >= max_retries:
                raise
            attempt += 1
            logger.warning(
                "Upstream closed connection mid-response (%s); retrying on a "
                "fresh connection (attempt %d/%d)%s",
                exc,
                attempt,
                max_retries,
                f" [{request_id}]" if request_id else "",
            )


# Image compression availability (do not retain a global compressor instance)
_image_compressor_available: bool | None = None


def _get_image_compressor():
    """Create a short-lived image compressor on demand."""
    global _image_compressor_available
    if _image_compressor_available is False:
        return None

    try:
        from headroom.image import ImageCompressor

        # Callers own closing the compressor; this helper only memoizes whether
        # the optional image stack is importable.
        compressor = ImageCompressor()
        if _image_compressor_available is None:
            logger.info("Image compression enabled (model: chopratejas/technique-router)")
        _image_compressor_available = True
        return compressor
    except ImportError as e:
        if _image_compressor_available is not False:
            logger.warning(f"Image compression not available: {e}")
        _image_compressor_available = False
        return None


# Always-on file logging to the workspace logs directory for `headroom perf` analysis.
# Resolved lazily so HEADROOM_WORKSPACE_DIR env-var changes are honored.


def _headroom_log_dir() -> Path:
    return _paths.log_dir()


def _setup_file_logging() -> None:
    """Add a RotatingFileHandler to the headroom root logger.

    Writes to ~/.headroom/logs/proxy.log with automatic rotation:
    - Rotates at 10 MB
    - Keeps 5 backups (~50 MB max)
    """
    from logging.handlers import RotatingFileHandler

    try:
        log_dir = _headroom_log_dir()
        log_dir.mkdir(parents=True, exist_ok=True)
        log_path = log_dir / "proxy.log"
        handler = RotatingFileHandler(
            log_path,
            maxBytes=10 * 1024 * 1024,  # 10 MB
            backupCount=5,
            encoding="utf-8",
        )
        handler.setLevel(logging.INFO)
        handler.setFormatter(
            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
        )
        # Attach to the headroom root logger so all sub-loggers are captured.
        # Disable propagation to root to avoid duplicate writes when
        # wrap.py redirects stderr to the same log file.
        headroom_logger = logging.getLogger("headroom")
        headroom_logger.setLevel(logging.INFO)
        if not any(isinstance(h, RotatingFileHandler) for h in headroom_logger.handlers):
            headroom_logger.addHandler(handler)
        headroom_logger.propagate = False
    except OSError:
        # Non-fatal: can't write logs (read-only fs, permissions, etc.)
        pass


def _selected_context_tool() -> str:
    raw = os.environ.get(_CONTEXT_TOOL_ENV, _CONTEXT_TOOL_RTK).strip().lower()
    normalized = raw.replace("_", "-")
    if normalized in ("leanctx", _CONTEXT_TOOL_LEAN_CTX):
        return _CONTEXT_TOOL_LEAN_CTX
    return _CONTEXT_TOOL_RTK


def _context_tool_label(tool: str) -> str:
    if tool == _CONTEXT_TOOL_LEAN_CTX:
        return "lean-ctx"
    return "RTK"


def _context_tool_default_scope(tool: str) -> str:
    if tool == _CONTEXT_TOOL_LEAN_CTX:
        return "local"
    return _RTK_GAIN_SCOPE_GLOBAL


def _rtk_gain_scope() -> str:
    raw = os.environ.get(_RTK_GAIN_SCOPE_ENV, "").strip().lower()
    if not raw:
        return _RTK_GAIN_SCOPE_GLOBAL
    if raw in _RTK_GAIN_SCOPES:
        return raw

    logger.warning(
        "event=rtk_gain_scope_invalid env=%s value=%r default=%s",
        _RTK_GAIN_SCOPE_ENV,
        raw,
        _RTK_GAIN_SCOPE_GLOBAL,
    )
    return _RTK_GAIN_SCOPE_GLOBAL


def _rtk_gain_command(rtk_path: Any, scope: str) -> list[str]:
    command = [str(rtk_path), "gain"]
    if scope == _RTK_GAIN_SCOPE_PROJECT:
        command.append("--project")
    command.extend(["--format", "json"])
    return command


def _coerce_int(value: Any, default: int = 0) -> int:
    try:
        return int(value or 0)
    except (TypeError, ValueError):
        return default


def _coerce_float(value: Any, default: float = 0.0) -> float:
    try:
        return float(value or 0.0)
    except (TypeError, ValueError):
        return default


def _first_value(mapping: dict[str, Any], keys: tuple[str, ...], default: Any = 0) -> Any:
    for key in keys:
        if key in mapping and mapping[key] is not None:
            return mapping[key]
    return default


def _context_tool_summary_payload(
    *,
    tool: str,
    installed: bool,
    scope: str | None = None,
    summary: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Normalize RTK/lean-ctx lifetime gain output into one schema.

    Both tools expose cumulative counters, but field names vary slightly.
    Headroom computes session values by subtracting a startup baseline, so
    keeping raw input/output counters is necessary for a truthful session
    savings percentage.
    """

    summary = summary or {}
    input_tokens = _coerce_int(
        _first_value(
            summary,
            (
                "total_input",
                "total_input_tokens",
                "input_tokens",
                "tokens_input",
                "totalBefore",
            ),
        )
    )
    output_tokens = _coerce_int(
        _first_value(
            summary,
            (
                "total_output",
                "total_output_tokens",
                "output_tokens",
                "tokens_output",
                "totalAfter",
            ),
        )
    )
    tokens_saved = _coerce_int(
        _first_value(
            summary,
            (
                "total_saved",
                "tokens_saved",
                "total_tokens_saved",
                "saved_tokens",
                "totalSaved",
            ),
        )
    )
    if tokens_saved <= 0 and input_tokens > 0 and output_tokens >= 0:
        tokens_saved = max(input_tokens - output_tokens, 0)
    if input_tokens <= 0 and tokens_saved > 0 and output_tokens >= 0:
        input_tokens = tokens_saved + output_tokens

    lifetime_savings_pct = _coerce_float(
        _first_value(
            summary,
            (
                "avg_savings_pct",
                "average_savings_pct",
                "savings_pct",
                "savings_percent",
                "avgSavingsPct",
            ),
            0.0,
        )
    )
    if lifetime_savings_pct <= 0 and input_tokens > 0:
        lifetime_savings_pct = (tokens_saved / input_tokens) * 100.0

    return {
        "tool": tool,
        "label": _context_tool_label(tool),
        "installed": installed,
        "scope": scope or _context_tool_default_scope(tool),
        "total_commands": _coerce_int(
            _first_value(
                summary,
                (
                    "total_commands",
                    "commands",
                    "command_count",
                    "totalCommandCount",
                ),
            )
        ),
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "tokens_saved": tokens_saved,
        # Backward-compatible name. See `lifetime_avg_savings_pct` and
        # `session_savings_pct` below for explicit scopes.
        "avg_savings_pct": lifetime_savings_pct,
        "lifetime_avg_savings_pct": lifetime_savings_pct,
        "total_time_ms": _coerce_int(
            _first_value(summary, ("total_time_ms", "time_ms", "totalTimeMs"))
        ),
    }


def _context_tool_zero_payload(
    *,
    tool: str,
    installed: bool,
    scope: str | None = None,
) -> dict[str, Any]:
    return _context_tool_summary_payload(
        tool=tool,
        installed=installed,
        scope=scope,
        summary={},
    )


def _read_rtk_lifetime_stats() -> dict[str, Any] | None:
    """Read rtk's lifetime stats using the configured gain scope."""

    from headroom.rtk import get_rtk_path

    scope = _rtk_gain_scope()
    rtk_path = get_rtk_path()
    if not rtk_path:
        return _context_tool_zero_payload(
            tool=_CONTEXT_TOOL_RTK,
            installed=False,
            scope=scope,
        )

    try:
        result = run(
            _rtk_gain_command(rtk_path, scope),
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            data = json.loads(result.stdout)
            summary = data.get("summary", {})
            payload = _context_tool_summary_payload(
                tool=_CONTEXT_TOOL_RTK,
                installed=True,
                scope=scope,
                summary=summary if isinstance(summary, dict) else {},
            )
        else:
            # A failed read is "no data", never a zero counter — a synthetic
            # zero here re-pins the session baseline and inflates session
            # savings by the tool's whole lifetime on recovery.
            stderr_excerpt = (result.stderr or "")[:200]
            logger.warning(
                "event=rtk_stats_subprocess_failed reason=non_zero_exit rc=%s stderr=%r",
                result.returncode,
                stderr_excerpt,
            )
            return None
    except Exception as exc:
        # Reason is the exception class name (without payload — RTK
        # exceptions can carry filesystem paths).
        logger.warning(
            "event=rtk_stats_subprocess_failed reason=%s error=%s",
            type(exc).__name__,
            exc,
        )
        return None

    return payload


def _read_lean_ctx_lifetime_stats() -> dict[str, Any] | None:
    """Read lean-ctx's current project-level lifetime stats."""

    from headroom.lean_ctx import get_lean_ctx_path

    lean_ctx_path = get_lean_ctx_path()
    if not lean_ctx_path:
        return _context_tool_zero_payload(tool=_CONTEXT_TOOL_LEAN_CTX, installed=False)

    try:
        result = run(
            [str(lean_ctx_path), "gain", "--json"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        # Failed reads return None ("no data") — mirrors the rtk reader so
        # the baseline logic never sees synthetic zeros from either tool.
        if result.returncode != 0 or not result.stdout.strip():
            logger.warning(
                "event=lean_ctx_stats_subprocess_failed reason=non_zero_exit rc=%s",
                result.returncode,
            )
            return None

        data = json.loads(result.stdout)
        summary = data.get("summary", data) if isinstance(data, dict) else {}
        if not isinstance(summary, dict):
            logger.warning("event=lean_ctx_stats_subprocess_failed reason=bad_payload")
            return None

        return _context_tool_summary_payload(
            tool=_CONTEXT_TOOL_LEAN_CTX,
            installed=True,
            summary=summary,
        )
    except Exception as exc:
        logger.warning(
            "event=lean_ctx_stats_subprocess_failed reason=%s",
            type(exc).__name__,
        )
        return None


def _read_context_tool_lifetime_stats(tool: str) -> dict[str, Any] | None:
    if tool == _CONTEXT_TOOL_LEAN_CTX:
        return _read_lean_ctx_lifetime_stats()
    return _read_rtk_lifetime_stats()


async def initialize_context_tool_session_baseline() -> None:
    """Pin the current context-tool counters as the proxy-session baseline."""

    tool = _selected_context_tool()
    payload = await asyncio.to_thread(_read_context_tool_lifetime_stats, tool)
    with _context_tool_stats_cache_lock:
        if payload is None or not payload.get("installed", False):
            # Failed or tool-absent read: defer the pin to the first
            # successful read (guarded lazy-init) — pinning zeros here would
            # inflate session savings by the tool's whole lifetime once it
            # recovers or gets installed.
            _context_tool_session_baseline.update(
                {
                    "initialized": False,
                    "tool": tool,
                    "total_commands": 0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "tokens_saved": 0,
                    "total_time_ms": 0,
                    "captured_at": time.time(),
                }
            )
        else:
            _context_tool_session_baseline.update(
                {
                    "initialized": True,
                    "tool": tool,
                    "total_commands": int(payload.get("total_commands", 0) or 0),
                    "input_tokens": int(payload.get("input_tokens", 0) or 0),
                    "output_tokens": int(payload.get("output_tokens", 0) or 0),
                    "tokens_saved": int(payload.get("tokens_saved", 0) or 0),
                    "total_time_ms": int(payload.get("total_time_ms", 0) or 0),
                    "captured_at": time.time(),
                }
            )
        _context_tool_stats_cache.update(
            {
                "expires_at": 0.0,
                "has_value": False,
                "tool": None,
                "value": None,
            }
        )


async def initialize_rtk_session_baseline() -> None:
    """Backward-compatible alias for initialize_context_tool_session_baseline."""

    await initialize_context_tool_session_baseline()


def _get_context_tool_stats() -> dict[str, Any] | None:
    """Get context-tool savings for the current Headroom proxy session.

    RTK and lean-ctx persist project-level lifetime counters. Dashboard stats
    should be session-local, so we subtract the counter snapshot captured at
    proxy startup instead of resetting the tool's own history.
    """

    tool = _selected_context_tool()
    now = time.monotonic()
    with _context_tool_stats_cache_lock:
        cached_value = cast(dict[str, Any] | None, _context_tool_stats_cache["value"])
        if (
            _context_tool_stats_cache["has_value"]
            and now < float(_context_tool_stats_cache["expires_at"])
            and _context_tool_stats_cache.get("tool") == tool
        ):
            return cached_value

    payload = _read_context_tool_lifetime_stats(tool)
    with _context_tool_stats_cache_lock:
        # Baseline mutations only happen on successful reads from an
        # installed tool — a failed read (None) or a tool-absent zero payload
        # must never pin or re-pin, or session deltas inflate by the whole
        # lifetime when the tool comes back.
        tool_installed = payload is not None and bool(payload.get("installed", False))
        if (
            payload is not None
            and tool_installed
            and (
                not _context_tool_session_baseline["initialized"]
                or _context_tool_session_baseline.get("tool") != tool
            )
        ):
            _context_tool_session_baseline.update(
                {
                    "initialized": True,
                    "tool": tool,
                    "total_commands": int(payload.get("total_commands", 0) or 0),
                    "input_tokens": int(payload.get("input_tokens", 0) or 0),
                    "output_tokens": int(payload.get("output_tokens", 0) or 0),
                    "tokens_saved": int(payload.get("tokens_saved", 0) or 0),
                    "total_time_ms": int(payload.get("total_time_ms", 0) or 0),
                    "captured_at": time.time(),
                }
            )

        if payload is not None:
            lifetime_total_commands = int(payload.get("total_commands", 0) or 0)
            lifetime_input_tokens = int(payload.get("input_tokens", 0) or 0)
            lifetime_output_tokens = int(payload.get("output_tokens", 0) or 0)
            lifetime_tokens_saved = int(payload.get("tokens_saved", 0) or 0)
            lifetime_total_time_ms = int(payload.get("total_time_ms", 0) or 0)
            baseline_total_commands = int(_context_tool_session_baseline["total_commands"])
            baseline_input_tokens = int(_context_tool_session_baseline["input_tokens"])
            baseline_output_tokens = int(_context_tool_session_baseline["output_tokens"])
            baseline_tokens_saved = int(_context_tool_session_baseline["tokens_saved"])
            baseline_total_time_ms = int(_context_tool_session_baseline["total_time_ms"])
            # A tool-absent payload carries zero counters that are not a
            # genuine external reset — only successful installed reads may
            # re-pin the baseline.
            counter_reset_detected = tool_installed and (
                lifetime_total_commands < baseline_total_commands
                or lifetime_input_tokens < baseline_input_tokens
                or lifetime_output_tokens < baseline_output_tokens
                or lifetime_tokens_saved < baseline_tokens_saved
                or lifetime_total_time_ms < baseline_total_time_ms
            )
            if counter_reset_detected:
                baseline_total_commands = lifetime_total_commands
                baseline_input_tokens = lifetime_input_tokens
                baseline_output_tokens = lifetime_output_tokens
                baseline_tokens_saved = lifetime_tokens_saved
                baseline_total_time_ms = lifetime_total_time_ms
                _context_tool_session_baseline.update(
                    {
                        "total_commands": baseline_total_commands,
                        "input_tokens": baseline_input_tokens,
                        "output_tokens": baseline_output_tokens,
                        "tokens_saved": baseline_tokens_saved,
                        "total_time_ms": baseline_total_time_ms,
                        "captured_at": time.time(),
                    }
                )

            session_total_commands = max(lifetime_total_commands - baseline_total_commands, 0)
            session_input_tokens = max(lifetime_input_tokens - baseline_input_tokens, 0)
            session_output_tokens = max(lifetime_output_tokens - baseline_output_tokens, 0)
            session_tokens_saved = max(lifetime_tokens_saved - baseline_tokens_saved, 0)
            session_total_time_ms = max(lifetime_total_time_ms - baseline_total_time_ms, 0)
            session_savings_pct = (
                round(session_tokens_saved / session_input_tokens * 100.0, 4)
                if session_input_tokens > 0
                else None
            )
            session_avg_time_ms = (
                round(session_total_time_ms / session_total_commands, 2)
                if session_total_commands > 0 and session_total_time_ms > 0
                else None
            )
            lifetime_savings_pct = float(payload.get("lifetime_avg_savings_pct", 0.0) or 0.0)

            payload = {
                **payload,
                "tool": tool,
                "label": _context_tool_label(tool),
                # Backward-compatible session-delta fields.
                "total_commands": session_total_commands,
                "input_tokens": session_input_tokens,
                "output_tokens": session_output_tokens,
                "tokens_saved": session_tokens_saved,
                "total_time_ms": session_total_time_ms,
                "session_savings_pct": session_savings_pct,
                "session_avg_time_ms": session_avg_time_ms,
                # Keep old field for compatibility, but declare its scope.
                "avg_savings_pct": lifetime_savings_pct,
                "avg_savings_pct_scope": "lifetime",
                "lifetime_avg_savings_pct": lifetime_savings_pct,
                "lifetime_total_commands": lifetime_total_commands,
                "lifetime_input_tokens": lifetime_input_tokens,
                "lifetime_output_tokens": lifetime_output_tokens,
                "lifetime_tokens_saved": lifetime_tokens_saved,
                "lifetime_total_time_ms": lifetime_total_time_ms,
                "session_baseline_total_commands": baseline_total_commands,
                "session_baseline_input_tokens": baseline_input_tokens,
                "session_baseline_output_tokens": baseline_output_tokens,
                "session_baseline_tokens_saved": baseline_tokens_saved,
                "session_baseline_total_time_ms": baseline_total_time_ms,
                "session_baseline_captured_at": _context_tool_session_baseline.get(
                    "captured_at", 0.0
                ),
                "session": {
                    "commands": session_total_commands,
                    "input_tokens": session_input_tokens,
                    "output_tokens": session_output_tokens,
                    "tokens_saved": session_tokens_saved,
                    "savings_pct": session_savings_pct,
                    "total_time_ms": session_total_time_ms,
                    "avg_time_ms": session_avg_time_ms,
                },
                "lifetime": {
                    "commands": lifetime_total_commands,
                    "input_tokens": lifetime_input_tokens,
                    "output_tokens": lifetime_output_tokens,
                    "tokens_saved": lifetime_tokens_saved,
                    "savings_pct": lifetime_savings_pct,
                    "total_time_ms": lifetime_total_time_ms,
                },
                "baseline": {
                    "commands": baseline_total_commands,
                    "input_tokens": baseline_input_tokens,
                    "output_tokens": baseline_output_tokens,
                    "tokens_saved": baseline_tokens_saved,
                    "total_time_ms": baseline_total_time_ms,
                    "captured_at": _context_tool_session_baseline.get("captured_at", 0.0),
                },
                "sampled_at": time.time(),
                "sample_ttl_seconds": CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
                "refresh_interval_seconds": CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
                "counter_reset_detected": counter_reset_detected,
            }

        _context_tool_stats_cache.update(
            {
                "expires_at": time.monotonic() + CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
                "has_value": True,
                "tool": tool,
                "value": payload,
            }
        )
    return payload


def _get_rtk_stats() -> dict[str, Any] | None:
    """Backward-compatible alias for selected context-tool stats."""

    return _get_context_tool_stats()


def is_anthropic_auth(headers: dict[str, str]) -> bool:
    """Detect Anthropic auth signals in request headers."""
    if headers.get("x-api-key") or headers.get("anthropic-version"):
        return True
    auth = headers.get("authorization", "")
    if auth.startswith("Bearer sk-ant-"):
        return True
    return False


# ---------------------------------------------------------------------------
# Internal-header stripping (PR-A5 — fixes P5-49).
# ---------------------------------------------------------------------------
#
# `x-headroom-*` request headers (e.g. ``x-headroom-bypass``,
# ``x-headroom-mode``, ``x-headroom-user-id``, ``x-headroom-stack``,
# ``x-headroom-base-url``) are internal control flags consumed by the
# proxy itself. They MUST NOT leak upstream — leaking them would (a)
# fingerprint the proxy to subscription enforcers and (b) expose the
# user-id/stack/base-url internals to whichever vendor terminates the
# request.
#
# Inbound read paths (bypass gating, ``_extract_tags`` reading
# ``x-headroom-*``, memory ``x-headroom-user-id`` lookup) keep using
# the original dict / ``request.headers``. The stripped copy is what
# every upstream-bound forwarder receives.
#
# Note: response-side ``X-Headroom-*`` injection (e.g.
# ``x-headroom-tokens-saved``) is unrelated — the proxy is allowed to
# tell its client about its own work. This helper only filters
# request-side headers.

_INTERNAL_HEADER_PREFIX = INTERNAL_HEADER_PREFIX
_STRIP_INTERNAL_HEADERS_ENV = STRIP_INTERNAL_HEADERS_ENV
_STRIP_INTERNAL_HEADERS_DEFAULT = STRIP_INTERNAL_HEADERS_DEFAULT


def get_strip_internal_headers_mode() -> StripInternalHeadersMode:
    """Return the active internal-header strip mode.

    Read at request time so operators can flip behaviour without a
    restart. Unknown values raise loudly per the no-silent-fallback
    build constraint.
    """
    return resolve_strip_internal_headers_mode(os.environ.get(_STRIP_INTERNAL_HEADERS_ENV))


def _strip_internal_headers(headers: dict[str, str]) -> dict[str, str]:
    """Return a copy of ``headers`` with internal ``x-headroom-*`` keys stripped.

    Used at every upstream call site to prevent fingerprinting / leakage of
    internal flags like ``x-headroom-bypass``, ``x-headroom-mode``,
    ``x-headroom-user-id``, ``x-headroom-stack``, ``x-headroom-base-url``.
    Case-insensitive on the prefix. Returns a NEW dict; never mutates the
    caller's mapping. Pure function. No regex.

    When the operator opt-in ``HEADROOM_STRIP_INTERNAL_HEADERS=disabled``
    is set, returns a shallow copy unchanged. That mode is for diagnostic
    shadow tracing only and is documented as a per-deploy choice.
    """
    return strip_internal_headers(headers, mode=get_strip_internal_headers_mode())


def log_outbound_headers(
    *,
    forwarder: str,
    stripped_count: int,
    request_id: str | None,
) -> None:
    """Structured log line for every upstream forwarder header strip.

    Emitted once per outbound request (paired with ``log_outbound_request``).
    Per realignment build constraint #8 we log every cache-affecting
    decision; per #8/#11 we never log header values, only the count of
    stripped internal headers.
    """
    logger.info(
        "event=outbound_headers forwarder=%s stripped_count=%d request_id=%s",
        forwarder,
        stripped_count,
        request_id or "",
    )


# ---------------------------------------------------------------------------
# Beta-header merge + per-session stickiness (PR-A6 — fixes P5-50; preps P0-6).
# ---------------------------------------------------------------------------
#
# Anthropic's `anthropic-beta` and OpenAI's `OpenAI-Beta` request headers
# carry a comma-separated list of opt-in beta tokens. Two cache-killer
# patterns motivated PR-A6:
#
#   1. Mid-session mutation: when memory is enabled the proxy historically
#      did an ad-hoc concat of `context-management-2025-06-27` onto the
#      client value (anthropic.py:1244-1248) — every variant produced a
#      different byte sequence and the order was undefined when the same
#      client value already contained a Headroom-required token.
#
#   2. Token drop-out across turns: clients (Claude Code, Codex CLI) MAY
#      drop a beta token between turn N and turn N+1 even when the proxy
#      mutated turn N to add it. The cache hot zone is positional, so the
#      next turn's prefix bytes hash differently and the prefix-cache
#      read misses.
#
# PR-A6 introduces:
#   * `merge_anthropic_beta` / `merge_openai_beta`: deterministic, pure,
#     order-preserving merge. Client tokens first (in their original order),
#     then Headroom-required tokens (in the order passed). Dedupe is
#     case-insensitive but preserves original casing of first occurrence.
#     Per Anthropic guide §6.3 #6: sticky-on means we add but never reorder.
#
#   * `SessionBetaTracker`: bounded LRU cache keyed by `(provider,
#     session_id)` tracking every beta token observed for that session.
#     On every request we union the client value with previously-seen
#     tokens and update the seen set — so a beta seen in turn N is
#     present in turn N+1 even if the client drops it. LRU bound (default
#     1000 sessions) prevents unbounded growth. Reentrant lock so future
#     callers from inside another locked method don't self-deadlock.
#
# Operator opt-in `HEADROOM_BETA_HEADER_STICKY=disabled` short-circuits
# the tracker (returns the client value verbatim). That mode is loud and
# explicit per realignment build constraint #4 — NOT a silent fallback.

_BETA_HEADER_STICKY_ENV = BETA_HEADER_STICKY_ENV
_BETA_HEADER_STICKY_DEFAULT = BETA_HEADER_STICKY_DEFAULT

_BETA_TRACKER_MAX_SESSIONS_ENV = BETA_TRACKER_MAX_SESSIONS_ENV
_BETA_TRACKER_MAX_SESSIONS_DEFAULT = BETA_TRACKER_MAX_SESSIONS_DEFAULT


def get_beta_header_sticky_mode() -> BetaHeaderStickyMode:
    """Return the active beta-header stickiness mode.

    Read at request time so operators can flip behaviour without a
    restart. Unknown values raise loudly per the no-silent-fallback
    build constraint.
    """
    return resolve_beta_header_sticky_mode(os.environ.get(_BETA_HEADER_STICKY_ENV))


def get_beta_tracker_max_sessions() -> int:
    """Return the LRU bound for `SessionBetaTracker` (sessions cap)."""
    return resolve_beta_tracker_max_sessions(os.environ.get(_BETA_TRACKER_MAX_SESSIONS_ENV))


def _split_beta_tokens(value: str | None) -> list[str]:
    """Split a comma-separated beta-header value into trimmed tokens.

    Empty/whitespace-only entries are dropped. Pure function, no regex.
    """
    if not value:
        return []
    out: list[str] = []
    for raw in value.split(","):
        token = raw.strip()
        if token:
            out.append(token)
    return out


def _merge_beta_tokens(client_value: str | None, headroom_required: list[str]) -> str:
    """Shared deterministic merge for `anthropic-beta` / `OpenAI-Beta` tokens.

    Rules (per Anthropic guide §6.3 #6 "sticky-on; add but never reorder"):

    * Client tokens come first, in their original order.
    * Headroom-required tokens append in the order given, skipping any
      token already present (case-insensitive).
    * Dedupe is case-insensitive but the FIRST occurrence's casing wins
      (prevents drift when client uses one casing across turns).
    * Returns ``""`` when both inputs are empty.

    Pure function. No regex. No global state.
    """
    seen_lower: set[str] = set()
    out: list[str] = []
    for token in _split_beta_tokens(client_value):
        lower = token.lower()
        if lower in seen_lower:
            continue
        seen_lower.add(lower)
        out.append(token)
    for token in headroom_required:
        if not token:
            continue
        token = token.strip()
        if not token:
            continue
        lower = token.lower()
        if lower in seen_lower:
            continue
        seen_lower.add(lower)
        out.append(token)
    return ",".join(out)


def merge_anthropic_beta(client_value: str | None, headroom_required: list[str]) -> str:
    """Merge client `anthropic-beta` value with Headroom-required tokens.

    See `_merge_beta_tokens` for full semantics. Order is deterministic:
    client tokens first (in their original order), then headroom tokens
    (in the order passed). No sorting — sticky-on per Anthropic guide
    §6.3 #6 means we add but never reorder. Dedupe is case-insensitive
    but preserves the original casing of the first occurrence.

    Returns ``""`` when both inputs are empty.
    """
    return _merge_beta_tokens(client_value, headroom_required)


def merge_openai_beta(client_value: str | None, headroom_required: list[str]) -> str:
    """Merge client `OpenAI-Beta` value with Headroom-required tokens.

    Mirror of `merge_anthropic_beta`. Same semantics — the OpenAI header
    follows the same comma-separated convention and the same cache-stable
    rules apply.
    """
    return _merge_beta_tokens(client_value, headroom_required)


class SessionBetaTracker:
    """Bounded LRU tracker of beta-header tokens observed per (provider, session).

    On every request:
      * Read the client's beta-header value.
      * Union with previously-seen tokens for this session (sticky-on).
      * Update the session's seen set.
      * Return the union (preserving first-seen order).

    Bounded by `max_sessions` (default 1000) via `OrderedDict` LRU
    eviction: hits move-to-end; overflow pops oldest. Reentrant lock so
    future callers from inside another locked method don't self-deadlock
    (mirrors `CompressionCache` pattern).

    The tracker is provider-aware: the same `session_id` for Anthropic
    and OpenAI keeps independent token sets (clients/upstreams differ on
    which tokens are valid).
    """

    def __init__(self, max_sessions: int | None = None) -> None:
        if max_sessions is None:
            max_sessions = get_beta_tracker_max_sessions()
        if max_sessions <= 0:
            raise ValueError("max_sessions must be > 0")
        self._max_sessions: int = max_sessions
        # OrderedDict per `compression_cache.py` LRU pattern. Entries
        # store the per-session ordered token list (preserving first-seen
        # order). RLock allows future callers from inside another locked
        # method to enter without self-deadlock.
        self._lock = threading.RLock()
        self._sessions: OrderedDict[tuple[str, str], list[str]] = OrderedDict()

    @property
    def active_sessions(self) -> int:
        with self._lock:
            return len(self._sessions)

    def _key(self, provider: str, session_id: str) -> tuple[str, str]:
        return (provider, session_id)

    def record_and_get_sticky_betas(
        self,
        provider: str,
        session_id: str,
        client_value: str | None,
    ) -> str:
        """Union client tokens with session-seen tokens; update; return.

        ``provider`` is the upstream identifier (``anthropic`` /
        ``openai``). ``session_id`` is the proxy's per-conversation ID
        (e.g. `SessionTrackerStore.compute_session_id` output for the
        HTTP path; the WS handler's per-connection UUID for the WS
        path — note WS sessions are short-lived and won't accumulate
        cross-turn).

        When `HEADROOM_BETA_HEADER_STICKY=disabled` returns the client
        value verbatim (operator diagnostic opt-in; documented as a
        per-deploy choice, NOT a silent fallback).

        Returns the merged comma-separated value (possibly empty).
        """
        if not provider:
            raise ValueError("provider must be non-empty")
        if not session_id:
            raise ValueError("session_id must be non-empty")

        if get_beta_header_sticky_mode() == "disabled":
            # Diagnostic mode — return the client value verbatim, do not
            # touch tracker state. This is loud (operators read the env
            # var) and per-deploy.
            return (client_value or "").strip()

        client_tokens = _split_beta_tokens(client_value)
        key = self._key(provider, session_id)

        with self._lock:
            previous = self._sessions.get(key)
            if previous is None:
                merged_list: list[str] = []
                seen_lower: set[str] = set()
            else:
                # Move-to-end on hit (LRU touch).
                self._sessions.move_to_end(key)
                merged_list = list(previous)
                seen_lower = {t.lower() for t in merged_list}

            # Append client tokens preserving order; first-seen casing wins.
            for token in client_tokens:
                lower = token.lower()
                if lower in seen_lower:
                    continue
                seen_lower.add(lower)
                merged_list.append(token)

            self._sessions[key] = merged_list
            self._sessions.move_to_end(key)

            # Bound: evict oldest until at-or-below cap.
            while len(self._sessions) > self._max_sessions:
                self._sessions.popitem(last=False)

            return ",".join(merged_list)

    def reset(self) -> None:
        """Clear all session state (test helper)."""
        with self._lock:
            self._sessions.clear()


# Process-wide singleton. Lazily replaced by tests via `reset` /
# `_reset_session_beta_tracker_for_test`. One tracker for both providers
# — the (provider, session_id) key keeps namespaces independent.
_session_beta_tracker_lock = threading.Lock()
_session_beta_tracker: SessionBetaTracker | None = None


def get_session_beta_tracker() -> SessionBetaTracker:
    """Return the process-wide `SessionBetaTracker` singleton.

    Lazily constructed so the env-var bound (`HEADROOM_BETA_TRACKER_MAX_SESSIONS`)
    is honored at first use. Tests use `_reset_session_beta_tracker_for_test`.
    """
    global _session_beta_tracker
    with _session_beta_tracker_lock:
        if _session_beta_tracker is None:
            _session_beta_tracker = SessionBetaTracker()
        return _session_beta_tracker


def _reset_session_beta_tracker_for_test() -> None:
    """Clear the process-wide tracker (test-only)."""
    global _session_beta_tracker
    with _session_beta_tracker_lock:
        _session_beta_tracker = None


def log_beta_header_merge(
    *,
    provider: str,
    session_id: str | None,
    client_betas_count: int,
    sticky_betas_count: int,
    headroom_added: list[str],
    request_id: str | None,
) -> None:
    """Structured log for every cache-affecting beta-header merge.

    `headroom_added` is a list of public, documented beta tokens
    (e.g. ``context-management-2025-06-27``,
    ``responses_websockets=2026-02-06``) — safe to log. We intentionally
    do NOT log the raw client value because beta tokens, while public,
    can carry experiment IDs the user has not opted to share with
    Headroom logs. Emitting counts only makes the decision auditable.
    """
    logger.info(
        "event=beta_header_merge provider=%s session_id=%s "
        "client_betas=%d sticky_betas=%d headroom_added=%s request_id=%s",
        provider,
        session_id or "",
        client_betas_count,
        sticky_betas_count,
        ",".join(headroom_added) if headroom_added else "",
        request_id or "",
    )


# ---------------------------------------------------------------------------
# Memory-tool injection session-stickiness (PR-A7 — closes P0-6).
# ---------------------------------------------------------------------------
#
# Memory adds `memory_save` / `memory_search` tool definitions to
# `body["tools"]` when memory is enabled for a request. The cache-killer
# pattern motivated by guide §6.3 #2 ("tool list change → cache bust"):
#
#   * Mid-session toggle: memory is enabled in turn N (tool definitions
#     injected) and disabled in turn N+1 (tool list shrinks). The next
#     turn's prefix bytes hash differently, prefix-cache misses, and the
#     full prompt re-runs at provider cost.
#
#   * Tool definition drift: memory adds the SAME logical tool but the
#     bytes differ across turns (insertion order, dict key order, schema
#     drift between deploys, etc.). Even with the tool list intact the
#     prefix bytes change.
#
# PR-A7 introduces:
#
#   * `SessionToolTracker`: bounded LRU keyed by (provider, session_id)
#     storing the GOLDEN tool-definition bytes injected on the first
#     turn. Subsequent turns of that session always inject the same
#     bytes — even if memory is disabled mid-session (sticky-on per
#     guide §6.3 #2). Provider-aware so the same `session_id` under
#     two providers keeps independent state.
#
# The golden bytes are produced by `serialize_body_canonical` of the
# tool definition object so they are deterministic across deploys
# regardless of dict insertion ordering quirks.
#
# Operator opt-in `HEADROOM_TOOL_INJECTION_STICKY=disabled` short-
# circuits the tracker; per-turn decision flows through unchanged. That
# mode is loud and explicit per realignment build constraint #4 — NOT a
# silent fallback. It exists for diagnostic shadow tracing / emergency
# rollback only.


def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode:
    """Return the active memory-tool stickiness mode.

    Read at request time so operators can flip behaviour without a
    restart. Unknown values raise loudly per the no-silent-fallback
    build constraint.
    """
    return _get_tool_injection_sticky_mode()


def get_tool_tracker_max_sessions() -> int:
    """Return the LRU bound for `SessionToolTracker` (sessions cap)."""
    return _get_tool_tracker_max_sessions()


def serialize_tool_definition_canonical(tool_definition: dict[str, Any]) -> bytes:
    """Deterministic byte serialization of a single memory tool definition.

    Uses ``serialize_body_canonical`` semantics (compact separators, UTF-8,
    no ASCII escaping). Python 3.7+ dict insertion order is preserved by
    ``json.dumps`` so callers must construct the tool definition with a
    stable key order — which the static schemas in
    ``headroom/proxy/memory_handler.py`` and
    ``headroom/proxy/memory_tool_adapter.py`` already do.

    Returned bytes pin the golden tool definition for a session: every
    follow-up turn must inject byte-equal output to keep the prefix
    cache hot.
    """
    return serialize_body_canonical(tool_definition)


class SessionToolTracker(_SessionToolTracker):
    """Env-aware compatibility wrapper for the pure session tool tracker."""

    def __init__(self, max_sessions: int | None = None) -> None:
        if max_sessions is None:
            max_sessions = get_tool_tracker_max_sessions()
        super().__init__(max_sessions=max_sessions)


# Process-wide singleton. Lazily replaced by tests via
# `_reset_session_tool_tracker_for_test`.
_session_tool_tracker_lock = threading.Lock()
_session_tool_tracker: SessionToolTracker | None = None


def get_session_tool_tracker() -> SessionToolTracker:
    """Return the process-wide `SessionToolTracker` singleton.

    Lazily constructed so the env-var bound
    (`HEADROOM_TOOL_TRACKER_MAX_SESSIONS`) is honored at first use.
    Tests use ``_reset_session_tool_tracker_for_test``.
    """
    global _session_tool_tracker
    with _session_tool_tracker_lock:
        if _session_tool_tracker is None:
            _session_tool_tracker = SessionToolTracker()
        return _session_tool_tracker


def _reset_session_tool_tracker_for_test() -> None:
    """Clear the process-wide tracker (test-only)."""
    global _session_tool_tracker
    with _session_tool_tracker_lock:
        _session_tool_tracker = None


def log_tool_injection_decision(
    *,
    provider: str,
    session_id: str | None,
    decision: ToolInjectionDecision,
    tool_definition_bytes_count: int,
    request_id: str | None,
) -> None:
    """Structured log for every cache-affecting tool-injection decision.

    Per realignment build constraint #8 we log every cache-affecting
    decision. ``tool_definition_bytes_count`` is the per-tool byte count
    summed across all memory tools injected this turn. We do NOT log the
    tool definition contents (might contain user-specific schemas) per
    constraint #11.
    """
    _log_tool_injection_decision(
        logger=logger,
        provider=provider,
        session_id=session_id,
        decision=decision,
        tool_definition_bytes_count=tool_definition_bytes_count,
        request_id=request_id,
    )


def _extract_tool_name(tool_definition: dict[str, Any]) -> str | None:
    """Extract a stable tool name from a memory tool definition.

    Handles three formats:
      * Anthropic custom: ``{"name": "memory_save", ...}``
      * Anthropic native: ``{"type": "memory_20250818", "name": "memory"}``
      * OpenAI function: ``{"type": "function", "function": {"name": "memory_save", ...}}``
    """
    return extract_tool_name(tool_definition)


def apply_session_sticky_memory_tools(
    *,
    provider: Literal["anthropic", "openai"],
    session_id: str | None,
    request_id: str | None,
    existing_tools: list[dict[str, Any]] | None,
    memory_tools_to_inject: list[dict[str, Any]],
    inject_this_turn: bool,
) -> tuple[list[dict[str, Any]], bool]:
    """Apply sticky-on memory tool injection per `SessionToolTracker`.

    The single coordination point for all memory-tool injection sites
    (Anthropic custom tools, Anthropic native tool, OpenAI function tools).

    Logic (guide §6.3 #2):

      * If ``HEADROOM_TOOL_INJECTION_STICKY=disabled``: bypass tracker,
        inject only when ``inject_this_turn`` is True. Diagnostic mode.

      * If session previously injected and tracker has golden bytes:
        ALWAYS inject the golden bytes verbatim (sticky-on). Memory-this-
        turn flag is irrelevant — once injected, always injected.

      * If session has NOT previously injected:
          - ``inject_this_turn=True``: serialize ``memory_tools_to_inject``,
            record golden bytes, append to tools list.
          - ``inject_this_turn=False``: skip; no future replay obligation.

    Memory tools whose names already appear in ``existing_tools`` are
    NOT re-appended (the client owns the canonical definition then).

    ``session_id`` may be ``None`` (e.g. WS path with no per-turn
    session); in that case the tracker is bypassed and the caller's
    ``inject_this_turn`` flag drives the decision verbatim. We log the
    bypass once so operators can see it.

    Returns ``(updated_tools, was_injected)``. The returned list is a
    fresh list (caller-safe). ``was_injected`` is True iff at least one
    memory tool was added to the list.
    """
    if provider not in ("anthropic", "openai"):
        raise ValueError(f"unsupported provider: {provider!r}")

    tools_out: list[dict[str, Any]] = list(existing_tools) if existing_tools else []
    existing_names: set[str] = set()
    for t in tools_out:
        n = _extract_tool_name(t)
        if n:
            existing_names.add(n)

    # Diagnostic / rollback path.
    if get_tool_injection_sticky_mode() == "disabled":
        if not inject_this_turn:
            log_tool_injection_decision(
                provider=provider,
                session_id=session_id,
                decision="skip_disabled_via_env",
                tool_definition_bytes_count=0,
                request_id=request_id,
            )
            return tools_out, False
        # Disabled mode + inject_this_turn=True: append the definitions
        # verbatim without recording golden bytes (per-turn decision
        # passes through as the broken behavior — explicit operator
        # opt-in only). Skip names already in the list.
        added_bytes = 0
        for tool_def in memory_tools_to_inject:
            tn = _extract_tool_name(tool_def)
            if tn is None or tn in existing_names:
                continue
            tools_out.append(tool_def)
            existing_names.add(tn)
            added_bytes += len(serialize_memory_tool_definition_canonical(tool_def))
        log_tool_injection_decision(
            provider=provider,
            session_id=session_id,
            decision="skip_disabled_via_env",
            tool_definition_bytes_count=added_bytes,
            request_id=request_id,
        )
        return tools_out, added_bytes > 0

    # Sticky path requires a session_id. None means we cannot track —
    # fall back to the caller's per-turn decision (loud, single log line)
    # so WS handlers / pre-session paths remain functional.
    if not session_id:
        if not inject_this_turn:
            log_tool_injection_decision(
                provider=provider,
                session_id=None,
                decision="skip",
                tool_definition_bytes_count=0,
                request_id=request_id,
            )
            return tools_out, False
        added_bytes = 0
        for tool_def in memory_tools_to_inject:
            tn = _extract_tool_name(tool_def)
            if tn is None or tn in existing_names:
                continue
            tools_out.append(tool_def)
            existing_names.add(tn)
            added_bytes += len(serialize_memory_tool_definition_canonical(tool_def))
        log_tool_injection_decision(
            provider=provider,
            session_id=None,
            decision="inject_first_time",
            tool_definition_bytes_count=added_bytes,
            request_id=request_id,
        )
        return tools_out, added_bytes > 0

    tracker = get_session_tool_tracker()
    previously_injected = tracker.should_inject(provider, session_id)

    if previously_injected:
        # Sticky replay: always inject the golden bytes. inject_this_turn
        # flag is intentionally ignored (memory may be disabled this turn
        # but the cache prefix demands the same tool list as before).
        golden = tracker.get_golden_definitions(provider, session_id) or []
        replay_bytes = 0
        for tool_name, golden_bytes in golden:
            if tool_name in existing_names:
                # Client also has a tool by this name — don't double up.
                # Their bytes win (the client's choice, not ours to gate).
                continue
            try:
                replay = replay_golden_memory_tool_definition(
                    tool_name=tool_name,
                    golden_tool_bytes=golden_bytes,
                )
            except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                logger.error(
                    "corrupt golden tool bytes for session %s tool %s: %s — skipping tool injection",
                    session_id,
                    tool_name,
                    exc,
                    exc_info=True,
                )
                continue
            tools_out.append(replay.tool_definition)
            existing_names.add(replay.tool_name)
            replay_bytes += len(replay.canonical_bytes)
        log_tool_injection_decision(
            provider=provider,
            session_id=session_id,
            decision="inject_sticky_replay",
            tool_definition_bytes_count=replay_bytes,
            request_id=request_id,
        )
        return tools_out, replay_bytes > 0

    # Fresh session.
    if not inject_this_turn:
        log_tool_injection_decision(
            provider=provider,
            session_id=session_id,
            decision="skip",
            tool_definition_bytes_count=0,
            request_id=request_id,
        )
        return tools_out, False

    # First-time inject: serialize, record, append.
    added_bytes = 0
    for tool_def in memory_tools_to_inject:
        tn = _extract_tool_name(tool_def)
        if tn is None or tn in existing_names:
            continue
        golden_bytes = serialize_memory_tool_definition_canonical(tool_def)
        tracker.record_injection(
            provider=provider,
            session_id=session_id,
            tool_name=tn,
            tool_definition_bytes=golden_bytes,
        )
        tools_out.append(tool_def)
        existing_names.add(tn)
        added_bytes += len(golden_bytes)
    log_tool_injection_decision(
        provider=provider,
        session_id=session_id,
        decision="inject_first_time",
        tool_definition_bytes_count=added_bytes,
        request_id=request_id,
    )
    return tools_out, added_bytes > 0


# ─── Session-sticky CCR tool injection (PR-B7) ─────────────────────────
#
# Per realignment plan PR-B7 (`REALIGNMENT/04-phase-B-live-zone.md`):
# once a session has performed any CCR compression, the
# `headroom_retrieve` tool stays registered in `body["tools"]` for every
# subsequent request in that session — never toggled off.
#
# The legacy `CCRToolInjector.has_compressed_content` flips on/off based
# on whether the *latest request* contained compression markers, which
# bust the prompt cache every time the flag flips. Sticky-on means the
# tool list bytes stay byte-stable across turns once injected.


class SessionCcrTracker(_SessionCcrTracker):
    """Env-aware compatibility wrapper for the pure CCR session tracker."""

    def __init__(self, max_sessions: int | None = None) -> None:
        if max_sessions is None:
            max_sessions = get_tool_tracker_max_sessions()
        super().__init__(max_sessions=max_sessions)


# Process-wide singleton.
_session_ccr_tracker_lock = threading.Lock()
_session_ccr_tracker: SessionCcrTracker | None = None


def get_session_ccr_tracker() -> SessionCcrTracker:
    """Return the process-wide :class:`SessionCcrTracker` singleton."""
    global _session_ccr_tracker
    with _session_ccr_tracker_lock:
        if _session_ccr_tracker is None:
            _session_ccr_tracker = SessionCcrTracker()
        return _session_ccr_tracker


def _reset_session_ccr_tracker_for_test() -> None:
    """Clear the process-wide CCR tracker (test-only)."""
    global _session_ccr_tracker
    with _session_ccr_tracker_lock:
        _session_ccr_tracker = None


def has_new_ccr_markers(
    *,
    current_detected_hashes: list[str],
    previous_forwarded_messages: list[dict[str, Any]] | None,
    provider: Literal["anthropic", "openai", "google"],
) -> bool:
    """Whether the about-to-forward content carries CCR markers NOT already forwarded.

    ``overlay_cached_prefix`` (#1850) replays the previously-forwarded (compressed)
    prefix byte-identical to keep the prompt cache warm — which reintroduces the
    ``hash=…`` markers that prefix already carried. Those markers are *historical*:
    the agent saw them last turn and the retrieve-tool state was already settled
    for them. Only markers that are genuinely NEW this turn justify overriding the
    tool-injection deferral (#1006); counting the replayed ones would re-inject the
    tool on every frozen turn and bust the *tools* cache segment (undoing the very
    cache-safety the overlay provides).

    Returns True iff ``current_detected_hashes`` contains a hash that is not present
    in ``previous_forwarded_messages``.
    """
    return _has_new_ccr_markers(
        current_detected_hashes=current_detected_hashes,
        previous_forwarded_messages=previous_forwarded_messages,
        provider=provider,
    )


def should_inject_ccr_tool(
    *,
    configured_inject_tool: bool,
    frozen_message_count: int,
    has_compressed_content: bool,
) -> tuple[bool, bool]:
    """Decide whether the ``headroom_retrieve`` tool must be injected this turn.

    This is the decision the Anthropic handler used to inline. It is extracted
    so the #1006 regression can be pinned at the decision point itself.

    Tool injection is normally deferred when there is a frozen message prefix
    (``frozen_message_count > 0``) to preserve the prompt cache. But if
    compression emitted fresh markers this turn, deferring would hand the agent
    a ``<<ccr:hash>>`` marker with no tool to redeem it — silent data loss. In
    that case we override the deferral and inject anyway (one cache miss is
    cheaper than dropped content).

    Returns ``(should_inject, is_marker_override)``. ``is_marker_override`` is
    True only when injection happens *because* of new markers despite a deferral,
    so the caller can log the override distinctly.
    """
    return _should_inject_ccr_tool(
        configured_inject_tool=configured_inject_tool,
        frozen_message_count=frozen_message_count,
        has_compressed_content=has_compressed_content,
    )


def apply_session_sticky_ccr_tool(
    *,
    provider: Literal["anthropic", "openai", "google"],
    session_id: str | None,
    request_id: str | None,
    existing_tools: list[dict[str, Any]] | None,
    has_compressed_content_this_turn: bool,
) -> tuple[list[dict[str, Any]], bool]:
    """Apply sticky-on CCR retrieval-tool injection per :class:`SessionCcrTracker`.

    Coordination point for both Anthropic and OpenAI handlers — replaces
    the legacy ``CCRToolInjector.inject_tool_definition`` "flip on, flip
    off" behaviour.

    Logic:

      * If ``session_id`` is None: tracker is bypassed and the per-turn
        ``has_compressed_content_this_turn`` flag drives the decision
        verbatim (matching legacy behaviour for WS / pre-session paths).
      * If the session has previously done CCR (``has_done_ccr``):
        ALWAYS inject the recorded golden bytes — even if this turn has
        no fresh compression. That is the load-bearing PR-B7 fix.
      * Otherwise, inject only when this turn produced compressed content.
        The first injection records the golden bytes for future turns.

    Tools whose name already equals ``CCR_TOOL_NAME`` (e.g. the client
    pre-registered it via MCP) are not re-appended; the client's bytes
    win.

    Returns ``(updated_tools, was_injected)``. ``updated_tools`` is a
    fresh list (caller-safe).
    """
    from headroom.ccr.tool_injection import CCR_TOOL_NAME

    if provider not in ("anthropic", "openai", "google"):
        raise ValueError(f"unsupported provider: {provider!r}")

    tools_out: list[dict[str, Any]] = list(existing_tools) if existing_tools else []
    existing_names: set[str] = set()
    for t in tools_out:
        n = _extract_tool_name(t)
        if n:
            existing_names.add(n)

    # Client (or MCP) already provided a tool by this name — don't double up.
    if CCR_TOOL_NAME in existing_names:
        log_tool_injection_decision(
            provider=provider,
            session_id=session_id,
            decision="skip",
            tool_definition_bytes_count=0,
            request_id=request_id,
        )
        return tools_out, False

    # No session_id (e.g. WS path): per-turn decision drives directly.
    if not session_id:
        if not has_compressed_content_this_turn:
            log_tool_injection_decision(
                provider=provider,
                session_id=None,
                decision="skip",
                tool_definition_bytes_count=0,
                request_id=request_id,
            )
            return tools_out, False
        replay = create_fresh_ccr_tool_definition(provider)
        tools_out.append(replay.tool_definition)
        log_tool_injection_decision(
            provider=provider,
            session_id=None,
            decision="inject_first_time",
            tool_definition_bytes_count=len(replay.canonical_bytes),
            request_id=request_id,
        )
        return tools_out, True

    tracker = get_session_ccr_tracker()
    previously_done = tracker.has_done_ccr(provider, session_id)

    if previously_done:
        # Sticky replay path. Always inject — even if this turn had no
        # fresh CCR compression. Prefer the recorded golden bytes; fall
        # back to a freshly serialized definition if (somehow) the
        # tracker lost them. Loud per build constraint #4: we log the
        # path taken either way.
        golden = tracker.get_golden_tool_bytes(provider, session_id)
        if golden is not None:
            try:
                replay = replay_golden_ccr_tool_definition(golden)
                tools_out.append(replay.tool_definition)
                log_tool_injection_decision(
                    provider=provider,
                    session_id=session_id,
                    decision="inject_sticky_replay",
                    tool_definition_bytes_count=len(replay.canonical_bytes),
                    request_id=request_id,
                )
                return tools_out, True
            except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                logger.error(
                    "corrupt golden CCR tool bytes for session %s: %s — regenerating fresh definition",
                    session_id,
                    exc,
                    exc_info=True,
                )
                # Fall through to fresh creation below
        # Tracker says "done CCR" but has no golden bytes (or they were corrupt). Pin
        # them now so future turns are stable.
        replay = create_fresh_ccr_tool_definition(provider)
        tracker.record_ccr_done(provider, session_id, replay.canonical_bytes)
        tools_out.append(replay.tool_definition)
        log_tool_injection_decision(
            provider=provider,
            session_id=session_id,
            decision="inject_sticky_replay",
            tool_definition_bytes_count=len(replay.canonical_bytes),
            request_id=request_id,
        )
        return tools_out, True

    # Fresh session — only inject when this turn produced compressed content.
    if not has_compressed_content_this_turn:
        log_tool_injection_decision(
            provider=provider,
            session_id=session_id,
            decision="skip",
            tool_definition_bytes_count=0,
            request_id=request_id,
        )
        return tools_out, False

    replay = create_fresh_ccr_tool_definition(provider)
    tracker.record_ccr_done(provider, session_id, replay.canonical_bytes)
    tools_out.append(replay.tool_definition)
    log_tool_injection_decision(
        provider=provider,
        session_id=session_id,
        decision="inject_first_time",
        tool_definition_bytes_count=len(replay.canonical_bytes),
        request_id=request_id,
    )
    return tools_out, True


async def _read_request_body_bytes(request: Request) -> bytes:
    """Read and (if needed) decompress the request body, returning raw UTF-8 bytes.

    Mirrors ``_read_request_json`` but returns the bytes pre-parse so
    forwarders can implement byte-faithful passthrough (PR-A3, fixes P0-2).
    Raises ``ValueError`` on any decompression failure.
    """
    encoding = (request.headers.get("content-encoding") or "").lower().strip()
    raw = await request.body()

    if encoding in ("zstd", "zstandard"):
        try:
            import zstandard

            dctx = zstandard.ZstdDecompressor()
            reader = dctx.stream_reader(raw)
            raw = reader.read()
            reader.close()
        except ImportError:
            raise ValueError(
                "Request body is zstd-compressed but the 'zstandard' package is not installed. "
                "Install it with: pip install zstandard"
            ) from None
        except Exception as exc:
            raise ValueError(f"Failed to decompress zstd request body: {exc}") from exc
    elif encoding == "gzip":
        import gzip as _gzip

        try:
            raw = _gzip.decompress(raw)
        except Exception as exc:
            raise ValueError(f"Failed to decompress gzip request body: {exc}") from exc
    elif encoding == "deflate":
        import zlib

        try:
            raw = zlib.decompress(raw)
        except Exception as exc:
            raise ValueError(f"Failed to decompress deflate request body: {exc}") from exc
    elif encoding == "br":
        try:
            import brotli

            raw = brotli.decompress(raw)
        except ImportError:
            raise ValueError(
                "Request body is brotli-compressed but the 'brotli' package is not installed."
            ) from None
        except Exception as exc:
            raise ValueError(f"Failed to decompress brotli request body: {exc}") from exc
    elif encoding and encoding != "identity":
        raise ValueError(f"Unsupported Content-Encoding: {encoding}")

    return cast(bytes, raw)


async def _read_request_json(request: Request) -> dict[str, Any]:
    """Read and parse JSON from a request, handling compressed bodies.

    Clients like OpenAI Codex may send zstd, gzip, or deflate-compressed
    request bodies.  Starlette's ``request.json()`` does not decompress
    automatically, causing a UnicodeDecodeError on compressed bytes.

    This helper inspects ``Content-Encoding``, decompresses if needed,
    then JSON-decodes the result.  It raises ``ValueError`` on any
    decompression or parse failure so callers can return a clean 400.
    """
    raw = await _read_request_body_bytes(request)

    # Decode and parse JSON
    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ValueError(f"Request body is not valid UTF-8 (possibly compressed?): {exc}") from exc

    result = json.loads(text)
    if not isinstance(result, dict):
        raise ValueError("Request body must be a JSON object, not " + type(result).__name__)
    return result


async def read_request_json_with_bytes(request: Request) -> tuple[dict[str, Any], bytes]:
    """Read JSON body AND return the original (decompressed) bytes.

    Returned bytes are post-content-decoding (zstd/gzip/deflate/br are
    decompressed) so they represent the body as the upstream API will
    receive it. Forwarders pair this with a ``BodyMutationTracker`` to
    decide between passthrough and canonical re-serialization.
    """
    raw = await _read_request_body_bytes(request)

    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ValueError(f"Request body is not valid UTF-8 (possibly compressed?): {exc}") from exc

    result = json.loads(text)
    if not isinstance(result, dict):
        raise ValueError("Request body must be a JSON object, not " + type(result).__name__)
    return result, raw


def _strip_per_call_annotations(obj: Any) -> Any:
    """Remove annotations that clients mutate between calls in one agent loop.

    ``cache_control`` is the main offender: clients (notably Claude Code)
    move the cache breakpoint to the newest message on each call, which
    means the exact same user-text message carries ``cache_control`` on
    call 1 and not on call 2. Hashing the raw message dicts therefore
    produces a different turn_id for every iteration of a single agent
    loop, collapsing ``turn_id`` to effectively ``request_id`` and
    breaking prompt-level aggregation downstream.
    """
    if isinstance(obj, dict):
        return {k: _strip_per_call_annotations(v) for k, v in obj.items() if k != "cache_control"}
    if isinstance(obj, list):
        return [_strip_per_call_annotations(item) for item in obj]
    return obj


def compute_turn_id(
    model: str,
    system: Any,
    messages: list[dict[str, Any]] | None,
) -> str | None:
    """Group all agent-loop API calls triggered by a single user prompt.

    A turn spans the user's text prompt plus every assistant tool-use and
    user tool-result message the agent appends while executing that prompt.
    Hashing the prefix up to and including the last user *text* message yields
    an id that is stable across the turn but rolls over when the user sends a
    new prompt.

    Returns None when no user-text message is present (nothing to identify).
    """
    if not messages:
        return None

    last_text_user_idx: int | None = None
    for i in range(len(messages) - 1, -1, -1):
        msg = messages[i]
        if not isinstance(msg, dict) or msg.get("role") != "user":
            continue
        content = msg.get("content")
        if isinstance(content, str) and content:
            last_text_user_idx = i
            break
        if isinstance(content, list):
            has_text = any(
                isinstance(block, dict) and block.get("type") == "text" for block in content
            )
            has_tool_result = any(
                isinstance(block, dict) and block.get("type") == "tool_result" for block in content
            )
            # An agent-loop continuation carries tool_result blocks; only a
            # fresh user turn is text-only.
            if has_text and not has_tool_result:
                last_text_user_idx = i
                break

    if last_text_user_idx is None:
        return None

    prefix = _strip_per_call_annotations(messages[: last_text_user_idx + 1])
    try:
        prefix_json = json.dumps(prefix, sort_keys=True, default=str)
    except (TypeError, ValueError):
        return None

    h = hashlib.sha256()
    h.update(model.encode("utf-8", errors="replace"))
    h.update(b"\0")
    if isinstance(system, str):
        h.update(system.encode("utf-8", errors="replace"))
    elif system is not None:
        try:
            normalized_system = _strip_per_call_annotations(system)
            h.update(json.dumps(normalized_system, sort_keys=True, default=str).encode("utf-8"))
        except (TypeError, ValueError):
            pass
    h.update(b"\0")
    h.update(prefix_json.encode("utf-8", errors="replace"))
    return h.hexdigest()[:16]


# ---------------------------------------------------------------------------
# Issue #746: Claude Code on-demand tool loading (deferral) detection
#
# When Claude Code points at a custom ``ANTHROPIC_BASE_URL`` (the proxy) with
# ``ENABLE_TOOL_SEARCH`` unset, it stops deferring MCP/system tool schemas
# behind the server-side Tool Search Tool and materializes them all into its
# local context window — tens of K tokens. That decision is made client-side
# before the request reaches us, so the proxy cannot reverse it; the only
# remedy is the ``ENABLE_TOOL_SEARCH`` env var (set automatically by
# ``headroom wrap claude``). For users who run ``claude`` manually we cannot
# touch their environment, so the proxy emits a single actionable hint.
# ---------------------------------------------------------------------------

_TOOL_SEARCH_TOOL_TYPE_PREFIX = "tool_search_tool_"
# Substrings of the ``anthropic-beta`` tokens that gate tool search:
# ``advanced-tool-use-2025-11-20`` (firstParty/foundry) and
# ``tool-search-tool-2025-10-19`` (vertex/bedrock/mantle/gateway).
_TOOL_SEARCH_BETA_MARKERS = ("advanced-tool-use", "tool-search-tool")

_tool_search_hint_lock = threading.Lock()
_tool_search_hint_emitted = False


def claude_code_tool_search_inactive(
    *,
    client: str | None,
    tools: Any,
    anthropic_beta: str | None,
) -> bool:
    """Return ``True`` when a Claude Code request is *not* deferring tools.

    Detected from request shape alone — no token thresholds, so it scales to
    any tool surface:

    * the request is from Claude Code (``client == "claude-code"``),
    * it carries one or more tool definitions, yet
    * it includes neither a ``tool_search_tool_*`` tool nor a tool-search
      ``anthropic-beta`` token.

    In that combination Claude Code has eagerly materialized every tool schema
    into its local context window (issue #746).
    """
    if client != "claude-code":
        return False
    if not isinstance(tools, list) or not tools:
        return False
    for tool in tools:
        if isinstance(tool, dict) and str(tool.get("type", "")).startswith(
            _TOOL_SEARCH_TOOL_TYPE_PREFIX
        ):
            return False
    beta = (anthropic_beta or "").lower()
    return not any(marker in beta for marker in _TOOL_SEARCH_BETA_MARKERS)


def format_tool_search_disabled_hint(tools: list[Any]) -> str:
    """Build the one-time, actionable hint for issue #746.

    Reports factual, directional numbers (tool count and serialized schema
    size) rather than a derived token estimate, which avoids implying a
    precision the proxy cannot measure for the client's tokenizer.
    """
    try:
        schema_kb = len(json.dumps(tools, separators=(",", ":"), default=str)) / 1024
    except (TypeError, ValueError):
        schema_kb = 0.0
    return (
        f"Claude Code is sending all {len(tools)} tool definitions eagerly "
        f"(~{schema_kb:.0f} KB of tool schema in local context) because "
        "ENABLE_TOOL_SEARCH is unset with a custom ANTHROPIC_BASE_URL. Set "
        "ENABLE_TOOL_SEARCH=true (or auto) to keep on-demand tool loading active, "
        "or launch via `headroom wrap claude` (which sets it automatically). "
        "See https://github.com/chopratejas/headroom/issues/746"
    )


def tool_search_hint_pending() -> bool:
    """Cheap, lock-free check of whether the one-time hint may still fire.

    Lets the request hot path skip the (O(number-of-tools)) detection scan on
    every request once the hint has already been emitted. A benign race here
    only costs one extra detection scan, never a duplicate warning — the
    actual one-shot guarantee lives in :func:`take_tool_search_hint_slot`.
    """
    return not _tool_search_hint_emitted


def take_tool_search_hint_slot() -> bool:
    """Return ``True`` exactly once per process, gating the one-time hint.

    Thread-safe so concurrent requests cannot each emit the warning.
    """
    global _tool_search_hint_emitted
    if _tool_search_hint_emitted:
        return False
    with _tool_search_hint_lock:
        if _tool_search_hint_emitted:
            return False
        _tool_search_hint_emitted = True
        return True


def reset_tool_search_hint_state() -> None:
    """Reset the one-time hint guard. Test helper only."""
    global _tool_search_hint_emitted
    with _tool_search_hint_lock:
        _tool_search_hint_emitted = False


# ---------------------------------------------------------------------------
# Server-side Tool Search injection (opencode / non-Claude-Code clients).
#
# Clients that eagerly materialize every tool schema (opencode ships ~135 tool
# defs ≈ 28k tokens on EVERY request) never opt into Anthropic's Tool Search
# Tool themselves. Unlike the Claude Code case above — where the schemas are
# already in the client's own context and the proxy can't reverse it — a plain
# API client's tools live only in the request body, so the proxy CAN defer them:
# mark the non-core tools ``defer_loading: true`` and inject a tool_search tool.
# Anthropic then excludes deferred tools from the context window (they stop
# counting as input tokens until the model searches for one), while every tool
# stays callable. Deterministic output → the tools prefix still prompt-caches.
# ---------------------------------------------------------------------------

# Core coding tools kept non-deferred so routine edit/read/run loops never pay a
# search round-trip. Everything else (Slack/Linear/Sentry/Notion/Snowflake/…) is
# deferred and loaded on demand. Anthropic recommends keeping the 3–5 (here a few
# more) most frequent tools resident.
_TOOL_SEARCH_CORE_TOOLS = frozenset(
    {
        "bash",
        "bash_background",
        "bash_background_output",
        "bash_background_wait",
        "bash_background_kill",
        "read",
        "write",
        "edit",
        "multiedit",
        "apply_patch",
        "glob",
        "grep",
        "task",
        "todowrite",
        "todoread",
        "webfetch",
        "question",
        "skill",
    }
)
_TOOL_SEARCH_DEFAULT_TYPE = "tool_search_tool_regex_20251119"
_TOOL_SEARCH_DEFAULT_NAME = "tool_search_tool_regex"
# Below this many tools the ~search round-trip isn't worth it (Anthropic's own
# guidance: standard calling is better under ~10 tools).
_TOOL_SEARCH_MIN_TOOLS = 12


def inject_tool_search_deferral(
    tools: Any,
    *,
    core_tools: frozenset[str] = _TOOL_SEARCH_CORE_TOOLS,
    search_type: str = _TOOL_SEARCH_DEFAULT_TYPE,
    search_name: str = _TOOL_SEARCH_DEFAULT_NAME,
) -> Any:
    """Return a new ``tools`` list with non-core tools deferred + a search tool
    injected, or the original list unchanged when injection doesn't apply.

    No-op when: not a list, fewer than ``_TOOL_SEARCH_MIN_TOOLS``, a tool_search
    tool is already present (client already defers), or nothing would be deferred.

    Invariants enforced (else Anthropic 400s): the search tool is never deferred;
    at least one tool stays non-deferred; a deferred tool never carries
    ``cache_control`` — if the client's tools cache breakpoint sat on a now-deferred
    tool, it is moved to the last non-deferred real tool so the (smaller) tools
    prefix still caches.
    """
    if not isinstance(tools, list) or len(tools) < _TOOL_SEARCH_MIN_TOOLS:
        return tools
    for tool in tools:
        if isinstance(tool, dict) and str(tool.get("type", "")).startswith(
            _TOOL_SEARCH_TOOL_TYPE_PREFIX
        ):
            return tools  # client already uses tool search — leave it alone

    search_tool = {"type": search_type, "name": search_name}
    out: list[Any] = [search_tool]
    deferred = 0
    dropped_cache_control = False
    last_resident_real: dict[str, Any] | None = None
    resident_has_cache_control = False

    for tool in tools:
        if not isinstance(tool, dict) or tool.get("type") or tool.get("name") in core_tools:
            # Non-dict, server/typed tools (web_search, computer, …), and core
            # tools stay resident and unchanged.
            out.append(tool)
            if isinstance(tool, dict) and not tool.get("type"):
                last_resident_real = tool
                resident_has_cache_control = resident_has_cache_control or bool(
                    tool.get("cache_control")
                )
            continue
        new_tool = dict(tool)
        new_tool["defer_loading"] = True
        if new_tool.pop("cache_control", None) is not None:
            dropped_cache_control = True
        out.append(new_tool)
        deferred += 1

    if deferred == 0:
        return tools  # nothing to defer → don't perturb the cache prefix
    # Preserve a tools cache breakpoint: if we stripped cache_control off a
    # deferred tool and no resident tool carries one, move it to the last
    # resident real tool (never the search tool, to keep its shape canonical).
    if dropped_cache_control and not resident_has_cache_control and last_resident_real is not None:
        last_resident_real["cache_control"] = {"type": "ephemeral"}
    return out


# ---------------------------------------------------------------------------
# Server-side Tool Search injection — OpenAI Responses API (gpt-5.4+).
#
# The OpenAI-side analogue of inject_tool_search_deferral above. OpenAI shipped
# the same idea for the Responses API on gpt-5.4+: mark a function/MCP tool
# ``defer_loading: true`` and add a ``{"type": "tool_search"}`` tool, and OpenAI
# keeps the deferred tools' heavy parameter schemas OUT of the model's context
# (only name+description remain) until the model searches for one — while every
# tool stays callable and the prompt cache is preserved. Same win as Anthropic
# (~15-25k tool-schema tokens -> ~200) for clients that ship a big tool surface
# and never opt into tool search themselves (opencode, plain API clients).
#
# Differences from the Anthropic path that require a separate function:
#   * Responses function tools carry ``type: "function"`` (Anthropic real tools
#     have no ``type``), so the resident/defer test is inverted — we defer
#     ``function`` (non-core) and ``mcp`` tools and keep OTHER typed/hosted tools
#     (web_search, file_search, code_interpreter, computer, image_generation, and
#     the search tool itself) resident.
#   * Model-gated: only gpt-5.4+ support it; older models 400 on the fields.
#   * No ``cache_control`` (OpenAI caches automatically), so no breakpoint move.
# ---------------------------------------------------------------------------

_OPENAI_TOOL_SEARCH_TYPE = "tool_search"
_OPENAI_TOOL_SEARCH_MIN_TOOLS = 12
_OPENAI_TOOL_SEARCH_RESIDENT_NAMES = frozenset({"terminal"})
# gpt-5.4 is the first model with Responses tool_search (OpenAI docs). Version-
# gated by default; overridable per deployment via a regex in
# HEADROOM_OPENAI_TOOL_SEARCH_MODELS (matched against the model name) so new
# model families can be enabled without a code edit + release.
_OPENAI_TOOL_SEARCH_MIN_VERSION = (5, 4)


def _model_supports_openai_tool_search(model: str | None) -> bool:
    """True when an OpenAI model supports the Responses ``tool_search`` feature.

    Default gate: ``gpt-<major>.<minor>`` >= 5.4. A regex in
    ``HEADROOM_OPENAI_TOOL_SEARCH_MODELS`` (matched against the model name) wins
    when set; a malformed pattern falls back to the version gate rather than
    crashing.
    """
    if not model:
        return False
    override = os.environ.get("HEADROOM_OPENAI_TOOL_SEARCH_MODELS", "").strip()
    if override:
        try:
            return re.search(override, model) is not None
        except re.error:
            pass  # malformed override → fall back to the version gate
    match = re.match(r"gpt-(\d+)(?:\.(\d+))?", model.strip().lower())
    if not match:
        return False
    major, minor = int(match.group(1)), int(match.group(2) or 0)
    return (major, minor) >= _OPENAI_TOOL_SEARCH_MIN_VERSION


def inject_tool_search_deferral_openai(
    tools: Any,
    model: str | None,
    *,
    core_tools: frozenset[str] = _TOOL_SEARCH_CORE_TOOLS,
) -> Any:
    """Return a new Responses ``tools`` list with non-core function/MCP tools
    deferred + a ``{"type": "tool_search"}`` tool injected, or the original list
    unchanged when injection doesn't apply.

    No-op when: the model doesn't support tool search (gpt-5.4+ only), ``tools``
    is not a list, there are fewer than ``_OPENAI_TOOL_SEARCH_MIN_TOOLS``, a
    tool_search tool is already present (client already defers), or nothing would
    be deferred. Core coding tools and hosted/typed tools (web_search,
    file_search, code_interpreter, computer, …) stay resident and unchanged, so
    routine edit/read/run loops never pay a search round-trip and the request
    stays valid; the injected search tool is itself resident.
    """
    if not _model_supports_openai_tool_search(model):
        return tools
    if not isinstance(tools, list) or len(tools) < _OPENAI_TOOL_SEARCH_MIN_TOOLS:
        return tools
    for tool in tools:
        if isinstance(tool, dict) and tool.get("type") == _OPENAI_TOOL_SEARCH_TYPE:
            return tools  # client already uses tool search — leave it alone

    out: list[Any] = [{"type": _OPENAI_TOOL_SEARCH_TYPE}]
    deferred = 0
    for tool in tools:
        if not isinstance(tool, dict):
            out.append(tool)
            continue
        ttype = tool.get("type")
        # Deferrable: a non-core function, or an MCP server (OpenAI models are
        # trained to search namespaces / MCP servers). Everything else — core
        # coding tools and other hosted tools — stays resident.
        deferrable = (
            ttype == "function"
            and tool.get("name") not in core_tools
            and tool.get("name") not in _OPENAI_TOOL_SEARCH_RESIDENT_NAMES
        ) or ttype == "mcp"
        if deferrable and not tool.get("defer_loading"):
            new_tool = dict(tool)
            new_tool["defer_loading"] = True
            out.append(new_tool)
            deferred += 1
        else:
            out.append(tool)

    if deferred == 0:
        return tools  # nothing to defer → don't perturb the request / cache prefix
    return out