akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(proxy): model-aware cold-prefix hook — reasoning compaction (Kimi/GLM) + cold recompaction (CC) (#2555)

Browse source 
## Description

Adds a **model-aware cold-prefix cache-miss hook** plus **plain-text
reasoning compaction**, both off by default behind flags. Motivation:
prior-turn reasoning and stale prefix content are re-sent and (for some
models) re-billed every turn; when the prompt cache has lapsed,
rewriting the prefix is free. What we do depends on the model's
reasoning shape.

| | plain-text reasoning (Kimi/GLM/DeepSeek) | encrypted reasoning
(Claude/Codex) |
|---|---|---|
| **warm turn** | Kompress reasoning (deterministic → cache-stable) |
leave it (encrypted; can't shrink) |
| **cold turn** | drop the full reasoning block | dedupe + drop
superseded reads (recompact whole prefix) |

Closes #

## Type of Change

- [x] New feature (non-breaking change that adds functionality)
- [x] Performance improvement

## Changes Made

- `headroom/transforms/thinking_compactor.py` (new): shape-driven
reasoning compaction for the OpenAI-chat path — Kimi `reasoning_content`
field + GLM/DeepSeek inline `<think>` spans; deterministic memoized
Kompress (warm) or drop (cold); `keep_last_turns` protects the active
reasoning; no-ops on encrypted-reasoning models.
- `headroom/transforms/cold_prefix.py` (new): the cold-decision surface
— `is_cold_prefix` (idle > TTL + margin), `has_plaintext_reasoning`,
`cold_recompact_messages` (lossless whole-prefix dedupe/superseded), and
`anthropic_cache_ttl_seconds` (reads CC's **real** cache TTL from
request `cache_control.ttl` + `DISABLE_/ENABLE_/FORCE_PROMPT_CACHING_*`
env controls instead of a hardcoded 300s guess).
- `headroom/proxy/handlers/openai.py`: PRE_SEND reasoning compaction
(warm Kompress / cold drop).
- `headroom/proxy/handlers/anthropic.py`: cold-prefix recompaction —
token mode via `frozen_message_count=0`, and **cache mode** via a
whole-prefix lossless recompaction that skips the byte-identical
splice/overlay on a confirmed-cold turn; cold decision uses CC's real
TTL.
- Flags (all off by default): `HEADROOM_THINKING_COMPACT` (+
`HEADROOM_THINKING_COMPACT_KEEP_LAST`), `HEADROOM_COLD_RECOMPACT`.

## Testing

- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`)
- [x] New tests added for new functionality (module self-checks)
- [x] Manual testing performed (live provider calls)

### Test Output

```text
$ ruff check headroom/transforms/cold_prefix.py headroom/transforms/thinking_compactor.py \
      headroom/proxy/handlers/openai.py headroom/proxy/handlers/anthropic.py
All checks passed!

$ mypy <same 4 files>
Success: no issues found in 4 source files

$ python headroom/transforms/cold_prefix.py
cold_prefix self-check OK
$ python headroom/transforms/thinking_compactor.py
thinking_compactor self-check OK
```

## Real Behavior Proof

- Environment: live Kimi K2.7 via Fireworks
(`accounts/fireworks/models/kimi-k2p7-code`) + real Modal Kompress
endpoint; Claude models via Anthropic API.
- Exact command / steps: 2-turn replay — turn 1 produces reasoning; turn
2 re-sends it through the transform; compare `usage.prompt_tokens`.
- Observed result:
- Kimi reasoning resend is real, billable plain text: WITH reasoning =
2,643 vs WITHOUT = 1,085 input tokens (+1,558/block).
  - Warm Kompress (real Modal endpoint): 2,427 → 2,190 prompt_tokens.
  - **Cold drop: 2,330 → 714** (= none-baseline; full block removed).
- opencode confirmed to resend `reasoning_content` across turns (real
`opencode run` trace).
- Cold recompaction (dedupe/superseded) on a real 3.4M-token Claude Code
prefix: ~3.7%.
- CC TTL detection self-check pins the bug: `is_cold_prefix` at
idle=400s is `False` under the real 1h TTL (safe) but `True` under the
old 300s guess (would bust a warm cache).
- Not tested: Codex/Responses API path (deferred — no prefix tracker
there, encrypted reasoning); cross-provider live cache-mode cold turn on
a real >TTL idle gap (measured on captured prefixes instead).

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] My changes generate no new warnings
- [x] I have added tests that prove my feature works (module
self-checks)
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`

## Additional Notes

- All behavior is flag-gated and off by default — zero change for
existing users.
- Follow-ups: (1) Codex/Responses API wiring (gap #2, deferred); (2) a
cross-provider cache-TTL learner (estimate real TTL per provider from
JSONL cache-bust observations) so Kimi/OpenAI cold detection is
empirical rather than the 300s fallback — candidate for an enterprise
plugin.
This commit is contained in:
Tejas Chopra 2026-07-25 09:53:25 -07:00 • committed by GitHub
parent a6d4921e82
commit cb8f4b6436
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1061 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

149
headroom/cache/ttl_observations.py vendored Normal file
View file

@ -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()

99
headroom/proxy/handlers/anthropic.py
View file

@ -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} "

124
headroom/proxy/handlers/openai.py
View file

@ -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
# `<think>…</think>` (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,

282
headroom/transforms/cold_prefix.py Normal file
View file

@ -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 ``<think>``). 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_<FAMILY>``). 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
``<think>…</think>`` 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 "<think>" in c and "</think>" 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": "<think>x</think> 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()

417
headroom/transforms/thinking_compactor.py Normal file
View file

@ -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 ``<think>…</think>`` 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 = "<think>", "</think>"
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 ``<think>…</think>`` 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 <think>…</think> in string content
c = m.get("content")
if isinstance(c, str) and "<think>" 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 <think>) ---
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 <think> on an OLD assistant turn -> inner compacted, wrapper kept
{"role": "assistant", "content": f"<think>{long}</think> 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"] == "<think>short summary</think> 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] # <think> 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()