headroom/tests/test_org_scale_limits.py
Tejas Chopra 4fa88026d9
feat(compress): session-aware /v1/compress (sidecar mode) + /v1/usage relay (#3270)
> Replaces #3262 (same changeset, squashed to one conventional commit —
the stacked branch's history could not pass commitlint after #3261's
squash-merge broke ancestry, and force-pushing the original branch was
not permitted). All review findings from the two max-effort reviews are
already incorporated; #3261 is merged.

## Why

Gateways that own routing (e.g. Kong as the upstream caller) can't use
Headroom's proxy path, and the stateless `/v1/compress` pushes all
byte-replay bookkeeping onto the caller. This PR moves that state into
the endpoint: **the caller sends the raw conversation + a session id
every turn, forwards the returned bytes verbatim, and gets a
byte-identical prefix — provider prompt cache preserved, no forwarding
through Headroom.**

## Design

- **Session pre-work** mirrors the proxy's Zone 1: content-addressed
swap of previously-computed compressed bytes, then freeze the **entire
locally-replayable prefix** (`compute_frozen_count`).
- **Freeze posture deliberately differs from the proxy's `min(tracker,
cache)`**: in sidecar mode, whatever this endpoint previously returned
*is* the provider's cache contract — recompressing an already-returned
message (even into a smaller form) is a bust. Over-freezing only forgoes
tail compression; it can never bust. (A test caught exactly this:
recompression drift produced a smaller form, and
`overlay_cached_prefix`'s non-inflation guard then couldn't repair it.)
- **`PrefixCacheTracker.record_returned()`** — the sidecar equivalent of
"last forwarded", captured at return time because whatever is returned
is what the caller forwards.
- **`POST /v1/usage`** (same loopback exposure policy): the caller
relays the provider's usage block; `update_from_response` makes freeze
decisions provider-confirmed. Optional — skipping it degrades freeze
precision, never correctness.
- Sessions are NUL-namespaced (`compress\x00<id>`, unspoofable via HTTP
headers); the registry's TTL/LRU lifecycle from #3261 applies
automatically. No session id ⇒ stateless contract byte-for-byte
unchanged.

## Hardening (from two max-effort code reviews, all applied)

- `compress_user_messages` + session_id → 400 (user-message rewrites are
not content-addressed → guaranteed later bust).
- Session-mode timeout / lock-busy → 503 `compression_timeout` /
`session_busy` with retry semantics, instead of failing open with raw
bytes (desync bust).
- Header-based session ids gated behind
`HEADROOM_COMPRESS_SESSION_FROM_HEADER` (default off).
- `/v1/usage` validation: unknown/expired session → 404; both cache
fields absent → 400; single-present-zero → `{"applied": false, "reason":
"no_cache_signal"}` (never wipes freeze state).
- Warm-turn savings recomputed from the raw payload (honest
`tokens_saved`), all CPU work in the executor under a per-session turn
lock.

## Caller contract (Kong)

1. Send raw history + `config.session_id` (or `x-headroom-session-id`
with the env gate on) every turn.
2. Forward the returned `messages` to the provider **verbatim**.
3. Optionally relay the provider's usage block to `/v1/usage`.

## Testing

20 cases in `tests/test_compress_session_mode.py`: stateless regression
+ no state leakage, invalid-id rejection, 2-turn and 3-turn whole-prefix
byte-stability, tracker-loss stability, spoof-resistance, header gating,
lock-busy 503s, usage validation and no-signal handling,
unknown/expired-session 404, TTL-eviction fail-open, explicit
`frozen_message_count` precedence. Plus the full local suite green (11k+
tests).

## Phase 2 (follow-up)

#3263 migrates the proxy request path onto this same session engine so
both modes share one compression/state codepath.

🤖 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:34:23 +05:30

165 lines
6.3 KiB
Python

"""Org-scale sizing knobs: shared-process stores must be tunable and safe.
One Headroom process shared by many users (gateway sidecar/pool) stresses
stores that were sized for a single user's workload:
* the per-session compression-cache entry cap
(``HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES``),
* the process-wide frozen-verdicts store
(``HEADROOM_FROZEN_VERDICTS_MAX``), and
* the session registry under churn (active sessions must survive a flood
of transient ones — the LRU property at scale).
Registry TTL/LRU mechanics live in ``test_compression_cache_registry.py``.
"""
from __future__ import annotations
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
# --------------------------------------------------------------------------- #
# Per-session entry cap is plumbed through and env-tunable. #
# --------------------------------------------------------------------------- #
def test_compression_cache_entry_cap_is_plumbed(monkeypatch) -> None:
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_MAX_ENTRIES", 123)
proxy = _make_proxy()
assert proxy._get_compression_cache("s").max_entries == 123
def test_compression_cache_entry_cap_env_parsing(monkeypatch) -> None:
import importlib
import headroom.proxy.helpers as helpers_mod
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "50000")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 50000
# Floor: an absurdly small value cannot disable the cache.
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "1")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 100
# Garbage falls back to the default.
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "banana")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 10000
monkeypatch.delenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 10000
def test_compression_cache_ttl_env_rejects_non_finite(monkeypatch) -> None:
"""'nan'/'inf' parse as floats but poison every idle comparison — they
must fall back to the default like any other unparseable value."""
import importlib
import headroom.proxy.helpers as helpers_mod
for bad in ("nan", "inf", "-inf"):
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", bad)
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 3900.0, bad
# Below the 600s floor clamps up; above it passes through.
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "60")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 600.0
monkeypatch.delenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 3900.0
# --------------------------------------------------------------------------- #
# Frozen-verdicts store: process-wide, so it must be sizeable per deployment. #
# --------------------------------------------------------------------------- #
def test_frozen_verdicts_cap_env(monkeypatch) -> None:
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "65536")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 65536
# Floor: cannot be sized below 256.
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "1")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 256
# Garbage falls back to the default.
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "banana")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 4096
monkeypatch.delenv("HEADROOM_FROZEN_VERDICTS_MAX")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 4096
def test_frozen_verdicts_eviction_honors_configured_cap(monkeypatch) -> None:
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "256")
router = ContentRouter(ContentRouterConfig())
for key in range(300):
router._record_frozen_verdict(key, True)
assert len(router._frozen_verdicts) == 256
# FIFO: the oldest keys were evicted, the newest survive.
assert 0 not in router._frozen_verdicts
assert 299 in router._frozen_verdicts
# --------------------------------------------------------------------------- #
# Session registry under org-scale churn: active sessions always survive a #
# flood of transient ones (the property that keeps busts away at capacity). #
# --------------------------------------------------------------------------- #
def test_active_sessions_survive_transient_flood(monkeypatch) -> None:
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 100)
proxy = _make_proxy()
active = [f"active-{i}" for i in range(40)]
active_caches = {sid: proxy._get_compression_cache(sid) for sid in active}
# 400 transient sessions arrive interleaved with active-session traffic —
# 4x the cap, forcing repeated capacity evictions along the way.
for i in range(400):
proxy._get_compression_cache(f"transient-{i}")
if i % 5 == 0: # active sessions keep making requests
for sid in active:
proxy._get_compression_cache(sid)
# Every active session survived with its instance (and therefore its
# byte-replay state) intact; evictions only ever hit transient sessions.
for sid in active:
assert proxy._get_compression_cache(sid) is active_caches[sid], (
f"active session {sid} lost its cache to transient churn"
)
assert len(proxy._compression_caches) <= 100 + len(active)