From d12ea501222b92f371b5427595c14ec2466a3f39 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 26 Aug 2026 15:57:31 +0530 Subject: [PATCH] feat(proxy): unify proxy and sidecar compression on one session engine (#3271) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Replaces #3263 (same changeset, squashed to one conventional commit — after the base PRs squash-merged, the stacked branch's commit history could not pass the commitlint gate against main, and force-pushing the original branch was not permitted). #3261 and #3270 (which replaced #3262) are merged; this is the last piece of the stack. ## Goal One brain. The cache-management tier — freeze computation, Zone-1 byte swap, cached-prefix overlay — previously existed twice: inline in the proxy request handlers, and (as of #3270) in the `/v1/compress` sidecar path. This PR extracts it into **`headroom/proxy/session_engine.py`**, invoked by BOTH. Every future cache-management fix lands in both modes by construction. ## Design **`prepare_turn(...)` → `TurnPrep`** — freeze + `mark_stable_from_messages` + `apply_cached`, with two *deliberately different, documented* freeze policies: - `FREEZE_POLICY_CONFIRMED_CLAMP`: `min(tracker_frozen, cache_count)` — never freeze past provider-confirmed (the #327 posture). The Anthropic proxy passes its already-composed tracker/strict-override value, reproducing the previous `min()` byte-for-byte. - `FREEZE_POLICY_REPLAYABLE`: `max(cache_count, explicit)` — freeze everything locally replayable, because whatever was previously returned *is* the provider's cache contract; recompressing it (even "better") busts. **`finalize_turn(...)` → `TurnFinal`** — the byte-identical cached-prefix replay (`overlay_cached_prefix`) + conditional token recount hook. Run as a **strictly behavior-preserving extraction**: the bar was every pre-existing test passing *unmodified*, and it held. ## What migrated | Path | Status | |---|---| | `/v1/compress` sidecar turn | ✅ engine (REPLAYABLE); lock, executor offload, savings accounting, record_returned unchanged | | `anthropic.py` token-mode pre-block + overlay | ✅ engine (CONFIRMED_CLAMP); background compression, cold-start fast pass, `_cold_recompact_active` skip preserved | | `openai.py` proxy token-mode pre-block + overlay | ✅ engine (REPLAYABLE — formula-identical to the old bare `compute_frozen_count`); the added `mark_stable` call means the freeze now survives entry-level LRU eviction (test-pinned); the router's `_frozen_verdicts` remains the boundary-message protection | | `openai.py` cache-mode branch | ⏸ keeps bare `apply_cached` — cache mode keeps the latest observation mutable by design | Also fixed for BOTH handlers: overlay replay now runs under backpressure (shedding it busted every gated session's prompt cache exactly at peak load), and the inflation guard exempts replayed prefixes. ## Hardening (max-effort review, all applied) `/v1/usage` applies on the executor under the per-session turn lock with a timed acquire (503 `session_busy`); registry eviction skips sessions mid-turn; `peek()` is expiry-aware; silent fallbacks log warnings; RequestOutcome recorded on session 503s. ## Testing - `tests/test_session_engine.py`: 13 direct unit tests — both policies, explicit-pin precedence, REPLAYABLE-without-pin ≡ bare `compute_frozen_count`, overlay fires/doesn't, recount only on replay, freeze-survives-entry-eviction. - Parity bar: full pre-existing suites pass unmodified — cache-stability (Anthropic + OpenAI), overlay, backpressure (incl. replay-under-saturation regression), cold-start fast pass, cache-mode, session-mode byte-stability, compress-API, org-scale, registry. Full local suite: 11k+ green. - ruff check/format clean (CI's ruff 0.16.3). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE Co-authored-by: Claude Fable 5 --- headroom/cache/prefix_tracker.py | 14 +- headroom/proxy/handlers/anthropic.py | 70 +++--- headroom/proxy/handlers/openai.py | 289 ++++++++++++++++++----- headroom/proxy/server.py | 46 +++- headroom/proxy/session_engine.py | 185 +++++++++++++++ tests/test_compress_session_mode.py | 141 +++++++++++ tests/test_compression_cache_registry.py | 51 ++++ tests/test_session_engine.py | 225 ++++++++++++++++++ 8 files changed, 923 insertions(+), 98 deletions(-) create mode 100644 headroom/proxy/session_engine.py create mode 100644 tests/test_session_engine.py diff --git a/headroom/cache/prefix_tracker.py b/headroom/cache/prefix_tracker.py index 9b145c333..ea47417e3 100644 --- a/headroom/cache/prefix_tracker.py +++ b/headroom/cache/prefix_tracker.py @@ -1270,14 +1270,24 @@ class SessionTrackerStore: self._lineage_counter = itertools.count(1) def peek(self, session_id: str) -> PrefixCacheTracker | None: - """Return the tracker for ``session_id`` if one exists, else 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. """ - return self._trackers.get(session_id) + 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.""" diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 00c719d2e..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,10 +2071,8 @@ 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 @@ -2087,16 +2091,18 @@ class AnthropicHandlerMixin: 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: logger.debug( "[%s] Cached-prefix replay skipped: reason=%s", diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 18b6bfb33..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" @@ -9815,24 +9860,47 @@ class OpenAIHandlerMixin: # 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. - from headroom.cache.prefix_tracker import overlay_cached_prefix + # 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, + ) - with comp_cache.session_turn_lock: + # 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() - derived_frozen = comp_cache.compute_frozen_count(messages) - session_frozen = max(derived_frozen, frozen_message_count or 0) - comp_cache.mark_stable_from_messages(messages, session_frozen) - pipeline_input = comp_cache.apply_cached(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=pipeline_input, model=model, **pipeline_kwargs) + 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. - final = overlay_cached_prefix( - result.messages, messages, prev_original, prev_returned - ) - replayed = final != result.messages + 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 @@ -9843,7 +9911,21 @@ class OpenAIHandlerMixin: _tok = get_tokenizer(model_name) raw_tokens_before = _tok.count_messages(messages) final_tokens_after = _tok.count_messages(final) - except Exception: # nosec B110 - fall back to pipeline counts + 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) @@ -9855,9 +9937,11 @@ class OpenAIHandlerMixin: info = { "id": session_id, "frozen_message_count": session_frozen, - "cached_prefix_replayed": replayed, + "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 @@ -9942,6 +10026,32 @@ class OpenAIHandlerMixin: 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={ @@ -10081,45 +10191,112 @@ class OpenAIHandlerMixin: # 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 session is - # answered without leaving a footprint, and the session keeps the - # provider its compress call inferred rather than a default from here. - tracker = self.session_tracker_store.peek(f"compress\x00{session_id}") - last_returned = tracker.get_last_forwarded_messages() if tracker is not None else [] - if not last_returned: - # 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 — it only means this telemetry landed - # nowhere. + # 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( - status_code=404, + { + "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": "unknown_session", + "type": "session_busy", "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." + "A compress turn for this session is in flight; retry the usage relay." ), } }, ) - - tracker.update_from_response( - cache_read_tokens=cache_read, - cache_write_tokens=cache_write, - messages=last_returned, - original_messages=tracker.get_last_original_messages(), - ) + if frozen_count is None: + return self._compress_usage_unknown_session(session_id) return JSONResponse( { "session_id": session_id, - "frozen_message_count": tracker.get_frozen_message_count(), + "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/server.py b/headroom/proxy/server.py index dac2fe6e5..233a0367c 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1617,10 +1617,18 @@ class HeadroomProxy( ): 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) @@ -1633,6 +1641,16 @@ class HeadroomProxy( 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. @@ -1657,20 +1675,32 @@ class HeadroomProxy( # 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: evict_count = min( max(1, MAX_COMPRESSION_CACHE_SESSIONS // 4), len(self._compression_caches), ) - for _ in range(evict_count): - sid, _evicted = self._compression_caches.popitem(last=False) + 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) - logger.info( - "Evicted %d least-recently-used compression caches " - "(exceeded %d max sessions)", - evict_count, - MAX_COMPRESSION_CACHE_SESSIONS, - ) + if evictable: + logger.info( + "Evicted %d least-recently-used compression caches " + "(exceeded %d max sessions)", + len(evictable), + MAX_COMPRESSION_CACHE_SESSIONS, + ) cache = CompressionCache(max_entries=COMPRESSION_CACHE_MAX_ENTRIES) self._compression_caches[session_id] = cache 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/tests/test_compress_session_mode.py b/tests/test_compress_session_mode.py index 01380fa32..cc3a23563 100644 --- a/tests/test_compress_session_mode.py +++ b/tests/test_compress_session_mode.py @@ -425,3 +425,144 @@ def test_explicit_frozen_count_still_wins_when_larger() -> None: 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 index f510ff8d8..da3aba4e8 100644 --- a/tests/test_compression_cache_registry.py +++ b/tests/test_compression_cache_registry.py @@ -132,3 +132,54 @@ def test_sweep_is_rate_limited(monkeypatch) -> None: 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_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