mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(proxy): unify proxy and sidecar compression on one session engine (#3271)
> 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 <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra
GitHub
parent
4fa88026d9
commit
d12ea50122
8 changed files with 923 additions and 98 deletions
14
headroom/cache/prefix_tracker.py
vendored
14
headroom/cache/prefix_tracker.py
vendored
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
185
headroom/proxy/session_engine.py
Normal file
185
headroom/proxy/session_engine.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
"""Session-turn engine — the single cache-management brain for both modes.
|
||||
|
||||
One conversation turn, from the cache's point of view, is always the same
|
||||
three-step dance regardless of who owns the upstream call:
|
||||
|
||||
1. **Prepare** (:func:`prepare_turn`): decide how many leading messages are
|
||||
frozen, mark the stable prefix, and swap previously-computed compressed
|
||||
bytes into the working copy (``apply_cached`` — "Zone 1").
|
||||
2. Run the compression pipeline over the prepared input (owned by the
|
||||
caller: the proxy handlers wrap it in background/cold-start/backpressure
|
||||
orchestration, the sidecar path runs it inline on the executor).
|
||||
3. **Finalize** (:func:`finalize_turn`): replay last turn's exact
|
||||
previously-forwarded/returned prefix over any residual drift the pipeline
|
||||
introduced (``overlay_cached_prefix``), so the bytes that leave the
|
||||
process are byte-identical to what the provider already cached.
|
||||
|
||||
Historically the proxy request handlers (anthropic + openai token mode) and
|
||||
the sidecar ``/v1/compress`` session path each carried their own inline copy
|
||||
of steps 1 and 3. This module is the shared implementation: a
|
||||
cache-management fix landed here reaches BOTH modes at once.
|
||||
|
||||
Freeze policies
|
||||
---------------
|
||||
|
||||
The one deliberate behavioural difference between the modes lives in step 1,
|
||||
and it is a *policy parameter*, not a fork of the code:
|
||||
|
||||
``FREEZE_POLICY_CONFIRMED_CLAMP`` — ``min(tracker_frozen, cache_count)``.
|
||||
The proxy sees the provider's responses, so ``tracker_frozen`` is the
|
||||
provider-confirmed cached prefix (from ``cache_read_input_tokens``).
|
||||
Freezing is clamped by BOTH bounds: never past what the provider
|
||||
actually has cached (freezing more would forgo compression of content
|
||||
that is not yet cache-protected — the #327 posture), and never past what
|
||||
the local cache can byte-replay (freezing a message whose entry was
|
||||
evicted would pass through raw original bytes).
|
||||
|
||||
``FREEZE_POLICY_REPLAYABLE`` — ``max(cache_count, explicit_frozen or 0)``.
|
||||
Freeze everything the local cache can byte-replay. Used by callers with
|
||||
no provider-confirmed count to clamp against: the sidecar ``/v1/compress``
|
||||
endpoint (it never sees the provider's response — whatever it previously
|
||||
RETURNED is the provider's cache contract, so every already-returned
|
||||
message must come back byte-identical), and the OpenAI proxy token path
|
||||
(its tracker feeds cache mode, not token mode). Recompressing an
|
||||
already-returned message — even into a *smaller* form — is a bust: the
|
||||
drift was observed in practice, and ``overlay_cached_prefix``'s
|
||||
non-inflation guard cannot repair a shrunken form (replaying the larger
|
||||
original bytes would "inflate" the candidate). Freezing the entire
|
||||
locally-replayable prefix eliminates that recompression outright.
|
||||
Over-freezing relative to the provider's real cache only forgoes tail
|
||||
compression; it can never bust. An explicit ``frozen_message_count``
|
||||
from the caller still wins when larger — the caller may know more about
|
||||
the provider cache than local state does.
|
||||
|
||||
Why the Anthropic proxy path cannot simply adopt the replayable posture: its
|
||||
provider-confirmed clamp deliberately KEEPS not-yet-cached content
|
||||
compressible, and its overlay inputs (tracker snapshots) are refreshed on
|
||||
every response, so drift repair is reliable there. Each posture is correct
|
||||
for the information its mode actually has.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from headroom.cache.prefix_tracker import overlay_cached_prefix
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FREEZE_POLICY_CONFIRMED_CLAMP = "confirmed_clamp"
|
||||
FREEZE_POLICY_REPLAYABLE = "replayable"
|
||||
|
||||
_FREEZE_POLICIES = (FREEZE_POLICY_CONFIRMED_CLAMP, FREEZE_POLICY_REPLAYABLE)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TurnPrep:
|
||||
"""Result of :func:`prepare_turn`.
|
||||
|
||||
``frozen_message_count`` is what the pipeline must be told to skip;
|
||||
``pipeline_input`` is the working copy with previously-compressed bytes
|
||||
swapped in (never the caller's list — ``apply_cached`` copies).
|
||||
"""
|
||||
|
||||
frozen_message_count: int
|
||||
pipeline_input: list[dict[str, Any]]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TurnFinal:
|
||||
"""Result of :func:`finalize_turn`.
|
||||
|
||||
``messages`` are the bytes to forward/return; ``replayed`` says whether
|
||||
the overlay restored last turn's prefix over pipeline drift; ``tokens``
|
||||
is the recount of ``messages`` when a ``count_tokens`` hook was supplied
|
||||
and the overlay actually fired (None otherwise — the pipeline's own
|
||||
count is still valid when nothing was replaced).
|
||||
"""
|
||||
|
||||
messages: list[dict[str, Any]]
|
||||
replayed: bool
|
||||
tokens: int | None = None
|
||||
|
||||
|
||||
def prepare_turn(
|
||||
comp_cache: Any,
|
||||
messages: list[dict[str, Any]],
|
||||
*,
|
||||
policy: str,
|
||||
tracker_frozen: int | None = None,
|
||||
explicit_frozen: int | None = None,
|
||||
) -> TurnPrep:
|
||||
"""Freeze decision + stable marking + cached-byte swap for one turn.
|
||||
|
||||
Args:
|
||||
comp_cache: the session's ``CompressionCache``.
|
||||
messages: the caller's RAW message list (never mutated).
|
||||
policy: ``FREEZE_POLICY_CONFIRMED_CLAMP`` or ``FREEZE_POLICY_REPLAYABLE`` —
|
||||
see the module docstring for why they differ.
|
||||
tracker_frozen: provider-confirmed frozen count (proxy policy only;
|
||||
``None`` means "nothing confirmed" and freezes 0 there).
|
||||
explicit_frozen: caller-pinned frozen count (sidecar policy only;
|
||||
wins when larger than the locally-derived bound).
|
||||
"""
|
||||
if policy not in _FREEZE_POLICIES:
|
||||
raise ValueError(f"unknown freeze policy: {policy!r}")
|
||||
|
||||
cache_count = comp_cache.compute_frozen_count(messages)
|
||||
if policy == FREEZE_POLICY_CONFIRMED_CLAMP:
|
||||
# Never freeze past the provider-confirmed prefix, and never past
|
||||
# what local state can byte-replay.
|
||||
frozen = min(tracker_frozen or 0, cache_count)
|
||||
else:
|
||||
# Freeze the entire locally-replayable prefix; an explicit caller
|
||||
# pin may extend it (the caller vouches the provider cached those
|
||||
# exact raw bytes, so passing them through untouched is correct).
|
||||
frozen = max(cache_count, explicit_frozen or 0)
|
||||
|
||||
comp_cache.mark_stable_from_messages(messages, frozen)
|
||||
return TurnPrep(
|
||||
frozen_message_count=frozen,
|
||||
pipeline_input=comp_cache.apply_cached(messages),
|
||||
)
|
||||
|
||||
|
||||
def finalize_turn(
|
||||
result_messages: list[dict[str, Any]],
|
||||
original_messages: list[dict[str, Any]],
|
||||
prev_original: list[dict[str, Any]] | None,
|
||||
prev_returned: list[dict[str, Any]] | None,
|
||||
*,
|
||||
count_tokens: Callable[[list[dict[str, Any]]], int] | None = None,
|
||||
) -> TurnFinal:
|
||||
"""Replay last turn's exact forwarded/returned prefix over pipeline drift.
|
||||
|
||||
``overlay_cached_prefix`` self-guards (positional alignment, append-only
|
||||
shape, non-inflation), so calling this is always safe: when replay is not
|
||||
provably correct it returns the pipeline's own output unchanged.
|
||||
|
||||
``count_tokens`` is invoked only when the overlay actually replaced
|
||||
bytes — the pipeline's own token count is still accurate otherwise. A
|
||||
failing hook falls back to "no recount" rather than failing the turn.
|
||||
"""
|
||||
final = overlay_cached_prefix(result_messages, original_messages, prev_original, prev_returned)
|
||||
replayed = final != result_messages
|
||||
tokens: int | None = None
|
||||
if replayed and count_tokens is not None:
|
||||
try:
|
||||
tokens = count_tokens(final)
|
||||
except Exception as e:
|
||||
# Fail-open: the turn still forwards, but the caller keeps the
|
||||
# pipeline's count of messages that are NOT being forwarded —
|
||||
# tokens_saved accounting is stale for this turn. Loud, not
|
||||
# silent: a tokenizer that cannot count the replayed form is a
|
||||
# bug worth surfacing even though it must not fail the request.
|
||||
logger.warning(
|
||||
"finalize_turn: token recount of replayed prefix failed "
|
||||
"(%s: %s); keeping the pipeline's pre-overlay count",
|
||||
type(e).__name__,
|
||||
e,
|
||||
)
|
||||
tokens = None
|
||||
return TurnFinal(messages=final, replayed=replayed, tokens=tokens)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
225
tests/test_session_engine.py
Normal file
225
tests/test_session_engine.py
Normal file
|
|
@ -0,0 +1,225 @@
|
|||
"""Unit tests for the shared session-turn engine (headroom/proxy/session_engine).
|
||||
|
||||
The engine is the single cache-management brain for the proxy request paths
|
||||
and the sidecar /v1/compress path; these tests pin its two freeze policies
|
||||
and the overlay finalization directly, without an HTTP harness.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.cache.compression_cache import CompressionCache
|
||||
from headroom.proxy.session_engine import (
|
||||
FREEZE_POLICY_CONFIRMED_CLAMP,
|
||||
FREEZE_POLICY_REPLAYABLE,
|
||||
finalize_turn,
|
||||
prepare_turn,
|
||||
)
|
||||
|
||||
|
||||
def _tool_msg(content: str, call_id: str = "c1") -> dict:
|
||||
return {"role": "tool", "tool_call_id": call_id, "content": content}
|
||||
|
||||
|
||||
def _history_with_cached_tool(
|
||||
cache: CompressionCache, original: str, compressed: str
|
||||
) -> list[dict]:
|
||||
"""A 3-message history whose tool result has a cached compressed form."""
|
||||
cache.store_compressed(cache.content_hash(original), compressed, tokens_saved=10)
|
||||
return [
|
||||
{"role": "user", "content": "get items"},
|
||||
{"role": "assistant", "content": "calling"},
|
||||
_tool_msg(original),
|
||||
]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# prepare_turn: freeze policies #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_sidecar_policy_freezes_full_replayable_prefix() -> None:
|
||||
cache = CompressionCache()
|
||||
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
|
||||
# user, assistant, cached tool are all stable; the trailing message is
|
||||
# always excluded by compute_frozen_count.
|
||||
assert prep.frozen_message_count == 3
|
||||
# The swap replaced the tool result with its cached compressed form.
|
||||
assert prep.pipeline_input[2]["content"] == "[compressed]"
|
||||
# The caller's list is never mutated.
|
||||
assert messages[2]["content"].startswith("ORIGINAL")
|
||||
|
||||
|
||||
def test_sidecar_policy_explicit_pin_wins_when_larger() -> None:
|
||||
cache = CompressionCache()
|
||||
messages = [
|
||||
{"role": "user", "content": "a"},
|
||||
_tool_msg("never seen before " * 50), # not in cache -> derived stops here
|
||||
{"role": "user", "content": "next"},
|
||||
]
|
||||
derived = cache.compute_frozen_count(messages)
|
||||
assert derived == 1 # only the leading plain message
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE, explicit_frozen=2)
|
||||
assert prep.frozen_message_count == 2
|
||||
|
||||
|
||||
def test_sidecar_policy_derived_wins_when_explicit_smaller() -> None:
|
||||
cache = CompressionCache()
|
||||
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE, explicit_frozen=1)
|
||||
assert prep.frozen_message_count == 3
|
||||
|
||||
|
||||
def test_proxy_policy_clamps_by_cache_count() -> None:
|
||||
"""Provider says 5 messages are cached, but local state can only replay 3:
|
||||
freezing past the replayable bound would forward raw bytes."""
|
||||
cache = CompressionCache()
|
||||
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
|
||||
messages.append(_tool_msg("uncached " * 50, "c2"))
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=5)
|
||||
assert prep.frozen_message_count == 3
|
||||
|
||||
|
||||
def test_proxy_policy_clamps_by_tracker() -> None:
|
||||
"""Local state could replay 3, but the provider only confirmed 1: content
|
||||
past the confirmed prefix stays compressible (the #327 posture)."""
|
||||
cache = CompressionCache()
|
||||
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=1)
|
||||
assert prep.frozen_message_count == 1
|
||||
|
||||
|
||||
def test_proxy_policy_none_tracker_freezes_nothing() -> None:
|
||||
cache = CompressionCache()
|
||||
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=None)
|
||||
assert prep.frozen_message_count == 0
|
||||
|
||||
|
||||
def test_unknown_policy_rejected() -> None:
|
||||
cache = CompressionCache()
|
||||
with pytest.raises(ValueError):
|
||||
prepare_turn(cache, [], policy="wat")
|
||||
|
||||
|
||||
def test_prepare_marks_frozen_tool_results_stable() -> None:
|
||||
cache = CompressionCache()
|
||||
original = "ORIGINAL " * 100
|
||||
messages = _history_with_cached_tool(cache, original, "[compressed]")
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
|
||||
assert cache.content_hash(original) in cache._stable_hashes
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# finalize_turn: overlay + recount hook #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _prev_pair() -> tuple[list[dict], list[dict]]:
|
||||
prev_original = [
|
||||
{"role": "user", "content": "ORIGINAL " * 100},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
prev_returned = [
|
||||
{"role": "user", "content": "[returned-form]"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
return prev_original, prev_returned
|
||||
|
||||
|
||||
def test_finalize_replays_previous_returned_prefix() -> None:
|
||||
prev_original, prev_returned = _prev_pair()
|
||||
current = prev_original + [{"role": "user", "content": "next"}]
|
||||
# The pipeline "drifted": it emitted the raw original for message 0.
|
||||
drifted = [dict(m) for m in current]
|
||||
|
||||
counted: list[int] = []
|
||||
|
||||
def _count(msgs: list[dict]) -> int:
|
||||
counted.append(len(json.dumps(msgs)))
|
||||
return 42
|
||||
|
||||
turn = finalize_turn(drifted, current, prev_original, prev_returned, count_tokens=_count)
|
||||
assert turn.replayed
|
||||
assert turn.messages[0]["content"] == "[returned-form]"
|
||||
assert turn.messages[-1]["content"] == "next"
|
||||
assert turn.tokens == 42
|
||||
assert len(counted) == 1
|
||||
|
||||
|
||||
def test_finalize_noop_without_prev_snapshots() -> None:
|
||||
current = [{"role": "user", "content": "hi"}]
|
||||
calls: list[int] = []
|
||||
turn = finalize_turn(current, current, [], [], count_tokens=lambda m: calls.append(1) or 1)
|
||||
assert not turn.replayed
|
||||
assert turn.messages == current
|
||||
assert turn.tokens is None
|
||||
assert not calls # count_tokens only runs when the overlay fired
|
||||
|
||||
|
||||
def test_finalize_count_hook_failure_falls_back() -> None:
|
||||
prev_original, prev_returned = _prev_pair()
|
||||
current = prev_original + [{"role": "user", "content": "next"}]
|
||||
|
||||
def _boom(_msgs: list[dict]) -> int:
|
||||
raise RuntimeError("tokenizer down")
|
||||
|
||||
turn = finalize_turn(
|
||||
[dict(m) for m in current], current, prev_original, prev_returned, count_tokens=_boom
|
||||
)
|
||||
assert turn.replayed
|
||||
assert turn.tokens is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# OpenAI proxy token-path migration: formula identity + marking benefit. #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_replayable_without_pin_equals_bare_cache_count() -> None:
|
||||
"""The OpenAI proxy token path historically froze on compute_frozen_count
|
||||
alone; REPLAYABLE with no explicit pin must be formula-identical, so its
|
||||
migration onto the engine is a pure extraction."""
|
||||
cache = CompressionCache()
|
||||
messages = _history_with_cached_tool(cache, "AAAA " * 50, "[c1]")
|
||||
messages.append(_tool_msg("uncached content", call_id="c2"))
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
|
||||
assert prep.frozen_message_count == cache.compute_frozen_count(messages)
|
||||
# And that count stops at the uncached tool_result (index 3).
|
||||
assert prep.frozen_message_count == 3
|
||||
|
||||
|
||||
def test_marking_preserves_freeze_across_entry_eviction() -> None:
|
||||
"""The one real benefit mark_stable_from_messages adds on the migrated
|
||||
path: an in-prefix tool_result stays stable via `_stable_hashes` even
|
||||
after its compressed ENTRY is evicted by the per-cache LRU, so the frozen
|
||||
count does not collapse at that position on the next turn."""
|
||||
cache = CompressionCache(max_entries=100)
|
||||
original = "BBBB " * 50
|
||||
messages = _history_with_cached_tool(cache, original, "[c1]")
|
||||
messages.append({"role": "user", "content": "next"})
|
||||
|
||||
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
|
||||
assert prep.frozen_message_count == 3 # tool in prefix, marked stable
|
||||
|
||||
# Simulate entry LRU turnover: the compressed entry disappears.
|
||||
h = cache.content_hash(original)
|
||||
with cache._lock:
|
||||
cache._cache.pop(h, None)
|
||||
|
||||
# Without marking, the frozen count would collapse to 2 here; the
|
||||
# stable-hash record keeps the position frozen.
|
||||
assert cache.compute_frozen_count(messages) == 3
|
||||
Loading…
Add table
Add a link
Reference in a new issue