akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into fix/holdout-conversation-key

Browse source
This commit is contained in:
JD Davis 2026-08-26 12:42:22 -05:00 • committed by GitHub
commit 4c75bf3cf3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 2327 additions and 124 deletions
Download patch file Download diff file Expand all files Collapse all files

6
headroom/cache/compression_cache.py vendored
View file

@ -123,6 +123,12 @@ class CompressionCache:
# `RLock` (not `Lock`) so future code can call locked methods from
# inside another locked method without self-deadlock.
self._lock = threading.RLock()
# Serializes one sidecar-mode compress turn per session (pre-work,
# pipeline, post-work run as one block on an executor thread). The
# sidecar contract is sequential turns per conversation; this lock
# keeps a contract-violating concurrent pair from interleaving and
# tearing the tracker's prev-original/prev-returned snapshots.
self.session_turn_lock = threading.Lock()
self._cache: OrderedDict[str, _CacheEntry] = OrderedDict()
# `_stable_hashes` is CONTENT-KEYED, not positional. It records "we
# have seen this content before and it is known not to compress

40
headroom/cache/prefix_tracker.py vendored
View file

@ -965,6 +965,26 @@ class PrefixCacheTracker:
def get_last_forwarded_messages(self) -> list[dict[str, Any]]:
return copy.deepcopy(self._last_forwarded_messages)
def record_returned(
self,
original_messages: list[dict[str, Any]],
returned_messages: list[dict[str, Any]],
) -> None:
"""Record the compressed form handed back to a compress-only caller.
Sidecar mode (session-aware ``/v1/compress``): Headroom does not
forward upstream, but whatever it RETURNS is what the caller forwards
— the same fact ``update_from_response`` records in proxy mode, just
captured at return time instead of send time. Only the transcript
snapshots and the activity clock move here; frozen-prefix counts are
left untouched because no provider response has confirmed anything
yet — they advance when the caller relays usage via ``/v1/usage``
(``update_from_response``), or stay at their conservative local value.
"""
self._last_activity = time.time()
self._last_original_messages = copy.deepcopy(original_messages)
self._last_forwarded_messages = copy.deepcopy(returned_messages)
def resolved_cache_ttl_seconds(self) -> int:
"""Effective prompt-cache lifetime for this session's provider."""
if self.config.cache_ttl_seconds is not None:
@ -1249,6 +1269,26 @@ class SessionTrackerStore:
self._lineage_affinities: dict[str, str | None] = {}
self._lineage_counter = itertools.count(1)
def peek(self, session_id: str) -> PrefixCacheTracker | None:
"""Return the live tracker for ``session_id``, else None.
Never creates: lookup paths that must not leave a footprint (e.g. the
``/v1/usage`` unknown-session check, where ``get_or_create`` would let
a flood of novel ids grow the store unboundedly within each TTL
window) use this instead of :meth:`get_or_create`.
A TTL-expired-but-unswept tracker answers None too: the sweep runs
lazily from get_or_create at 60s granularity, so without this check an
expired session would keep answering with stale pre-expiry state — and
a caller that then touched it (``update_from_response`` stamps
``_last_activity``) would resurrect the dead tracker indefinitely,
making the documented 404-on-expired contract nondeterministic.
"""
tracker = self._trackers.get(session_id)
if tracker is None or tracker.is_expired:
return None
return tracker
def get_or_create(self, session_id: str, provider: str) -> PrefixCacheTracker:
"""Get existing tracker or create a new one for this session."""
self._maybe_cleanup()

22
headroom/providers/registry.py
View file

@ -162,9 +162,27 @@ def resolve_extra_headers(
def resolve_api_targets(overrides: ProviderApiOverrides) -> ProviderApiTargets:
"""Resolve normalized upstream provider targets from configured overrides."""
from headroom.copilot_auth import is_copilot_upstream_url
openai = _normalize_api_url(overrides.openai, default=DEFAULT_OPENAI_API_URL)
# GitHub Copilot serves BOTH its OpenAI surface (``/chat/completions``,
# ``/responses``) and its Anthropic surface (``/v1/messages``, for Claude
# models) from the same host. When the OpenAI target is a Copilot host
# (``wrap copilot --subscription`` / ``wrap vscode`` both point it there so
# GPT models work) but no Anthropic target was set, Claude-model requests
# fell back to ``DEFAULT_ANTHROPIC_API_URL`` (api.anthropic.com) and 401'd
# with the Copilot bearer — "Invalid bearer token" (#3247). Default the
# Anthropic target to the same Copilot host so those requests reach the
# surface that actually serves them. An explicit ``ANTHROPIC_TARGET_API_URL``
# still wins (only a ``None`` override is filled in here).
anthropic_override = overrides.anthropic
if anthropic_override is None and is_copilot_upstream_url(openai):
anthropic_override = openai
return ProviderApiTargets(
anthropic=_normalize_api_url(overrides.anthropic, default=DEFAULT_ANTHROPIC_API_URL),
openai=_normalize_api_url(overrides.openai, default=DEFAULT_OPENAI_API_URL),
anthropic=_normalize_api_url(anthropic_override, default=DEFAULT_ANTHROPIC_API_URL),
openai=openai,
gemini=_normalize_api_url(overrides.gemini, default=DEFAULT_GEMINI_API_URL),
cloudcode=_normalize_api_url(overrides.cloudcode, default=DEFAULT_CLOUDCODE_API_URL),
vertex=_normalize_api_url(overrides.vertex, default=DEFAULT_VERTEX_API_URL),

88
headroom/proxy/handlers/anthropic.py
View file

@ -1685,36 +1685,42 @@ class AnthropicHandlerMixin:
if is_token_mode(self.config.mode):
comp_cache = self._get_compression_cache(session_id)
# Re-freeze boundary: consecutive stable messages from start.
# Safety: never freeze beyond provider-confirmed cached prefix.
# `prefix_tracker.frozen_message_count` (set above) is the
# AUTHORITATIVE positional truth — derived from Anthropic's
# `cache_read_input_tokens` response. `compute_frozen_count`
# provides a defensive lower bound from local cache state.
# Use the smaller; never extend past what Anthropic actually
# has cached.
# Freeze + stable marking + Zone-1 swap now live in the
# shared session engine (PROXY policy: clamp by BOTH the
# provider-confirmed count and the locally-replayable
# bound — see session_engine.py's module docstring).
# `frozen_message_count` here has already been through
# tracker + strict-override logic above, so it is the
# AUTHORITATIVE positional truth derived from Anthropic's
# `cache_read_input_tokens` response.
#
# Issue #327: a previous version walked past
# `prefix_tracker.frozen_message_count` whenever an upcoming
# tool_result's content-hash matched `_stable_hashes` or
# `should_defer_compression` returned True. That conflated
# content equality with positional cache membership: the
# prefix cache is positional (bytes 0..K cached, anything
# past K is fresh), but `_stable_hashes` is content-keyed
# and grows unbounded. On long Claude Code sessions where
# tool_result content rhymes across turns (repeated system
# prompts, repeated file reads, etc.), the walker advanced
# Issue #327 (history kept at the call site): a previous
# version walked past `prefix_tracker.frozen_message_count`
# whenever an upcoming tool_result's content-hash matched
# `_stable_hashes` or `should_defer_compression` returned
# True. That conflated content equality with positional
# cache membership: the prefix cache is positional (bytes
# 0..K cached, anything past K is fresh), but
# `_stable_hashes` is content-keyed and grows unbounded.
# On long Claude Code sessions where tool_result content
# rhymes across turns, the walker advanced
# `frozen_message_count` to `len(messages)` and the
# pipeline produced `transforms_applied=[]` on 73% of
# requests. The walker has been removed; trust
# `prefix_tracker` clamped by `compute_frozen_count`.
cache_frozen_count = comp_cache.compute_frozen_count(messages)
frozen_message_count = min(frozen_message_count, cache_frozen_count)
# Record all tool_results in the verified frozen prefix as stable
comp_cache.mark_stable_from_messages(messages, frozen_message_count)
from headroom.proxy.session_engine import (
FREEZE_POLICY_CONFIRMED_CLAMP,
prepare_turn,
)
# Zone 1: Swap cached compressed versions into working copy
working_messages = comp_cache.apply_cached(messages)
_prep = prepare_turn(
comp_cache,
messages,
policy=FREEZE_POLICY_CONFIRMED_CLAMP,
tracker_frozen=frozen_message_count,
)
frozen_message_count = _prep.frozen_message_count
working_messages = _prep.pipeline_input
if (
getattr(self, "_background_compression_enabled", False)
and frozen_message_count == 0
@ -2065,39 +2071,43 @@ class AnthropicHandlerMixin:
# previously-forwarded prefix keeps it byte-identical → cache hits.
# Append-only-guarded and idempotent (cache mode already replays), so
# it is safe to run unconditionally here.
from headroom.cache.prefix_tracker import (
normalize_message_cache_control,
overlay_cached_prefix,
)
from headroom.cache.prefix_tracker import normalize_message_cache_control
from headroom.proxy.session_engine import finalize_turn
_overlay_replayed = False
# 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:
_ov = overlay_cached_prefix(
_final = finalize_turn(
optimized_messages,
original_client_messages,
previous_original_messages,
previous_forwarded_messages,
count_tokens=tokenizer.count_messages,
)
_overlay_replayed = _ov != optimized_messages
_overlay_replayed = _final.replayed
if _overlay_replayed:
optimized_messages = _ov
optimized_tokens = tokenizer.count_messages(optimized_messages)
optimized_messages = _final.messages
if _final.tokens is not None:
optimized_tokens = _final.tokens
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

575
headroom/proxy/handlers/openai.py
View file

@ -1593,6 +1593,14 @@ WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0
# "lossless" would otherwise look like it worked).
COMPRESS_MODES = ("ccr", "lossy_inline", "lossless_then_lossy")
# Max wait for a sidecar session's turn lock, on the executor. MUST stay
# well below COMPRESSION_TIMEOUT_SECONDS: with an untimed acquire, a slow
# turn's 503-driven retries would park executor workers blocked on the lock
# doing no work, each recording timeout debt toward the compression
# quarantine. Failing the acquire raises TimeoutError, which maps to the
# session-mode 503 retry path.
_SESSION_TURN_LOCK_TIMEOUT_SECONDS = 10.0
def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]:
"""Return the ``x-codex-*`` headers from an upstream WS handshake response.
@ -3753,17 +3761,45 @@ class OpenAIHandlerMixin:
if is_token_mode(self.config.mode):
comp_cache = self._get_compression_cache(openai_session_id)
# Zone 1: Swap cached compressed versions
working_messages = comp_cache.apply_cached(messages)
# Re-freeze boundary. Token mode can use the compression
# cache's positional frozen count. Cache mode must keep the
# latest observation mutable even when the compression
# cache has no compressible entry for it yet; otherwise
# OpenAI-compatible tool-call clients freeze the entire
# conversation and report near-zero savings.
if not is_cache_mode(self.config.mode):
openai_frozen_count = comp_cache.compute_frozen_count(messages)
# Token mode: shared engine, REPLAYABLE policy — its
# formula with no explicit pin is exactly this path's
# historical freeze (compute_frozen_count alone; the
# tracker count feeds cache mode below, never token
# mode). The engine also runs
# mark_stable_from_messages, which this path skipped:
# that marks tool_results INSIDE the frozen prefix as
# stable — redundant in the common case (an in-prefix
# tool_result is already stable via its cache entry)
# but it keeps `_stable_hashes` bookkeeping identical
# across all three paths, e.g. preserving stability
# across cache-entry LRU turnover. Note it can never
# mark the BOUNDARY tool_result that stopped the
# count (it sits outside messages[:frozen]) — the
# protection against re-compressing a passthrough
# boundary tool_result under rising context pressure
# is the router-level `_frozen_verdicts` pin, on every
# path, unchanged by this migration.
from headroom.proxy.session_engine import (
FREEZE_POLICY_REPLAYABLE,
prepare_turn,
)
_prep = prepare_turn(
comp_cache,
messages,
policy=FREEZE_POLICY_REPLAYABLE,
)
working_messages = _prep.pipeline_input
openai_frozen_count = _prep.frozen_message_count
else:
# Cache mode: Zone-1 swap only. The latest observation
# must stay mutable even when the compression cache
# has no entry for it yet (otherwise OpenAI-compatible
# tool-call clients freeze the entire conversation and
# report near-zero savings), so the freeze comes from
# the tracker (set above), never from the cache count.
working_messages = comp_cache.apply_cached(messages)
result = await self._run_compression_in_executor(
lambda: self.openai_pipeline.apply(
@ -3854,21 +3890,30 @@ class OpenAIHandlerMixin:
# Cache-safety (ALL modes): forward the previously-cached (compressed)
# prefix byte-identical, so freezing can't bust the prompt cache. See the
# matching guard in the Anthropic handler for the full rationale. Append-
# only-guarded and idempotent (cache mode already replays).
from headroom.cache.prefix_tracker import overlay_cached_prefix
# only-guarded and idempotent (cache mode already replays). Shared
# implementation: session_engine.finalize_turn.
from headroom.proxy.session_engine import finalize_turn
_ov = overlay_cached_prefix(
_final = finalize_turn(
optimized_messages,
original_client_messages,
openai_prefix_tracker.get_last_original_messages(),
openai_prefix_tracker.get_last_forwarded_messages(),
count_tokens=tokenizer.count_messages,
)
if _ov != optimized_messages:
optimized_messages = _ov
optimized_tokens = tokenizer.count_messages(optimized_messages)
if _final.replayed:
optimized_messages = _final.messages
if _final.tokens is not None:
optimized_tokens = _final.tokens
# Guard: if "optimization" inflated tokens, revert to originals
if optimized_tokens > original_tokens:
# Guard: if "optimization" inflated tokens, revert to originals.
# NEVER after the overlay replayed (same exemption as the Anthropic
# handler): the replayed prefix is the exact bytes the provider
# cached, and reverting to raw originals re-forwards the uncompressed
# prefix — trading a 90% read discount for a full cache re-write. The
# nominal "inflation" there is an artifact of comparing the cached
# (compressed) forwarding against the raw original count.
if optimized_tokens > original_tokens and not _final.replayed:
logger.warning(
f"[{request_id}] Optimization inflated tokens "
f"({original_tokens} -> {optimized_tokens}), reverting to original messages"
@ -9591,6 +9636,9 @@ class OpenAIHandlerMixin:
headers = dict(request.headers)
tags = extract_tags(headers)
client = classify_client(headers)
# Initialized before the try so the TimeoutError handler can branch on
# it even if the failure happened before session parsing.
session_id = None
try:
# Use OpenAI pipeline (messages are in OpenAI format from TS SDK)
@ -9655,6 +9703,64 @@ class OpenAIHandlerMixin:
}
},
)
# Session-aware sidecar mode (opt-in): with a session id the
# endpoint keeps the byte-replay state ITSELF — the same
# per-session machinery the proxy path uses (compression cache +
# prefix tracker, with the registry's TTL/LRU lifecycle) — so a
# gateway that owns routing (e.g. Kong) can resend the RAW
# conversation every turn and still get a byte-identical prefix
# back. Contract: the caller forwards the returned messages
# verbatim, and may relay provider usage via POST /v1/usage for
# telemetry/attribution. Without a session id, behaviour is the
# stateless contract, unchanged.
session_id = compress_config.get("session_id")
# The x-headroom-session-id header is honored only behind an
# explicit env opt-in: deployments whose gateways already stamp
# that header on ALL traffic (it is the documented proxy-path
# session key) would otherwise silently flip stateless callers
# into session mode on upgrade — and a header value shared across
# conversations (Claude Code subagents do exactly this) would
# blend unrelated conversations into one replay state.
if session_id is None and os.environ.get(
"HEADROOM_COMPRESS_SESSION_FROM_HEADER", ""
).lower() in ("1", "true"):
session_id = request.headers.get("x-headroom-session-id")
if session_id is not None and (
not isinstance(session_id, str) or not session_id.strip() or len(session_id) > 256
):
return JSONResponse(
status_code=400,
content={
"error": {
"type": "invalid_request",
"message": (
f"Invalid config.session_id: {session_id!r}. "
"Expected a non-empty string of at most 256 characters."
),
}
},
)
if session_id is not None and compress_user_messages:
# User/assistant rewrites are not content-addressed (the
# session cache replays tool_result content only), so once the
# tracker's overlay snapshots expire a rewritten user message
# would come back in RAW form — a guaranteed prefix bust inside
# the tracker-TTL/cache-TTL window. Refuse the combination
# rather than bust later.
return JSONResponse(
status_code=400,
content={
"error": {
"type": "invalid_request",
"message": (
"config.compress_user_messages is not supported with "
"config.session_id: user-message rewrites cannot be "
"byte-replayed across turns, which would bust the "
"provider prompt cache."
),
}
},
)
# Mode selection. Default is marker-free (see _no_ccr_pipeline):
# no caller of this route can resolve a CCR marker unless it opts in
# with mode="ccr", which restores the full marker + store behaviour.
@ -9701,23 +9807,160 @@ class OpenAIHandlerMixin:
if frozen_message_count is not None:
pipeline_kwargs["frozen_message_count"] = frozen_message_count
# Offload the CPU-bound pipeline to the bounded compression executor
# (mirrors the request handlers above). Running apply() inline blocked
# the single event loop on a large payload, so even GET /health stalled
# until it finished (#718). The executor also enforces a timeout so a
# too-large body fails fast instead of hanging forever.
result = await self._run_compression_in_executor(
lambda: pipeline.apply(
messages=messages,
model=model,
**pipeline_kwargs,
),
# Sidecar session pre-work: swap in previously-computed compressed
# bytes (Zone 1), then freeze the ENTIRE locally-replayable prefix
# (`compute_frozen_count`). This deliberately differs from the
# proxy path's `min(tracker, cache)` posture: in sidecar mode,
# whatever this endpoint previously RETURNED is the provider's
# cache contract, so every already-returned message must come back
# byte-identical — recompressing it (even "better") is a bust.
# Over-freezing relative to the provider's actual cache only
# forgoes tail compression; it can never bust. The tracker's
# /v1/usage-fed freeze count is deliberately NOT a freeze floor —
# freezing a message whose cache entry was evicted would forward
# raw original bytes. An explicit config.frozen_message_count
# still wins when larger: the caller may know more about the
# provider cache than local state does.
comp_cache = None
session_tracker = None
if session_id:
# Namespaced with a NUL separator so sidecar sessions can
# never collide with proxy-path session ids: NUL cannot
# appear in an HTTP header value, so no client-supplied
# x-headroom-session-id on the proxy path can spoof its way
# into a sidecar session's tracker or replay cache (the same
# trick SessionTrackerStore uses for its synthetic lineage
# keys). A plain "compress:" string prefix was spoofable.
_session_key = f"compress\x00{session_id}"
_tracker_provider = (
"anthropic"
if ("claude" in model_name.lower() or "anthropic" in model_name.lower())
else "openai"
)
comp_cache = self._get_compression_cache(_session_key)
session_tracker = self.session_tracker_store.get_or_create(
_session_key, _tracker_provider
)
def _run_stateless():
result = pipeline.apply(messages=messages, model=model, **pipeline_kwargs)
return (
result,
result.messages,
result.tokens_before,
result.tokens_after,
None,
)
def _run_session_turn():
# One sidecar turn as a single executor-side block: every step
# here is CPU-bound (content hashing, deep compares, token
# counts, full-transcript deepcopies) and must stay off the
# event loop for the same reason pipeline.apply does (#718).
# The per-session lock serializes contract-violating
# concurrent turns so an older in-flight turn cannot tear or
# overwrite a newer turn's tracker snapshots mid-flight.
# Cache management (freeze + swap + overlay) lives in the
# shared session engine — one brain for this path and the
# proxy request paths.
from headroom.proxy.session_engine import (
FREEZE_POLICY_REPLAYABLE,
finalize_turn,
prepare_turn,
)
# TIMED acquire, strictly shorter than the executor timeout:
# an untimed `with lock:` here lets one slow session's
# 503-driven retries park executor workers doing no work —
# each blocked worker records timeout debt and can arm the
# compression quarantine for ALL traffic. Failing fast maps
# to the same TimeoutError → session-mode 503 → retry path.
if not comp_cache.session_turn_lock.acquire(
timeout=_SESSION_TURN_LOCK_TIMEOUT_SECONDS
):
raise TimeoutError(
f"session turn lock busy for {session_id!r} "
"(a previous turn for this session is still running)"
)
try:
prev_original = session_tracker.get_last_original_messages()
prev_returned = session_tracker.get_last_forwarded_messages()
prep = prepare_turn(
comp_cache,
messages,
policy=FREEZE_POLICY_REPLAYABLE,
explicit_frozen=frozen_message_count,
)
session_frozen = prep.frozen_message_count
pipeline_kwargs["frozen_message_count"] = session_frozen
result = pipeline.apply(
messages=prep.pipeline_input, model=model, **pipeline_kwargs
)
# Replay last turn's exact returned prefix over any drift
# the pipeline introduced — byte-identical is the contract
# the caller forwards on.
turn = finalize_turn(result.messages, messages, prev_original, prev_returned)
final = turn.messages
# Savings are reported against the caller's RAW payload,
# not the cache-swapped pipeline input: on a warm turn the
# swap has already shrunk the input before the pipeline
# counts it, which made every warm turn report ~0 saved.
try:
from headroom.tokenizers import get_tokenizer
_tok = get_tokenizer(model_name)
raw_tokens_before = _tok.count_messages(messages)
final_tokens_after = _tok.count_messages(final)
except Exception as e:
# Fail-open, but LOUD: this fallback reverts to the
# pipeline's counts of the cache-swapped input, which
# silently resurrects the ~0-saved warm-turn bug the
# raw recount exists to fix — per-model, so it can
# hide indefinitely without this log.
logger.warning(
"[compress:%s] raw-payload token recount failed for "
"model %s (%s: %s); savings for this turn are "
"reported against the cache-swapped input",
session_id,
model_name,
type(e).__name__,
e,
)
raw_tokens_before = result.tokens_before
final_tokens_after = result.tokens_after
comp_cache.update_from_result(messages, final)
# Record this turn's result as the new "last returned" —
# the sidecar equivalent of "last forwarded", captured at
# return time because whatever we hand back IS what the
# caller sends upstream.
session_tracker.record_returned(messages, final)
info = {
"id": session_id,
"frozen_message_count": session_frozen,
"cached_prefix_replayed": turn.replayed,
}
return result, final, raw_tokens_before, final_tokens_after, info
finally:
comp_cache.session_turn_lock.release()
# Offload the CPU-bound work to the bounded compression executor
# (mirrors the request handlers above). Running it inline blocked
# the single event loop on a large payload, so even GET /health
# stalled until it finished (#718). The executor also enforces a
# timeout so a too-large body fails fast instead of hanging.
(
result,
final_messages,
tokens_before,
tokens_after,
session_info,
) = await self._run_compression_in_executor(
_run_session_turn if session_id else _run_stateless,
timeout=COMPRESSION_TIMEOUT_SECONDS,
)
ccr_hashes = _response_ccr_hashes(result.messages, result.markers_inserted)
tokens_before = result.tokens_before
tokens_after = result.tokens_after
ccr_hashes = _response_ccr_hashes(final_messages, result.markers_inserted)
tokens_saved = max(0, tokens_before - tokens_after)
latency_ms = (time.time() - start_time) * 1000
await self._record_request_outcome(
@ -9749,28 +9992,82 @@ class OpenAIHandlerMixin:
)
)
return JSONResponse(
{
"messages": result.messages,
"tokens_before": result.tokens_before,
"tokens_after": result.tokens_after,
"tokens_saved": result.tokens_before - result.tokens_after,
"compression_ratio": (
result.tokens_after / result.tokens_before
if result.tokens_before > 0
else 1.0
),
"transforms_applied": result.transforms_applied,
"transforms_summary": result.transforms_summary,
"ccr_hashes": ccr_hashes,
}
)
_payload = {
"messages": final_messages,
"tokens_before": tokens_before,
"tokens_after": tokens_after,
# Clamped like the telemetry above: the overlay's byte-replay
# can legitimately return a slightly larger prefix than the
# pipeline's best effort, and a negative "saved" here while
# telemetry records 0 would be two answers for one number.
"tokens_saved": tokens_saved,
"compression_ratio": (tokens_after / tokens_before if tokens_before > 0 else 1.0),
"transforms_applied": result.transforms_applied,
"transforms_summary": result.transforms_summary,
"ccr_hashes": ccr_hashes,
}
if session_info is not None:
_payload["session"] = session_info
return JSONResponse(_payload)
except TimeoutError:
self.metrics.record_compression_failed("timeout")
if session_id:
# Fail-open-with-originals is WRONG for a session call: the
# timed-out worker cannot be cancelled and may still finish
# and record its compressed result as "last returned" — while
# the caller, handed the originals, forwards those instead.
# The desynced snapshot then busts the next turn. A 503 tells
# the gateway to retry; the retry lands on whatever state the
# straggler recorded and replays it consistently.
logger.warning(
"Compression timed out after %.0fs for session %r; "
"returning 503 (session mode cannot fail open without "
"desyncing replay state)",
COMPRESSION_TIMEOUT_SECONDS,
session_id,
)
# Same outcome recording as the stateless timeout path below:
# session timeouts hit the largest transcripts, and skipping
# the RequestOutcome here under-counts exactly those requests
# when dashboards reconcile failure counters against outcomes.
_timeout_latency_ms = (time.time() - start_time) * 1000
await self._record_request_outcome(
RequestOutcome(
request_id=(
await self._next_request_id()
if hasattr(self, "_next_request_id")
else f"compress_{int(time.time())}"
),
provider="compress",
model=model if isinstance(model, str) else str(model),
original_tokens=0,
optimized_tokens=0,
output_tokens=0,
tokens_saved=0,
attempted_input_tokens=0,
total_latency_ms=_timeout_latency_ms,
overhead_ms=_timeout_latency_ms,
num_messages=len(messages) if isinstance(messages, list) else 0,
tags=tags,
client=client,
)
)
return JSONResponse(
status_code=503,
content={
"error": {
"type": "compression_timeout",
"message": (
"Compression timed out; retry this turn. "
"Session replay state remains consistent."
),
}
},
)
logger.warning(
"Compression timed out after %.0fs; failing open with original messages",
COMPRESSION_TIMEOUT_SECONDS,
)
self.metrics.record_compression_failed("timeout")
latency_ms = (time.time() - start_time) * 1000
await self._record_request_outcome(
RequestOutcome(
@ -9820,6 +10117,186 @@ class OpenAIHandlerMixin:
},
)
async def handle_compress_usage(self, request: Request) -> JSONResponse:
"""Relay of the provider's usage block for a sidecar compress session.
POST /v1/usage
Body: {"session_id": "...",
"usage": {"cache_read_input_tokens": N,
"cache_creation_input_tokens": N}}
The session-aware ``/v1/compress`` never sees the provider's response
(the caller owns routing). This relay feeds the provider-confirmed
numbers into the session's tracker — the same signal the proxy path
reads from the response itself — powering cache-hit/miss attribution,
idle-vs-prefix-change classification, and savings accounting for
sidecar sessions.
Deliberately NOT a freeze input: the compress path freezes exactly the
locally-replayable prefix (``compute_frozen_count``), and raising that
to a provider-confirmed count could freeze a message whose cache entry
was evicted — which would forward raw original bytes and bust the very
prefix the count vouched for. Optional: skipping this call costs
telemetry fidelity, never correctness.
"""
from fastapi.responses import JSONResponse
from headroom.proxy.helpers import _read_request_json
def _invalid(message: str) -> JSONResponse:
return JSONResponse(
status_code=400,
content={"error": {"type": "invalid_request", "message": message}},
)
try:
body = await _read_request_json(request)
except Exception:
return _invalid("Invalid JSON in request body.")
session_id = body.get("session_id")
if not isinstance(session_id, str) or not session_id.strip() or len(session_id) > 256:
return _invalid(
"Missing or invalid session_id: expected a non-empty string "
"of at most 256 characters."
)
usage = body.get("usage")
if not isinstance(usage, dict):
return _invalid("Missing or invalid usage: expected an object.")
# A usage block carrying NEITHER cache field is a no-signal relay (an
# OpenAI-style {"prompt_tokens": N} forwarded verbatim, for example).
# Defaulting the absent fields to 0 would make update_from_response
# treat it as a provider-confirmed fully-cold turn and wipe the
# tracker's cached-prefix state — so absence of both is a 400, not 0.
if "cache_read_input_tokens" not in usage and "cache_creation_input_tokens" not in usage:
return _invalid(
"usage must carry cache_read_input_tokens and/or "
"cache_creation_input_tokens; a block with neither carries no "
"cache signal and is not accepted."
)
def _token_field(name: str) -> int | None:
value = usage.get(name, 0)
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
return None
return value
cache_read = _token_field("cache_read_input_tokens")
cache_write = _token_field("cache_creation_input_tokens")
if cache_read is None or cache_write is None:
return _invalid(
"usage.cache_read_input_tokens and usage.cache_creation_input_tokens "
"must be non-negative integers when present."
)
# Same NUL-separated namespace as handle_compress: unspoofable from
# any HTTP header. peek() (never get_or_create) so a flood of novel
# session ids cannot grow the tracker store — an unknown or expired
# session is answered without leaving a footprint, and the session
# keeps the provider its compress call inferred rather than a default
# from here.
_session_key = f"compress\x00{session_id}"
tracker = self.session_tracker_store.peek(_session_key)
if tracker is None:
return self._compress_usage_unknown_session(session_id)
# No create, no LRU bump: the cache is only needed for its turn lock.
comp_cache = self._peek_compression_cache(_session_key)
# A relay whose only present field is zero carries no positive cache
# signal (an OpenAI-mapped gateway naturally sends
# {"cache_read_input_tokens": 0} with no write field — OpenAI has no
# write signal). Applying it would hit update_from_response's
# total_cached == 0 branch and wipe the tracker's cached-prefix
# state — a "provider-confirmed fully cold" reset the relay never
# actually asserted. Only a relay with BOTH fields present may claim
# a genuine fully-cold turn.
_both_present = (
"cache_read_input_tokens" in usage and "cache_creation_input_tokens" in usage
)
if cache_read + cache_write == 0 and not _both_present:
return JSONResponse(
{
"session_id": session_id,
"frozen_message_count": tracker.get_frozen_message_count(),
"applied": False,
"reason": "no_cache_signal",
}
)
def _apply_usage():
# Off the event loop (full-transcript deepcopies + per-message
# token estimation live in update_from_response), and under the
# session turn lock: an unlocked update here races the
# executor-side compress turn — record_returned installs turn
# N+1's snapshots, then this write would roll them back to turn
# N's copies and the next overlay would refuse to replay.
lock = comp_cache.session_turn_lock if comp_cache is not None else None
if lock is not None and not lock.acquire(timeout=_SESSION_TURN_LOCK_TIMEOUT_SECONDS):
raise TimeoutError(f"session turn lock busy for {session_id!r}")
try:
last_returned = tracker.get_last_forwarded_messages()
if not last_returned:
return None
tracker.update_from_response(
cache_read_tokens=cache_read,
cache_write_tokens=cache_write,
messages=last_returned,
original_messages=tracker.get_last_original_messages(),
)
return tracker.get_frozen_message_count()
finally:
if lock is not None:
lock.release()
try:
frozen_count = await self._run_compression_in_executor(
_apply_usage, timeout=COMPRESSION_TIMEOUT_SECONDS
)
except TimeoutError:
return JSONResponse(
status_code=503,
content={
"error": {
"type": "session_busy",
"message": (
"A compress turn for this session is in flight; retry the usage relay."
),
}
},
)
if frozen_count is None:
return self._compress_usage_unknown_session(session_id)
return JSONResponse(
{
"session_id": session_id,
"frozen_message_count": frozen_count,
"applied": True,
}
)
@staticmethod
def _compress_usage_unknown_session(session_id: str):
from fastapi.responses import JSONResponse
# No compress state for this session: never seen, or the tracker's
# session TTL reclaimed it. Note the byte-replay cache lives longer
# than the tracker, so a 404 here does NOT mean the next /v1/compress
# loses replay — only this telemetry relay landed nowhere.
return JSONResponse(
status_code=404,
content={
"error": {
"type": "unknown_session",
"message": (
f"No usage-tracking state for session {session_id!r} "
"(never seen, or expired). Compression replay for the "
"session may still be active; only this telemetry "
"relay landed nowhere."
),
}
},
)
async def _maybe_compress_passthrough_responses(
self, body: bytes, *, client: str | None = None
) -> bytes:

50
headroom/proxy/helpers.py
View file

@ -1245,8 +1245,54 @@ 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).
# Non-finite floats ("nan"/"inf") parse but poison every idle comparison,
# so they are rejected like any other unparseable value.
_ttl_env = float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900"))
if _ttl_env != _ttl_env or _ttl_env in (float("inf"), float("-inf")):
raise ValueError("non-finite TTL")
COMPRESSION_CACHE_TTL_SECONDS = max(600.0, _ttl_env)
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
# ---------------------------------------------------------------------------

136
headroom/proxy/server.py
View file

@ -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,67 @@ 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
# Skip sessions with a turn in flight (session_turn_lock held): popping
# one would hand its retry a FRESH cache with a NEW lock — straggler
# and retry then run unserialized against the same tracker, and the
# retry's empty cache recompresses previously-returned content into
# different bytes. An in-flight session is by definition not idle; it
# will be swept on a later pass once genuinely quiet.
expired = [
sid
for sid, seen in self._compression_cache_last_seen.items()
if now - seen > COMPRESSION_CACHE_TTL_SECONDS
and (cache := self._compression_caches.get(sid)) is not None
and not cache.session_turn_lock.locked()
]
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 _peek_compression_cache(self, session_id: str) -> CompressionCache | None:
"""Return the session's cache if one exists — no create, no LRU bump.
For lookup paths that must not leave a footprint or distort access
recency (e.g. /v1/usage taking the session turn lock): an unknown
session answers None instead of allocating an empty cache.
"""
with self._compression_caches_lock:
return self._compression_caches.get(session_id)
def _get_compression_cache(self, session_id: str) -> CompressionCache:
"""Get or create a CompressionCache for a session.
@ -1588,27 +1659,55 @@ 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.
# Sessions with a turn in flight (session_turn_lock held) are
# skipped: popping one splits its lock across two cache
# instances and desyncs the straggler from its retry (see the
# TTL sweep's comment). If every candidate is mid-turn, no
# eviction happens this round — briefly exceeding the cap is
# cheaper than a guaranteed prefix bust.
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]
logger.info(
"Evicted %d compression caches (exceeded %d max sessions)",
len(oldest_keys),
MAX_COMPRESSION_CACHE_SESSIONS,
evict_count = min(
max(1, MAX_COMPRESSION_CACHE_SESSIONS // 4),
len(self._compression_caches),
)
evictable = [
sid
for sid, c in self._compression_caches.items()
if not c.session_turn_lock.locked()
][:evict_count]
for sid in evictable:
del self._compression_caches[sid]
self._compression_cache_last_seen.pop(sid, None)
if evictable:
logger.info(
"Evicted %d least-recently-used compression caches "
"(exceeded %d max sessions)",
len(evictable),
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.
@ -5183,6 +5282,13 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
async def compress_messages(request: Request):
return await proxy.handle_compress(request)
# Sidecar-mode usage relay: same exposure policy as /v1/compress — the two
# form one contract (compress returns the bytes, usage reports what the
# provider said about them), so they must be reachable from the same place.
@app.post("/v1/usage", dependencies=_compress_dependencies)
async def compress_usage(request: Request):
return await proxy.handle_compress_usage(request)
register_provider_routes(app, proxy)
return app

185
headroom/proxy/session_engine.py Normal file
View file

@ -0,0 +1,185 @@
"""Session-turn engine — the single cache-management brain for both modes.
One conversation turn, from the cache's point of view, is always the same
three-step dance regardless of who owns the upstream call:
1. **Prepare** (:func:`prepare_turn`): decide how many leading messages are
frozen, mark the stable prefix, and swap previously-computed compressed
bytes into the working copy (``apply_cached`` — "Zone 1").
2. Run the compression pipeline over the prepared input (owned by the
caller: the proxy handlers wrap it in background/cold-start/backpressure
orchestration, the sidecar path runs it inline on the executor).
3. **Finalize** (:func:`finalize_turn`): replay last turn's exact
previously-forwarded/returned prefix over any residual drift the pipeline
introduced (``overlay_cached_prefix``), so the bytes that leave the
process are byte-identical to what the provider already cached.
Historically the proxy request handlers (anthropic + openai token mode) and
the sidecar ``/v1/compress`` session path each carried their own inline copy
of steps 1 and 3. This module is the shared implementation: a
cache-management fix landed here reaches BOTH modes at once.
Freeze policies
---------------
The one deliberate behavioural difference between the modes lives in step 1,
and it is a *policy parameter*, not a fork of the code:
``FREEZE_POLICY_CONFIRMED_CLAMP`` — ``min(tracker_frozen, cache_count)``.
The proxy sees the provider's responses, so ``tracker_frozen`` is the
provider-confirmed cached prefix (from ``cache_read_input_tokens``).
Freezing is clamped by BOTH bounds: never past what the provider
actually has cached (freezing more would forgo compression of content
that is not yet cache-protected — the #327 posture), and never past what
the local cache can byte-replay (freezing a message whose entry was
evicted would pass through raw original bytes).
``FREEZE_POLICY_REPLAYABLE`` — ``max(cache_count, explicit_frozen or 0)``.
Freeze everything the local cache can byte-replay. Used by callers with
no provider-confirmed count to clamp against: the sidecar ``/v1/compress``
endpoint (it never sees the provider's response — whatever it previously
RETURNED is the provider's cache contract, so every already-returned
message must come back byte-identical), and the OpenAI proxy token path
(its tracker feeds cache mode, not token mode). Recompressing an
already-returned message — even into a *smaller* form — is a bust: the
drift was observed in practice, and ``overlay_cached_prefix``'s
non-inflation guard cannot repair a shrunken form (replaying the larger
original bytes would "inflate" the candidate). Freezing the entire
locally-replayable prefix eliminates that recompression outright.
Over-freezing relative to the provider's real cache only forgoes tail
compression; it can never bust. An explicit ``frozen_message_count``
from the caller still wins when larger — the caller may know more about
the provider cache than local state does.
Why the Anthropic proxy path cannot simply adopt the replayable posture: its
provider-confirmed clamp deliberately KEEPS not-yet-cached content
compressible, and its overlay inputs (tracker snapshots) are refreshed on
every response, so drift repair is reliable there. Each posture is correct
for the information its mode actually has.
"""
from __future__ import annotations
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
from headroom.cache.prefix_tracker import overlay_cached_prefix
logger = logging.getLogger(__name__)
FREEZE_POLICY_CONFIRMED_CLAMP = "confirmed_clamp"
FREEZE_POLICY_REPLAYABLE = "replayable"
_FREEZE_POLICIES = (FREEZE_POLICY_CONFIRMED_CLAMP, FREEZE_POLICY_REPLAYABLE)
@dataclass(frozen=True)
class TurnPrep:
"""Result of :func:`prepare_turn`.
``frozen_message_count`` is what the pipeline must be told to skip;
``pipeline_input`` is the working copy with previously-compressed bytes
swapped in (never the caller's list — ``apply_cached`` copies).
"""
frozen_message_count: int
pipeline_input: list[dict[str, Any]]
@dataclass(frozen=True)
class TurnFinal:
"""Result of :func:`finalize_turn`.
``messages`` are the bytes to forward/return; ``replayed`` says whether
the overlay restored last turn's prefix over pipeline drift; ``tokens``
is the recount of ``messages`` when a ``count_tokens`` hook was supplied
and the overlay actually fired (None otherwise — the pipeline's own
count is still valid when nothing was replaced).
"""
messages: list[dict[str, Any]]
replayed: bool
tokens: int | None = None
def prepare_turn(
comp_cache: Any,
messages: list[dict[str, Any]],
*,
policy: str,
tracker_frozen: int | None = None,
explicit_frozen: int | None = None,
) -> TurnPrep:
"""Freeze decision + stable marking + cached-byte swap for one turn.
Args:
comp_cache: the session's ``CompressionCache``.
messages: the caller's RAW message list (never mutated).
policy: ``FREEZE_POLICY_CONFIRMED_CLAMP`` or ``FREEZE_POLICY_REPLAYABLE`` —
see the module docstring for why they differ.
tracker_frozen: provider-confirmed frozen count (proxy policy only;
``None`` means "nothing confirmed" and freezes 0 there).
explicit_frozen: caller-pinned frozen count (sidecar policy only;
wins when larger than the locally-derived bound).
"""
if policy not in _FREEZE_POLICIES:
raise ValueError(f"unknown freeze policy: {policy!r}")
cache_count = comp_cache.compute_frozen_count(messages)
if policy == FREEZE_POLICY_CONFIRMED_CLAMP:
# Never freeze past the provider-confirmed prefix, and never past
# what local state can byte-replay.
frozen = min(tracker_frozen or 0, cache_count)
else:
# Freeze the entire locally-replayable prefix; an explicit caller
# pin may extend it (the caller vouches the provider cached those
# exact raw bytes, so passing them through untouched is correct).
frozen = max(cache_count, explicit_frozen or 0)
comp_cache.mark_stable_from_messages(messages, frozen)
return TurnPrep(
frozen_message_count=frozen,
pipeline_input=comp_cache.apply_cached(messages),
)
def finalize_turn(
result_messages: list[dict[str, Any]],
original_messages: list[dict[str, Any]],
prev_original: list[dict[str, Any]] | None,
prev_returned: list[dict[str, Any]] | None,
*,
count_tokens: Callable[[list[dict[str, Any]]], int] | None = None,
) -> TurnFinal:
"""Replay last turn's exact forwarded/returned prefix over pipeline drift.
``overlay_cached_prefix`` self-guards (positional alignment, append-only
shape, non-inflation), so calling this is always safe: when replay is not
provably correct it returns the pipeline's own output unchanged.
``count_tokens`` is invoked only when the overlay actually replaced
bytes — the pipeline's own token count is still accurate otherwise. A
failing hook falls back to "no recount" rather than failing the turn.
"""
final = overlay_cached_prefix(result_messages, original_messages, prev_original, prev_returned)
replayed = final != result_messages
tokens: int | None = None
if replayed and count_tokens is not None:
try:
tokens = count_tokens(final)
except Exception as e:
# Fail-open: the turn still forwards, but the caller keeps the
# pipeline's count of messages that are NOT being forwarded —
# tokens_saved accounting is stale for this turn. Loud, not
# silent: a tokenizer that cannot count the replayed form is a
# bug worth surfacing even though it must not fail the request.
logger.warning(
"finalize_turn: token recount of replayed prefix failed "
"(%s: %s); keeping the pipeline's pre-overlay count",
type(e).__name__,
e,
)
tokens = None
return TurnFinal(messages=final, replayed=replayed, tokens=tokens)

14
headroom/transforms/content_router.py
View file

@ -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)

71
tests/test_anthropic_pre_upstream_backpressure.py
View file

@ -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)

568
tests/test_compress_session_mode.py Normal file
View file

@ -0,0 +1,568 @@
"""Session-aware /v1/compress (sidecar mode) + the /v1/usage relay.
Contract under test: a gateway that owns routing (e.g. Kong) sends the RAW
conversation plus a session id every turn; Headroom keeps the byte-replay
state itself and returns a byte-identical prefix; the gateway forwards the
result verbatim and may relay provider usage via POST /v1/usage to make
freeze decisions exact.
The critical property is byte-stability: content already returned for a
session must come back byte-for-byte identical on later turns, or the
provider prompt cache busts.
"""
from __future__ import annotations
import json
import pytest
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
def _make_client() -> TestClient:
config = ProxyConfig(
optimize=True,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
image_optimize=False,
)
app = create_app(config)
client = TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345))
return client
def _big_tool_history() -> list[dict]:
"""A conversation whose tool result is large enough to be compressed."""
items = [
{
"id": i,
"score": 0.99 if i % 30 == 0 else 0.6,
"msg": f"Result {i:03d}{' error' if i % 30 == 0 else ' ok'}",
"blob": f"payload-{i:04d}-" + "".join(chr(97 + (i * 7 + j) % 26) for j in range(240)),
}
for i in range(200)
]
return [
{"role": "user", "content": "Get items"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "c1", "type": "function", "function": {"name": "get", "arguments": "{}"}}
],
},
{"role": "tool", "tool_call_id": "c1", "content": json.dumps(items)},
]
def _compress(client: TestClient, messages: list[dict], **config) -> dict:
resp = client.post(
"/v1/compress",
json={"model": "gpt-4o", "messages": messages, "config": config},
)
assert resp.status_code == 200, resp.text
return resp.json()
# --------------------------------------------------------------------------- #
# Stateless behaviour is unchanged (regression guard). #
# --------------------------------------------------------------------------- #
# The NUL separator makes the namespace unspoofable from any HTTP header.
SESSION_KEY_PREFIX = "compress\x00"
def test_no_session_id_stays_stateless() -> None:
with _make_client() as client:
body = _compress(client, _big_tool_history())
assert "session" not in body
# And nothing session-shaped leaked into the registry.
proxy = client.app.state.proxy
assert not any(k.startswith(SESSION_KEY_PREFIX) for k in proxy._compression_caches)
def test_invalid_session_id_is_rejected() -> None:
with _make_client() as client:
for bad in ["", " ", "x" * 300, 42]:
resp = client.post(
"/v1/compress",
json={
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"config": {"session_id": bad},
},
)
assert resp.status_code == 400, f"session_id {bad!r} was not rejected"
def test_compress_user_messages_rejected_with_session() -> None:
"""User-message rewrites are not content-addressed, so they cannot be
byte-replayed after tracker state expires — the combination is a latent
prefix-cache bust and must be refused up front."""
with _make_client() as client:
resp = client.post(
"/v1/compress",
json={
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"config": {"session_id": "conv-x", "compress_user_messages": True},
},
)
assert resp.status_code == 400
assert "compress_user_messages" in resp.json()["error"]["message"]
def test_session_key_is_not_spoofable_via_string_prefix() -> None:
"""A caller passing 'compress:...' (or similar) as its session id must
land on a key that no proxy-path header value can also produce."""
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="compress:sneaky")
proxy = client.app.state.proxy
keys = [k for k in proxy._compression_caches if "sneaky" in k]
assert keys == [f"{SESSION_KEY_PREFIX}compress:sneaky"]
# NUL cannot appear in an HTTP header value, so no x-headroom-session-id
# on the proxy path can collide with this key.
assert all("\x00" in k for k in keys)
# --------------------------------------------------------------------------- #
# The core sidecar property: turn 2 replays turn 1's exact bytes. #
# --------------------------------------------------------------------------- #
def test_second_turn_replays_first_turn_bytes() -> None:
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-1")
assert turn1["session"]["id"] == "conv-1"
# The tool result must actually have been compressed, otherwise the
# byte-stability assertion below is vacuous.
t1_tool_content = turn1["messages"][2]["content"]
assert t1_tool_content != history[2]["content"]
assert turn1["tokens_saved"] > 0
# Turn 2: the caller resends the RAW history (as real clients do) plus
# the new turns. Headroom must return the OLD prefix byte-identical to
# what it handed back on turn 1 — that is what the provider cached.
turn2_history = history + [
{"role": "assistant", "content": "The top items are listed above."},
{"role": "user", "content": "Now sort them by score."},
]
turn2 = _compress(client, turn2_history, session_id="conv-1")
assert turn2["messages"][2]["content"] == t1_tool_content
# The WHOLE turn-1 prefix, not just the tool result: any drifted byte
# anywhere in the leading messages is a provider-cache bust.
assert turn2["messages"][: len(turn1["messages"])] == turn1["messages"]
assert turn2["messages"][-1]["content"] == "Now sort them by score."
assert turn2["session"]["id"] == "conv-1"
# Savings must be reported against the RAW payload the caller sent —
# the warm turn still saved the caller ~everything turn 1 saved, even
# though the pipeline itself only saw the already-swapped input.
assert turn2["tokens_saved"] > 0
assert turn2["tokens_before"] > turn2["tokens_after"]
def test_third_turn_still_byte_stable() -> None:
"""The WHOLE returned prefix — every message, byte for byte — must be
stable across N turns. Checking only the tool result would let drift in
any other message (a mutated plain message, a moved marker) bust the
provider cache while the test stayed green.
"""
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-multi")
history2 = history + [{"role": "user", "content": "next"}]
turn2 = _compress(client, history2, session_id="conv-multi")
# Turn 2's leading messages must be exactly turn 1's returned bytes.
assert turn2["messages"][: len(turn1["messages"])] == turn1["messages"]
history3 = history2 + [
{"role": "assistant", "content": "ok"},
{"role": "user", "content": "and again"},
]
turn3 = _compress(client, history3, session_id="conv-multi")
# And turn 3's leading messages must be exactly turn 2's.
assert turn3["messages"][: len(turn2["messages"])] == turn2["messages"]
def test_prefix_stable_even_after_tracker_state_loss() -> None:
"""The overlay's tracker snapshots live shorter (600s session TTL) than
the compression cache (3900s). In that window the frozen+swap path is the
ONLY protection — this test kills the tracker between turns and demands
whole-prefix byte stability from frozen+swap alone.
"""
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-trackerloss")
proxy = client.app.state.proxy
# Simulate the tracker registry's TTL sweep reclaiming the session
# while the compression cache (longer TTL) survives.
store = proxy.session_tracker_store
removed = [k for k in list(store._trackers) if "conv-trackerloss" in k]
for k in removed:
del store._trackers[k]
assert removed, "tracker was never created for the session"
assert any("conv-trackerloss" in k for k in proxy._compression_caches)
turn2 = _compress(
client,
history + [{"role": "user", "content": "after tracker loss"}],
session_id="conv-trackerloss",
)
assert turn2["messages"][: len(turn1["messages"])] == turn1["messages"]
def test_header_session_id_ignored_by_default() -> None:
"""Deployments whose gateways stamp x-headroom-session-id on ALL traffic
must not silently flip stateless /v1/compress callers into session mode
(or blend conversations sharing one header value into one replay state)."""
with _make_client() as client:
resp = client.post(
"/v1/compress",
json={"model": "gpt-4o", "messages": _big_tool_history(), "config": {}},
headers={"x-headroom-session-id": "conv-header"},
)
assert resp.status_code == 200, resp.text
assert "session" not in resp.json()
def test_header_session_id_works_with_env_opt_in(monkeypatch) -> None:
monkeypatch.setenv("HEADROOM_COMPRESS_SESSION_FROM_HEADER", "1")
with _make_client() as client:
resp = client.post(
"/v1/compress",
json={"model": "gpt-4o", "messages": _big_tool_history(), "config": {}},
headers={"x-headroom-session-id": "conv-header"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["session"]["id"] == "conv-header"
def test_sessions_are_isolated() -> None:
with _make_client() as client:
history = _big_tool_history()
a1 = _compress(client, history, session_id="conv-a")
b1 = _compress(client, history, session_id="conv-b")
# Same content in, same compressed form out — but through separate
# session state. Interleave new turns and re-check both replay.
a2 = _compress(
client,
history + [{"role": "user", "content": "a follow-up"}],
session_id="conv-a",
)
b2 = _compress(
client,
history + [{"role": "user", "content": "b follow-up"}],
session_id="conv-b",
)
assert a2["messages"][2]["content"] == a1["messages"][2]["content"]
assert b2["messages"][2]["content"] == b1["messages"][2]["content"]
assert a2["messages"][-1]["content"] == "a follow-up"
assert b2["messages"][-1]["content"] == "b follow-up"
# --------------------------------------------------------------------------- #
# /v1/usage: telemetry relay for sidecar sessions. Deliberately NOT a freeze #
# input — freeze stays the locally-replayable bound (see handler docstring). #
# --------------------------------------------------------------------------- #
def test_usage_relay_is_recorded_and_freeze_stays_local() -> None:
with _make_client() as client:
history = _big_tool_history()
_compress(client, history, session_id="conv-usage")
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-usage",
"usage": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 50_000,
},
},
)
assert resp.status_code == 200, resp.text
# The tracker recorded the provider-confirmed prefix (telemetry).
assert resp.json()["frozen_message_count"] >= 1
# The next compress freezes from the LOCAL replayable bound, which
# covers the whole previously-returned prefix here.
turn2 = _compress(
client,
history + [{"role": "user", "content": "next"}],
session_id="conv-usage",
)
assert turn2["session"]["frozen_message_count"] >= 1
# An absurdly large confirmed count must never drag freezing past
# what local state can actually replay (that would forward raw bytes
# for evicted entries — the bust this design refuses).
resp2 = client.post(
"/v1/usage",
json={
"session_id": "conv-usage",
"usage": {"cache_read_input_tokens": 10_000_000},
},
)
assert resp2.status_code == 200
turn3 = _compress(
client,
history
+ [
{"role": "user", "content": "next"},
{"role": "assistant", "content": "done"},
{"role": "user", "content": "more"},
],
session_id="conv-usage",
)
# Freeze is capped by message count minus the trailing message — it
# can never exceed what exists, regardless of relayed numbers.
assert turn3["session"]["frozen_message_count"] < 6
def test_usage_unknown_session_is_404_and_leaves_no_footprint() -> None:
with _make_client() as client:
proxy = client.app.state.proxy
before = len(proxy.session_tracker_store._trackers)
for i in range(20):
resp = client.post(
"/v1/usage",
json={
"session_id": f"never-seen-{i}",
"usage": {"cache_read_input_tokens": 100},
},
)
assert resp.status_code == 404
assert resp.json()["error"]["type"] == "unknown_session"
# A flood of novel ids must not grow the tracker store (peek, never
# get_or_create): each ghost tracker would otherwise live a full TTL.
assert len(proxy.session_tracker_store._trackers) == before
def test_usage_without_cache_fields_is_rejected_not_treated_as_cold() -> None:
"""A usage block with NEITHER cache field (e.g. an OpenAI-style
{'prompt_tokens': N} relayed verbatim) carries no cache signal. Treating
the absent fields as 0 would tell the tracker 'provider confirmed fully
cold' and wipe its cached-prefix state on every signal-free relay."""
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-nosignal")
resp = client.post(
"/v1/usage",
json={"session_id": "conv-nosignal", "usage": {"prompt_tokens": 12345}},
)
assert resp.status_code == 400
assert "cache" in resp.json()["error"]["message"]
def test_usage_validation() -> None:
with _make_client() as client:
cases = [
{}, # no session_id
{"session_id": "s"}, # no usage
{"session_id": "s", "usage": "nope"}, # usage not a dict
{"session_id": "s", "usage": {"cache_read_input_tokens": -1}},
{"session_id": "s", "usage": {"cache_read_input_tokens": True}},
]
for body in cases:
resp = client.post("/v1/usage", json=body)
assert resp.status_code == 400, f"body {body!r} was not rejected"
# --------------------------------------------------------------------------- #
# Lifecycle: sidecar sessions ride the registry's TTL/LRU machinery. #
# --------------------------------------------------------------------------- #
def test_session_state_lives_in_registry_and_survives_eviction() -> None:
import time as _time
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-ttl")
proxy = client.app.state.proxy
_key = f"{SESSION_KEY_PREFIX}conv-ttl"
assert _key in proxy._compression_caches
# Simulate the idle-TTL sweep reclaiming the session.
now = _time.time()
proxy._compression_cache_last_seen[_key] = now - 999_999
proxy._compression_caches_last_cleanup = now - 61
proxy._get_compression_cache("unrelated")
assert _key not in proxy._compression_caches
# A post-eviction turn is fail-open: fresh state, valid response, and
# the compressed form is reproducible (deterministic pipeline), even
# though the replay guarantee had to restart from scratch.
turn2 = _compress(
client,
history + [{"role": "user", "content": "after the gap"}],
session_id="conv-ttl",
)
assert turn2["session"]["id"] == "conv-ttl"
assert turn2["messages"][-1]["content"] == "after the gap"
assert isinstance(turn1["messages"][2]["content"], str)
def test_explicit_frozen_count_still_wins_when_larger() -> None:
with _make_client() as client:
history = _big_tool_history()
# First turn with an explicit pin covering the whole tool result: the
# caller asserts the provider already cached it, so it must come back
# byte-for-byte untouched even though no session state exists yet.
turn1 = _compress(client, history, session_id="conv-pin", frozen_message_count=3)
assert turn1["messages"][2]["content"] == history[2]["content"]
assert turn1["session"]["frozen_message_count"] == 3
# --------------------------------------------------------------------------- #
# Review fixes: turn-lock contention, no-signal usage, expired trackers. #
# --------------------------------------------------------------------------- #
def test_compress_503_when_turn_lock_busy(monkeypatch) -> None:
"""A concurrent turn for the same session must fail fast with a 503,
not park an executor worker on an untimed lock acquire."""
import headroom.proxy.handlers.openai as openai_mod
monkeypatch.setattr(openai_mod, "_SESSION_TURN_LOCK_TIMEOUT_SECONDS", 0.05)
with _make_client() as client:
history = _big_tool_history()
_compress(client, history, session_id="conv-lock")
proxy = client.app.state.proxy
lock = proxy._compression_caches[f"{SESSION_KEY_PREFIX}conv-lock"].session_turn_lock
assert lock.acquire(timeout=1), "test could not take the turn lock"
try:
resp = client.post(
"/v1/compress",
json={
"model": "gpt-4o",
"messages": history + [{"role": "user", "content": "blocked"}],
"config": {"session_id": "conv-lock"},
},
)
assert resp.status_code == 503, resp.text
finally:
lock.release()
# With the lock free again the same turn succeeds.
after = _compress(
client,
history + [{"role": "user", "content": "blocked"}],
session_id="conv-lock",
)
assert after["session"]["id"] == "conv-lock"
def test_usage_503_when_turn_lock_busy(monkeypatch) -> None:
"""/v1/usage must take the same turn lock as the compress turn — an
unlocked update races the executor and rolls tracker snapshots back."""
import headroom.proxy.handlers.openai as openai_mod
monkeypatch.setattr(openai_mod, "_SESSION_TURN_LOCK_TIMEOUT_SECONDS", 0.05)
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-ulock")
proxy = client.app.state.proxy
lock = proxy._compression_caches[f"{SESSION_KEY_PREFIX}conv-ulock"].session_turn_lock
assert lock.acquire(timeout=1)
try:
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-ulock",
"usage": {
"cache_read_input_tokens": 100,
"cache_creation_input_tokens": 0,
},
},
)
assert resp.status_code == 503, resp.text
assert resp.json()["error"]["type"] == "session_busy"
finally:
lock.release()
def test_usage_single_zero_field_does_not_wipe_state() -> None:
"""{"cache_read_input_tokens": 0} with no write field (the natural
OpenAI-mapped relay on a cold turn) carries no cache signal — it must
not reset the tracker's provider-confirmed prefix state."""
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-zero")
# Establish real provider-confirmed state (both fields present).
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-zero",
"usage": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 50_000,
},
},
)
assert resp.status_code == 200
assert resp.json()["applied"] is True
established = resp.json()["frozen_message_count"]
assert established >= 1
# The no-signal relay is acknowledged but NOT applied.
resp2 = client.post(
"/v1/usage",
json={"session_id": "conv-zero", "usage": {"cache_read_input_tokens": 0}},
)
assert resp2.status_code == 200
body = resp2.json()
assert body["applied"] is False
assert body["reason"] == "no_cache_signal"
assert body["frozen_message_count"] == established # state intact
# A relay with BOTH fields zero is a genuine fully-cold assertion
# and IS applied.
resp3 = client.post(
"/v1/usage",
json={
"session_id": "conv-zero",
"usage": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
},
)
assert resp3.status_code == 200
assert resp3.json()["applied"] is True
def test_usage_404_for_ttl_expired_tracker() -> None:
"""peek() must treat a TTL-expired-but-unswept tracker as gone — a 200
here would resurrect the dead tracker on every relay."""
import time as _time
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-expired")
proxy = client.app.state.proxy
tracker = proxy.session_tracker_store._trackers[f"{SESSION_KEY_PREFIX}conv-expired"]
tracker._last_activity = _time.time() - 999_999
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-expired",
"usage": {"cache_read_input_tokens": 100},
},
)
assert resp.status_code == 404
assert resp.json()["error"]["type"] == "unknown_session"

185
tests/test_compression_cache_registry.py Normal file
View file

@ -0,0 +1,185 @@
"""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()

5
tests/test_integrations/agno/test_model.py
View file

@ -19,11 +19,6 @@ try:
AGNO_AVAILABLE = True
except ImportError:
AGNO_AVAILABLE = False
else:
try: # agno < 3: the per-message usage dataclass lived at agno.models.metrics
from agno.models.metrics import Metrics
except ImportError: # agno >= 3 moved it to agno.metrics, renamed MessageMetrics
from agno.metrics import MessageMetrics as Metrics
from headroom import HeadroomConfig, HeadroomMode

165
tests/test_org_scale_limits.py Normal file
View file

@ -0,0 +1,165 @@
"""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)

49
tests/test_provider_registry.py
View file

@ -56,6 +56,55 @@ def test_resolve_api_targets_normalizes_trailing_v1() -> None:
assert targets.vertex == "https://vertex.example"
def test_copilot_openai_target_routes_anthropic_to_copilot() -> None:
"""When the OpenAI target is a Copilot host and no Anthropic override is set,
the Anthropic target must default to the same Copilot host.
Copilot serves Claude models via its Anthropic surface (``/v1/messages``) on
the same host. Without this, Claude requests fell back to api.anthropic.com
and 401'd with the Copilot bearer ("Invalid bearer token", #3247).
"""
targets = resolve_api_targets(
ProviderApiOverrides(
anthropic=None,
openai="https://api.githubcopilot.com",
gemini=None,
cloudcode=None,
vertex=None,
)
)
assert targets.openai == "https://api.githubcopilot.com"
assert targets.anthropic == "https://api.githubcopilot.com"
def test_explicit_anthropic_override_wins_over_copilot_default() -> None:
"""An explicit Anthropic target is never overridden by the Copilot default."""
targets = resolve_api_targets(
ProviderApiOverrides(
anthropic="https://api.anthropic.com",
openai="https://api.githubcopilot.com",
gemini=None,
cloudcode=None,
vertex=None,
)
)
assert targets.anthropic == "https://api.anthropic.com"
def test_non_copilot_openai_target_leaves_anthropic_default() -> None:
"""A non-Copilot OpenAI target must not touch the Anthropic default."""
targets = resolve_api_targets(
ProviderApiOverrides(
anthropic=None,
openai="https://api.openai.com",
gemini=None,
cloudcode=None,
vertex=None,
)
)
assert targets.anthropic == "https://api.anthropic.com"
def test_proxy_config_exposes_provider_api_overrides() -> None:
config = ProxyConfig(
anthropic_api_url="https://anthropic.example",

49
tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py
View file

@ -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():

18
tests/test_proxy_compress_endpoint.py
View file

@ -303,7 +303,12 @@ class TestCompressEndpointCompression:
transforms_summary={"test_transform": 1},
markers_inserted=[],
)
run_compression = AsyncMock(return_value=result)
# The executor callable returns the 5-tuple contract of
# _run_stateless/_run_session_turn:
# (result, final_messages, tokens_before, tokens_after, session_info).
run_compression = AsyncMock(
return_value=(result, result.messages, result.tokens_before, result.tokens_after, None)
)
record_outcome = AsyncMock()
monkeypatch.setattr(proxy, "_run_compression_in_executor", run_compression)
monkeypatch.setattr(proxy, "_record_request_outcome", record_outcome)
@ -349,7 +354,16 @@ class TestCompressEndpointCompression:
monkeypatch.setattr(
proxy,
"_run_compression_in_executor",
AsyncMock(return_value=result),
# Same 5-tuple contract as _run_stateless (see above).
AsyncMock(
return_value=(
result,
result.messages,
result.tokens_before,
result.tokens_after,
None,
)
),
)
monkeypatch.setattr(proxy, "_record_request_outcome", AsyncMock())

225
tests/test_session_engine.py Normal file
View file

@ -0,0 +1,225 @@
"""Unit tests for the shared session-turn engine (headroom/proxy/session_engine).
The engine is the single cache-management brain for the proxy request paths
and the sidecar /v1/compress path; these tests pin its two freeze policies
and the overlay finalization directly, without an HTTP harness.
"""
from __future__ import annotations
import json
import pytest
from headroom.cache.compression_cache import CompressionCache
from headroom.proxy.session_engine import (
FREEZE_POLICY_CONFIRMED_CLAMP,
FREEZE_POLICY_REPLAYABLE,
finalize_turn,
prepare_turn,
)
def _tool_msg(content: str, call_id: str = "c1") -> dict:
return {"role": "tool", "tool_call_id": call_id, "content": content}
def _history_with_cached_tool(
cache: CompressionCache, original: str, compressed: str
) -> list[dict]:
"""A 3-message history whose tool result has a cached compressed form."""
cache.store_compressed(cache.content_hash(original), compressed, tokens_saved=10)
return [
{"role": "user", "content": "get items"},
{"role": "assistant", "content": "calling"},
_tool_msg(original),
]
# --------------------------------------------------------------------------- #
# prepare_turn: freeze policies #
# --------------------------------------------------------------------------- #
def test_sidecar_policy_freezes_full_replayable_prefix() -> None:
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
# user, assistant, cached tool are all stable; the trailing message is
# always excluded by compute_frozen_count.
assert prep.frozen_message_count == 3
# The swap replaced the tool result with its cached compressed form.
assert prep.pipeline_input[2]["content"] == "[compressed]"
# The caller's list is never mutated.
assert messages[2]["content"].startswith("ORIGINAL")
def test_sidecar_policy_explicit_pin_wins_when_larger() -> None:
cache = CompressionCache()
messages = [
{"role": "user", "content": "a"},
_tool_msg("never seen before " * 50), # not in cache -> derived stops here
{"role": "user", "content": "next"},
]
derived = cache.compute_frozen_count(messages)
assert derived == 1 # only the leading plain message
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE, explicit_frozen=2)
assert prep.frozen_message_count == 2
def test_sidecar_policy_derived_wins_when_explicit_smaller() -> None:
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE, explicit_frozen=1)
assert prep.frozen_message_count == 3
def test_proxy_policy_clamps_by_cache_count() -> None:
"""Provider says 5 messages are cached, but local state can only replay 3:
freezing past the replayable bound would forward raw bytes."""
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append(_tool_msg("uncached " * 50, "c2"))
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=5)
assert prep.frozen_message_count == 3
def test_proxy_policy_clamps_by_tracker() -> None:
"""Local state could replay 3, but the provider only confirmed 1: content
past the confirmed prefix stays compressible (the #327 posture)."""
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=1)
assert prep.frozen_message_count == 1
def test_proxy_policy_none_tracker_freezes_nothing() -> None:
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=None)
assert prep.frozen_message_count == 0
def test_unknown_policy_rejected() -> None:
cache = CompressionCache()
with pytest.raises(ValueError):
prepare_turn(cache, [], policy="wat")
def test_prepare_marks_frozen_tool_results_stable() -> None:
cache = CompressionCache()
original = "ORIGINAL " * 100
messages = _history_with_cached_tool(cache, original, "[compressed]")
messages.append({"role": "user", "content": "next"})
prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
assert cache.content_hash(original) in cache._stable_hashes
# --------------------------------------------------------------------------- #
# finalize_turn: overlay + recount hook #
# --------------------------------------------------------------------------- #
def _prev_pair() -> tuple[list[dict], list[dict]]:
prev_original = [
{"role": "user", "content": "ORIGINAL " * 100},
{"role": "assistant", "content": "ok"},
]
prev_returned = [
{"role": "user", "content": "[returned-form]"},
{"role": "assistant", "content": "ok"},
]
return prev_original, prev_returned
def test_finalize_replays_previous_returned_prefix() -> None:
prev_original, prev_returned = _prev_pair()
current = prev_original + [{"role": "user", "content": "next"}]
# The pipeline "drifted": it emitted the raw original for message 0.
drifted = [dict(m) for m in current]
counted: list[int] = []
def _count(msgs: list[dict]) -> int:
counted.append(len(json.dumps(msgs)))
return 42
turn = finalize_turn(drifted, current, prev_original, prev_returned, count_tokens=_count)
assert turn.replayed
assert turn.messages[0]["content"] == "[returned-form]"
assert turn.messages[-1]["content"] == "next"
assert turn.tokens == 42
assert len(counted) == 1
def test_finalize_noop_without_prev_snapshots() -> None:
current = [{"role": "user", "content": "hi"}]
calls: list[int] = []
turn = finalize_turn(current, current, [], [], count_tokens=lambda m: calls.append(1) or 1)
assert not turn.replayed
assert turn.messages == current
assert turn.tokens is None
assert not calls # count_tokens only runs when the overlay fired
def test_finalize_count_hook_failure_falls_back() -> None:
prev_original, prev_returned = _prev_pair()
current = prev_original + [{"role": "user", "content": "next"}]
def _boom(_msgs: list[dict]) -> int:
raise RuntimeError("tokenizer down")
turn = finalize_turn(
[dict(m) for m in current], current, prev_original, prev_returned, count_tokens=_boom
)
assert turn.replayed
assert turn.tokens is None
# --------------------------------------------------------------------------- #
# OpenAI proxy token-path migration: formula identity + marking benefit. #
# --------------------------------------------------------------------------- #
def test_replayable_without_pin_equals_bare_cache_count() -> None:
"""The OpenAI proxy token path historically froze on compute_frozen_count
alone; REPLAYABLE with no explicit pin must be formula-identical, so its
migration onto the engine is a pure extraction."""
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "AAAA " * 50, "[c1]")
messages.append(_tool_msg("uncached content", call_id="c2"))
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
assert prep.frozen_message_count == cache.compute_frozen_count(messages)
# And that count stops at the uncached tool_result (index 3).
assert prep.frozen_message_count == 3
def test_marking_preserves_freeze_across_entry_eviction() -> None:
"""The one real benefit mark_stable_from_messages adds on the migrated
path: an in-prefix tool_result stays stable via `_stable_hashes` even
after its compressed ENTRY is evicted by the per-cache LRU, so the frozen
count does not collapse at that position on the next turn."""
cache = CompressionCache(max_entries=100)
original = "BBBB " * 50
messages = _history_with_cached_tool(cache, original, "[c1]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
assert prep.frozen_message_count == 3 # tool in prefix, marked stable
# Simulate entry LRU turnover: the compressed entry disappears.
h = cache.content_hash(original)
with cache._lock:
cache._cache.pop(h, None)
# Without marking, the frozen count would collapse to 2 here; the
# stable-hash record keeps the position frozen.
assert cache.compute_frozen_count(messages) == 3