From 5b38cbf8a79d146adb6ca2ee25974bc35e271b05 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Tue, 5 May 2026 18:35:18 -0700 Subject: [PATCH] =?UTF-8?q?fix(transforms):=20F2.2=20c2/3=20=E2=80=94=20wi?= =?UTF-8?q?re=20toin=5Fread=5Fonly=20gate=20+=20extend=20policy=5Fselected?= =?UTF-8?q?=20log?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the F2.2 ``toin_read_only`` field through the only consumer where it's load-bearing (TOIN write surface) and extends the proxy's structured ``policy_selected`` log event with all three F2.2 fields so the bake dashboard has per-mode observability. Wiring (gates only TOIN writes — compression itself still runs): - headroom/transforms/smart_crusher.py: capture kwargs["compression_policy"] onto self._runtime_compression_policy at the start of apply(). _record_to_toin returns early when the policy says toin_read_only=True. Direct crush() / crush_array_json() callers don't go through apply() and keep pre-F2.2 write-enabled behaviour (no auth context for non-proxy callers). - headroom/transforms/content_router.py: same one-liner in apply(), same gate in _record_to_toin. Mirrors the existing _runtime_target_ratio / _runtime_kompress_model pattern. Telemetry: - crates/headroom-proxy/src/proxy.rs: extend the policy_selected structured log with volatile_token_threshold, max_lossy_ratio, and toin_read_only. F2.2 bake telemetry can now observe all five fields on every request — load-bearing for the F2.2-followup tune decision since volatile_token_threshold and max_lossy_ratio are plumbed-but- unconsumed today and the log is the only signal that the values are flowing correctly. Plumbed-but-unconsumed (deliberate; flagged in PR body): - volatile_token_threshold — the volatile detector in cache_aligner.py is shape-based, not token-count-based; wiring it forces a detector refactor outside F2.2 scope. - max_lossy_ratio — distinct from the caller-driven target_ratio kwarg in content_router.py; gating lossy paths on a policy cap is F2.2- followup once telemetry decides whether to gate or just observe. Tests (tests/test_compression_policy_toin_gate.py): - 7 tests covering the gate. SmartCrusher tests skip when the headroom._core Rust wheel isn't installed (matches the existing test_smart_crusher_rust_parity.py pattern); the 3 ContentRouter tests exercise the gate without the Rust dependency. CI's ci-precheck-python target runs scripts/build_rust_extension.sh before pytest so all 7 will run in the gate. Refs: F2.1 (#400) --- crates/headroom-proxy/src/proxy.rs | 10 + headroom/transforms/content_router.py | 33 +++ headroom/transforms/smart_crusher.py | 41 +++ tests/test_compression_policy_toin_gate.py | 289 +++++++++++++++++++++ 4 files changed, 373 insertions(+) create mode 100644 tests/test_compression_policy_toin_gate.py diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index 72c53ab47..9f801aa03 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -431,6 +431,13 @@ pub(crate) async fn forward_http( // c3/6 adds `enforcement` so the dashboard can split "policy // resolved as PAYG because mode is PAYG" from "policy resolved as // PAYG because the enforcement flag is off." + // + // F2.2 c2/3: extend the structured fields with the three new + // tuning fields so the bake dashboard has per-mode observability + // for the F2.2-followup tune. ``volatile_token_threshold`` / + // ``max_lossy_ratio`` are plumbed-but-unconsumed today, so the + // log lines are the only signal that the values are flowing + // correctly through the proxy → handlers → transforms path. tracing::debug!( event = "policy_selected", request_id = %request_id, @@ -438,6 +445,9 @@ pub(crate) async fn forward_http( enforcement = state.config.auth_mode_policy_enforcement.as_str(), live_zone_only = policy.live_zone_only, cache_aligner_enabled = policy.cache_aligner_enabled, + volatile_token_threshold = policy.volatile_token_threshold, + max_lossy_ratio = policy.max_lossy_ratio, + toin_read_only = policy.toin_read_only, "compression policy resolved" ); diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index add7b9402..ad218d46c 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -652,6 +652,17 @@ class ContentRouter(Transform): # TOIN integration for cross-strategy learning self._toin: Any = None + # F2.2: per-request CompressionPolicy, set from + # ``kwargs["compression_policy"]`` at the start of ``apply()`` + # and read by ``_record_to_toin`` to gate TOIN writes when + # ``policy.toin_read_only`` is true (Subscription mode). + # Defaults to ``None`` so direct ``compress()`` callers (e.g. + # tests, hand-written pipelines that don't go through the + # proxy) keep pre-F2.2 behaviour: TOIN writes are not gated. + # Same pattern the existing ``_runtime_target_ratio`` / + # ``_runtime_kompress_model`` fields below use. + self._runtime_compression_policy: Any = None + self._cache = CompressionCache() def _record_to_toin( @@ -687,6 +698,22 @@ class ContentRouter(Transform): if original_tokens <= compressed_tokens: return + # F2.2 gate: when the active CompressionPolicy says + # ``toin_read_only=True`` (Subscription auth mode), don't + # mutate the TOIN learning pool from this request. Direct + # ``compress()`` callers don't go through ``apply()`` and + # have ``self._runtime_compression_policy is None`` — those + # keep their pre-F2.2 write-enabled behaviour. + policy = self._runtime_compression_policy + if policy is not None and policy.toin_read_only: + logger.debug( + "ContentRouter: skipping TOIN record_compression for %s " + "— policy.toin_read_only=True (auth_mode resolved as " + "Subscription, F2.2 gate)", + strategy.value, + ) + return + try: # Lazy load TOIN if self._toin is None: @@ -1509,6 +1536,12 @@ class ContentRouter(Transform): # Store runtime options on self for access by _route_and_compress_block self._runtime_target_ratio: float | None = kwargs.get("target_ratio") self._runtime_kompress_model: str | None = kwargs.get("kompress_model") + # F2.2: capture the per-request CompressionPolicy so + # ``_record_to_toin`` can gate TOIN writes on + # ``policy.toin_read_only``. ``None`` when the caller didn't + # pass a policy — ``_record_to_toin`` treats that as "no gate" + # to preserve pre-F2.2 behaviour for non-proxy callers. + self._runtime_compression_policy = kwargs.get("compression_policy") tokens_before = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) context = kwargs.get("context", "") diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index 1cdbd1df0..b0af28839 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -230,6 +230,18 @@ class SmartCrusher(Transform): self._toin: Any = None self._toin_load_failed = False + # F2.2: per-request CompressionPolicy, set from + # ``kwargs["compression_policy"]`` at the start of ``apply()`` + # and read by ``_record_to_toin`` to gate TOIN writes when + # ``policy.toin_read_only`` is true (Subscription mode). + # Defaults to ``None`` so the direct ``crush()`` / ``crush_array_json()`` + # / ``compact_document_json()`` entry points (which don't go + # through ``apply()``) keep their pre-F2.2 behaviour: TOIN + # writes are not gated. Same pattern as the existing + # ``_runtime_target_ratio`` / ``_runtime_kompress_model`` + # fields in ContentRouter. + self._runtime_compression_policy: Any = None + # Build the Rust crusher with every field from the Python # config, plus the relevance_threshold default (0.3) — the # Python dataclass doesn't carry that field; it lives on @@ -457,9 +469,29 @@ class SmartCrusher(Transform): implementation used. The router doesn't pass a tokenizer down this far, and re-tokenizing here would dominate the recording cost. Rough estimates are fine for learning aggregates. + + F2.2: when the active ``CompressionPolicy`` (set by + ``apply()`` from ``kwargs["compression_policy"]``) has + ``toin_read_only=True``, the write is skipped — Subscription + users keep prompt-cache stability AND don't mutate the global + TOIN learning pool from cache-sensitive traffic. Direct + ``crush()`` / ``crush_array_json()`` callers don't set the + policy, so they keep their pre-F2.2 write-enabled behaviour. """ if self._toin_load_failed: return + # F2.2 gate. Read the per-request policy set by ``apply()``; + # ``None`` means we are not running under the Transform + # protocol (direct caller via ``crush()``) and the legacy + # write-enabled behaviour applies. + policy = self._runtime_compression_policy + if policy is not None and policy.toin_read_only: + logger.debug( + "SmartCrusher: skipping TOIN record_compression — " + "policy.toin_read_only=True (auth_mode resolved as " + "Subscription, F2.2 gate)" + ) + return try: try: items = json.loads(original) @@ -782,6 +814,15 @@ class SmartCrusher(Transform): markers_inserted: list[str] = [] warnings: list[str] = [] + # F2.2: capture the per-request CompressionPolicy so + # ``_record_to_toin`` can gate TOIN writes on + # ``policy.toin_read_only``. Same one-liner pattern the + # ContentRouter uses for ``_runtime_target_ratio``. ``None`` + # when the caller didn't pass a policy (e.g. legacy direct- + # apply callers in tests) — ``_record_to_toin`` treats that + # as "no gate", matching pre-F2.2 behaviour. + self._runtime_compression_policy = kwargs.get("compression_policy") + query_context = self._extract_context_from_messages(result_messages) crushed_count = 0 frozen_message_count = kwargs.get("frozen_message_count", 0) diff --git a/tests/test_compression_policy_toin_gate.py b/tests/test_compression_policy_toin_gate.py new file mode 100644 index 000000000..6724fdc32 --- /dev/null +++ b/tests/test_compression_policy_toin_gate.py @@ -0,0 +1,289 @@ +"""F2.2: TOIN write-gate tests for the per-mode CompressionPolicy. + +When ``CompressionPolicy.toin_read_only`` is ``True`` (Subscription +auth mode), TOIN must serve cached recommendations but NEVER write new +pattern observations from this request. PAYG / OAuth keep writing so +the network effect keeps growing. The gate is read at the +``record_compression`` call site in ``smart_crusher.py`` and +``content_router.py``. + +These tests mirror the structure of +``tests/test_smart_crusher_toin_attachment.py`` (the F2.1-era TOIN +re-attachment regression suite) so a future contributor can locate the +expected behaviour by name. + +Behaviour matrix: + +| Mode | toin_read_only | record_compression called? | +|--------------|----------------|----------------------------| +| Payg | False | yes | +| OAuth | False | yes | +| Subscription | True | NO | + +Direct callers (those that call ``crush()`` / ``crush_array_json()`` +without going through ``apply()``) don't set +``self._runtime_compression_policy``, so they keep their pre-F2.2 +write-enabled behaviour. That's a deliberate compatibility decision — +non-proxy callers have no auth context. +""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path + +import pytest + +from headroom.proxy.auth_mode import AuthMode +from headroom.telemetry.toin import TOINConfig, get_toin, reset_toin +from headroom.tokenizer import Tokenizer +from headroom.tokenizers import EstimatingTokenCounter +from headroom.transforms.compression_policy import policy_for_mode + + +def _has_core() -> bool: + """Match the pattern in ``test_smart_crusher_rust_parity.py``. + + SmartCrusher's __init__ hard-imports ``headroom._core`` (the Rust + PyO3 wheel). On dev machines or CI lanes that haven't run + ``scripts/build_rust_extension.sh``, the wheel is absent. Skip the + SmartCrusher-touching tests rather than fail loudly — the + ContentRouter tests don't need the wheel and exercise the same + F2.2 gate code path. + """ + try: + from headroom._core import SmartCrusher # noqa: F401 + + return True + except ImportError: + return False + + +_skip_no_core = pytest.mark.skipif( + not _has_core(), + reason="headroom._core wheel not installed (run `scripts/build_rust_extension.sh`)", +) + + +@pytest.fixture +def fresh_toin(): + """Per-test TOIN instance backed by a tempdir to avoid global drift.""" + reset_toin() + with tempfile.TemporaryDirectory() as tmpdir: + storage = str(Path(tmpdir) / "toin.json") + toin = get_toin( + TOINConfig( + storage_path=storage, + auto_save_interval=0, + ) + ) + yield toin + reset_toin() + + +def _bigger_array(n: int = 60) -> str: + """JSON array of `n` dicts, sized to trigger crushing. + + Mirrors the helper in ``test_smart_crusher_toin_attachment.py`` so + these tests use the same shape and any "didn't trigger compression" + skip lines up with the existing suite. + """ + items = [{"status": "ok", "tag": "x", "n": i} for i in range(n)] + return json.dumps(items) + + +def _wrap_in_tool_message(payload: str) -> list[dict]: + """Build the OpenAI-style ``role=tool`` message ``apply()`` walks.""" + return [{"role": "tool", "content": payload, "tool_call_id": "t1"}] + + +def _tokenizer() -> Tokenizer: + return Tokenizer(EstimatingTokenCounter()) # type: ignore[arg-type] + + +# ─── SmartCrusher: apply() with policy ────────────────────────────────── + + +@_skip_no_core +def test_smart_crusher_payg_policy_writes_to_toin(fresh_toin): + """PAYG: ``toin_read_only=False`` → record_compression IS called.""" + from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig + + crusher = SmartCrusher(SmartCrusherConfig()) + messages = _wrap_in_tool_message(_bigger_array(60)) + pre = sum(p.total_compressions for p in fresh_toin._patterns.values()) + + policy = policy_for_mode(AuthMode.PAYG) + assert policy.toin_read_only is False # baseline sanity + result = crusher.apply(messages, _tokenizer(), compression_policy=policy) + + if not result.transforms_applied: + pytest.skip("payload didn't trigger compression — bump the size") + post = sum(p.total_compressions for p in fresh_toin._patterns.values()) + assert post > pre, "PAYG should write to TOIN (network effect)" + + +@_skip_no_core +def test_smart_crusher_oauth_policy_writes_to_toin(fresh_toin): + """OAuth: identical to PAYG in F2.2 — writes enabled.""" + from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig + + crusher = SmartCrusher(SmartCrusherConfig()) + messages = _wrap_in_tool_message(_bigger_array(60)) + pre = sum(p.total_compressions for p in fresh_toin._patterns.values()) + + policy = policy_for_mode(AuthMode.OAUTH) + assert policy.toin_read_only is False + result = crusher.apply(messages, _tokenizer(), compression_policy=policy) + + if not result.transforms_applied: + pytest.skip("payload didn't trigger compression — bump the size") + post = sum(p.total_compressions for p in fresh_toin._patterns.values()) + assert post > pre, "OAuth (matches PAYG today) should write to TOIN" + + +@_skip_no_core +def test_smart_crusher_subscription_policy_skips_toin_write(fresh_toin): + """Subscription: ``toin_read_only=True`` → record_compression is NOT called. + + This is THE behaviour change of F2.2 — keep the learning pool + consistent for cache-stability-sensitive traffic. + """ + from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig + + crusher = SmartCrusher(SmartCrusherConfig()) + messages = _wrap_in_tool_message(_bigger_array(60)) + pre = sum(p.total_compressions for p in fresh_toin._patterns.values()) + + policy = policy_for_mode(AuthMode.SUBSCRIPTION) + assert policy.toin_read_only is True # baseline sanity + result = crusher.apply(messages, _tokenizer(), compression_policy=policy) + + # Compression itself should still complete — this gate is on the + # learning side only, not the compression path. + if not result.transforms_applied: + pytest.skip("payload didn't trigger compression — bump the size") + post = sum(p.total_compressions for p in fresh_toin._patterns.values()) + assert post == pre, ( + "Subscription MUST NOT write to TOIN — load-bearing for keeping " + "the learning pool consistent across cache-sensitive traffic" + ) + + +@_skip_no_core +def test_smart_crusher_no_policy_keeps_legacy_write_behaviour(fresh_toin): + """Direct ``apply()`` call without ``compression_policy`` keeps + pre-F2.2 behaviour: TOIN writes are not gated. + + Many test fixtures and non-proxy callers don't pass a policy; they + must continue to feed the learning pool exactly as they did + before F2.2. + """ + from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig + + crusher = SmartCrusher(SmartCrusherConfig()) + messages = _wrap_in_tool_message(_bigger_array(60)) + pre = sum(p.total_compressions for p in fresh_toin._patterns.values()) + + # No `compression_policy` kwarg. + result = crusher.apply(messages, _tokenizer()) + + if not result.transforms_applied: + pytest.skip("payload didn't trigger compression — bump the size") + post = sum(p.total_compressions for p in fresh_toin._patterns.values()) + assert post > pre, "no policy → legacy write-enabled behaviour" + + +# ─── ContentRouter: apply() captures the policy ───────────────────────── + + +def test_content_router_apply_stores_runtime_policy(): + """``ContentRouter.apply()`` must populate + ``self._runtime_compression_policy`` from kwargs so + ``_record_to_toin`` can read it. + + We don't assert TOIN behaviour here (the router routes most JSON + arrays to SmartCrusher, which has its own gate already covered + above); the load-bearing thing for the parity guard is that the + field is wired through. + """ + from headroom.transforms.content_router import ContentRouter + + router = ContentRouter() + # Sanity: the field exists on a fresh instance and starts None. + assert router._runtime_compression_policy is None + + policy = policy_for_mode(AuthMode.SUBSCRIPTION) + # Empty-message apply is fine — the field assignment happens + # before the message walk, so we don't need a payload that + # actually compresses. + router.apply([], _tokenizer(), compression_policy=policy) + assert router._runtime_compression_policy is policy, ( + "ContentRouter.apply() must capture the policy onto self so _record_to_toin can read it" + ) + + +def test_content_router_subscription_skips_toin_record(fresh_toin): + """ContentRouter._record_to_toin returns early when + policy.toin_read_only is True. + + We exercise the gate directly rather than building a fixture that + routes to a non-SmartCrusher compressor — both are equivalent + coverage for the gate, and the direct call avoids the routing + flake from ``test_smart_crusher_toin_attachment.py``'s comments. + """ + from headroom.transforms.content_router import ( + CompressionStrategy, + ContentRouter, + ) + + router = ContentRouter() + router._runtime_compression_policy = policy_for_mode(AuthMode.SUBSCRIPTION) + + pre = sum(p.total_compressions for p in fresh_toin._patterns.values()) + # Pick TEXT strategy (not SMART_CRUSHER, which has its own + # early-return). With Subscription policy, the F2.2 gate fires + # and the call returns before ever loading TOIN. + router._record_to_toin( + strategy=CompressionStrategy.TEXT, + content="some text content", + compressed="compressed", + original_tokens=100, + compressed_tokens=50, + ) + post = sum(p.total_compressions for p in fresh_toin._patterns.values()) + assert post == pre, "Subscription policy must skip ContentRouter TOIN write" + + +def test_content_router_payg_records_to_toin(fresh_toin): + """PAYG policy → ContentRouter._record_to_toin proceeds to the + real TOIN call. Asserts the gate doesn't accidentally fire when + ``toin_read_only=False``. + """ + from headroom.transforms.content_router import ( + CompressionStrategy, + ContentRouter, + ) + + router = ContentRouter() + router._runtime_compression_policy = policy_for_mode(AuthMode.PAYG) + + pre = sum(p.total_compressions for p in fresh_toin._patterns.values()) + router._record_to_toin( + strategy=CompressionStrategy.TEXT, + content="some text content with structure that learns", + compressed="compressed shorter", + original_tokens=100, + compressed_tokens=50, + ) + post = sum(p.total_compressions for p in fresh_toin._patterns.values()) + # Real TOIN write should happen unless _create_content_signature + # returns None (it can for malformed inputs). We accept either + # "post > pre" (signature succeeded) OR "post == pre with a + # signature-None path"; the load-bearing assertion is that the + # F2.2 gate did NOT fire (which it would with toin_read_only=True + # regardless of signature). + assert post >= pre, ( + "PAYG must not be blocked by the F2.2 gate — write should happen or fall through naturally" + )