"""Prefix Cache Tracker — session-scoped state for cache-aware compression. Tracks provider prefix cache state between turns so the transform pipeline can freeze already-cached messages and only compress new content. Problem: Clients like Claude Code already manage prefix caching (up to 4 cache_control breakpoints, growing-prefix strategy). If Headroom compresses or modifies messages in the cached prefix, it invalidates the cache — replacing a 90% read discount (Anthropic) or 50% (OpenAI) with a 25% write penalty. Solution: After each API response, record how many tokens the provider cached. On the next turn, freeze that many messages so the transform pipeline skips them entirely. """ from __future__ import annotations import copy import hashlib import json import logging import time from dataclasses import dataclass from typing import Any logger = logging.getLogger(__name__) # Provider cache economics for cost comparisons _PROVIDER_READ_DISCOUNT = { "anthropic": 0.9, # 90% discount on reads "openai": 0.5, # 50% discount on reads "gemini": 0.9, "bedrock": 0.9, } _PROVIDER_WRITE_PENALTY = { "anthropic": 0.25, # 25% surcharge on writes "openai": 0.0, # No write penalty "gemini": 0.0, "bedrock": 0.25, } # Default prompt-cache lifetime per provider, in seconds. Used by # `classify_cache_miss` to decide whether a miss is most likely a TTL # lapse (idle longer than this) versus a prefix-content change. Anthropic's # default ephemeral cache is 5 minutes (matches # headroom.cache.anthropic.ANTHROPIC_CACHE_TTL_SECONDS); the others are best- # effort defaults and only matter once those providers are wired in. A # session that opts into Anthropic's 1h cache breakpoint can override this # via the tracker config (see PrefixFreezeConfig.cache_ttl_seconds). _PROVIDER_CACHE_TTL_SECONDS = { "anthropic": 300, # 5 minutes (default ephemeral cache) "openai": 300, # automatic prefix cache, ~5-10 min; conservative floor "gemini": 300, "bedrock": 300, } @dataclass class PrefixFreezeConfig: """Configuration for cache-aware prefix freezing.""" enabled: bool = True min_cached_tokens: int = 1024 # Min cached tokens to activate freeze session_ttl_seconds: int = 600 # Session tracker cleanup TTL force_compress_threshold: float = 0.5 # Bust cache if compression saves > this fraction # Provider prompt-cache lifetime used by `classify_cache_miss` to tell a # TTL lapse from a prefix change. `None` falls back to the per-provider # default in `_PROVIDER_CACHE_TTL_SECONDS`. Set to 3600 for a session that # uses Anthropic's 1h cache breakpoint so idle-gap attribution stays honest. cache_ttl_seconds: int | None = None @dataclass class FreezeStats: """Statistics from prefix freezing for metrics/dashboard.""" busts_avoided: int = 0 tokens_preserved: int = 0 compression_foregone_tokens: int = 0 net_benefit_tokens: int = 0 # tokens_preserved - compression_foregone frozen_message_count: int = 0 turn_number: int = 0 # Cache-miss attribution verdicts. `reason` is one of these literals so # metrics/dashboard can bucket without re-deriving the logic. See # PrefixCacheTracker.classify_cache_miss. MISS_TTL_EXPIRY = "ttl_expiry" MISS_PREFIX_CHANGE = "prefix_change" MISS_COLD_START = "cold_start" # no prior cached prefix to miss against MISS_UNKNOWN = "unknown" # expected a hit, content stable, idle within TTL @dataclass class CacheMissAttribution: """Why a turn that expected a prompt-cache hit missed instead. Produced by :meth:`PrefixCacheTracker.classify_cache_miss`. ``is_miss`` is False when the turn actually hit cache (or there was nothing to hit), in which case ``reason`` is informational only. """ is_miss: bool reason: str # one of the MISS_* literals idle_seconds: float = 0.0 cache_ttl_seconds: int = 0 expected_cached_tokens: int = 0 cache_read_tokens: int = 0 prefix_changed: bool = False ttl_exceeded: bool = False def _strip_cache_control(obj: Any) -> Any: """Recursively drop ``cache_control`` for content-only equality checks. Clients (notably Claude Code) move the cache_control breakpoint to the newest message on every call, so the exact same message carries cache_control on one turn and not the next. That per-call annotation must be ignored when deciding whether this turn append-only-extends the previous one — otherwise a moved marker spuriously fails the check and we skip the byte-identical replay, busting the cache.""" if isinstance(obj, dict): return {k: _strip_cache_control(v) for k, v in obj.items() if k != "cache_control"} if isinstance(obj, list): return [_strip_cache_control(v) for v in obj] return obj # Keys that carry NO semantic payload for the model — transport / caching-directive # / telemetry / client-routing annotations that clients attach and vary turn-to-turn. # Grounded in provider API docs (Anthropic Messages, OpenAI Chat+Responses, Bedrock # Converse) + client-library field inventories (litellm, Vercel AI SDK, opencode, # Claude Code, Cline). Dropped from the cross-turn prefix-equality key ONLY. # # NOTE ON SAFETY: this projection is a COMPARISON KEY, never a source to rebuild # forwarded bytes — the cache-stable-delta path always forwards the previously # forwarded bytes + the raw appended delta. So dropping these can't deprive the # model. What we must NOT do is drop a *semantic* field (that would mask a real # divergence and replay a stale prefix), which is why: (1) reasoning SIGNATURES are # NOT in this set (Anthropic 400s if a thinking block is altered/missing, and a # present/absent flip is a real divergence we want to detect); (2) tool inputs / # arguments / json payloads are treated as OPAQUE and compared verbatim (see # _OPAQUE_PAYLOAD_KEYS) so a user key that happens to be named "index"/"state" is # never stripped from inside a tool call. _NON_SEMANTIC_KEYS = frozenset( { # cache-breakpoint markers (moved to the newest block every turn) "cache_control", # Anthropic (per-block) "cachePoint", # Bedrock (per-block content block) # litellm unified-message / tool annotations "caller", # litellm programmatic-tool tag on tool_use "provider_specific_fields", "reasoning_content", # litellm display echo (the paired signature is separate) "reasoning_items", "annotations", # citation/display metadata # OpenAI response echoes that can ride on assistant messages "system_fingerprint", "service_tier", # Vercel AI SDK / opencode part transport "providerMetadata", "providerOptions", "callProviderMetadata", "state", "providerExecuted", "synthetic", "ignored", # streaming-assembly artifact "index", } ) # Values under these keys are opaque semantic payloads (tool-call input, OpenAI # stringified arguments, Bedrock tool_result json). They are compared VERBATIM — we # never recurse into them to strip "noise" keys, because arbitrary user data there # may legitimately contain keys that collide with _NON_SEMANTIC_KEYS (e.g. an # `input` of {"state": "CA", "index": 3}). Recursing would corrupt the comparison. _OPAQUE_PAYLOAD_KEYS = frozenset({"input", "arguments", "json"}) def _canonicalize_for_prefix_compare(obj: Any) -> Any: """Representation-agnostic canonical form for cross-turn prefix equality. Providers accept several *equivalent* encodings for the same message, and real clients vary them turn-to-turn; a raw-dict prefix compare then fails spuriously and drops cache mode to raw (uncompressed) forwarding. This normalizes ONLY representation: * drops non-semantic annotation / cache-directive / telemetry keys (_NON_SEMANTIC_KEYS) at any message/block level; * wraps a bare string ``content`` into ``[{"type": "text", "text": ...}]`` (Anthropic's string sugar, which litellm flips per turn); * leaves tool ``input`` / ``arguments`` / ``json`` payloads verbatim (_OPAQUE_PAYLOAD_KEYS) so user data is never corrupted; * KEEPS all real content (text, tool name/input, tool_result content, reasoning signatures, ids) so two messages canonicalize-equal iff they are semantically identical. Used ONLY as a comparison key for the cache-stable delta path; the original, unmodified messages are always what gets forwarded. """ if isinstance(obj, dict): out: dict[str, Any] = {} for key, value in obj.items(): if key in _NON_SEMANTIC_KEYS: continue if key in _OPAQUE_PAYLOAD_KEYS: out[key] = value # verbatim — do not recurse into user payloads elif key == "content" and isinstance(value, str): out[key] = [{"type": "text", "text": value}] else: out[key] = _canonicalize_for_prefix_compare(value) return out if isinstance(obj, list): canon = [_canonicalize_for_prefix_compare(value) for value in obj] # Drop blocks that projected to {} — a pure cache-directive content block # (e.g. Bedrock {"cachePoint": {...}}) whose only key was non-semantic. Left # in place it would be an empty-dict entry, so a directive block moving # position across turns would spuriously fail the length/order compare. return [value for value in canon if value != {}] return obj def extract_cache_stable_delta( current_messages: list[dict[str, Any]], previous_original_messages: list[dict[str, Any]] | None, previous_forwarded_messages: list[dict[str, Any]] | None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: """Return ``(stable_forwarded_prefix, appended_delta_messages)`` when the current request append-only-extends the previous one, else ``None``. Provider-agnostic delta engine for cache mode. "Append-only" is decided by comparing the *canonicalized* prefix (:func:`_canonicalize_for_prefix_compare`, which ignores per-turn transport / cache-directive / client-annotation noise across Anthropic / OpenAI / Bedrock and the common clients), so a moved cache marker or shape churn does not spuriously collapse cache mode to raw forwarding. On a match the caller replays the byte-identical previously-forwarded prefix and compresses ONLY the appended delta. This is a COMPARISON + slice only: the returned prefix is the previously-forwarded bytes verbatim and the delta is the raw appended messages — never a rebuild from the canonical projection — so the projection dropping non-semantic fields is safe. """ if not previous_original_messages or previous_forwarded_messages is None: return None prefix_len = len(previous_original_messages) if len(current_messages) < prefix_len: return None if _canonicalize_for_prefix_compare( current_messages[:prefix_len] ) != _canonicalize_for_prefix_compare(previous_original_messages): return None return ( copy.deepcopy(previous_forwarded_messages), copy.deepcopy(current_messages[prefix_len:]), ) def overlay_cached_prefix( optimized_messages: list[dict[str, Any]], current_original_messages: list[dict[str, Any]], previous_original_messages: list[dict[str, Any]] | None, previous_forwarded_messages: list[dict[str, Any]] | None, ) -> list[dict[str, Any]]: """Replay the previously-forwarded (cached, compressed) prefix byte-identical. Provider-agnostic cache-safety guard for the freeze path. When a message is "frozen", the compression pipeline may emit the agent's ORIGINAL bytes for it — but the provider cached whatever we FORWARDED last turn (the compressed form). Forwarding the original then mismatches the cached prefix and busts the prompt cache from that point (100% of observed misses were this ``prefix_change``). This overlays the exact previously-forwarded prefix onto the corresponding leading messages so the forwarded prefix stays byte-for-byte what the provider hashed for its cache key. Safe only when this turn append-only-extends the previous turn (the standard growing-conversation shape): the previous ORIGINAL messages must be an exact prefix of the current ORIGINAL messages, and there is exactly one forwarded message per original. Otherwise the previous forwarded bytes may not correspond to the same positions, so we return ``optimized_messages`` unchanged (accept a possible bust rather than forward wrong content). This makes freezing byte-identical in BOTH proxy modes, so the only remaining difference between them is how large a mutable (still-compressible) tail each leaves — not whether the frozen prefix busts the cache. """ prev_orig = previous_original_messages prev_fwd = previous_forwarded_messages if not prev_orig or not prev_fwd: return optimized_messages n = len(prev_orig) # Positional 1:1 correspondence between prev_orig[i] and prev_fwd[i] holds # only when last turn forwarded exactly one message per original (the # append-only, no-injection shape update_from_response records). If the # counts differ, an injected / dropped / merged message shifted the # mapping, so replaying prev_fwd[i] at position i could forward the wrong # content — bail (leave this turn's output untouched) rather than risk it. if len(prev_fwd) != n: logger.debug( "overlay: forwarded/original count mismatch (prev_fwd=%d, prev_orig=%d) " "— skipping cached-prefix replay (possible bust)", len(prev_fwd), n, ) return optimized_messages # Append-only guard on CONTENT ONLY, message-by-message. Replay the # previously-forwarded (cached, compressed) bytes for the longest LEADING # run of messages that is byte-for-byte (content-canonical) identical to # what we forwarded last turn, and stop at the first divergence. # # This is the cache-safety centerpiece for token mode (which relies solely # on this replay; cache mode is already byte-stable by construction). The # prior all-or-nothing guard busted the ENTIRE cached prefix the moment any # single leading message failed to canonicalize-equal last turn — most # commonly the just-added assistant turn, whose client-resent form can # differ trivially from the copy we reconstructed and recorded. Stopping at # the first divergence instead keeps the (much larger) cache-hit region # up to that point and only re-forwards from the changed message onward. # # Comparison uses the shared canonicalizer (not just cache_control # stripping) so it is robust to ALL per-turn transport / annotation churn — # cache_control movement (Anthropic), litellm `caller`, # provider_specific_fields, streaming `index`, string<->block content shape, # etc. Content stability is what the provider's prefix cache actually keys # on. Safe by construction: we only replay prev_fwd[k] where # current_original[k] canonicalize-equals prev_orig[k], and prev_fwd[k] # positionally corresponds to prev_orig[k] (guaranteed by the count check # above), so no wrong bytes are ever forwarded. limit = min(n, len(current_original_messages), len(optimized_messages)) k = 0 while k < limit and _canonicalize_for_prefix_compare( current_original_messages[k] ) == _canonicalize_for_prefix_compare(prev_orig[k]): k += 1 if k == 0: logger.debug( "overlay: prefix diverged at message 0 — no cached-prefix replay " "(cold prefix or client rewrote history head)" ) return optimized_messages if k < n: logger.debug( "overlay: cached-prefix replay for %d/%d leading messages " "(diverged at %d — re-forwarding tail fresh)", k, n, k, ) # Replay the cached (compressed) prefix byte-identical up to the first # divergence; keep this turn's freshly-produced output for the rest. return list(prev_fwd[:k]) + list(optimized_messages[k:]) def normalize_message_cache_control( messages: list[dict[str, Any]], ) -> list[dict[str, Any]]: """Own message-level cache_control placement so breakpoints stay bounded. Two forces pile up cache_control markers turn over turn: clients move the breakpoint to the newest message each call, and ``overlay_cached_prefix`` replays the markers that rode on each turn's then-newest message. Anthropic hard-errors at **>4 cache_control blocks total** (system + tools + messages), so on a long conversation the accumulation eventually 400s. Fix: strip EVERY message-level cache_control and re-place a **single** ephemeral breakpoint on the last block of the last block-style message. One breakpoint caches the whole message prefix up to it, and — because the provider's cache key is message CONTENT, not marker presence (moving the breakpoint forward is the documented client pattern and it hits) — stripping and re-placing markers never busts. system/tools breakpoints live outside ``messages`` and are left untouched (they still count toward the 4 limit, so holding messages to one breakpoint leaves room for them). Only block-style (list) content can carry cache_control; string content is left as-is. Returns the input unchanged when there is nothing to normalize. """ changed = False out: list[dict[str, Any]] = [] last_block_idx = -1 for i, msg in enumerate(messages): content = msg.get("content") if isinstance(msg, dict) else None if isinstance(content, list): had = any(isinstance(b, dict) and "cache_control" in b for b in content) stripped = [ {k: v for k, v in b.items() if k != "cache_control"} if isinstance(b, dict) else b for b in content ] out.append({**msg, "content": stripped} if had else msg) changed = changed or had if stripped and isinstance(stripped[-1], dict): last_block_idx = i else: out.append(msg) # Re-place exactly one breakpoint on the last block-style message. if last_block_idx >= 0: msg = out[last_block_idx] content = list(msg["content"]) content[-1] = {**content[-1], "cache_control": {"type": "ephemeral"}} out[last_block_idx] = {**msg, "content": content} changed = True return out if changed else messages class PrefixCacheTracker: """Tracks provider prefix cache state across turns in a session. Usage: tracker = PrefixCacheTracker("anthropic") # Before compression (turn 2+): frozen = tracker.get_frozen_message_count() result = pipeline.apply(messages, model, frozen_message_count=frozen) # After API response: tracker.update_from_response( cache_read_tokens=usage["cache_read_input_tokens"], cache_write_tokens=usage["cache_creation_input_tokens"], messages=optimized_messages, tokenizer=tokenizer, ) """ def __init__(self, provider: str, config: PrefixFreezeConfig | None = None): self.provider = provider self.config = config or PrefixFreezeConfig() self._cached_token_count: int = 0 self._cached_message_count: int = 0 self._turn_number: int = 0 self._last_activity: float = time.time() self._last_original_messages: list[dict[str, Any]] = [] self._last_forwarded_messages: list[dict[str, Any]] = [] # Idle gap (seconds) since the PREVIOUS turn's response, captured by # SessionTrackerStore.get_or_create at fetch time — BEFORE it refreshes # _last_activity. Without this snapshot, seconds_since_activity() reads # ~0 on every request (the fetch itself bumps the clock), so the # net-cost/TTL P_alive gate could never see idle time. The handler reads # this and forwards it to the pipeline as `idle_seconds`. self._idle_seconds_at_fetch: float = 0.0 # Session-scoped ReadMaturationManager (Mechanism B), created # lazily by the handler when read maturation is enabled. Rides # here so it shares the session's affinity and TTL cleanup. self.read_maturation_manager: Any = None # Stats self._busts_avoided: int = 0 self._tokens_preserved: int = 0 self._compression_foregone_tokens: int = 0 def get_frozen_message_count(self) -> int: """How many leading messages to skip compression on the next turn. Returns 0 on turn 0 (cold start) or if caching is disabled/below threshold. """ if not self.config.enabled: return 0 if self._turn_number == 0: return 0 if self._cached_token_count < self.config.min_cached_tokens: return 0 return self._cached_message_count def update_from_response( self, cache_read_tokens: int, cache_write_tokens: int, messages: list[dict[str, Any]], message_token_counts: list[int] | None = None, original_messages: list[dict[str, Any]] | None = None, ) -> None: """Update tracker with cache metrics from the API response. Called after every API call. Computes how many messages to freeze on the next turn based on the cache_read_tokens reported. Args: cache_read_tokens: Tokens read from cache (cache hit portion). cache_write_tokens: Tokens written to cache (new cache entries). messages: The messages that were sent to the API. message_token_counts: Pre-computed token counts per message. If None, estimates from content length. """ self._last_activity = time.time() self._turn_number += 1 self._last_original_messages = copy.deepcopy(original_messages or messages) self._last_forwarded_messages = copy.deepcopy(messages) # Compute total cached tokens (read + write = what's in cache now) total_cached = cache_read_tokens + cache_write_tokens if total_cached == 0: self._cached_token_count = 0 self._cached_message_count = 0 return # Estimate per-message token counts if not provided if message_token_counts is None: message_token_counts = self._estimate_message_tokens(messages) # Walk messages from the start, accumulating tokens until we exceed # the cached amount. All messages within the cached prefix are frozen. accumulated = 0 frozen_count = 0 for i, tok_count in enumerate(message_token_counts): accumulated += tok_count if accumulated <= total_cached: frozen_count = i + 1 else: break self._cached_token_count = total_cached self._cached_message_count = frozen_count logger.debug( "PrefixCacheTracker[%s]: turn=%d, cached=%d tokens, " "frozen=%d/%d messages (read=%d, write=%d)", self.provider, self._turn_number, total_cached, frozen_count, len(messages), cache_read_tokens, cache_write_tokens, ) def get_last_original_messages(self) -> list[dict[str, Any]]: return copy.deepcopy(self._last_original_messages) def get_last_forwarded_messages(self) -> list[dict[str, Any]]: return copy.deepcopy(self._last_forwarded_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: return self.config.cache_ttl_seconds return _PROVIDER_CACHE_TTL_SECONDS.get(self.provider, 300) def classify_cache_miss( self, cache_read_tokens: int, current_forwarded_messages: list[dict[str, Any]], idle_seconds: float | None = None, ) -> CacheMissAttribution: """Attribute *this turn's* cache outcome: hit, TTL lapse, or prefix change. Call this BEFORE :meth:`update_from_response` — it reads the state captured from the *previous* turn (``_cached_token_count``, ``_last_forwarded_messages``, ``_last_activity``), all of which ``update_from_response`` overwrites. Attribution only fires when the previous turn left a cacheable prefix (``_cached_token_count > 0``); the very first warm turn has nothing to miss against, so it is reported as ``cold_start`` with ``is_miss=False``. When a hit was expected but ``cache_read_tokens == 0``: * If the idle gap since the last turn exceeded the provider cache TTL, the cache entry had already lapsed — ``ttl_expiry``. **TTL wins ties:** once the entry expired, a coincident prefix change is moot (the issue asks "should I move 5m → 1h?", which only the TTL signal answers). * Otherwise, if the forwarded prefix changed versus last turn, the new bytes couldn't match the cached prefix — ``prefix_change``. * If neither signal fires (stable prefix, within TTL) we can't explain it from local state — ``unknown`` (e.g. provider-side eviction). A partial read (``0 < cache_read_tokens``) counts as a hit here; the existing model-aware bust detection in PrometheusMetrics already covers partial-invalidation accounting, and double-counting it as a "miss" would muddy the 5m-vs-1h signal this method exists to provide. Returns a :class:`CacheMissAttribution`; ``is_miss`` is False for hits and cold starts. """ if idle_seconds is None: idle_seconds = self.seconds_since_activity() ttl = self.resolved_cache_ttl_seconds() expected = self._cached_token_count # Nothing was cached last turn → cold start, not a miss. if expected <= 0: return CacheMissAttribution( is_miss=False, reason=MISS_COLD_START, idle_seconds=idle_seconds, cache_ttl_seconds=ttl, expected_cached_tokens=expected, cache_read_tokens=cache_read_tokens, ) # We expected a hit. A non-zero read means the prefix cache worked. if cache_read_tokens > 0: return CacheMissAttribution( is_miss=False, reason="hit", idle_seconds=idle_seconds, cache_ttl_seconds=ttl, expected_cached_tokens=expected, cache_read_tokens=cache_read_tokens, ) # Full miss on a prefix we expected cached. Attribute it. ttl_exceeded = idle_seconds > ttl prefix_changed = not self._forwarded_prefix_stable(current_forwarded_messages) if ttl_exceeded: reason = MISS_TTL_EXPIRY # TTL wins ties (see docstring) elif prefix_changed: reason = MISS_PREFIX_CHANGE else: reason = MISS_UNKNOWN return CacheMissAttribution( is_miss=True, reason=reason, idle_seconds=idle_seconds, cache_ttl_seconds=ttl, expected_cached_tokens=expected, cache_read_tokens=cache_read_tokens, prefix_changed=prefix_changed, ttl_exceeded=ttl_exceeded, ) def _forwarded_prefix_stable(self, current_forwarded_messages: list[dict[str, Any]]) -> bool: """True if last turn's forwarded prefix is still an exact prefix of this turn's. The cached prefix is whatever we forwarded last turn. If those exact messages still lead the current forwarded list, the bytes the provider hashed for its cache key are unchanged, so a miss can't be blamed on content. Anything else (a frozen message rewritten, the prefix reordered, the list now shorter) counts as a prefix change. """ prev = self._last_forwarded_messages if not prev: # No recorded prefix to compare — can't claim it changed. return True if len(current_forwarded_messages) < len(prev): return False return current_forwarded_messages[: len(prev)] == prev def record_bust_avoided(self, tokens_preserved: int, compression_foregone: int) -> None: """Record when we chose to preserve cache over compressing.""" self._busts_avoided += 1 self._tokens_preserved += tokens_preserved self._compression_foregone_tokens += compression_foregone def should_force_compress( self, message_index: int, message_tokens: int, estimated_compressed_tokens: int, ) -> bool: """Check if compression savings outweigh cache preservation. Returns True if we should bust the cache and compress anyway. This happens when compression would save a large fraction of tokens AND the savings exceed the cache read discount. """ if message_index >= self._cached_message_count: return True # Not in frozen prefix, always compress if message_tokens == 0: return False savings_fraction = (message_tokens - estimated_compressed_tokens) / message_tokens # Would compression savings exceed the cache read discount? read_discount = _PROVIDER_READ_DISCOUNT.get(self.provider, 0.5) return savings_fraction > read_discount @property def is_expired(self) -> bool: """Check if this tracker has been idle beyond TTL.""" return (time.time() - self._last_activity) > self.config.session_ttl_seconds def seconds_since_activity(self) -> float: """Wall-clock seconds since this tracker last saw activity. #856 P3b feeds this to the net-cost gate as an idle signal: as it approaches the provider's prompt-cache TTL (~300s for Anthropic), P_alive decays toward 0 and deep edits near cache lapse become free. Distinct from :attr:`is_expired`, which uses the much longer session-tracker *cleanup* TTL (``session_ttl_seconds``), not the cache TTL. Wiring caveat: ``SessionTrackerStore.get_or_create`` refreshes ``_last_activity`` on access, so a caller that wants the idle gap since the *previous turn's response* must read this before fetching the tracker for the current request (or the store must capture it at fetch time). ``update_from_response`` is the per-turn activity stamp. """ return max(0.0, time.time() - self._last_activity) @property def stats(self) -> FreezeStats: """Return stats for dashboard/metrics.""" return FreezeStats( busts_avoided=self._busts_avoided, tokens_preserved=self._tokens_preserved, compression_foregone_tokens=self._compression_foregone_tokens, net_benefit_tokens=self._tokens_preserved - self._compression_foregone_tokens, frozen_message_count=self._cached_message_count, turn_number=self._turn_number, ) @staticmethod def _estimate_message_tokens(messages: list[dict[str, Any]]) -> list[int]: """Rough token count per message (chars / 3.5). Counts text, tool_result content, and tool_use input fields for accurate Anthropic-format estimation. """ counts = [] for msg in messages: content = msg.get("content", "") if isinstance(content, str): chars = len(content) elif isinstance(content, list): chars = 0 for block in content: if not isinstance(block, dict): continue block_type = block.get("type", "") if block_type == "text": chars += len(block.get("text", "")) elif block_type == "tool_result": inner = block.get("content", "") if isinstance(inner, str): chars += len(inner) elif isinstance(inner, list): chars += sum( len(b.get("text", "")) for b in inner if isinstance(b, dict) ) elif block_type == "tool_use": inp = block.get("input") if isinstance(inp, str): chars += len(inp) elif isinstance(inp, dict): chars += len(json.dumps(inp, separators=(",", ":"))) else: text = block.get("text", "") if text: chars += len(text) else: chars = 0 # OpenAI function-calling: the assistant's command lives in the # top-level `tool_calls` (or legacy `function_call`) field, NOT in # `content` (which is empty/None on a tool-call turn). Anthropic puts # the equivalent in a `tool_use` content BLOCK (counted above), but # the OpenAI shape was never counted here. That under-counted every # tool-based assistant turn to ~0, so the frozen-prefix estimate # overshot the real cache boundary and froze the NEWEST delta — which # is why OpenAI/Kimi (fireworks) tool harnesses got ~zero compression # while text/back-tick harnesses (command in `content`) compressed. for tc in msg.get("tool_calls") or []: if isinstance(tc, dict): fn = tc.get("function") or {} chars += len(str(fn.get("name", ""))) + len(str(fn.get("arguments", ""))) fc = msg.get("function_call") if isinstance(fc, dict): chars += len(str(fc.get("name", ""))) + len(str(fc.get("arguments", ""))) # Add overhead for role, block structure, etc. chars += 20 counts.append(max(1, int(chars / 3.5))) return counts class SessionTrackerStore: """Manages PrefixCacheTracker instances across sessions. Keyed by session ID (from x-headroom-session-id header or computed hash). Automatically cleans up expired sessions. """ def __init__(self, default_config: PrefixFreezeConfig | None = None): self._trackers: dict[str, PrefixCacheTracker] = {} self._default_config = default_config or PrefixFreezeConfig() self._last_cleanup: float = time.time() self._cleanup_interval: float = 60.0 # Cleanup every 60s 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() if session_id in self._trackers: tracker = self._trackers[session_id] # Snapshot idle-since-last-response BEFORE bumping the access clock, # so the net-cost/TTL gate sees the true gap (see the attribute's # docstring in PrefixCacheTracker.__init__). tracker._idle_seconds_at_fetch = max(0.0, time.time() - tracker._last_activity) tracker._last_activity = time.time() return tracker tracker = PrefixCacheTracker(provider, self._default_config) tracker._idle_seconds_at_fetch = 0.0 # cold start: nothing cached to lapse self._trackers[session_id] = tracker return tracker def compute_session_id( self, request: Any, model: str, messages: list[dict[str, Any]], ) -> str: """Compute a session ID from the request. Priority: 1. x-headroom-session-id header (explicit) 2. Hash of (model + system prompt) — stable per conversation The system prompt is harvested from ``role:"system"`` entries in ``messages``. Anthropic carries the system prompt as a top-level ``body["system"]`` field instead, so its handler prepends that as a synthetic ``role:"system"`` message before calling this — otherwise every Anthropic conversation on the same model would collapse to one session id and their session-sticky state (CCR/memory tools, beta headers, frozen prefix) would cross-contaminate. """ # Check for explicit session header if hasattr(request, "headers"): session_header = request.headers.get("x-headroom-session-id") if session_header: return str(session_header) # Fall back to hashing model + all system-text content. system_parts: list[str] = [] for msg in messages: if msg.get("role") == "system": content = msg.get("content", "") if isinstance(content, str): system_parts.append(content) elif isinstance(content, list): for block in content: if isinstance(block, dict) and block.get("type") == "text": system_parts.append(block.get("text", "")) system_content = json.dumps(system_parts, ensure_ascii=False, separators=(",", ":")) key = f"{model}:{system_content}" return hashlib.md5(key.encode()).hexdigest()[:16] # nosec B324 def _maybe_cleanup(self) -> None: """Remove expired trackers periodically.""" now = time.time() if now - self._last_cleanup < self._cleanup_interval: return expired = [sid for sid, tracker in self._trackers.items() if tracker.is_expired] for sid in expired: del self._trackers[sid] if expired: logger.debug("SessionTrackerStore: cleaned up %d expired sessions", len(expired)) self._last_cleanup = now @property def active_sessions(self) -> int: """Number of active session trackers.""" return len(self._trackers)