headroom/tests/test_compression_cache_registry.py
Tejas Chopra d12ea50122
feat(proxy): unify proxy and sidecar compression on one session engine (#3271)
> Replaces #3263 (same changeset, squashed to one conventional commit —
after the base PRs squash-merged, the stacked branch's commit history
could not pass the commitlint gate against main, and force-pushing the
original branch was not permitted). #3261 and #3270 (which replaced
#3262) are merged; this is the last piece of the stack.

## Goal

One brain. The cache-management tier — freeze computation, Zone-1 byte
swap, cached-prefix overlay — previously existed twice: inline in the
proxy request handlers, and (as of #3270) in the `/v1/compress` sidecar
path. This PR extracts it into **`headroom/proxy/session_engine.py`**,
invoked by BOTH. Every future cache-management fix lands in both modes
by construction.

## Design

**`prepare_turn(...)` → `TurnPrep`** — freeze +
`mark_stable_from_messages` + `apply_cached`, with two *deliberately
different, documented* freeze policies:
- `FREEZE_POLICY_CONFIRMED_CLAMP`: `min(tracker_frozen, cache_count)` —
never freeze past provider-confirmed (the #327 posture). The Anthropic
proxy passes its already-composed tracker/strict-override value,
reproducing the previous `min()` byte-for-byte.
- `FREEZE_POLICY_REPLAYABLE`: `max(cache_count, explicit)` — freeze
everything locally replayable, because whatever was previously returned
*is* the provider's cache contract; recompressing it (even "better")
busts.

**`finalize_turn(...)` → `TurnFinal`** — the byte-identical
cached-prefix replay (`overlay_cached_prefix`) + conditional token
recount hook.

Run as a **strictly behavior-preserving extraction**: the bar was every
pre-existing test passing *unmodified*, and it held.

## What migrated

| Path | Status |
|---|---|
| `/v1/compress` sidecar turn |  engine (REPLAYABLE); lock, executor
offload, savings accounting, record_returned unchanged |
| `anthropic.py` token-mode pre-block + overlay |  engine
(CONFIRMED_CLAMP); background compression, cold-start fast pass,
`_cold_recompact_active` skip preserved |
| `openai.py` proxy token-mode pre-block + overlay |  engine
(REPLAYABLE — formula-identical to the old bare `compute_frozen_count`);
the added `mark_stable` call means the freeze now survives entry-level
LRU eviction (test-pinned); the router's `_frozen_verdicts` remains the
boundary-message protection |
| `openai.py` cache-mode branch | ⏸ keeps bare `apply_cached` — cache
mode keeps the latest observation mutable by design |

Also fixed for BOTH handlers: overlay replay now runs under backpressure
(shedding it busted every gated session's prompt cache exactly at peak
load), and the inflation guard exempts replayed prefixes.

## Hardening (max-effort review, all applied)

`/v1/usage` applies on the executor under the per-session turn lock with
a timed acquire (503 `session_busy`); registry eviction skips sessions
mid-turn; `peek()` is expiry-aware; silent fallbacks log warnings;
RequestOutcome recorded on session 503s.

## Testing

- `tests/test_session_engine.py`: 13 direct unit tests — both policies,
explicit-pin precedence, REPLAYABLE-without-pin ≡ bare
`compute_frozen_count`, overlay fires/doesn't, recount only on replay,
freeze-survives-entry-eviction.
- Parity bar: full pre-existing suites pass unmodified — cache-stability
(Anthropic + OpenAI), overlay, backpressure (incl.
replay-under-saturation regression), cold-start fast pass, cache-mode,
session-mode byte-stability, compress-API, org-scale, registry. Full
local suite: 11k+ green.
- ruff check/format clean (CI's ruff 0.16.3).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-26 15:57:31 +05:30

185 lines
6.3 KiB
Python

"""Session-level lifecycle of the compression-cache registry.
Covers the two eviction paths on ``HeadroomProxy._get_compression_cache``:
* capacity eviction must be LRU by *access* (a busy long-lived session
survives; the idlest session goes), not FIFO by creation, and
* the lazy idle-TTL sweep must reclaim sessions whose provider prompt
cache has lapsed, while an access refreshes the clock.
Entry-level LRU/limits inside a single ``CompressionCache`` live in
``test_compression_cache.py``.
"""
from __future__ import annotations
import time
import pytest
pytest.importorskip("fastapi")
def _make_proxy():
from headroom.proxy.server import ProxyConfig, create_app
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
return app.state.proxy
def test_capacity_eviction_is_lru_not_fifo(monkeypatch) -> None:
"""At capacity, the idlest session is evicted — not the oldest-created."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 4)
proxy = _make_proxy()
for sid in ("a", "b", "c", "d"):
proxy._get_compression_cache(sid)
# "a" is the oldest-created; touch it so "b" becomes the LRU.
cache_a = proxy._get_compression_cache("a")
proxy._get_compression_cache("e")
assert "b" not in proxy._compression_caches
assert proxy._get_compression_cache("a") is cache_a
assert "b" not in proxy._compression_cache_last_seen
def test_capacity_eviction_count_respects_small_caps(monkeypatch) -> None:
"""A cap below 4 still evicts at least one session instead of looping."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 2)
proxy = _make_proxy()
proxy._get_compression_cache("a")
proxy._get_compression_cache("b")
proxy._get_compression_cache("c")
assert len(proxy._compression_caches) == 2
assert "a" not in proxy._compression_caches
def test_idle_ttl_sweep_evicts_expired_sessions(monkeypatch) -> None:
"""A session idle past the TTL is reclaimed by the lazy sweep."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
proxy._get_compression_cache("stale")
proxy._get_compression_cache("fresh")
now = time.time()
# Backdate "stale" past the TTL and allow the sweep to run again.
proxy._compression_cache_last_seen["stale"] = now - 101.0
proxy._compression_caches_last_cleanup = (
now - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger")
assert "stale" not in proxy._compression_caches
assert "stale" not in proxy._compression_cache_last_seen
assert "fresh" in proxy._compression_caches
def test_access_refreshes_ttl_clock(monkeypatch) -> None:
"""Accessing a session resets its idle clock, so it survives the sweep."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
proxy._get_compression_cache("busy")
now = time.time()
proxy._compression_cache_last_seen["busy"] = now - 101.0
# Access refreshes last_seen before any sweep can see it as expired.
cache = proxy._get_compression_cache("busy")
proxy._compression_caches_last_cleanup = (
now - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger")
assert proxy._get_compression_cache("busy") is cache
def test_sweep_is_rate_limited(monkeypatch) -> None:
"""Within the cleanup interval, even an expired session is not swept."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
proxy._get_compression_cache("stale")
proxy._compression_cache_last_seen["stale"] = time.time() - 101.0
# _compression_caches_last_cleanup is recent (set in __init__), so the
# sweep must not run yet.
proxy._get_compression_cache("trigger")
assert "stale" in proxy._compression_caches
def test_ttl_sweep_never_evicts_a_session_mid_turn(monkeypatch) -> None:
"""Popping a session whose turn lock is held splits the lock across two
cache instances: the straggler and its retry then run unserialized and
the retry's empty cache recompresses previously-returned bytes."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
cache = proxy._get_compression_cache("mid-turn")
now = time.time()
proxy._compression_cache_last_seen["mid-turn"] = now - 999.0
assert cache.session_turn_lock.acquire(timeout=1)
try:
proxy._compression_caches_last_cleanup = (
now - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger-1")
# In-flight: must survive the sweep despite being far past TTL.
assert proxy._compression_caches.get("mid-turn") is cache
finally:
cache.session_turn_lock.release()
# Turn finished: the next sweep may reclaim it.
proxy._compression_caches_last_cleanup = (
time.time() - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger-2")
assert "mid-turn" not in proxy._compression_caches
def test_capacity_eviction_skips_locked_sessions(monkeypatch) -> None:
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 2)
proxy = _make_proxy()
cache_a = proxy._get_compression_cache("a")
proxy._get_compression_cache("b")
assert cache_a.session_turn_lock.acquire(timeout=1)
try:
# "a" is the LRU but mid-turn — capacity pressure must evict "b".
proxy._get_compression_cache("c")
assert proxy._compression_caches.get("a") is cache_a
assert "b" not in proxy._compression_caches
finally:
cache_a.session_turn_lock.release()