mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Closes #856 **P1 of the #856 phased plan** — pure functions, zero behavior change. (Closing keyword links the issue; if P2 hasn't started when this merges, reopen #856 or it remains the design record for the P2/P3 follow-up PRs.) ## What Adds the break-even decision rule for deep (pre-cache-marker) edits to `CompressionPolicy`: ``` gain = ΔT · (w + r·(R−1)) − P_alive · (w − r) · S ``` - `net_mutation_gain()`, `should_mutate_deep()` (gain > 0), `break_even_reads()` (R = ((w−r)/r)·(S/ΔT−1) ≈ 11.5·S/ΔT) on the Rust struct (source of truth) and the Python hand-mirror, following the existing F2.1/F2.2 parity pattern. - `CACHE_WRITE_MULTIPLIER = 1.25` / `CACHE_READ_MULTIPLIER = 0.1` public constants (Anthropic 5-minute tier). - Inputs clamped (`expected_reads ≥ 0`, `p_alive ∈ [0,1]`); methods take `&self`/`self` so a follow-up can add per-mode margins. - The formula derives the existing Subscription live-zone policy as its S=0 special case rather than contradicting it. **No callers yet.** P2 (consuming this in `TransformPipeline` behind `HEADROOM_NET_COST_POLICY`, replacing the binary `live_zone_only` gate, with decision telemetry) is specified in #856 and awaits maintainer direction — this PR just lands the audited arithmetic both dispatchers will share. ## Tests Golden-value parity: 6 new Rust unit tests and 7 new Python tests assert the **identical scenario numbers** (loss −53 200 for a 2K shave under a 50K warm suffix at R=10; win +61 000 for a 50K shave under a 10K suffix at R=3; S=0 always profitable; P_alive=0 always profitable — the idle-timer window; clamping; break-even 276 reads for the 2K/50K anchor). A drift on either side trips the pair loudly, same contract as the existing field-map parity test. - `cargo test -p headroom-core --lib compression_policy`: 12 passed (6 existing + 6 new) - `pytest tests/test_compression_policy.py`: 17 passed (10 existing + 7 new) - `cargo fmt --check`, `cargo clippy -p headroom-core` clean; `ruff check` + `ruff format --check` clean ## Real behavior proof Not applicable in the runtime sense — this PR intentionally adds **no runtime behavior** (pure functions, no call sites). The arithmetic is validated against the research anchors above in both languages' test suites; live decision telemetry arrives with P2 where the formula first gates real traffic. ## Out of scope P2 (flag-gated pipeline consumption + telemetry), P3 (deep-edit batching, idle-timer compaction near TTL lapse), retiring the deprecated `volatile_token_threshold`/`max_lossy_ratio` fields — all tracked in #856. --------- Co-authored-by: Ash Rhodes <ashley.rhodes@king.com>
237 lines
10 KiB
Python
237 lines
10 KiB
Python
"""Tests for the Python ``CompressionPolicy`` and its parity with Rust.
|
|
|
|
The Python module is a hand-mirror of
|
|
``headroom_core::compression_policy::CompressionPolicy``. These tests
|
|
pin both halves: that the per-mode values are right, and that the
|
|
Python and Rust sides agree on the field map. F2.2 will likely retire
|
|
the hand-mirror via PyO3 — until then, this file is the canary.
|
|
|
|
F2.2 extends the F2.1 surface with three tuning fields:
|
|
``volatile_token_threshold``, ``max_lossy_ratio``, ``toin_read_only``.
|
|
Per-mode value tests below mirror the Rust unit tests in
|
|
``crates/headroom-core/src/compression_policy.rs``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from headroom.proxy.auth_mode import AuthMode
|
|
from headroom.transforms.compression_policy import (
|
|
CompressionPolicy,
|
|
policy_default_payg,
|
|
policy_for_mode,
|
|
)
|
|
|
|
|
|
class TestCompressionPolicyForMode:
|
|
"""Per-mode field assertions. Mirrors the Rust unit tests in
|
|
`crates/headroom-core/src/compression_policy.rs`.
|
|
"""
|
|
|
|
def test_payg_is_aggressive(self):
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
assert p.live_zone_only is False, "PAYG can touch outside live zone"
|
|
assert p.cache_aligner_enabled is True, "PAYG runs cache aligner"
|
|
|
|
def test_payg_tuning_fields_aggressive(self):
|
|
# F2.2: per-mode tuning fields. Values are the conservative
|
|
# defaults pending bake telemetry (see PR body and module
|
|
# docstring).
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
assert p.volatile_token_threshold == 128, (
|
|
"PAYG volatile threshold is the relaxed default; F2.2-followup will tune"
|
|
)
|
|
assert p.max_lossy_ratio == pytest.approx(0.45), (
|
|
"PAYG max_lossy_ratio caps lossy paths at 0.45; F2.2-followup will tune"
|
|
)
|
|
assert p.toin_read_only is False, (
|
|
"PAYG keeps TOIN write-enabled — network effect feeds on PAYG traffic"
|
|
)
|
|
|
|
def test_oauth_matches_payg_today(self):
|
|
# Canary: when F2.2-followup diverges OAuth from PAYG, this test
|
|
# fails and forces a deliberate update on BOTH sides (Rust +
|
|
# Python). Covers ALL fields (F2.1 + F2.2) so a future field-
|
|
# level divergence trips the assertion just as loudly as a flag
|
|
# flip.
|
|
oauth = policy_for_mode(AuthMode.OAUTH)
|
|
payg = policy_for_mode(AuthMode.PAYG)
|
|
assert oauth == payg, (
|
|
"F2.1+F2.2 ship OAuth=PAYG; F2.2-followup will diverge based on telemetry. "
|
|
"If you are reading this assertion failure: also update "
|
|
"crates/headroom-core/src/compression_policy.rs "
|
|
"::oauth_matches_payg_today, otherwise the Rust + Python "
|
|
"parities silently drift apart."
|
|
)
|
|
|
|
def test_subscription_disables_cache_aligner(self):
|
|
p = policy_for_mode(AuthMode.SUBSCRIPTION)
|
|
assert p.live_zone_only is True, "Subscription is live-zone-only"
|
|
assert p.cache_aligner_enabled is False, (
|
|
"Subscription MUST skip cache aligner — load-bearing for issues #327 / #388"
|
|
)
|
|
|
|
def test_subscription_tuning_fields_conservative(self):
|
|
# F2.2: per-mode tuning fields. Subscription is the conservative
|
|
# end — tighter threshold, lower lossy cap, TOIN read-only — so
|
|
# cache prefixes stay stable and the learning pool isn't
|
|
# mutated from cache-stability-sensitive traffic.
|
|
p = policy_for_mode(AuthMode.SUBSCRIPTION)
|
|
assert p.volatile_token_threshold == 32, (
|
|
"Subscription volatile threshold flags content earlier (cache stability)"
|
|
)
|
|
assert p.max_lossy_ratio == pytest.approx(0.25), (
|
|
"Subscription max_lossy_ratio caps lossy paths at 0.25 (conservative)"
|
|
)
|
|
assert p.toin_read_only is True, (
|
|
"Subscription MUST be TOIN read-only — load-bearing for keeping the "
|
|
"learning pool consistent across cache-sensitive traffic"
|
|
)
|
|
|
|
def test_max_lossy_ratio_in_unit_interval(self):
|
|
# Defensive: every per-mode `max_lossy_ratio` MUST be in
|
|
# ``[0.0, 1.0]`` because it expresses a fraction. A tune that
|
|
# drifts outside the unit interval is a bug — catch it cheaply
|
|
# here rather than at the eventual consumer site.
|
|
for mode in (AuthMode.PAYG, AuthMode.OAUTH, AuthMode.SUBSCRIPTION):
|
|
r = policy_for_mode(mode).max_lossy_ratio
|
|
assert 0.0 <= r <= 1.0, f"max_lossy_ratio for {mode!r} = {r} is outside [0.0, 1.0]"
|
|
|
|
|
|
class TestPolicyDefaultPayg:
|
|
"""The constant used when the enforcement flag is disabled."""
|
|
|
|
def test_default_payg_equals_for_mode_payg(self):
|
|
assert policy_default_payg() == policy_for_mode(AuthMode.PAYG)
|
|
|
|
|
|
class TestImmutability:
|
|
"""The struct is `frozen=True`; mutation must raise."""
|
|
|
|
def test_policy_is_frozen(self):
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
with pytest.raises((AttributeError, Exception)):
|
|
# Attempting to mutate a frozen dataclass raises
|
|
# FrozenInstanceError (subclass of AttributeError on
|
|
# CPython 3.10+). Catch both for compatibility.
|
|
p.live_zone_only = True # type: ignore[misc]
|
|
|
|
def test_f22_tuning_fields_also_frozen(self):
|
|
# Each F2.2 field gets its own immutability assertion — a
|
|
# future refactor that accidentally drops `frozen=True` on the
|
|
# dataclass would silently allow per-request mutation. The
|
|
# F2.1 test only covered ``live_zone_only``; explicit per-
|
|
# field coverage prevents quiet regressions.
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
for attr_name in ("volatile_token_threshold", "max_lossy_ratio", "toin_read_only"):
|
|
with pytest.raises((AttributeError, Exception)):
|
|
setattr(p, attr_name, 0) # type: ignore[misc]
|
|
|
|
|
|
class TestRustParityFieldMap:
|
|
"""The Python policy must have the same fields as the Rust struct.
|
|
|
|
The canonical Rust struct lives at
|
|
``crates/headroom-core/src/compression_policy.rs``. When you add a
|
|
field there for a future PR, add it here AND update this test.
|
|
Otherwise the parity silently drifts.
|
|
"""
|
|
|
|
def test_field_set_matches_rust(self):
|
|
# Hard-coded set — when Rust grows fields, this test fails until
|
|
# Python catches up. F2.2 added three: volatile_token_threshold,
|
|
# max_lossy_ratio, toin_read_only.
|
|
expected_fields = {
|
|
"live_zone_only",
|
|
"cache_aligner_enabled",
|
|
"volatile_token_threshold",
|
|
"max_lossy_ratio",
|
|
"toin_read_only",
|
|
}
|
|
actual_fields = {f.name for f in CompressionPolicy.__dataclass_fields__.values()}
|
|
assert actual_fields == expected_fields, (
|
|
f"Python CompressionPolicy fields drifted from Rust. "
|
|
f"Expected exactly {expected_fields}, got {actual_fields}. "
|
|
f"Update both `headroom/transforms/compression_policy.py` "
|
|
f"and `crates/headroom-core/src/compression_policy.rs` in "
|
|
f"the same commit."
|
|
)
|
|
|
|
|
|
class TestNetCostFormula:
|
|
"""Net-cost mutation formula (#856) — Rust parity.
|
|
|
|
Scenario values are golden: the Rust unit tests in
|
|
``crates/headroom-core/src/compression_policy.rs`` assert the
|
|
identical numbers, so a drift in either side trips the parity pair
|
|
loudly.
|
|
"""
|
|
|
|
def test_small_shave_deep_suffix_is_loss(self):
|
|
# 2000*(1.25 + 0.1*9) - 1.0*1.15*50000 = 4300 - 57500 = -53200.
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
gain = p.net_mutation_gain(2_000, 50_000, 10.0, 1.0)
|
|
assert abs(gain - (-53_200.0)) < 1.0
|
|
assert not p.should_mutate_deep(2_000, 50_000, 10.0, 1.0)
|
|
|
|
def test_big_shave_shallow_suffix_is_win(self):
|
|
# 50000*(1.25 + 0.1*2) - 1.0*1.15*10000 = 72500 - 11500 = 61000.
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
gain = p.net_mutation_gain(50_000, 10_000, 3.0, 1.0)
|
|
assert abs(gain - 61_000.0) < 1.0
|
|
assert p.should_mutate_deep(50_000, 10_000, 3.0, 1.0)
|
|
|
|
def test_live_zone_edit_always_profitable(self):
|
|
# S = 0 derives the existing Subscription live-zone policy as a
|
|
# special case of the formula.
|
|
p = policy_for_mode(AuthMode.SUBSCRIPTION)
|
|
assert p.should_mutate_deep(1, 0, 0.0, 1.0)
|
|
assert p.should_mutate_deep(2_000, 0, 0.0, 1.0)
|
|
|
|
def test_cold_cache_ignores_suffix(self):
|
|
# P_alive = 0 (TTL lapsed): the idle-timer compaction window.
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
assert p.should_mutate_deep(2_000, 50_000, 0.0, 0.0)
|
|
|
|
def test_clamps_out_of_range_inputs(self):
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
clamped = p.net_mutation_gain(2_000, 50_000, -5.0, 7.0)
|
|
reference = p.net_mutation_gain(2_000, 50_000, 0.0, 1.0)
|
|
assert abs(clamped - reference) < 1e-6
|
|
|
|
def test_nan_inputs_guarded(self):
|
|
# NaN reads -> 0, NaN p_alive -> 1 (same as Rust): the gain stays
|
|
# finite instead of poisoning the mutate decision.
|
|
import math
|
|
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
guarded = p.net_mutation_gain(2_000, 50_000, float("nan"), float("nan"))
|
|
assert math.isfinite(guarded)
|
|
reference = p.net_mutation_gain(2_000, 50_000, 0.0, 1.0)
|
|
assert abs(guarded - reference) < 1e-6
|
|
|
|
def test_negative_int_inputs_clamped(self):
|
|
# Rust takes u32 — negative Python ints must not flip the sign of
|
|
# the result; they clamp to 0.
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
assert p.net_mutation_gain(-2_000, -50_000, 5.0, 1.0) == p.net_mutation_gain(0, 0, 5.0, 1.0)
|
|
assert p.break_even_reads(-5, 10_000) == 0.0
|
|
assert p.net_mutation_gain(2_000, -1, 5.0, 1.0) == p.net_mutation_gain(2_000, 0, 5.0, 1.0)
|
|
|
|
def test_break_even_reads_matches_research_anchor(self):
|
|
# R = 11.5*(S/dT - 1): 2K/50K -> 276; 50K/10K -> negative
|
|
# (profitable from the first read); dT=0 -> 0.
|
|
p = policy_for_mode(AuthMode.PAYG)
|
|
assert abs(p.break_even_reads(2_000, 50_000) - 276.0) < 0.5
|
|
assert p.break_even_reads(50_000, 10_000) < 0.0
|
|
assert p.break_even_reads(0, 10_000) == 0.0
|
|
|
|
def test_constants_match_rust(self):
|
|
from headroom.transforms.compression_policy import (
|
|
CACHE_READ_MULTIPLIER,
|
|
CACHE_WRITE_MULTIPLIER,
|
|
)
|
|
|
|
assert CACHE_WRITE_MULTIPLIER == 1.25
|
|
assert CACHE_READ_MULTIPLIER == 0.1
|