chopratejas--headroom/headroom/transforms/compression_policy.py

"""Per-auth-mode compression policy — Phase F PR-F2.1, extended in F2.2 (Python parity).

Hand-mirrored port of `headroom_core::compression_policy::CompressionPolicy`
(Rust). The Rust crate is the source of truth; this module exists so
the Python proxy's `TransformPipeline` (which still runs `CacheAligner`
and other detector-only transforms) can read the same per-mode flags
the Rust dispatcher does.

A parity test (`tests/test_compression_policy.py`) instantiates one of
each variant and asserts the field map matches what the Rust unit
tests assert. F2.2 should consider exposing the Rust struct via PyO3
to retire this hand-mirror — that's deliberately out of scope here so
F2.1/F2.2 can ship.

See `crates/headroom-core/src/compression_policy.rs` for the canonical
docstring (per-mode rationale, why-a-struct, etc.).
"""

from __future__ import annotations

import math
import os
from dataclasses import dataclass

from headroom.proxy.auth_mode import AuthMode

# ── F2.2 per-mode default values (CONSERVATIVE pending bake telemetry) ──
# Mirrors the Rust ``pub(crate) const`` block in
# ``crates/headroom-core/src/compression_policy.rs``. Centralised here so
# a follow-up tune lands in one place per language. The Rust tests assert
# the values directly; the Python parity tests assert the *fields* match
# (not the values — that would double-pin against drift in a way that
# masks real divergence).
#
# Per the realignment build constraints (project memory
# ``feedback_realignment_build_constraints.md``): "configurable / no
# hardcoded values". The configuration *is* the per-mode default — we
# deliberately do NOT add a separate env var per field. Operators tune
# by editing these constants and shipping a new build, mirroring the
# Rust pattern exactly.

#: PAYG: aggressive — let volatile content noise up to ~128 tokens slip
#: before flagging. Higher than Subscription because PAYG users opt in
#: to aggressive compression.
_VOLATILE_TOKEN_THRESHOLD_PAYG: int = 128

#: Subscription: conservative — flag volatile content earlier (32
#: tokens) so cache prefixes stay stable.
_VOLATILE_TOKEN_THRESHOLD_SUBSCRIPTION: int = 32

#: PAYG: cap lossy compression at 45% of original tokens. Aggressive
#: but bounded — F2.1 had no cap (effectively ``1.0``), F2.2 introduces one.
_MAX_LOSSY_RATIO_PAYG: float = 0.45

#: Subscription: conservative cap at 25%. Cache stability over savings.
_MAX_LOSSY_RATIO_SUBSCRIPTION: float = 0.25

#: Anthropic prompt-cache write multiplier: a ``cache_creation`` token
#: costs 1.25x a plain input token (5-minute TTL tier). Input to the
#: net-cost mutation formula (#856). Mirrors the Rust ``pub const``.
CACHE_WRITE_MULTIPLIER: float = 1.25

#: Anthropic prompt-cache read multiplier: a ``cache_read`` token costs
#: 0.1x a plain input token. Input to the net-cost mutation formula
#: (#856). Mirrors the Rust ``pub const``.
CACHE_READ_MULTIPLIER: float = 0.1


@dataclass(frozen=True, slots=True)
class CompressionPolicy:
    """Per-auth-mode policy that downstream compression stages consult.

    Five fields after F2.2:

    - ``live_zone_only`` (F2.1) — gate for the Python TransformPipeline
      to skip pre-cache-marker mutation.
    - ``cache_aligner_enabled`` (F2.1) — gate for the Python
      ``CacheAligner`` transform.
    - ``volatile_token_threshold`` (F2.2) — per-mode token threshold
      below which content is treated as cache-stable. Plumbed through
      the struct; no current detector consumes it (intentional — the
      detector refactor is a follow-up).
    - ``max_lossy_ratio`` (F2.2) — per-mode upper bound on lossy
      compression aggressiveness (fraction in ``[0.0, 1.0]``). Plumbed
      through the struct; no current compressor consumes it.
    - ``toin_read_only`` (F2.2) — TOIN learning gate. ``True`` =
      serve cached patterns but never write new observations from
      this request (Subscription).
    """

    live_zone_only: bool
    """When True, transforms MUST NOT modify bytes outside the post-
    cache-marker live zone. The Rust live-zone dispatcher is already
    live-zone-only by construction, so this flag is effectively a
    no-op on the Rust path; it exists for the Python TransformPipeline
    where transforms like CacheAligner inspect the cached prefix."""

    cache_aligner_enabled: bool
    """When False, the CacheAligner transform's `should_apply` MUST
    return False — that's the load-bearing F2.1 gate for the cache-
    instability complaints in #327 / #388."""

    volatile_token_threshold: int
    """F2.2: per-mode token-count threshold below which content is
    treated as cache-stable. Subscription is conservative
    (low → flag aggressively → keep prompts stable); PAYG aggressive
    (high → tolerate more volatile noise). Plumbed but unconsumed in
    F2.2 — the volatile detector in ``cache_aligner.py`` is shape-
    based, not token-count-based; wiring it is a follow-up."""

    max_lossy_ratio: float
    """F2.2: per-mode upper bound on lossy compression aggressiveness,
    expressed as the fraction of original tokens that may be dropped
    (``0.0`` = no lossy, ``1.0`` = unlimited). Subscription ``0.25``;
    PAYG ``0.45``. Plumbed but unconsumed in F2.2 — distinct from the
    caller-driven ``target_ratio`` kwarg in the Python ContentRouter."""

    toin_read_only: bool
    """F2.2: when True, TOIN serves cached recommendations but
    never writes new pattern observations from this request.
    Subscription True (consistency over learning); PAYG/OAuth False
    (network effect keeps growing). The gate is read by
    ``smart_crusher.py`` and ``content_router.py`` at the
    ``record_compression`` call sites."""

    def net_mutation_gain(
        self,
        delta_t: int,
        suffix_tokens: int,
        expected_reads: float,
        p_alive: float,
    ) -> float:
        """Net gain (in plain-input-token cost units) of a mutation that
        removes ``delta_t`` tokens from a message whose cached suffix is
        ``suffix_tokens`` long (#856).

        Mirrors ``CompressionPolicy::net_mutation_gain`` in the Rust
        crate (source of truth — see its docstring for the derivation)::

            gain = dT * (w + r*(R - 1)) - P_alive * (w - r) * (S + dT)

        The warm-case penalty covers ``S + dT``: with a live cache the
        ``dT`` tokens are already cache-written, so keeping them costs
        only reads — a mutation avoids at most ``dT*r*R``, not a fresh
        write.

        Inputs are clamped: ``delta_t``/``suffix_tokens`` to ``>= 0``
        (the Rust signature takes ``u32``), ``expected_reads`` to
        ``>= 0`` (NaN → 0), ``p_alive`` to ``[0, 1]`` (NaN → 1, the
        conservative full-penalty assumption — same as Rust).
        """
        w = CACHE_WRITE_MULTIPLIER
        r = CACHE_READ_MULTIPLIER
        dt = max(0, delta_t)
        suffix = max(0, suffix_tokens)
        # Python max()/min() propagate NaN from the first argument, unlike
        # f32::max in the Rust source of truth — guard explicitly.
        reads = 0.0 if math.isnan(expected_reads) else max(expected_reads, 0.0)
        alive = 1.0 if math.isnan(p_alive) else min(max(p_alive, 0.0), 1.0)
        return float(dt) * (w + r * (reads - 1.0)) - alive * (w - r) * float(suffix + dt)

    def should_mutate_deep(
        self,
        delta_t: int,
        suffix_tokens: int,
        expected_reads: float,
        p_alive: float,
    ) -> bool:
        """Decision form of :meth:`net_mutation_gain`: mutate iff the
        gain is strictly positive."""
        return self.net_mutation_gain(delta_t, suffix_tokens, expected_reads, p_alive) > 0.0

    def break_even_reads(self, delta_t: int, suffix_tokens: int) -> float:
        """Remaining-read count at which a warm-cache (``p_alive=1``)
        mutation breaks even::

            R = ((w - r) / r) * (S/dT)   = 11.5 * S/dT  (Anthropic 5-min)

        With the corrected penalty this reproduces the #856 anchors
        exactly: 2K/50K -> 287.5, 50K/10K -> 2.3.

        Returns 0 when ``delta_t`` is ``<= 0`` (no savings — callers
        gate on ``delta_t > 0``; the Rust signature takes ``u32``).
        Mirrors the Rust method.
        """
        if delta_t <= 0:
            return 0.0
        w = CACHE_WRITE_MULTIPLIER
        r = CACHE_READ_MULTIPLIER
        return ((w - r) / r) * (float(max(0, suffix_tokens)) / float(delta_t))


def policy_for_mode(mode: AuthMode) -> CompressionPolicy:
    """Resolve the F2.1+F2.2 policy for an auth mode.

    PAYG and OAuth are identical (aggressive: live-zone-not-only,
    cache-aligner on, relaxed thresholds, TOIN write-enabled).
    Subscription is the user-visible win: live-zone-only with cache
    aligner disabled, tighter thresholds, TOIN read-only.

    F2.2-followup may diverge OAuth from PAYG once telemetry is
    collected.
    """
    if mode == AuthMode.PAYG:
        return CompressionPolicy(
            live_zone_only=False,
            cache_aligner_enabled=True,
            volatile_token_threshold=_VOLATILE_TOKEN_THRESHOLD_PAYG,
            max_lossy_ratio=_MAX_LOSSY_RATIO_PAYG,
            toin_read_only=False,
        )
    if mode == AuthMode.OAUTH:
        # Identical to PAYG in F2.1/F2.2. The parity test in
        # ``tests/test_compression_policy.py`` is the canary that
        # catches a future divergence and forces a deliberate update
        # there + in the Rust crate.
        return CompressionPolicy(
            live_zone_only=False,
            cache_aligner_enabled=True,
            volatile_token_threshold=_VOLATILE_TOKEN_THRESHOLD_PAYG,
            max_lossy_ratio=_MAX_LOSSY_RATIO_PAYG,
            toin_read_only=False,
        )
    if mode == AuthMode.SUBSCRIPTION:
        return CompressionPolicy(
            live_zone_only=True,
            cache_aligner_enabled=False,
            volatile_token_threshold=_VOLATILE_TOKEN_THRESHOLD_SUBSCRIPTION,
            max_lossy_ratio=_MAX_LOSSY_RATIO_SUBSCRIPTION,
            toin_read_only=True,
        )
    raise ValueError(f"Unhandled AuthMode variant: {mode!r}")


def policy_default_payg() -> CompressionPolicy:
    """The PAYG-equivalent policy used when the
    ``HEADROOM_PROXY_AUTH_MODE_POLICY_ENFORCEMENT`` flag is disabled
    (default in F2.1 c1-c4; flipped to enabled in c5/5).

    Centralised so the proxy handlers do not duplicate the constant,
    and so a future change to PAYG semantics propagates to both the
    enforcement-on and enforcement-off paths.
    """
    return policy_for_mode(AuthMode.PAYG)


_ENFORCEMENT_ENV = "HEADROOM_PROXY_AUTH_MODE_POLICY_ENFORCEMENT"


def is_enforcement_enabled() -> bool:
    """Read the enforcement flag from the environment.

    Same env var the Rust proxy reads (``Config::auth_mode_policy_enforcement``)
    so the two paths stay in lockstep with one operator switch.
    Default (when unset): ``True`` from F2.1 c5/5 onward, matching
    the Rust default after the c5 flip.

    NOT cached — read every call so an operator can flip the env var
    in a hot-reload scenario. The cost is one ``dict.get`` per call,
    well below noise.
    """
    val = os.environ.get(_ENFORCEMENT_ENV, "enabled").strip().lower()
    # Same set of off-values the telemetry beacon honours
    # (`headroom/telemetry/beacon.py::_OFF_VALUES`) so operators don't
    # have to remember a different vocabulary per flag.
    return val not in ("disabled", "off", "false", "0", "no")


def resolve_policy(auth_mode: AuthMode | None) -> CompressionPolicy:
    """Resolve the effective ``CompressionPolicy`` for a request.

    - If the enforcement flag is off, returns PAYG-equivalent
      regardless of the classified auth mode.
    - If the enforcement flag is on and ``auth_mode`` is ``None``,
      returns PAYG-equivalent (defensive default for the unclassified
      / batch-row path).
    - Otherwise returns the per-mode policy.

    This is the single public entry point handlers should call when
    deriving the policy from a request's classification result.
    """
    if auth_mode is None or not is_enforcement_enabled():
        return policy_default_payg()
    return policy_for_mode(auth_mode)