feat(compress): session-aware /v1/compress (sidecar mode) + /v1/usage relay (#3270)

> Replaces #3262 (same changeset, squashed to one conventional commit —
the stacked branch's history could not pass commitlint after #3261's
squash-merge broke ancestry, and force-pushing the original branch was
not permitted). All review findings from the two max-effort reviews are
already incorporated; #3261 is merged.

## Why

Gateways that own routing (e.g. Kong as the upstream caller) can't use
Headroom's proxy path, and the stateless `/v1/compress` pushes all
byte-replay bookkeeping onto the caller. This PR moves that state into
the endpoint: **the caller sends the raw conversation + a session id
every turn, forwards the returned bytes verbatim, and gets a
byte-identical prefix — provider prompt cache preserved, no forwarding
through Headroom.**

## Design

- **Session pre-work** mirrors the proxy's Zone 1: content-addressed
swap of previously-computed compressed bytes, then freeze the **entire
locally-replayable prefix** (`compute_frozen_count`).
- **Freeze posture deliberately differs from the proxy's `min(tracker,
cache)`**: in sidecar mode, whatever this endpoint previously returned
*is* the provider's cache contract — recompressing an already-returned
message (even into a smaller form) is a bust. Over-freezing only forgoes
tail compression; it can never bust. (A test caught exactly this:
recompression drift produced a smaller form, and
`overlay_cached_prefix`'s non-inflation guard then couldn't repair it.)
- **`PrefixCacheTracker.record_returned()`** — the sidecar equivalent of
"last forwarded", captured at return time because whatever is returned
is what the caller forwards.
- **`POST /v1/usage`** (same loopback exposure policy): the caller
relays the provider's usage block; `update_from_response` makes freeze
decisions provider-confirmed. Optional — skipping it degrades freeze
precision, never correctness.
- Sessions are NUL-namespaced (`compress\x00<id>`, unspoofable via HTTP
headers); the registry's TTL/LRU lifecycle from #3261 applies
automatically. No session id ⇒ stateless contract byte-for-byte
unchanged.

## Hardening (from two max-effort code reviews, all applied)

- `compress_user_messages` + session_id → 400 (user-message rewrites are
not content-addressed → guaranteed later bust).
- Session-mode timeout / lock-busy → 503 `compression_timeout` /
`session_busy` with retry semantics, instead of failing open with raw
bytes (desync bust).
- Header-based session ids gated behind
`HEADROOM_COMPRESS_SESSION_FROM_HEADER` (default off).
- `/v1/usage` validation: unknown/expired session → 404; both cache
fields absent → 400; single-present-zero → `{"applied": false, "reason":
"no_cache_signal"}` (never wipes freeze state).
- Warm-turn savings recomputed from the raw payload (honest
`tokens_saved`), all CPU work in the executor under a per-session turn
lock.

## Caller contract (Kong)

1. Send raw history + `config.session_id` (or `x-headroom-session-id`
with the env gate on) every turn.
2. Forward the returned `messages` to the provider **verbatim**.
3. Optionally relay the provider's usage block to `/v1/usage`.

## Testing

20 cases in `tests/test_compress_session_mode.py`: stateless regression
+ no state leakage, invalid-id rejection, 2-turn and 3-turn whole-prefix
byte-stability, tracker-loss stability, spoof-resistance, header gating,
lock-busy 503s, usage validation and no-signal handling,
unknown/expired-session 404, TTL-eviction fail-open, explicit
`frozen_message_count` precedence. Plus the full local suite green (11k+
tests).

## Phase 2 (follow-up)

#3263 migrates the proxy request path onto this same session engine so
both modes share one compression/state codepath.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

headroom/cache/compression_cache.py

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

headroom/cache/prefix_tracker.py

@@ -965,6 +965,26 @@ class PrefixCacheTracker:
     def get_last_forwarded_messages(self) -> list[dict[str, Any]]:
         return copy.deepcopy(self._last_forwarded_messages)
 
+    def record_returned(
+        self,
+        original_messages: list[dict[str, Any]],
+        returned_messages: list[dict[str, Any]],
+    ) -> None:
+        """Record the compressed form handed back to a compress-only caller.
+
+        Sidecar mode (session-aware ``/v1/compress``): Headroom does not
+        forward upstream, but whatever it RETURNS is what the caller forwards
+        — the same fact ``update_from_response`` records in proxy mode, just
+        captured at return time instead of send time. Only the transcript
+        snapshots and the activity clock move here; frozen-prefix counts are
+        left untouched because no provider response has confirmed anything
+        yet — they advance when the caller relays usage via ``/v1/usage``
+        (``update_from_response``), or stay at their conservative local value.
+        """
+        self._last_activity = time.time()
+        self._last_original_messages = copy.deepcopy(original_messages)
+        self._last_forwarded_messages = copy.deepcopy(returned_messages)
+
     def resolved_cache_ttl_seconds(self) -> int:
         """Effective prompt-cache lifetime for this session's provider."""
         if self.config.cache_ttl_seconds is not None:
@@ -1249,6 +1269,16 @@ class SessionTrackerStore:
         self._lineage_affinities: dict[str, str | None] = {}
         self._lineage_counter = itertools.count(1)
 
+    def peek(self, session_id: str) -> PrefixCacheTracker | None:
+        """Return the tracker for ``session_id`` if one exists, 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`.
+        """
+        return self._trackers.get(session_id)
+
     def get_or_create(self, session_id: str, provider: str) -> PrefixCacheTracker:
         """Get existing tracker or create a new one for this session."""
         self._maybe_cleanup()

headroom/proxy/handlers/openai.py

@@ -9591,6 +9591,9 @@ class OpenAIHandlerMixin:
         headers = dict(request.headers)
         tags = extract_tags(headers)
         client = classify_client(headers)
+        # Initialized before the try so the TimeoutError handler can branch on
+        # it even if the failure happened before session parsing.
+        session_id = None
 
         try:
             # Use OpenAI pipeline (messages are in OpenAI format from TS SDK)
@@ -9655,6 +9658,64 @@ class OpenAIHandlerMixin:
                         }
                     },
                 )
+            # Session-aware sidecar mode (opt-in): with a session id the
+            # endpoint keeps the byte-replay state ITSELF — the same
+            # per-session machinery the proxy path uses (compression cache +
+            # prefix tracker, with the registry's TTL/LRU lifecycle) — so a
+            # gateway that owns routing (e.g. Kong) can resend the RAW
+            # conversation every turn and still get a byte-identical prefix
+            # back. Contract: the caller forwards the returned messages
+            # verbatim, and may relay provider usage via POST /v1/usage for
+            # telemetry/attribution. Without a session id, behaviour is the
+            # stateless contract, unchanged.
+            session_id = compress_config.get("session_id")
+            # The x-headroom-session-id header is honored only behind an
+            # explicit env opt-in: deployments whose gateways already stamp
+            # that header on ALL traffic (it is the documented proxy-path
+            # session key) would otherwise silently flip stateless callers
+            # into session mode on upgrade — and a header value shared across
+            # conversations (Claude Code subagents do exactly this) would
+            # blend unrelated conversations into one replay state.
+            if session_id is None and os.environ.get(
+                "HEADROOM_COMPRESS_SESSION_FROM_HEADER", ""
+            ).lower() in ("1", "true"):
+                session_id = request.headers.get("x-headroom-session-id")
+            if session_id is not None and (
+                not isinstance(session_id, str) or not session_id.strip() or len(session_id) > 256
+            ):
+                return JSONResponse(
+                    status_code=400,
+                    content={
+                        "error": {
+                            "type": "invalid_request",
+                            "message": (
+                                f"Invalid config.session_id: {session_id!r}. "
+                                "Expected a non-empty string of at most 256 characters."
+                            ),
+                        }
+                    },
+                )
+            if session_id is not None and compress_user_messages:
+                # User/assistant rewrites are not content-addressed (the
+                # session cache replays tool_result content only), so once the
+                # tracker's overlay snapshots expire a rewritten user message
+                # would come back in RAW form — a guaranteed prefix bust inside
+                # the tracker-TTL/cache-TTL window. Refuse the combination
+                # rather than bust later.
+                return JSONResponse(
+                    status_code=400,
+                    content={
+                        "error": {
+                            "type": "invalid_request",
+                            "message": (
+                                "config.compress_user_messages is not supported with "
+                                "config.session_id: user-message rewrites cannot be "
+                                "byte-replayed across turns, which would bust the "
+                                "provider prompt cache."
+                            ),
+                        }
+                    },
+                )
             # Mode selection. Default is marker-free (see _no_ccr_pipeline):
             # no caller of this route can resolve a CCR marker unless it opts in
             # with mode="ccr", which restores the full marker + store behaviour.
@@ -9701,23 +9762,121 @@ class OpenAIHandlerMixin:
             if frozen_message_count is not None:
                 pipeline_kwargs["frozen_message_count"] = frozen_message_count
 
-            # Offload the CPU-bound pipeline to the bounded compression executor
-            # (mirrors the request handlers above). Running apply() inline blocked
-            # the single event loop on a large payload, so even GET /health stalled
-            # until it finished (#718). The executor also enforces a timeout so a
-            # too-large body fails fast instead of hanging forever.
-            result = await self._run_compression_in_executor(
-                lambda: pipeline.apply(
-                    messages=messages,
-                    model=model,
-                    **pipeline_kwargs,
-                ),
+            # Sidecar session pre-work: swap in previously-computed compressed
+            # bytes (Zone 1), then freeze the ENTIRE locally-replayable prefix
+            # (`compute_frozen_count`). This deliberately differs from the
+            # proxy path's `min(tracker, cache)` posture: in sidecar mode,
+            # whatever this endpoint previously RETURNED is the provider's
+            # cache contract, so every already-returned message must come back
+            # byte-identical — recompressing it (even "better") is a bust.
+            # Over-freezing relative to the provider's actual cache only
+            # forgoes tail compression; it can never bust. The tracker's
+            # /v1/usage-fed freeze count is deliberately NOT a freeze floor —
+            # freezing a message whose cache entry was evicted would forward
+            # raw original bytes. An explicit config.frozen_message_count
+            # still wins when larger: the caller may know more about the
+            # provider cache than local state does.
+            comp_cache = None
+            session_tracker = None
+            if session_id:
+                # Namespaced with a NUL separator so sidecar sessions can
+                # never collide with proxy-path session ids: NUL cannot
+                # appear in an HTTP header value, so no client-supplied
+                # x-headroom-session-id on the proxy path can spoof its way
+                # into a sidecar session's tracker or replay cache (the same
+                # trick SessionTrackerStore uses for its synthetic lineage
+                # keys). A plain "compress:" string prefix was spoofable.
+                _session_key = f"compress\x00{session_id}"
+                _tracker_provider = (
+                    "anthropic"
+                    if ("claude" in model_name.lower() or "anthropic" in model_name.lower())
+                    else "openai"
+                )
+                comp_cache = self._get_compression_cache(_session_key)
+                session_tracker = self.session_tracker_store.get_or_create(
+                    _session_key, _tracker_provider
+                )
+
+            def _run_stateless():
+                result = pipeline.apply(messages=messages, model=model, **pipeline_kwargs)
+                return (
+                    result,
+                    result.messages,
+                    result.tokens_before,
+                    result.tokens_after,
+                    None,
+                )
+
+            def _run_session_turn():
+                # One sidecar turn as a single executor-side block: every step
+                # here is CPU-bound (content hashing, deep compares, token
+                # counts, full-transcript deepcopies) and must stay off the
+                # event loop for the same reason pipeline.apply does (#718).
+                # The per-session lock serializes contract-violating
+                # concurrent turns so an older in-flight turn cannot tear or
+                # overwrite a newer turn's tracker snapshots mid-flight.
+                from headroom.cache.prefix_tracker import overlay_cached_prefix
+
+                with comp_cache.session_turn_lock:
+                    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)
+                    pipeline_kwargs["frozen_message_count"] = session_frozen
+                    result = pipeline.apply(messages=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
+                    # Savings are reported against the caller's RAW payload,
+                    # not the cache-swapped pipeline input: on a warm turn the
+                    # swap has already shrunk the input before the pipeline
+                    # counts it, which made every warm turn report ~0 saved.
+                    try:
+                        from headroom.tokenizers import get_tokenizer
+
+                        _tok = get_tokenizer(model_name)
+                        raw_tokens_before = _tok.count_messages(messages)
+                        final_tokens_after = _tok.count_messages(final)
+                    except Exception:  # nosec B110 - fall back to pipeline counts
+                        raw_tokens_before = result.tokens_before
+                        final_tokens_after = result.tokens_after
+                    comp_cache.update_from_result(messages, final)
+                    # Record this turn's result as the new "last returned" —
+                    # the sidecar equivalent of "last forwarded", captured at
+                    # return time because whatever we hand back IS what the
+                    # caller sends upstream.
+                    session_tracker.record_returned(messages, final)
+                    info = {
+                        "id": session_id,
+                        "frozen_message_count": session_frozen,
+                        "cached_prefix_replayed": replayed,
+                    }
+                    return result, final, raw_tokens_before, final_tokens_after, info
+
+            # Offload the CPU-bound work to the bounded compression executor
+            # (mirrors the request handlers above). Running it inline blocked
+            # the single event loop on a large payload, so even GET /health
+            # stalled until it finished (#718). The executor also enforces a
+            # timeout so a too-large body fails fast instead of hanging.
+            (
+                result,
+                final_messages,
+                tokens_before,
+                tokens_after,
+                session_info,
+            ) = await self._run_compression_in_executor(
+                _run_session_turn if session_id else _run_stateless,
                 timeout=COMPRESSION_TIMEOUT_SECONDS,
             )
-            ccr_hashes = _response_ccr_hashes(result.messages, result.markers_inserted)
 
-            tokens_before = result.tokens_before
-            tokens_after = result.tokens_after
+            ccr_hashes = _response_ccr_hashes(final_messages, result.markers_inserted)
+
             tokens_saved = max(0, tokens_before - tokens_after)
             latency_ms = (time.time() - start_time) * 1000
             await self._record_request_outcome(
@@ -9749,28 +9908,56 @@ class OpenAIHandlerMixin:
                 )
             )
 
-            return JSONResponse(
-                {
-                    "messages": result.messages,
-                    "tokens_before": result.tokens_before,
-                    "tokens_after": result.tokens_after,
-                    "tokens_saved": result.tokens_before - result.tokens_after,
-                    "compression_ratio": (
-                        result.tokens_after / result.tokens_before
-                        if result.tokens_before > 0
-                        else 1.0
-                    ),
-                    "transforms_applied": result.transforms_applied,
-                    "transforms_summary": result.transforms_summary,
-                    "ccr_hashes": ccr_hashes,
-                }
-            )
+            _payload = {
+                "messages": final_messages,
+                "tokens_before": tokens_before,
+                "tokens_after": tokens_after,
+                # Clamped like the telemetry above: the overlay's byte-replay
+                # can legitimately return a slightly larger prefix than the
+                # pipeline's best effort, and a negative "saved" here while
+                # telemetry records 0 would be two answers for one number.
+                "tokens_saved": tokens_saved,
+                "compression_ratio": (tokens_after / tokens_before if tokens_before > 0 else 1.0),
+                "transforms_applied": result.transforms_applied,
+                "transforms_summary": result.transforms_summary,
+                "ccr_hashes": ccr_hashes,
+            }
+            if session_info is not None:
+                _payload["session"] = session_info
+            return JSONResponse(_payload)
         except TimeoutError:
+            self.metrics.record_compression_failed("timeout")
+            if session_id:
+                # Fail-open-with-originals is WRONG for a session call: the
+                # timed-out worker cannot be cancelled and may still finish
+                # and record its compressed result as "last returned" — while
+                # the caller, handed the originals, forwards those instead.
+                # The desynced snapshot then busts the next turn. A 503 tells
+                # the gateway to retry; the retry lands on whatever state the
+                # straggler recorded and replays it consistently.
+                logger.warning(
+                    "Compression timed out after %.0fs for session %r; "
+                    "returning 503 (session mode cannot fail open without "
+                    "desyncing replay state)",
+                    COMPRESSION_TIMEOUT_SECONDS,
+                    session_id,
+                )
+                return JSONResponse(
+                    status_code=503,
+                    content={
+                        "error": {
+                            "type": "compression_timeout",
+                            "message": (
+                                "Compression timed out; retry this turn. "
+                                "Session replay state remains consistent."
+                            ),
+                        }
+                    },
+                )
             logger.warning(
                 "Compression timed out after %.0fs; failing open with original messages",
                 COMPRESSION_TIMEOUT_SECONDS,
             )
-            self.metrics.record_compression_failed("timeout")
             latency_ms = (time.time() - start_time) * 1000
             await self._record_request_outcome(
                 RequestOutcome(
@@ -9820,6 +10007,119 @@ class OpenAIHandlerMixin:
                 },
             )
 
+    async def handle_compress_usage(self, request: Request) -> JSONResponse:
+        """Relay of the provider's usage block for a sidecar compress session.
+
+        POST /v1/usage
+        Body: {"session_id": "...",
+               "usage": {"cache_read_input_tokens": N,
+                         "cache_creation_input_tokens": N}}
+
+        The session-aware ``/v1/compress`` never sees the provider's response
+        (the caller owns routing). This relay feeds the provider-confirmed
+        numbers into the session's tracker — the same signal the proxy path
+        reads from the response itself — powering cache-hit/miss attribution,
+        idle-vs-prefix-change classification, and savings accounting for
+        sidecar sessions.
+
+        Deliberately NOT a freeze input: the compress path freezes exactly the
+        locally-replayable prefix (``compute_frozen_count``), and raising that
+        to a provider-confirmed count could freeze a message whose cache entry
+        was evicted — which would forward raw original bytes and bust the very
+        prefix the count vouched for. Optional: skipping this call costs
+        telemetry fidelity, never correctness.
+        """
+        from fastapi.responses import JSONResponse
+
+        from headroom.proxy.helpers import _read_request_json
+
+        def _invalid(message: str) -> JSONResponse:
+            return JSONResponse(
+                status_code=400,
+                content={"error": {"type": "invalid_request", "message": message}},
+            )
+
+        try:
+            body = await _read_request_json(request)
+        except Exception:
+            return _invalid("Invalid JSON in request body.")
+
+        session_id = body.get("session_id")
+        if not isinstance(session_id, str) or not session_id.strip() or len(session_id) > 256:
+            return _invalid(
+                "Missing or invalid session_id: expected a non-empty string "
+                "of at most 256 characters."
+            )
+        usage = body.get("usage")
+        if not isinstance(usage, dict):
+            return _invalid("Missing or invalid usage: expected an object.")
+        # A usage block carrying NEITHER cache field is a no-signal relay (an
+        # OpenAI-style {"prompt_tokens": N} forwarded verbatim, for example).
+        # Defaulting the absent fields to 0 would make update_from_response
+        # treat it as a provider-confirmed fully-cold turn and wipe the
+        # tracker's cached-prefix state — so absence of both is a 400, not 0.
+        if "cache_read_input_tokens" not in usage and "cache_creation_input_tokens" not in usage:
+            return _invalid(
+                "usage must carry cache_read_input_tokens and/or "
+                "cache_creation_input_tokens; a block with neither carries no "
+                "cache signal and is not accepted."
+            )
+
+        def _token_field(name: str) -> int | None:
+            value = usage.get(name, 0)
+            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+                return None
+            return value
+
+        cache_read = _token_field("cache_read_input_tokens")
+        cache_write = _token_field("cache_creation_input_tokens")
+        if cache_read is None or cache_write is None:
+            return _invalid(
+                "usage.cache_read_input_tokens and usage.cache_creation_input_tokens "
+                "must be non-negative integers when present."
+            )
+
+        # Same NUL-separated namespace as handle_compress: unspoofable from
+        # any HTTP header. peek() (never get_or_create) so a flood of novel
+        # session ids cannot grow the tracker store — an unknown 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.
+            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."
+                        ),
+                    }
+                },
+            )
+
+        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 JSONResponse(
+            {
+                "session_id": session_id,
+                "frozen_message_count": tracker.get_frozen_message_count(),
+            }
+        )
+
     async def _maybe_compress_passthrough_responses(
         self, body: bytes, *, client: str | None = None
     ) -> bytes:

headroom/proxy/helpers.py

@@ -1271,10 +1271,12 @@ try:
     # Floor of 600s: never below the prefix tracker's session TTL, or the
     # sweep could reclaim the byte-identical swap map while it is the only
     # remaining protection for a still-live provider prefix (see above).
-    COMPRESSION_CACHE_TTL_SECONDS = max(
-        600.0,
-        float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900")),
-    )
+    # Non-finite floats ("nan"/"inf") parse but poison every idle comparison,
+    # so they are rejected like any other unparseable value.
+    _ttl_env = float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900"))
+    if _ttl_env != _ttl_env or _ttl_env in (float("inf"), float("-inf")):
+        raise ValueError("non-finite TTL")
+    COMPRESSION_CACHE_TTL_SECONDS = max(600.0, _ttl_env)
 except ValueError:
     COMPRESSION_CACHE_TTL_SECONDS = 3900.0
 

headroom/proxy/server.py

@@ -5252,6 +5252,13 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
     async def compress_messages(request: Request):
         return await proxy.handle_compress(request)
 
+    # Sidecar-mode usage relay: same exposure policy as /v1/compress — the two
+    # form one contract (compress returns the bytes, usage reports what the
+    # provider said about them), so they must be reachable from the same place.
+    @app.post("/v1/usage", dependencies=_compress_dependencies)
+    async def compress_usage(request: Request):
+        return await proxy.handle_compress_usage(request)
+
     register_provider_routes(app, proxy)
 
     return app

tests/test_compress_session_mode.py

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

tests/test_org_scale_limits.py

@@ -75,6 +75,28 @@ def test_compression_cache_entry_cap_env_parsing(monkeypatch) -> None:
     assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 10000
 
 
+def test_compression_cache_ttl_env_rejects_non_finite(monkeypatch) -> None:
+    """'nan'/'inf' parse as floats but poison every idle comparison — they
+    must fall back to the default like any other unparseable value."""
+    import importlib
+
+    import headroom.proxy.helpers as helpers_mod
+
+    for bad in ("nan", "inf", "-inf"):
+        monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", bad)
+        importlib.reload(helpers_mod)
+        assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 3900.0, bad
+
+    # Below the 600s floor clamps up; above it passes through.
+    monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "60")
+    importlib.reload(helpers_mod)
+    assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 600.0
+
+    monkeypatch.delenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS")
+    importlib.reload(helpers_mod)
+    assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 3900.0
+
+
 # --------------------------------------------------------------------------- #
 # Frozen-verdicts store: process-wide, so it must be sizeable per deployment.  #
 # --------------------------------------------------------------------------- #

tests/test_proxy_compress_endpoint.py

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