diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index 00abf07b0..dcbe0c0ee 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -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 diff --git a/headroom/cache/prefix_tracker.py b/headroom/cache/prefix_tracker.py index 200adddb8..ea47417e3 100644 --- a/headroom/cache/prefix_tracker.py +++ b/headroom/cache/prefix_tracker.py @@ -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() diff --git a/headroom/providers/registry.py b/headroom/providers/registry.py index 04c212a4d..a8b816f3b 100644 --- a/headroom/providers/registry.py +++ b/headroom/providers/registry.py @@ -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), diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 9b0c9d502..538cc2b6d 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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 diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index fd928d110..42ac1a8c0 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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: diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 40efb6571..8d2a2f43c 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -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 # --------------------------------------------------------------------------- diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index a392e9fe1..233a0367c 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -36,6 +36,7 @@ import os import sys import threading import time +from collections import OrderedDict from collections.abc import Callable from dataclasses import fields, is_dataclass, replace from datetime import datetime, timezone @@ -128,6 +129,8 @@ from headroom.proxy.cost import ( merge_cost_stats, # noqa: F401 ) from headroom.proxy.helpers import ( + COMPRESSION_CACHE_MAX_ENTRIES, + COMPRESSION_CACHE_TTL_SECONDS, COMPRESSION_TIMEOUT_SECONDS, # noqa: F401 EAGER_PRELOAD_TIMEOUT_SECONDS, MAX_COMPRESSION_CACHE_SESSIONS, # noqa: F401 @@ -1028,7 +1031,14 @@ class HeadroomProxy( # `CompressionCache` instances have their own internal lock guarding # `_cache`/`_stable_hashes`/`_first_seen` against concurrent # async-dispatched requests for the same session. - self._compression_caches: dict[str, CompressionCache] = {} + # Ordered by last access: `_get_compression_cache` moves a session to + # the end on every hit, so capacity eviction drops the idlest sessions + # — whose provider prefix cache has lapsed anyway — never a busy + # long-lived one. `_compression_cache_last_seen` drives the idle-TTL + # sweep in `_maybe_cleanup_compression_caches`. + self._compression_caches: OrderedDict[str, CompressionCache] = OrderedDict() + self._compression_cache_last_seen: dict[str, float] = {} + self._compression_caches_last_cleanup: float = time.time() self._compression_caches_lock = threading.RLock() self.logger = ( @@ -1580,6 +1590,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 diff --git a/headroom/proxy/session_engine.py b/headroom/proxy/session_engine.py new file mode 100644 index 000000000..4864ae08a --- /dev/null +++ b/headroom/proxy/session_engine.py @@ -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) diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 090408ab5..fa445edc9 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -1947,7 +1947,19 @@ class ContentRouter(Transform): # we match that posture with a dedicated lock rather than relying on # GIL atomicity (which would not protect the read-then-evict sequence). self._frozen_verdicts: dict[int, bool] = {} - self._frozen_verdicts_max = 4096 + # The store is process-wide (one router per pipeline, shared by every + # session), so the cap must scale with the number of CONCURRENT + # sessions, not one user's workload: at org scale (many users behind + # one sidecar) 4096 churns in minutes and FIFO eviction lets tightened + # thresholds flip a still-cached block's verdict — a prefix bust. + # Read at construction so tests and multi-tenant deployments can size + # it via HEADROOM_FROZEN_VERDICTS_MAX without a module reload. + try: + self._frozen_verdicts_max = max( + 256, int(os.environ.get("HEADROOM_FROZEN_VERDICTS_MAX", "4096")) + ) + except ValueError: + self._frozen_verdicts_max = 4096 self._frozen_lock = threading.Lock() # Reset verdicts whenever the shadowed cache is cleared. self._cache.register_on_clear(self._clear_frozen_verdicts) diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index be51e48a0..383157bd5 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/tests/test_anthropic_pre_upstream_backpressure.py @@ -1094,3 +1094,74 @@ def test_response_cache_keys_on_lookup_messages_not_mutated(): assert cache.set_messages == cache.get_messages # And specifically the raw lookup messages, not the scanner's rewrite. assert cache.set_messages == [{"role": "user", "content": "hello"}] + + +# --------------------------------------------------------------------------- # +# Backpressure must not bust the provider prompt cache: the compression # +# pipeline is skipped under saturation, but the previously-forwarded # +# (compressed) prefix must still be replayed byte-identical. Forwarding raw # +# originals would mismatch the bytes the provider cached — busting every # +# gated session's prefix exactly when the proxy is busiest. # +# --------------------------------------------------------------------------- # + + +def test_backpressure_passthrough_replays_cached_prefix(stage_log_capture): + prev_original = [{"role": "user", "content": "ORIGINAL " * 6000}] + prev_forwarded = [{"role": "user", "content": "[compressed-form]"}] + + async def _run() -> None: + sem = asyncio.Semaphore(1) + await sem.acquire() # saturate: the request's acquire will time out + handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem) + handler.config.optimize = True + handler.config.anthropic_pre_upstream_acquire_timeout_seconds = 0.01 + handler.anthropic_pipeline = SimpleNamespace(apply=MagicMock()) + + tracker = SimpleNamespace( + _cached_token_count=0, + get_frozen_message_count=lambda: 0, + get_last_original_messages=lambda: copy.deepcopy(prev_original), + get_last_forwarded_messages=lambda: copy.deepcopy(prev_forwarded), + update_from_response=lambda *a, **k: None, + record_request=lambda *a, **k: None, + ) + handler.session_tracker_store = SimpleNamespace( + compute_session_id=lambda *a, **k: "sess-1", + get_or_create=lambda *a, **k: tracker, + resolve_tracker=lambda *a, **k: tracker, + ) + + forwarded_bodies: list[dict] = [] + orig_retry = handler._retry_request + + async def _capturing_retry(method, url, headers, body, **kw): + forwarded_bodies.append(copy.deepcopy(body)) + return await orig_retry(method, url, headers, body, **kw) + + handler._retry_request = _capturing_retry + + req = _build_request( + { + "model": "claude-3-5-sonnet-latest", + "messages": copy.deepcopy(prev_original) + + [{"role": "user", "content": "next turn"}], + }, + {"authorization": "Bearer sk-ant-api-test"}, + ) + try: + response = await handler.handle_anthropic_messages(req) + assert response.status_code == 200 + # Saturation must still skip the CPU-bound pipeline... + assert not handler.anthropic_pipeline.apply.called + finally: + sem.release() + + assert forwarded_bodies, "request never reached upstream" + sent = forwarded_bodies[-1]["messages"] + # ...but the forwarded prefix must be last turn's exact bytes, not the + # raw original (which the provider never cached). + assert sent[0]["content"] == "[compressed-form]" + assert sent[-1]["content"] == "next turn" + + with _tokenizer_patch(): + anyio.run(_run) diff --git a/tests/test_compress_session_mode.py b/tests/test_compress_session_mode.py new file mode 100644 index 000000000..cc3a23563 --- /dev/null +++ b/tests/test_compress_session_mode.py @@ -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" diff --git a/tests/test_compression_cache_registry.py b/tests/test_compression_cache_registry.py new file mode 100644 index 000000000..da3aba4e8 --- /dev/null +++ b/tests/test_compression_cache_registry.py @@ -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() diff --git a/tests/test_integrations/agno/test_model.py b/tests/test_integrations/agno/test_model.py index a3fe45c40..d1aac8982 100644 --- a/tests/test_integrations/agno/test_model.py +++ b/tests/test_integrations/agno/test_model.py @@ -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 diff --git a/tests/test_org_scale_limits.py b/tests/test_org_scale_limits.py new file mode 100644 index 000000000..d9ceb4879 --- /dev/null +++ b/tests/test_org_scale_limits.py @@ -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) diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py index 0b0ff5b04..07310cd41 100644 --- a/tests/test_provider_registry.py +++ b/tests/test_provider_registry.py @@ -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", diff --git a/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py b/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py index c79071b3b..1b79fdf40 100644 --- a/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py +++ b/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py @@ -199,7 +199,18 @@ def test_bypass_header_does_not_invoke_cached_prefix_replay(monkeypatch): assert captured[-1]["messages"] == messages -def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch): +def test_backpressure_still_invokes_cached_prefix_replay(monkeypatch): + """INVERTED from the pre-#3261 contract this test used to pin. + + Backpressure sheds the compression PIPELINE (the CPU-heavy stage), but + the byte-identical cached-prefix replay must STILL run: skipping it + forwarded raw originals over a compressed cached prefix, busting every + gated session's prompt cache exactly at peak load (the saturated path + previously emitted `Cached-prefix replay skipped: + reason=pre_upstream_backpressure` — that skip was the bug). The replay + self-guards and no-ops here (no previous turn), so the raw messages + still pass through unchanged. + """ app = create_app( _config( optimize=True, @@ -208,12 +219,29 @@ def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch): ) ) - def fail_if_called(*args, **kwargs): # noqa: ANN002, ANN003 - raise AssertionError("cached-prefix replay must be skipped under backpressure") + from headroom.cache import prefix_tracker as _pt - monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", fail_if_called) - debug = Mock() - monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.debug", debug) + real_overlay = _pt.overlay_cached_prefix + overlay_calls: list[int] = [] + + def spy(*args, **kwargs): # noqa: ANN002, ANN003 + overlay_calls.append(1) + return real_overlay(*args, **kwargs) + + # Patch every binding of overlay_cached_prefix: the handler historically + # imported it from prefix_tracker per-request, and the shared session + # engine (headroom.proxy.session_engine, later in this stack) binds it + # at module import — cover both so this test holds across the stack. + monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", spy) + try: + import headroom.proxy.session_engine as _se + + monkeypatch.setattr(_se, "overlay_cached_prefix", spy) + except ImportError: + pass + + info = Mock() + monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.info", info) proxy = app.state.proxy class _SaturatedSemaphore: @@ -236,10 +264,13 @@ def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch): proxy.anthropic_pre_upstream_sem.release() assert response.status_code == 200, response.text - assert captured[-1]["messages"] == messages - assert any( - call.args and call.args[-1] == "pre_upstream_backpressure" for call in debug.call_args_list + # Backpressure engaged (pipeline shed)... + assert any("pre_upstream_backpressure" in str(call) for call in info.call_args_list), ( + "backpressure did not engage — the test setup no longer saturates" ) + # ...but the replay ran (and, with no previous turn, no-op'd safely). + assert overlay_calls, "cached-prefix replay must run under backpressure" + assert captured[-1]["messages"] == messages def test_optimize_on_aligned_history_preserves_replay(): diff --git a/tests/test_proxy_compress_endpoint.py b/tests/test_proxy_compress_endpoint.py index c6434e4ed..9ab8084f5 100644 --- a/tests/test_proxy_compress_endpoint.py +++ b/tests/test_proxy_compress_endpoint.py @@ -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()) diff --git a/tests/test_session_engine.py b/tests/test_session_engine.py new file mode 100644 index 000000000..a29ae4c10 --- /dev/null +++ b/tests/test_session_engine.py @@ -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