diff --git a/headroom/cache/ttl_observations.py b/headroom/cache/ttl_observations.py new file mode 100644 index 000000000..148a2640f --- /dev/null +++ b/headroom/cache/ttl_observations.py @@ -0,0 +1,149 @@ +"""Cache-TTL learning seam — the OSS half of the cross-provider TTL learner. + +The proxy already attributes each turn's cache outcome +(:meth:`PrefixCacheTracker.classify_cache_miss`: ``hit`` / ``ttl_expiry`` / +``prefix_change`` / ``cold_start``). This module is the thin seam around that: + +1. :func:`record_cache_observation` appends each attribution to a local JSONL + (``cache_ttl_observations.jsonl``) when ``HEADROOM_CACHE_TTL_LEARN`` is set — + one row per turn: provider, model, idle, reason, hit/miss. +2. :func:`resolve_learned_ttl` reads a per-provider/model TTL table + (``cache_ttl_learned.json``) that an **offline** estimator (the + ``headroom-cache-ttl`` plugin) writes from those observations. + +``is_cold_prefix`` consults the learned TTL so cold detection is accurate for +providers that don't expose their TTL (Kimi / OpenAI / Codex). Claude Code reads +its TTL from config directly (see :func:`cold_prefix.anthropic_cache_ttl_seconds`), +so it doesn't need this — but its observations still feed the learner as ground +truth. The estimator lives out-of-process (offline batch over the JSONL) so there +is zero live-feedback risk on the hot path; here we only record + read. Every +function is best-effort and never raises. +""" + +from __future__ import annotations + +import json +import os +import time +from typing import Any + +_TRUTHY = ("1", "true", "yes", "on") +_LEARN_ENV = "HEADROOM_CACHE_TTL_LEARN" +_HEADROOM_DIR = os.path.expanduser("~/.headroom") +_OBS_DEFAULT = os.path.join(_HEADROOM_DIR, "cache_ttl_observations.jsonl") +_LEARNED_DEFAULT = os.path.join(_HEADROOM_DIR, "cache_ttl_learned.json") + + +def observations_enabled() -> bool: + return os.environ.get(_LEARN_ENV, "").strip().lower() in _TRUTHY + + +def _obs_path() -> str: + return os.environ.get("HEADROOM_CACHE_TTL_OBS_PATH") or _OBS_DEFAULT + + +def _learned_path() -> str: + return os.environ.get("HEADROOM_CACHE_TTL_LEARNED_PATH") or _LEARNED_DEFAULT + + +def record_cache_observation(*, provider: str, model: str, attribution: Any) -> None: + """Append one turn's cache outcome to the observation log (best-effort). + + Gated by ``HEADROOM_CACHE_TTL_LEARN`` (off by default → no file writes). + ``attribution`` is a ``CacheMissAttribution`` (duck-typed via getattr, so a + plain object works too). Records hits AND misses — the estimator needs the + idle gaps that still HIT to bound the TTL from below, and the ``ttl_expiry`` + misses to bound it from above. Never raises. + """ + if not observations_enabled(): + return + try: + row = { + "ts": round(time.time(), 3), + "provider": provider, + "model": model, + "reason": getattr(attribution, "reason", None), + "idle_seconds": round(float(getattr(attribution, "idle_seconds", 0.0) or 0.0), 1), + "ttl_assumed": int(getattr(attribution, "cache_ttl_seconds", 0) or 0), + "is_miss": bool(getattr(attribution, "is_miss", False)), + "cache_read": int(getattr(attribution, "cache_read_tokens", 0) or 0), + "expected_cached": int(getattr(attribution, "expected_cached_tokens", 0) or 0), + } + os.makedirs(os.path.dirname(_obs_path()), exist_ok=True) + with open(_obs_path(), "a", encoding="utf-8") as f: + f.write(json.dumps(row) + "\n") + except Exception: + pass # observability must never break a request + + +_learned_cache: dict[str, Any] | None = None +_learned_mtime: float = -1.0 + + +def resolve_learned_ttl(provider: str, model: str) -> int | None: + """Learned TTL (seconds) for ``(provider, model)``, or None if not learned yet. + + Reads ``cache_ttl_learned.json`` (written by the offline estimator), keyed by + ``"provider/model"`` first, then ``"provider"``. Cached by mtime so a hot-path + call is a dict lookup. Best-effort: returns None on any error, so the caller + falls back to its own default (config TTL for CC, static default elsewhere). + """ + global _learned_cache, _learned_mtime + path = _learned_path() + try: + mt = os.path.getmtime(path) + if _learned_cache is None or mt != _learned_mtime: + with open(path, encoding="utf-8") as f: + _learned_cache = json.load(f) + _learned_mtime = mt + except Exception: + return None + table = _learned_cache or {} + for key in (f"{provider}/{model}", provider): + v = table.get(key) + if isinstance(v, dict) and isinstance(v.get("ttl_seconds"), int | float): + return int(v["ttl_seconds"]) + if isinstance(v, int | float): + return int(v) + return None + + +def _demo() -> None: + import tempfile + + d = tempfile.mkdtemp() + os.environ["HEADROOM_CACHE_TTL_OBS_PATH"] = os.path.join(d, "obs.jsonl") + os.environ["HEADROOM_CACHE_TTL_LEARNED_PATH"] = os.path.join(d, "learned.json") + + class _Attr: + reason = "ttl_expiry" + idle_seconds = 420.0 + cache_ttl_seconds = 300 + is_miss = True + cache_read_tokens = 0 + expected_cached_tokens = 9000 + + # gated off by default -> no write + os.environ.pop("HEADROOM_CACHE_TTL_LEARN", None) + record_cache_observation(provider="openai", model="gpt-5.5", attribution=_Attr()) + assert not os.path.exists(os.environ["HEADROOM_CACHE_TTL_OBS_PATH"]), ( + "must not write when disabled" + ) + + os.environ["HEADROOM_CACHE_TTL_LEARN"] = "1" + record_cache_observation(provider="openai", model="gpt-5.5", attribution=_Attr()) + rows = [json.loads(x) for x in open(os.environ["HEADROOM_CACHE_TTL_OBS_PATH"])] + assert rows[0]["reason"] == "ttl_expiry" and rows[0]["idle_seconds"] == 420.0, rows[0] + + # learned-table read (exact model, then provider fallback) + with open(os.environ["HEADROOM_CACHE_TTL_LEARNED_PATH"], "w") as f: + json.dump({"openai/gpt-5.5": {"ttl_seconds": 1800}, "openai": 900}, f) + assert resolve_learned_ttl("openai", "gpt-5.5") == 1800 + assert resolve_learned_ttl("openai", "gpt-4o") == 900 # provider fallback + assert resolve_learned_ttl("anthropic", "claude-x") is None # not learned + + print("ttl_observations self-check OK") + + +if __name__ == "__main__": + _demo() diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 05dd6fdd8..e255ab886 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -1114,6 +1114,46 @@ class AnthropicHandlerMixin: original_client_messages, frozen_message_count, ) + # Cold-prefix cache-miss hook (HEADROOM_COLD_RECOMPACT). Claude's thinking is + # an encrypted handle we can't shrink, so when the prompt cache has lapsed + # (idle past TTL → dead, nothing to bust) we instead recompact the whole + # prefix — cross-turn dedupe (+HEADROOM_DEDUPE) + superseded-read drop + + # lossless folds. Decided once here, then applied per mode below: TOKEN mode + # sets frozen_message_count=0 (reuses the frozen==0 path); CACHE mode runs a + # lossless whole-prefix recompaction instead of the byte-identical splice + # (the splice preserves a dead cache) and skips the overlay replay. Both are + # deterministic → the recompacted prefix re-caches byte-stable on warm turns. + _cold_recompact_active = False + if os.environ.get("HEADROOM_COLD_RECOMPACT", "").strip().lower() in ( + "1", + "true", + "yes", + ): + from headroom.transforms.cold_prefix import ( + anthropic_cache_ttl_seconds, + is_cold_prefix, + ) + + # Read CC's ACTUAL prompt-cache TTL (request cache_control.ttl + the + # DISABLE_/ENABLE_/FORCE_PROMPT_CACHING_* env controls) instead of the + # static 300s guess — a wrong TTL is exactly what busts a warm cache. + # None ⇒ caching is OFF (no cache to bust) ⇒ recompact every turn. + _cc_ttl = anthropic_cache_ttl_seconds( + model, original_client_messages, system_prompt + ) + _cold_recompact_active = _cc_ttl is None or is_cold_prefix( + prefix_tracker, ttl_seconds=_cc_ttl + ) + if _cold_recompact_active: + logger.info( + "[%s] cold-prefix recompaction: cc_cache_ttl=%s idle=%.0fs — recompacting " + "whole prefix (dedupe/superseded-read/lossless)", + request_id, + "disabled" if _cc_ttl is None else f"{_cc_ttl}s", + idle_seconds, + ) + if is_token_mode(self.config.mode): + frozen_message_count = 0 # PR-A6 (P5-50, preps P0-6): session-sticky `anthropic-beta` merge. # Read the client's beta value (note: anthropic-beta is NOT @@ -1445,6 +1485,25 @@ class AnthropicHandlerMixin: pipeline_timing = result.timing original_tokens = result.tokens_before optimized_tokens = result.tokens_after + elif _cold_recompact_active: + # CACHE mode, cold turn: the prompt cache is dead, so the + # byte-identical splice preserves nothing. Recompact the whole + # prefix losslessly (dedupe + superseded-read drop + folds) and + # forward that — the overlay replay below is skipped on cold so + # this survives. Deterministic → re-caches byte-stable warm. + from headroom.transforms.cold_prefix import cold_recompact_messages + + recompacted, _cold_transforms = await self._run_compression_in_executor( + lambda: cold_recompact_messages( + original_client_messages, + tokenizer=tokenizer, + context=extract_user_query(original_client_messages), + ), + timeout=COMPRESSION_TIMEOUT_SECONDS, + ) + optimized_messages = recompacted + optimized_tokens = tokenizer.count_messages(optimized_messages) + transforms_applied = _cold_transforms else: previous_original_messages = prefix_tracker.get_last_original_messages() previous_forwarded_messages = prefix_tracker.get_last_forwarded_messages() @@ -1557,16 +1616,22 @@ class AnthropicHandlerMixin: overlay_cached_prefix, ) - _ov = overlay_cached_prefix( - optimized_messages, - original_client_messages, - prefix_tracker.get_last_original_messages(), - prefix_tracker.get_last_forwarded_messages(), - ) - _overlay_replayed = _ov != optimized_messages - if _overlay_replayed: - optimized_messages = _ov - optimized_tokens = tokenizer.count_messages(optimized_messages) + # On a confirmed-cold turn we deliberately do NOT replay the previously + # forwarded prefix: the cache is dead (nothing to keep byte-identical for) + # and the replay would clobber the whole-prefix recompaction we just did. + if _cold_recompact_active: + _overlay_replayed = False + else: + _ov = overlay_cached_prefix( + optimized_messages, + original_client_messages, + prefix_tracker.get_last_original_messages(), + prefix_tracker.get_last_forwarded_messages(), + ) + _overlay_replayed = _ov != optimized_messages + if _overlay_replayed: + optimized_messages = _ov + optimized_tokens = tokenizer.count_messages(optimized_messages) # Own cache_control placement: the client moves the breakpoint each # turn and the overlay replays past markers, so they accumulate ~1/turn @@ -2658,6 +2723,13 @@ class AnthropicHandlerMixin: cache_read_tokens=cr_tokens, current_forwarded_messages=optimized_messages, ) + from headroom.cache.ttl_observations import ( + record_cache_observation, + ) + + record_cache_observation( + provider="anthropic", model=model, attribution=miss + ) if miss.is_miss: logger.info( f"[{request_id}] CACHE-MISS-ATTRIBUTION: reason={miss.reason} " @@ -3289,6 +3361,13 @@ class AnthropicHandlerMixin: cache_read_tokens=cr_tokens, current_forwarded_messages=optimized_messages, ) + from headroom.cache.ttl_observations import ( + record_cache_observation, + ) + + record_cache_observation( + provider="anthropic", model=model, attribution=miss + ) if miss.is_miss: logger.info( f"[{request_id}] CACHE-MISS-ATTRIBUTION: reason={miss.reason} " diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 413382166..c3509bcb2 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3025,6 +3025,28 @@ class OpenAIHandlerMixin: messages, openai_frozen_count, ) + # Cold-prefix cache-miss hook, branch 2 (HEADROOM_COLD_RECOMPACT). When the + # prompt cache has lapsed (idle past TTL → dead, nothing to bust), unfreeze the + # whole prefix so the router recompacts it (cross-turn dedupe + superseded-read + # drop + lossless folds) — the lever for encrypted-reasoning models (Codex) + # whose reasoning we can't touch. Composes with branch 1 (plain-text reasoning + # drop at PRE_SEND for Kimi/GLM). Token mode only; deterministic → cache-stable. + if is_token_mode(self.config.mode) and os.environ.get( + "HEADROOM_COLD_RECOMPACT", "" + ).strip().lower() in ("1", "true", "yes"): + from headroom.cache.ttl_observations import resolve_learned_ttl + from headroom.transforms.cold_prefix import is_cold_prefix + + # OpenAI/Kimi don't expose their cache TTL, so use the offline-learned + # value if available (else is_cold_prefix falls back to the static guess). + _learned_ttl = resolve_learned_ttl("openai", model) + if is_cold_prefix(openai_prefix_tracker, ttl_seconds=_learned_ttl): + logger.info( + "[%s] cold-prefix recompaction: unfreezing prefix for " + "dedupe/superseded-read compaction", + request_id, + ) + openai_frozen_count = 0 _compression_failed = False original_messages = messages # Preserve for 400-retry fallback @@ -3434,6 +3456,68 @@ class OpenAIHandlerMixin: if tools or _original_tools is not None: body["tools"] = tools + # Reasoning compaction (HEADROOM_THINKING_COMPACT, off by default). Unlike + # Anthropic/OpenAI (encrypted reasoning handles, nothing to compress), + # Kimi/GLM/DeepSeek-R1 resend reasoning as PLAIN TEXT billed as input (Kimi + # K2.7: ~1,558 tok/block, kept every turn) — so Kompress can actually shrink + # it. Shape-driven: compacts the `reasoning_content` field (Kimi) and inline + # `` (GLM/DeepSeek); no-ops when neither is present (OpenAI's + # own encrypted models). Deterministic → forwarded prefix stays cache-stable; + # keep_last_turns protects the active reasoning. Never breaks the request. + if os.environ.get("HEADROOM_THINKING_COMPACT", "").strip().lower() in ( + "1", + "true", + "yes", + ): + try: + from headroom.cache.ttl_observations import resolve_learned_ttl + from headroom.transforms.cold_prefix import is_cold_prefix + from headroom.transforms.compression_units import find_content_router + from headroom.transforms.thinking_compactor import ( + compact_reasoning_openai_chat, + ) + + # Cold-prefix cache-miss hook: on a warm turn Kompress the reasoning + # (deterministic → cache-stable); on a COLD turn (idle past provider + # TTL, cache dead) DROP it outright — the full block, safe because the + # cold turn re-caches from scratch. See cold_prefix / thinking_compactor. + # Uses the offline-learned TTL (Kimi/OpenAI don't expose it). + _rc_cold = is_cold_prefix( + openai_prefix_tracker, ttl_seconds=resolve_learned_ttl("openai", model) + ) + _rc_router = find_content_router(self.openai_pipeline) + _rc_kompress = ( + (_rc_router._get_remote_kompress() or _rc_router._get_kompress()) + if _rc_router is not None + else None + ) + # Drop needs no compressor; Kompress (warm) does. + if _rc_kompress is not None or _rc_cold: + _rc_keep = int(os.environ.get("HEADROOM_THINKING_COMPACT_KEEP_LAST", "1")) + optimized_messages, _rc_stats = compact_reasoning_openai_chat( + optimized_messages, + kompress=_rc_kompress, + keep_last_turns=_rc_keep, + drop=_rc_cold, + ) + body["messages"] = optimized_messages + if _rc_stats["turns_compacted"]: + transforms_applied.append( + f"openai:reasoning_{'drop' if _rc_cold else 'compact'}:{_rc_stats['turns_compacted']}" + ) + logger.info( + "[%s] reasoning %s: %d turns, %d blocks, %d->%d words (cold=%s)", + request_id, + "drop" if _rc_cold else "compact", + _rc_stats["turns_compacted"], + _rc_stats["blocks"], + _rc_stats["words_before"], + _rc_stats["words_after"], + _rc_cold, + ) + except Exception as _rc_exc: # never break the request on compaction + logger.warning("[%s] reasoning compaction skipped: %s", request_id, _rc_exc) + presend_event = self.pipeline_extensions.emit( PipelineStage.PRE_SEND, operation="proxy.request", @@ -3826,6 +3910,27 @@ class OpenAIHandlerMixin: 0, total_input_tokens - cache_read_tokens - cache_write_tokens ) + # Cache-TTL learning seam: attribute this turn's cache outcome + # (hit / ttl_expiry / prefix_change) and record it BEFORE + # update_from_response overwrites the prior-turn state. Feeds the + # offline TTL learner (HEADROOM_CACHE_TTL_LEARN); best-effort. + try: + if hasattr(openai_prefix_tracker, "classify_cache_miss"): + from headroom.cache.ttl_observations import ( + record_cache_observation, + ) + + record_cache_observation( + provider="openai", + model=model, + attribution=openai_prefix_tracker.classify_cache_miss( + cache_read_tokens=cache_read_tokens, + current_forwarded_messages=optimized_messages, + ), + ) + except Exception: + pass + openai_prefix_tracker.update_from_response( cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, @@ -4128,6 +4233,25 @@ class OpenAIHandlerMixin: total_input_tokens, cache_read_tokens, ) + # Cache-TTL learning seam (see the /v1/chat path above): record the + # cache-outcome attribution before update_from_response. Best-effort. + try: + if hasattr(openai_prefix_tracker, "classify_cache_miss"): + from headroom.cache.ttl_observations import ( + record_cache_observation, + ) + + record_cache_observation( + provider="openai", + model=model, + attribution=openai_prefix_tracker.classify_cache_miss( + cache_read_tokens=cache_read_tokens, + current_forwarded_messages=optimized_messages, + ), + ) + except Exception: + pass + openai_prefix_tracker.update_from_response( cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, diff --git a/headroom/transforms/cold_prefix.py b/headroom/transforms/cold_prefix.py new file mode 100644 index 000000000..3777e1c1d --- /dev/null +++ b/headroom/transforms/cold_prefix.py @@ -0,0 +1,282 @@ +"""Cold-prefix cache-miss hook: decide what to rewrite when the prompt cache is dead. + +When the prefix cache has lapsed (idle since the last turn exceeded the provider +TTL), forwarding the byte-identical prefix buys nothing — the cache is gone — so +this is the safe moment for rewrites that would otherwise bust a warm cache. What +we rewrite depends on the model's reasoning shape: + +* **Plain-text reasoning (Kimi / GLM / DeepSeek-R1)** — reasoning is resent as + billable text (`reasoning_content` field or inline ````). On a cold turn + we can DROP the old reasoning outright (full block), not just Kompress it. +* **Encrypted reasoning (Claude / OpenAI Codex)** — the reasoning is an opaque + server-side handle billed free/light; touching it saves nothing. Instead, on a + cold turn, dedupe + drop superseded reads across the (now-unfreezable) prefix. + +This module is the *decision* surface (is-it-cold + which-shape); the handlers +apply the chosen rewrite. Never raises. + +Cache note: dropping/deduping the prefix is cache-safe here **because the cache is +already dead** — nothing to bust. The one cost is the cold turn re-caches a +smaller prefix, which then benefits every subsequent warm turn until the next +lapse. (For plain-text reasoning, deterministic Kompress-every-turn — see +``thinking_compactor`` — stays cache-stable on warm turns; the cold DROP is the +extra, aggressive step reserved for when the cache is confirmed gone.) +""" + +from __future__ import annotations + +import logging +from typing import Any + +log = logging.getLogger(__name__) + +_DEFAULT_MARGIN_SECONDS = 60.0 + + +def is_cold_prefix( + prefix_tracker: Any, + *, + margin_seconds: float = _DEFAULT_MARGIN_SECONDS, + ttl_seconds: float | None = None, +) -> bool: + """True when the prompt-cache prefix has (confidently) lapsed. + + Compares the idle gap captured at fetch (``_idle_seconds_at_fetch``) to the + provider cache TTL. Pass ``ttl_seconds`` to use a KNOWN TTL (e.g. from + :func:`anthropic_cache_ttl_seconds`, which reads CC's actual 5m/1h config) — + this is what makes cold detection reliable. When ``ttl_seconds`` is None it + falls back to the tracker's ``resolved_cache_ttl_seconds()`` (a static + per-provider guess), reliable only where that default is documented-correct. + + The margin makes us *confident* it's past TTL before treating it as cold — we + would rather miss a just-expired cache than rewrite a still-warm one (a wrong + TTL here is exactly what busts a warm cache). Returns False on any error + (conservative: never assume cold). + """ + try: + idle = float(getattr(prefix_tracker, "_idle_seconds_at_fetch", 0.0) or 0.0) + if ttl_seconds is not None: + ttl = float(ttl_seconds) + else: + ttl_fn = getattr(prefix_tracker, "resolved_cache_ttl_seconds", None) + if ttl_fn is None: + return False + ttl = float(ttl_fn()) + except Exception: + return False + return idle > ttl + margin_seconds + + +_ANTHROPIC_FAMILIES = ("opus", "sonnet", "haiku", "fable", "mythos") +_TRUTHY = ("1", "true", "yes", "on") + + +def _env_truthy(name: str) -> bool: + import os + + return os.environ.get(name, "").strip().lower() in _TRUTHY + + +def _anthropic_family(model: str) -> str | None: + m = model.lower() + for fam in _ANTHROPIC_FAMILIES: + if fam in m: + return fam + return None + + +def _cache_control_ttls(messages: list[dict[str, Any]], system: Any) -> set[str]: + """Collect explicit ``cache_control.ttl`` strings from a client request. + + Anthropic prompt caching: ``{"type":"ephemeral"}`` is the 5m default; + ``{"type":"ephemeral","ttl":"1h"}`` (with the extended-cache-ttl beta) is 1h. + We read whatever Claude Code actually sent — the authoritative TTL signal. + """ + ttls: set[str] = set() + + def _scan_blocks(blocks: Any) -> None: + if not isinstance(blocks, list): + return + for b in blocks: + if isinstance(b, dict): + cc = b.get("cache_control") + if isinstance(cc, dict) and isinstance(cc.get("ttl"), str): + ttls.add(cc["ttl"]) + + _scan_blocks(system) + for m in messages: + if not isinstance(m, dict): + continue + cc = m.get("cache_control") + if isinstance(cc, dict) and isinstance(cc.get("ttl"), str): + ttls.add(cc["ttl"]) + _scan_blocks(m.get("content")) + return ttls + + +def anthropic_cache_ttl_seconds( + model: str, messages: list[dict[str, Any]], system: Any = None +) -> int | None: + """The prompt-cache TTL Claude Code is actually using — not a guess. + + Returns: + * ``None`` — prompt caching is OFF (``DISABLE_PROMPT_CACHING`` or the + per-model ``DISABLE_PROMPT_CACHING_``). With no cache there is + nothing to bust, so the caller can recompact EVERY turn. + * ``3600`` — 1h caching (request ``cache_control.ttl == "1h"``, or + ``ENABLE_PROMPT_CACHING_1H``). + * ``300`` — 5m caching (default, ``FORCE_PROMPT_CACHING_5M``, or + ``cache_control.ttl == "5m"``). + + Priority: the request's ``cache_control.ttl`` is authoritative for 1h-vs-5m + (it already reflects CC's env config + overage checks and needs no env + sharing); the env vars are the OFF signal + a fallback. Reading a wrong TTL + is exactly what would bust a warm cache, so this replaces the hardcoded 300s + guess for CC. Never raises. + """ + try: + if _env_truthy("DISABLE_PROMPT_CACHING"): + return None + fam = _anthropic_family(model) + if fam and _env_truthy(f"DISABLE_PROMPT_CACHING_{fam.upper()}"): + return None + ttls = _cache_control_ttls(messages, system) + if "1h" in ttls: + return 3600 + if "5m" in ttls: + return 300 + # cache_control present without an explicit ttl (or none parsed): env hint, + # else Anthropic's 5m default. (We do NOT infer "off" from absence — a false + # off would recompact a warm cache every turn.) + if _env_truthy("FORCE_PROMPT_CACHING_5M"): + return 300 + if _env_truthy("ENABLE_PROMPT_CACHING_1H"): + return 3600 + return 300 + except Exception: + return 300 # safe default on any parse error + + +def has_plaintext_reasoning(messages: list[dict[str, Any]]) -> bool: + """True if any assistant turn carries reasoning as PLAIN TEXT we can drop/compress. + + Two shapes: a Kimi-style ``reasoning_content`` field, or an inline + ```` span in string content (GLM / DeepSeek-R1). Encrypted + reasoning (Claude signature / OpenAI ``encrypted_content``) never appears in + these forms, so this is False for those — which routes them to the dedupe / + superseded-read branch instead. + """ + for m in messages: + if m.get("role") != "assistant": + continue + rc = m.get("reasoning_content") + if isinstance(rc, str) and rc.strip(): + return True + c = m.get("content") + if isinstance(c, str) and "" in c and "" in c: + return True + return False + + +def cold_recompact_messages( + messages: list[dict[str, Any]], *, tokenizer: Any, context: str = "" +) -> tuple[list[dict[str, Any]], list[str]]: + """Lossless whole-prefix recompaction for a confirmed-cold turn. + + Runs a fresh lossless + cross-turn-dedup ContentRouter over the *entire* + conversation (``frozen_message_count=0``): superseded/stale-read drop + + verbatim dedupe + lossless folds — the safe, information-preserving rewrites + (never lossy, so old context the model relies on is not mangled). Used when + the prompt cache is dead (idle past TTL) and the byte-identical splice would + preserve nothing. Lossless + prefix-monotonic ⇒ deterministic per content ⇒ + the recompacted prefix re-caches and stays byte-stable on later warm turns. + + Returns (new_messages, transforms_applied). Fail-open: returns the input + unchanged on any error (never breaks the request). + """ + try: + from headroom.transforms.content_router import ( + ContentRouter, + ContentRouterConfig, + ) + + router = ContentRouter(ContentRouterConfig(lossless=True, enable_cross_turn_dedup=True)) + res = router.apply(list(messages), tokenizer, frozen_message_count=0, context=context) + return res.messages, list(res.transforms_applied) + except Exception as e: # never break the request + log.warning("cold-prefix recompaction failed (%s); leaving prefix unchanged", e) + return list(messages), [] + + +def _demo() -> None: + class _T: + def __init__(self, idle: float, ttl: float) -> None: + self._idle_seconds_at_fetch = idle + self._ttl = ttl + + def resolved_cache_ttl_seconds(self) -> float: + return self._ttl + + assert is_cold_prefix(_T(400, 300)) # 400 idle > 300 ttl + 60 margin? 400 > 360 ✓ + assert not is_cold_prefix(_T(350, 300)) # 350 < 360 → warm + assert not is_cold_prefix(_T(10, 300)) # back-to-back → warm + assert not is_cold_prefix(object()) # missing attrs → conservative False + + assert has_plaintext_reasoning([{"role": "assistant", "reasoning_content": "abc"}]) + assert has_plaintext_reasoning([{"role": "assistant", "content": "x y"}]) + assert not has_plaintext_reasoning([{"role": "assistant", "content": "plain"}]) + assert not has_plaintext_reasoning([{"role": "user", "reasoning_content": "x"}]) + + # --- CC cache-TTL detection (read the real TTL, don't guess 300s) --- + import os as _os + + _1h_msg = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "x", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + } + ] + _5m_msg = [ + { + "role": "user", + "content": [{"type": "text", "text": "x", "cache_control": {"type": "ephemeral"}}], + } + ] + assert anthropic_cache_ttl_seconds("claude-opus-4-6", _1h_msg) == 3600 # request says 1h + assert anthropic_cache_ttl_seconds("claude-opus-4-6", _5m_msg) == 300 # default 5m + + for _v in ( + "DISABLE_PROMPT_CACHING", + "DISABLE_PROMPT_CACHING_OPUS", + "ENABLE_PROMPT_CACHING_1H", + "FORCE_PROMPT_CACHING_5M", + ): + _os.environ.pop(_v, None) + _os.environ["DISABLE_PROMPT_CACHING"] = "1" + assert ( + anthropic_cache_ttl_seconds("claude-opus-4-6", _5m_msg) is None + ) # OFF -> recompact freely + del _os.environ["DISABLE_PROMPT_CACHING"] + _os.environ["DISABLE_PROMPT_CACHING_OPUS"] = "1" + assert anthropic_cache_ttl_seconds("claude-opus-4-6", []) is None # per-model OFF + assert anthropic_cache_ttl_seconds("claude-sonnet-4-6", []) == 300 # other family unaffected + del _os.environ["DISABLE_PROMPT_CACHING_OPUS"] + _os.environ["ENABLE_PROMPT_CACHING_1H"] = "1" + assert anthropic_cache_ttl_seconds("claude-opus-4-6", []) == 3600 # env 1h + _os.environ["FORCE_PROMPT_CACHING_5M"] = "1" + assert anthropic_cache_ttl_seconds("claude-opus-4-6", []) == 300 # 5m forced overrides 1h + del _os.environ["ENABLE_PROMPT_CACHING_1H"], _os.environ["FORCE_PROMPT_CACHING_5M"] + + # The exact bug: at 400s idle, the hardcoded 300s guess says "cold" (would bust a + # warm 1h cache); the real 1h TTL correctly says "warm". + assert is_cold_prefix(_T(400, 300), ttl_seconds=3600) is False # warm at 1h + assert is_cold_prefix(_T(400, 300)) is True # "cold" under the 300s guess (the bug) + assert is_cold_prefix(_T(3700, 300), ttl_seconds=3600) is True # genuinely cold past 1h + + print("cold_prefix self-check OK") + + +if __name__ == "__main__": + _demo() diff --git a/headroom/transforms/thinking_compactor.py b/headroom/transforms/thinking_compactor.py new file mode 100644 index 000000000..d2cb67abb --- /dev/null +++ b/headroom/transforms/thinking_compactor.py @@ -0,0 +1,417 @@ +"""Convert prior-turn extended-thinking blocks into Kompressed ``text`` blocks. + +On Claude 4.6+ models, prior-turn thinking is re-sent as input and **billed** +(verified live: opus-4-6 +995 tok/block, sonnet-4-6 +688; pre-4.6 models strip +it server-side, so this transform is a no-op there — gate on model generation at +the call site). Two findings dictate the mechanism: + +1. **Editing a thinking block in place is futile** — Anthropic pins the original + via the block ``signature`` and re-expands it server-side, ignoring whatever + text you send (verified: 835==835). The only ways to actually shrink thinking + are to *drop* the block or *convert it to a plain ``text`` block* (no + signature → the shorter text is billed as-is; verified 716<835). This + transform does the latter, running each block's text through Kompress. + +2. **Cache safety comes from determinism, not from "only touch the delta".** The + client re-sends the *original* thinking every turn, but the prompt cache holds + the *compacted* form we forwarded last turn. So we map original→compacted + deterministically (memoized by content hash) every turn — the forwarded prefix + is then byte-stable and the cache still hits. The last ``keep_last_turns`` + assistant turns keep their thinking intact (the active reasoning the model + needs); a turn aging out of that window is the only byte change, a bounded + recent-region re-write. + +Flag-gated at the call site (``HEADROOM_THINKING_COMPACT``). Fail-open: any +Kompress error leaves the original block untouched. ``keep_last_turns`` is the +quality knob. +""" + +from __future__ import annotations + +import hashlib +import logging +from collections import OrderedDict +from typing import Any + +log = logging.getLogger(__name__) + +# original-thinking-hash -> compacted text. Deterministic memo so the same +# thinking always yields the same bytes (cache stability), even if Kompress is +# nondeterministic. Bounded; ONNX Kompress is deterministic so eviction+recompute +# is byte-identical anyway. ponytail: crude LRU cap, fine given determinism. +_COMPACT_CACHE: OrderedDict[str, str] = OrderedDict() +_CACHE_CAP = 8192 + +# Prefix on the emitted text block so the compaction is legible to the model +# (and greppable in logs). Kept short; the token cost is negligible vs the block. +_MARKER = "[prior reasoning, compressed]" + + +def bills_prior_thinking(model: str) -> bool: + """True if ``model`` re-bills prior-turn thinking as input (so compaction pays). + + Claude 4.6+ (and the 5 family) keep prior-turn thinking in context and bill it; + pre-4.6 (sonnet-4-5, haiku-4-5, 3.x) strip it server-side. Verified live: + opus-4-6/sonnet-4-6 bill, sonnet-4-5/haiku-4-5 strip. **Conservative** — returns + False unless the version is confidently >= 4.6, because compacting on a stripping + model would turn free (stripped) thinking into billed text. (Opus 4.5 reportedly + bills too, but is excluded here pending verification — costs only missed savings.) + """ + nums: list[int] = [] + for part in model.lower().split("-"): + if part.isdigit(): + nums.append(int(part)) + elif nums: + break # version digits are contiguous; stop at the family/date boundary + if not nums: + return False + major = nums[0] + minor = nums[1] if len(nums) > 1 else 0 + return major >= 5 or (major, minor) >= (4, 6) + + +def _memo_compact(text: str, kompress: Any) -> str | None: + """Deterministically compact ``text`` via Kompress; None if no gain/failure.""" + key = hashlib.sha1(text.encode("utf-8", "replace")).hexdigest() + cached = _COMPACT_CACHE.get(key) + if cached is not None: + _COMPACT_CACHE.move_to_end(key) + return cached + try: + # allow_download=False: never block the request thread on a cold model + # (mirrors the proxy convention in ContentRouter._get_kompress callers). + result = kompress.compress(text, allow_download=False) + compacted = result.compressed + except Exception as e: # fail OPEN — never break the proxy on a bad compressor + log.warning("thinking compaction failed (%s); leaving block untouched", e) + return None + if not isinstance(compacted, str) or not compacted: + return None + _COMPACT_CACHE[key] = compacted + _COMPACT_CACHE.move_to_end(key) + while len(_COMPACT_CACHE) > _CACHE_CAP: + _COMPACT_CACHE.popitem(last=False) + return compacted + + +def compact_thinking_to_text( + messages: list[dict[str, Any]], + *, + kompress: Any, + keep_last_turns: int = 1, + min_words: int = 40, +) -> tuple[list[dict[str, Any]], dict[str, int]]: + """Replace ``thinking`` blocks with Kompressed ``text`` blocks. + + Every assistant turn except the last ``keep_last_turns`` (which keep their + thinking verbatim) has each of its ``thinking`` blocks converted to a + ``text`` block holding the Kompressed summary. Deterministic per content, so + the forwarded prefix stays cache-stable across turns. Never mutates the input + list or its message dicts in place. Never raises. + + Args: + messages: provider-native Anthropic messages. + kompress: an object with ``compress(text, allow_download=False) -> + KompressResult`` (local ``KompressCompressor`` or remote + ``RemoteKompressCompressor``; caller picks per HEADROOM_KOMPRESS_ENDPOINT). + keep_last_turns: number of most-recent assistant turns whose thinking is + left intact (the active reasoning). 0 compacts everything. + min_words: thinking blocks below this word count are left as-is (not worth + a Kompress call). + + Returns: + (new_messages, stats) where stats has ``turns_compacted``, ``blocks``, + ``words_before``, ``words_after``. + """ + stats = {"turns_compacted": 0, "blocks": 0, "words_before": 0, "words_after": 0} + if kompress is None: + return messages, stats + + asst_indices = [i for i, m in enumerate(messages) if m.get("role") == "assistant"] + keep = set(asst_indices[-keep_last_turns:]) if keep_last_turns > 0 else set() + + out: list[dict[str, Any]] = [] + for i, m in enumerate(messages): + content = m.get("content") + if ( + m.get("role") != "assistant" + or i in keep + or not isinstance(content, list) + or not any(isinstance(b, dict) and b.get("type") == "thinking" for b in content) + ): + out.append(m) + continue + + new_content: list[Any] = [] + turn_compacted = False + for block in content: + if not (isinstance(block, dict) and block.get("type") == "thinking"): + new_content.append(block) + continue + text = block.get("thinking", "") + words = len(text.split()) + if words < min_words: + new_content.append(block) + continue + compacted = _memo_compact(text, kompress) + # Skip if compaction failed or didn't actually shrink the block. + if compacted is None or len(compacted.split()) >= words: + new_content.append(block) + continue + text_block: dict[str, Any] = {"type": "text", "text": f"{_MARKER} {compacted}"} + # Preserve a cache breakpoint if one happened to sit on the thinking + # block (rare — breakpoints usually sit on the last block of a message), + # so we never silently drop a cache_control marker. + if "cache_control" in block: + text_block["cache_control"] = block["cache_control"] + new_content.append(text_block) + turn_compacted = True + stats["blocks"] += 1 + stats["words_before"] += words + stats["words_after"] += len(compacted.split()) + + if turn_compacted: + stats["turns_compacted"] += 1 + nm = dict(m) + nm["content"] = new_content + out.append(nm) + else: + out.append(m) + + return out, stats + + +def _compact_think_spans( + content: str, kompress: Any, min_words: int, *, drop: bool = False +) -> tuple[str, int, int, int]: + """Compact each ```` span (GLM / DeepSeek-R1 inline reasoning). + + ``drop=False`` (warm) Kompresses the inner text; ``drop=True`` (cold hook) + removes the whole span. Returns (new_content, blocks, words_before, + words_after). String-scan (the tags are a fixed literal delimiter, not a + heuristic). Leaves unmatched/short spans untouched. + """ + open_tag, close_tag = "", "" + blocks = wb = wa = 0 + parts: list[str] = [] + pos = 0 + while True: + start = content.find(open_tag, pos) + if start == -1: + parts.append(content[pos:]) + break + end = content.find(close_tag, start + len(open_tag)) + if end == -1: # unterminated — leave the remainder as-is + parts.append(content[pos:]) + break + parts.append(content[pos:start]) + inner = content[start + len(open_tag) : end] + words = len(inner.split()) + if words >= min_words and drop: + # Cold hook: drop the span entirely (append nothing). + blocks += 1 + wb += words + pos = end + len(close_tag) + continue + new_inner = inner + if words >= min_words: + comp = _memo_compact(inner, kompress) + if comp is not None and len(comp.split()) < words: + new_inner = comp + blocks += 1 + wb += words + wa += len(comp.split()) + parts.append(f"{open_tag}{new_inner}{close_tag}") + pos = end + len(close_tag) + return "".join(parts), blocks, wb, wa + + +def compact_reasoning_openai_chat( + messages: list[dict[str, Any]], + *, + kompress: Any, + keep_last_turns: int = 1, + min_words: int = 40, + drop: bool = False, +) -> tuple[list[dict[str, Any]], dict[str, int]]: + """Compact plain-text reasoning in OpenAI-chat messages (Kimi / GLM / DeepSeek-R1). + + Unlike Anthropic thinking / OpenAI reasoning (encrypted handles), these models + resend reasoning as PLAIN TEXT billed as input — so we can actually shrink it + (verified: Kimi K2.7 resends ``reasoning_content`` at +1,558 input tok/block). + Two shapes, both handled: + + * **Kimi:** the assistant message's ``reasoning_content`` field. + * **GLM / DeepSeek-R1:** inline ```` in string content. + + ``drop=False`` (warm): Kompress the reasoning (deterministic → cache-stable). + ``drop=True`` (cold-prefix hook): remove the reasoning outright — the full + block, not just ~15% — safe because the cold turn re-caches from scratch. + Shape-driven — no model gate; no-ops when no plain-text reasoning is present + (OpenAI's encrypted models, Kimi k2.6). Keeps the last ``keep_last_turns`` + assistant turns intact (the active reasoning the model uses). Never raises. + """ + stats = {"turns_compacted": 0, "blocks": 0, "words_before": 0, "words_after": 0} + if kompress is None and not drop: # drop needs no compressor + return messages, stats + asst_indices = [i for i, m in enumerate(messages) if m.get("role") == "assistant"] + keep = set(asst_indices[-keep_last_turns:]) if keep_last_turns > 0 else set() + + out: list[dict[str, Any]] = [] + for i, m in enumerate(messages): + if m.get("role") != "assistant" or i in keep: + out.append(m) + continue + changed = False + nm = dict(m) + # (1) Kimi: reasoning_content field (plain text, no signature → editable) + rc = m.get("reasoning_content") + if isinstance(rc, str) and len(rc.split()) >= min_words: + if drop: # cold hook: drop the whole reasoning block + nm["reasoning_content"] = "" + changed = True + stats["blocks"] += 1 + stats["words_before"] += len(rc.split()) + else: + comp = _memo_compact(rc, kompress) + if comp is not None and len(comp.split()) < len(rc.split()): + nm["reasoning_content"] = comp + changed = True + stats["blocks"] += 1 + stats["words_before"] += len(rc.split()) + stats["words_after"] += len(comp.split()) + # (2) GLM / DeepSeek-R1: inline in string content + c = m.get("content") + if isinstance(c, str) and "" in c: + new_c, b, wb, wa = _compact_think_spans(c, kompress, min_words, drop=drop) + if b: + nm["content"] = new_c + changed = True + stats["blocks"] += b + stats["words_before"] += wb + stats["words_after"] += wa + if changed: + stats["turns_compacted"] += 1 + out.append(nm) + else: + out.append(m) + return out, stats + + +def _demo() -> None: + """Self-check: assert conversion, keep-last, tool_use preservation, determinism.""" + + class _FakeKompress: + calls = 0 + + def compress(self, text: str, allow_download: bool = True) -> Any: # noqa: ARG002 + _FakeKompress.calls += 1 + return type("R", (), {"compressed": "short summary"})() + + k = _FakeKompress() + long = " ".join(["reasoning"] * 60) # 60 words > min_words + msgs: list[dict[str, Any]] = [ + {"role": "user", "content": "hi"}, + { # old assistant turn: thinking + tool_use -> thinking should become text + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": long, "signature": "sig1"}, + {"type": "tool_use", "id": "t1", "name": "calc", "input": {}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "ok"}], + }, + { # last assistant turn: must be KEPT full + "role": "assistant", + "content": [{"type": "thinking", "thinking": long, "signature": "sig2"}], + }, + ] + out, stats = compact_thinking_to_text(msgs, kompress=k, keep_last_turns=1) + + old = out[1]["content"] + assert old[0] == {"type": "text", "text": f"{_MARKER} short summary"}, old[0] + assert old[1]["type"] == "tool_use", "tool_use must be preserved" + assert out[3]["content"][0]["type"] == "thinking", "last turn must stay thinking" + assert stats["turns_compacted"] == 1 and stats["blocks"] == 1, stats + assert msgs[1]["content"][0]["type"] == "thinking", "input must not be mutated" + + # determinism: same thinking text -> memoized, no second Kompress call + calls_before = _FakeKompress.calls + out2, _ = compact_thinking_to_text(msgs, kompress=k, keep_last_turns=1) + assert out2[1]["content"][0] == old[0], "must be byte-identical across runs" + assert _FakeKompress.calls == calls_before, "identical thinking must hit the memo" + + # keep_last_turns=0 compacts the final turn too + out3, stats3 = compact_thinking_to_text(msgs, kompress=k, keep_last_turns=0) + assert out3[3]["content"][0]["type"] == "text", "keep_last_turns=0 compacts all" + assert stats3["turns_compacted"] == 2, stats3 + + # model gate: 4.6+ / 5.x bill (compact); pre-4.6 strip (skip) + assert bills_prior_thinking("claude-opus-4-6") + assert bills_prior_thinking("claude-sonnet-4-6") + assert bills_prior_thinking("claude-opus-4-8") + assert bills_prior_thinking("claude-sonnet-5") + assert not bills_prior_thinking("claude-sonnet-4-5-20250929") + assert not bills_prior_thinking("claude-haiku-4-5-20251001") + assert not bills_prior_thinking("claude-3-5-sonnet-20241022") + + # cache_control on a thinking block is carried to the emitted text block + msgs_cc: list[dict[str, Any]] = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": long, + "signature": "s", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": "next"}, + ] + out_cc, _ = compact_thinking_to_text(msgs_cc, kompress=k, keep_last_turns=0) + assert out_cc[1]["content"][0].get("cache_control") == {"type": "ephemeral"}, out_cc[1] + + # --- OpenAI-chat shapes (Kimi reasoning_content + GLM inline ) --- + oc_msgs: list[dict[str, Any]] = [ + {"role": "user", "content": "q"}, + # Kimi: reasoning_content field on an OLD assistant turn -> compacted + {"role": "assistant", "content": "kimi answer", "reasoning_content": long}, + {"role": "user", "content": "q2"}, + # GLM: inline on an OLD assistant turn -> inner compacted, wrapper kept + {"role": "assistant", "content": f"{long} glm answer"}, + {"role": "user", "content": "q3"}, + # OpenAI encrypted case: no plain-text reasoning -> must be a no-op + {"role": "assistant", "content": "plain answer, no reasoning"}, + ] + oc_out, oc_stats = compact_reasoning_openai_chat(oc_msgs, kompress=k, keep_last_turns=1) + assert oc_out[1]["reasoning_content"] == "short summary", oc_out[1] + assert oc_out[3]["content"] == "short summary glm answer", oc_out[3] + assert oc_out[5]["content"] == "plain answer, no reasoning", "no-reasoning must no-op" + assert oc_stats["turns_compacted"] == 2 and oc_stats["blocks"] == 2, oc_stats + # keep_last_turns protects the final assistant turn's reasoning + oc_msgs2: list[dict[str, Any]] = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a", "reasoning_content": long}, + ] + oc_out2, _ = compact_reasoning_openai_chat(oc_msgs2, kompress=k, keep_last_turns=1) + assert oc_out2[1]["reasoning_content"] == long, "last turn reasoning must be kept" + + # drop mode (cold-prefix hook): reasoning removed outright, no compressor needed + oc_drop, ds = compact_reasoning_openai_chat( + oc_msgs, kompress=None, keep_last_turns=1, drop=True + ) + assert oc_drop[1]["reasoning_content"] == "", "Kimi reasoning must be dropped" + assert oc_drop[3]["content"] == " glm answer", oc_drop[3] # span removed + assert oc_drop[5]["content"] == "plain answer, no reasoning", "no-op on encrypted" + assert ds["turns_compacted"] == 2 and ds["words_after"] == 0, ds + + print("thinking_compactor self-check OK") + + +if __name__ == "__main__": + _demo()