diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 9b0c9d502..00c719d2e 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2074,7 +2074,16 @@ class AnthropicHandlerMixin: # On a confirmed-cold turn we deliberately do NOT replay the previously # forwarded prefix: the cache is dead (nothing to keep byte-identical for) # and the replay would clobber the whole-prefix recompaction we just did. - if _decision.should_compress and not _skip_compression_for_backpressure: + # + # Backpressure skips the compression PIPELINE but must NOT skip this + # replay: on the saturated path `optimized_messages` is the raw + # originals, which mismatch the compressed prefix the provider cached + # — so every gated request busted its session's prompt cache exactly + # when traffic (and the re-write cost) peaked. The overlay itself is + # O(prefix) comparisons plus one token recount only when it actually + # replays, which is far cheaper than the whole-prefix cache re-write + # it prevents, so it stays on even under backpressure. + if _decision.should_compress: if _cold_recompact_active: _overlay_replayed = False else: @@ -2089,15 +2098,10 @@ class AnthropicHandlerMixin: optimized_messages = _ov optimized_tokens = tokenizer.count_messages(optimized_messages) else: - replay_skip_reason = ( - "pre_upstream_backpressure" - if _skip_compression_for_backpressure - else _decision.passthrough_reason - ) logger.debug( "[%s] Cached-prefix replay skipped: reason=%s", request_id, - replay_skip_reason, + _decision.passthrough_reason, ) # Own cache_control placement: the client moves the breakpoint each diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 40efb6571..76b2700d1 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1245,8 +1245,52 @@ try: except ValueError: EAGER_PRELOAD_TIMEOUT_SECONDS = 120.0 -# Maximum compression cache sessions (prevents unbounded memory growth) -MAX_COMPRESSION_CACHE_SESSIONS = 500 +# Maximum compression cache sessions (prevents unbounded memory growth). +# Overridable via HEADROOM_COMPRESSION_CACHE_MAX_SESSIONS for gateway +# deployments (e.g. Kong sidecars) that fan many concurrent sessions into one +# proxy process. Falls back to 500 on an unparseable value; floor of 1. +try: + MAX_COMPRESSION_CACHE_SESSIONS = max( + 1, int(os.environ.get("HEADROOM_COMPRESSION_CACHE_MAX_SESSIONS", "500")) + ) +except ValueError: + MAX_COMPRESSION_CACHE_SESSIONS = 500 + +# Idle TTL for per-session compression caches. Eviction is bust-free only +# once the provider's own prompt cache has lapsed, so this must exceed the +# LONGEST provider cache TTL Headroom serves — Anthropic's 1h extended +# breakpoint (3600s), not just the common 5m ephemeral cache. Evicting +# earlier would itself cause the bust this state exists to prevent: the +# session returns, the provider still holds the old bytes, but the map that +# replays them is gone. The cache must also outlive the prefix TRACKER's +# session TTL (600s): after the tracker expires, `apply_cached`'s +# byte-identical swap is the only thing still protecting the provider +# prefix. Default 3900s = 1h + 5m grace. Deployments that never opt into +# the 1h breakpoint can lower it via HEADROOM_COMPRESSION_CACHE_TTL_SECONDS. +try: + # Floor of 600s: never below the prefix tracker's session TTL, or the + # sweep could reclaim the byte-identical swap map while it is the only + # remaining protection for a still-live provider prefix (see above). + COMPRESSION_CACHE_TTL_SECONDS = max( + 600.0, + float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900")), + ) +except ValueError: + COMPRESSION_CACHE_TTL_SECONDS = 3900.0 + +# Entries per session compression cache. 10k covers a single conversation with +# ~2x headroom even at a 1M-token context (a compressible tool_result is at +# least a few hundred tokens, so at most ~5k can be live at once). Raise via +# HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES only for workloads that fan many +# concurrent conversations into ONE session id (shared fallback ids, heavy +# subagent fan-out) — entry LRU is hit-refreshed, so an undersized cap shows +# up as misses on still-live entries, i.e. prefix-cache busts. Floor of 100. +try: + COMPRESSION_CACHE_MAX_ENTRIES = max( + 100, int(os.environ.get("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "10000")) + ) +except ValueError: + COMPRESSION_CACHE_MAX_ENTRIES = 10000 # --------------------------------------------------------------------------- diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index a392e9fe1..b526493ad 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -36,6 +36,7 @@ import os import sys import threading import time +from collections import OrderedDict from collections.abc import Callable from dataclasses import fields, is_dataclass, replace from datetime import datetime, timezone @@ -128,6 +129,8 @@ from headroom.proxy.cost import ( merge_cost_stats, # noqa: F401 ) from headroom.proxy.helpers import ( + COMPRESSION_CACHE_MAX_ENTRIES, + COMPRESSION_CACHE_TTL_SECONDS, COMPRESSION_TIMEOUT_SECONDS, # noqa: F401 EAGER_PRELOAD_TIMEOUT_SECONDS, MAX_COMPRESSION_CACHE_SESSIONS, # noqa: F401 @@ -1028,7 +1031,14 @@ class HeadroomProxy( # `CompressionCache` instances have their own internal lock guarding # `_cache`/`_stable_hashes`/`_first_seen` against concurrent # async-dispatched requests for the same session. - self._compression_caches: dict[str, CompressionCache] = {} + # Ordered by last access: `_get_compression_cache` moves a session to + # the end on every hit, so capacity eviction drops the idlest sessions + # — whose provider prefix cache has lapsed anyway — never a busy + # long-lived one. `_compression_cache_last_seen` drives the idle-TTL + # sweep in `_maybe_cleanup_compression_caches`. + self._compression_caches: OrderedDict[str, CompressionCache] = OrderedDict() + self._compression_cache_last_seen: dict[str, float] = {} + self._compression_caches_last_cleanup: float = time.time() self._compression_caches_lock = threading.RLock() self.logger = ( @@ -1580,6 +1590,49 @@ class HeadroomProxy( loop = asyncio.get_running_loop() return await loop.run_in_executor(self._background_compression_executor, fn) + # How often the lazy TTL sweep in `_get_compression_cache` may run. + _COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS = 60.0 + + def _maybe_cleanup_compression_caches(self, now: float) -> None: + """Evict per-session compression caches idle past their TTL. + + Caller must hold `_compression_caches_lock`. Piggybacked on + `_get_compression_cache` (the same lazy-sweep pattern as + `PrefixCacheTrackerRegistry._maybe_cleanup`) so no background task is + needed: any traffic at all keeps memory tracking the active-session + window, and a fully idle process has no memory pressure worth a timer. + + A session idle longer than `COMPRESSION_CACHE_TTL_SECONDS` has + outlived the provider prompt cache its entries protect — the default + exceeds Anthropic's 1h extended breakpoint, the longest provider TTL + served — so evicting it cannot bust anything: the provider already + forgot the prefix. If the session does return, the cost is one + cache-write turn (fail-open), which it was going to pay regardless. + The TTL must never be set below the prefix tracker's session TTL: + after the tracker expires, this cache's byte-identical swap is the + only remaining protection for a still-live provider prefix. + """ + if now - self._compression_caches_last_cleanup < ( + self._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS + ): + return + self._compression_caches_last_cleanup = now + expired = [ + sid + for sid, seen in self._compression_cache_last_seen.items() + if now - seen > COMPRESSION_CACHE_TTL_SECONDS + ] + for sid in expired: + self._compression_caches.pop(sid, None) + self._compression_cache_last_seen.pop(sid, None) + if expired: + logger.info( + "Evicted %d compression caches idle > %.0fs (%d sessions remain)", + len(expired), + COMPRESSION_CACHE_TTL_SECONDS, + len(self._compression_caches), + ) + def _get_compression_cache(self, session_id: str) -> CompressionCache: """Get or create a CompressionCache for a session. @@ -1588,27 +1641,43 @@ class HeadroomProxy( for the same conversation) must return the **same** instance, otherwise the per-session cache state splits and the two halves diverge across requests. + + Every access refreshes both the LRU position and the idle-TTL clock, + so eviction — capacity or TTL — only ever hits sessions that have + gone quiet. Losing one costs at most a single cache-write turn + upstream; it never fails a request. """ with self._compression_caches_lock: - if session_id not in self._compression_caches: + now = time.time() + self._maybe_cleanup_compression_caches(now) + cache = self._compression_caches.get(session_id) + if cache is None: from headroom.cache.compression_cache import CompressionCache - # Evict oldest caches if at capacity + # Evict the least-recently-used quarter at capacity. The + # OrderedDict is maintained in access order, so the front is + # always the idlest session — never a busy long-lived one. if len(self._compression_caches) >= MAX_COMPRESSION_CACHE_SESSIONS: - # Remove oldest quarter to amortize cleanup cost - oldest_keys = list(self._compression_caches.keys())[ - : MAX_COMPRESSION_CACHE_SESSIONS // 4 - ] - for key in oldest_keys: - del self._compression_caches[key] + evict_count = min( + max(1, MAX_COMPRESSION_CACHE_SESSIONS // 4), + len(self._compression_caches), + ) + for _ in range(evict_count): + sid, _evicted = self._compression_caches.popitem(last=False) + self._compression_cache_last_seen.pop(sid, None) logger.info( - "Evicted %d compression caches (exceeded %d max sessions)", - len(oldest_keys), + "Evicted %d least-recently-used compression caches " + "(exceeded %d max sessions)", + evict_count, MAX_COMPRESSION_CACHE_SESSIONS, ) - self._compression_caches[session_id] = CompressionCache() - return self._compression_caches[session_id] + cache = CompressionCache(max_entries=COMPRESSION_CACHE_MAX_ENTRIES) + self._compression_caches[session_id] = cache + else: + self._compression_caches.move_to_end(session_id) + self._compression_cache_last_seen[session_id] = now + return cache def _setup_code_aware(self, config: ProxyConfig, transforms: list) -> str: """Set up code-aware compression if enabled. diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 090408ab5..fa445edc9 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -1947,7 +1947,19 @@ class ContentRouter(Transform): # we match that posture with a dedicated lock rather than relying on # GIL atomicity (which would not protect the read-then-evict sequence). self._frozen_verdicts: dict[int, bool] = {} - self._frozen_verdicts_max = 4096 + # The store is process-wide (one router per pipeline, shared by every + # session), so the cap must scale with the number of CONCURRENT + # sessions, not one user's workload: at org scale (many users behind + # one sidecar) 4096 churns in minutes and FIFO eviction lets tightened + # thresholds flip a still-cached block's verdict — a prefix bust. + # Read at construction so tests and multi-tenant deployments can size + # it via HEADROOM_FROZEN_VERDICTS_MAX without a module reload. + try: + self._frozen_verdicts_max = max( + 256, int(os.environ.get("HEADROOM_FROZEN_VERDICTS_MAX", "4096")) + ) + except ValueError: + self._frozen_verdicts_max = 4096 self._frozen_lock = threading.Lock() # Reset verdicts whenever the shadowed cache is cleared. self._cache.register_on_clear(self._clear_frozen_verdicts) diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index be51e48a0..383157bd5 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/tests/test_anthropic_pre_upstream_backpressure.py @@ -1094,3 +1094,74 @@ def test_response_cache_keys_on_lookup_messages_not_mutated(): assert cache.set_messages == cache.get_messages # And specifically the raw lookup messages, not the scanner's rewrite. assert cache.set_messages == [{"role": "user", "content": "hello"}] + + +# --------------------------------------------------------------------------- # +# Backpressure must not bust the provider prompt cache: the compression # +# pipeline is skipped under saturation, but the previously-forwarded # +# (compressed) prefix must still be replayed byte-identical. Forwarding raw # +# originals would mismatch the bytes the provider cached — busting every # +# gated session's prefix exactly when the proxy is busiest. # +# --------------------------------------------------------------------------- # + + +def test_backpressure_passthrough_replays_cached_prefix(stage_log_capture): + prev_original = [{"role": "user", "content": "ORIGINAL " * 6000}] + prev_forwarded = [{"role": "user", "content": "[compressed-form]"}] + + async def _run() -> None: + sem = asyncio.Semaphore(1) + await sem.acquire() # saturate: the request's acquire will time out + handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem) + handler.config.optimize = True + handler.config.anthropic_pre_upstream_acquire_timeout_seconds = 0.01 + handler.anthropic_pipeline = SimpleNamespace(apply=MagicMock()) + + tracker = SimpleNamespace( + _cached_token_count=0, + get_frozen_message_count=lambda: 0, + get_last_original_messages=lambda: copy.deepcopy(prev_original), + get_last_forwarded_messages=lambda: copy.deepcopy(prev_forwarded), + update_from_response=lambda *a, **k: None, + record_request=lambda *a, **k: None, + ) + handler.session_tracker_store = SimpleNamespace( + compute_session_id=lambda *a, **k: "sess-1", + get_or_create=lambda *a, **k: tracker, + resolve_tracker=lambda *a, **k: tracker, + ) + + forwarded_bodies: list[dict] = [] + orig_retry = handler._retry_request + + async def _capturing_retry(method, url, headers, body, **kw): + forwarded_bodies.append(copy.deepcopy(body)) + return await orig_retry(method, url, headers, body, **kw) + + handler._retry_request = _capturing_retry + + req = _build_request( + { + "model": "claude-3-5-sonnet-latest", + "messages": copy.deepcopy(prev_original) + + [{"role": "user", "content": "next turn"}], + }, + {"authorization": "Bearer sk-ant-api-test"}, + ) + try: + response = await handler.handle_anthropic_messages(req) + assert response.status_code == 200 + # Saturation must still skip the CPU-bound pipeline... + assert not handler.anthropic_pipeline.apply.called + finally: + sem.release() + + assert forwarded_bodies, "request never reached upstream" + sent = forwarded_bodies[-1]["messages"] + # ...but the forwarded prefix must be last turn's exact bytes, not the + # raw original (which the provider never cached). + assert sent[0]["content"] == "[compressed-form]" + assert sent[-1]["content"] == "next turn" + + with _tokenizer_patch(): + anyio.run(_run) diff --git a/tests/test_compression_cache_registry.py b/tests/test_compression_cache_registry.py new file mode 100644 index 000000000..f510ff8d8 --- /dev/null +++ b/tests/test_compression_cache_registry.py @@ -0,0 +1,134 @@ +"""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 diff --git a/tests/test_org_scale_limits.py b/tests/test_org_scale_limits.py new file mode 100644 index 000000000..716b1cc5e --- /dev/null +++ b/tests/test_org_scale_limits.py @@ -0,0 +1,143 @@ +"""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 + + +# --------------------------------------------------------------------------- # +# 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) diff --git a/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py b/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py index c79071b3b..1b79fdf40 100644 --- a/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py +++ b/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py @@ -199,7 +199,18 @@ def test_bypass_header_does_not_invoke_cached_prefix_replay(monkeypatch): assert captured[-1]["messages"] == messages -def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch): +def test_backpressure_still_invokes_cached_prefix_replay(monkeypatch): + """INVERTED from the pre-#3261 contract this test used to pin. + + Backpressure sheds the compression PIPELINE (the CPU-heavy stage), but + the byte-identical cached-prefix replay must STILL run: skipping it + forwarded raw originals over a compressed cached prefix, busting every + gated session's prompt cache exactly at peak load (the saturated path + previously emitted `Cached-prefix replay skipped: + reason=pre_upstream_backpressure` — that skip was the bug). The replay + self-guards and no-ops here (no previous turn), so the raw messages + still pass through unchanged. + """ app = create_app( _config( optimize=True, @@ -208,12 +219,29 @@ def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch): ) ) - def fail_if_called(*args, **kwargs): # noqa: ANN002, ANN003 - raise AssertionError("cached-prefix replay must be skipped under backpressure") + from headroom.cache import prefix_tracker as _pt - monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", fail_if_called) - debug = Mock() - monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.debug", debug) + real_overlay = _pt.overlay_cached_prefix + overlay_calls: list[int] = [] + + def spy(*args, **kwargs): # noqa: ANN002, ANN003 + overlay_calls.append(1) + return real_overlay(*args, **kwargs) + + # Patch every binding of overlay_cached_prefix: the handler historically + # imported it from prefix_tracker per-request, and the shared session + # engine (headroom.proxy.session_engine, later in this stack) binds it + # at module import — cover both so this test holds across the stack. + monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", spy) + try: + import headroom.proxy.session_engine as _se + + monkeypatch.setattr(_se, "overlay_cached_prefix", spy) + except ImportError: + pass + + info = Mock() + monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.info", info) proxy = app.state.proxy class _SaturatedSemaphore: @@ -236,10 +264,13 @@ def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch): proxy.anthropic_pre_upstream_sem.release() assert response.status_code == 200, response.text - assert captured[-1]["messages"] == messages - assert any( - call.args and call.args[-1] == "pre_upstream_backpressure" for call in debug.call_args_list + # Backpressure engaged (pipeline shed)... + assert any("pre_upstream_backpressure" in str(call) for call in info.call_args_list), ( + "backpressure did not engage — the test setup no longer saturates" ) + # ...but the replay ran (and, with no previous turn, no-op'd safely). + assert overlay_calls, "cached-prefix replay must run under backpressure" + assert captured[-1]["messages"] == messages def test_optimize_on_aligned_history_preserves_replay():