mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(read-maturation): activity-based hold-back Read maturation (Mechanism B) (#1068)
## What Splits **read maturation (Mechanism B)** out of #818 into its own PR, so #818 can stay focused on the other compression knobs + the SQLite CCR store. Audit-reads (traffic audits) stays in #818. Read maturation holds fresh large `Read` outputs **out of the provider prefix cache** while their file is still active, keeps them verbatim the whole time the model is working with them, and **matures** them into a CCR-backed marker once the file has been quiet for `quiesce_turns`. Only the final compressed form ever enters the cache, so **no cached byte is ever mutated — there is nothing to bust.** Activity-based rather than a fixed hold window: the `audit-reads` simulation showed next-touch gaps are fat-tailed (p50 = 4 turns, p90 = 81), so no fixed window covers the tail while a quiesce rule covers the activity cluster and lets the tail self-heal via partial-range re-reads. **Default OFF** — experimental, flag/env gated, validated in pilots first. ## Changes - `config.ReadMaturationConfig` (`enabled=False`, `quiesce_turns=5`, `max_hold_turns=25`, `min_size_bytes=2048`) - `ProxyConfig` fields mirroring the above - `ReadMaturationManager` transform + `relocate_cache_breakpoint` (`headroom/transforms/read_maturation.py`) - Session-scoped manager rides on `PrefixCacheTracker` — shares the session's cache affinity and TTL cleanup - Handler wiring in `anthropic.py`: runs **after** compression (so `read_lifecycle` markers are respected) and **before** body assembly; advisory — never fails the request - CLI flags + env vars (`--read-maturation*` / `HEADROOM_READ_MATURATION*`) - `_proxy_config_from_env` wiring for both the multi-worker and CLI server paths ## Bug fix included The `--read-maturation` flag was missing `envvar="HEADROOM_READ_MATURATION"`, so the env var was silently ignored on the CLI path (only the multi-worker `_proxy_config_from_env` path read it). Fixed here. ## Tests - `tests/test_read_maturation.py` — unit - `tests/test_read_maturation_handler_nobust.py` — handler never busts cache - `tests/test_live/test_live_maturation.py` — live harness `25 passed` locally; ruff check/format clean; mypy clean (only pre-existing `annotation-unchecked` notes).
This commit is contained in:
parent
f309244a77
commit
723b80c091
10 changed files with 1450 additions and 0 deletions
5
headroom/cache/prefix_tracker.py
vendored
5
headroom/cache/prefix_tracker.py
vendored
|
|
@ -93,6 +93,11 @@ class PrefixCacheTracker:
|
|||
self._last_original_messages: list[dict[str, Any]] = []
|
||||
self._last_forwarded_messages: list[dict[str, Any]] = []
|
||||
|
||||
# Session-scoped ReadMaturationManager (Mechanism B), created
|
||||
# lazily by the handler when read maturation is enabled. Rides
|
||||
# here so it shares the session's affinity and TTL cleanup.
|
||||
self.read_maturation_manager: Any = None
|
||||
|
||||
# Stats
|
||||
self._busts_avoided: int = 0
|
||||
self._tokens_preserved: int = 0
|
||||
|
|
|
|||
|
|
@ -474,6 +474,41 @@ def dashboard(port: int, no_open: bool) -> None:
|
|||
is_flag=True,
|
||||
help="Disable Read lifecycle management (stale/superseded Read compression)",
|
||||
)
|
||||
# Read maturation (Mechanism B) — experimental, OFF by default
|
||||
@click.option(
|
||||
"--read-maturation",
|
||||
is_flag=True,
|
||||
envvar="HEADROOM_READ_MATURATION",
|
||||
help=(
|
||||
"EXPERIMENTAL: activity-based read maturation — hold fresh Reads "
|
||||
"out of the provider prefix cache and compress them once their "
|
||||
"file quiesces (env: HEADROOM_READ_MATURATION=1)"
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--read-maturation-quiesce-turns",
|
||||
type=int,
|
||||
default=5,
|
||||
show_default=True,
|
||||
envvar="HEADROOM_READ_MATURATION_QUIESCE_TURNS",
|
||||
help="Read maturation: mature a held Read once its file is quiet this many assistant turns.",
|
||||
)
|
||||
@click.option(
|
||||
"--read-maturation-max-hold-turns",
|
||||
type=int,
|
||||
default=25,
|
||||
show_default=True,
|
||||
envvar="HEADROOM_READ_MATURATION_MAX_HOLD_TURNS",
|
||||
help="Read maturation: force-mature a Read held this many turns even if its file stays active.",
|
||||
)
|
||||
@click.option(
|
||||
"--read-maturation-min-size-bytes",
|
||||
type=int,
|
||||
default=2048,
|
||||
show_default=True,
|
||||
envvar="HEADROOM_READ_MATURATION_MIN_SIZE_BYTES",
|
||||
help="Read maturation: only hold/mature Read outputs at least this many bytes.",
|
||||
)
|
||||
# Memory System (Multi-Provider Support)
|
||||
@click.option(
|
||||
"--memory",
|
||||
|
|
@ -746,6 +781,10 @@ def proxy(
|
|||
disable_kompress_openai: bool | None,
|
||||
code_graph: bool,
|
||||
no_read_lifecycle: bool,
|
||||
read_maturation: bool,
|
||||
read_maturation_quiesce_turns: int,
|
||||
read_maturation_max_hold_turns: int,
|
||||
read_maturation_min_size_bytes: int,
|
||||
memory: bool,
|
||||
memory_db_path: str,
|
||||
memory_storage: str,
|
||||
|
|
@ -977,6 +1016,11 @@ def proxy(
|
|||
code_graph_watcher=code_graph,
|
||||
# Read lifecycle: ON by default (use --no-read-lifecycle to disable)
|
||||
read_lifecycle=not no_read_lifecycle,
|
||||
# Read maturation (Mechanism B): experimental, OFF by default
|
||||
read_maturation=read_maturation,
|
||||
read_maturation_quiesce_turns=read_maturation_quiesce_turns,
|
||||
read_maturation_max_hold_turns=read_maturation_max_hold_turns,
|
||||
read_maturation_min_size_bytes=read_maturation_min_size_bytes,
|
||||
# Memory System (Multi-Provider with auto-detection)
|
||||
# --learn implies --memory (need backend for storing patterns)
|
||||
# Stateless mode disables memory (requires SQLite on disk)
|
||||
|
|
|
|||
|
|
@ -277,6 +277,48 @@ class ReadLifecycleConfig:
|
|||
min_size_bytes: int = 512 # Skip tiny Read outputs (not worth the overhead)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReadMaturationConfig:
|
||||
"""Mechanism B: hold-back Read maturation (compress before cache entry).
|
||||
|
||||
Motivation (measured by `headroom audit-reads`): the median Read stays
|
||||
in context for ~118 assistant turns after it appears, billed at the
|
||||
provider's cache-read rate every request — a Read's lifetime cost is
|
||||
roughly 13x its size. The only cache-safe moment to shrink it is
|
||||
BEFORE it is ever cache-written.
|
||||
|
||||
Mechanics: a fresh large Read is held out of the provider prefix
|
||||
cache (the trailing cache breakpoint is relocated to just before it)
|
||||
while its file is ACTIVE, stays verbatim the whole time the model is
|
||||
working with it, and matures into a CCR-backed marker once the file
|
||||
has been quiet for `quiesce_turns`. Only that final compressed form
|
||||
ever enters the cache. No cached byte is ever mutated — there is
|
||||
nothing to bust.
|
||||
|
||||
Activity-based (not a fixed hold window) because the audit-reads
|
||||
simulation showed touch gaps are fat-tailed: next-touch p50 is 4
|
||||
turns but p90 is 81 — no fixed window covers the tail, while a
|
||||
quiesce rule covers the activity cluster and lets the tail self-heal
|
||||
via the model's observed habit of re-reading ranges from disk (95%
|
||||
of re-reads in real traffic are partial-range reads made while the
|
||||
full text was still in context).
|
||||
|
||||
Disabled by default while the mechanism is validated in pilots.
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
# Mature a held Read once its FILE has had no activity (reads or
|
||||
# edits) for this many assistant turns. Simulation: next-touch p50
|
||||
# is 4 turns, so 5 covers the median activity cluster.
|
||||
quiesce_turns: int = 5
|
||||
# Safety valve: mature regardless once held this many turns, bounding
|
||||
# the hold-out cost for files that stay active for long stretches.
|
||||
max_hold_turns: int = 25
|
||||
# Only hold/mature Reads at least this large; small Reads are cached
|
||||
# immediately as before (holding them costs more than it saves).
|
||||
min_size_bytes: int = 2048
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompressionProfile:
|
||||
"""Per-tool compression bias applied to statistically-determined K.
|
||||
|
|
|
|||
|
|
@ -1234,6 +1234,56 @@ class AnthropicHandlerMixin:
|
|||
optimized_tokens = tokenizer.count_messages(optimized_messages)
|
||||
tokens_saved = max(0, original_tokens - optimized_tokens)
|
||||
|
||||
# Mechanism B: activity-based read maturation (flag-gated,
|
||||
# default off). Runs after compression so read_lifecycle
|
||||
# markers are respected, and before body assembly so the
|
||||
# held-Read breakpoint relocation lands in the forwarded
|
||||
# request. Session state (matured markers) rides on the
|
||||
# prefix tracker — same affinity and TTL cleanup as the
|
||||
# freeze state. Advisory: must never fail the request.
|
||||
if self.config.read_maturation and not _bypass:
|
||||
try:
|
||||
from headroom.config import ReadMaturationConfig
|
||||
from headroom.transforms.read_maturation import (
|
||||
ReadMaturationManager,
|
||||
relocate_cache_breakpoint,
|
||||
)
|
||||
|
||||
maturation_mgr = prefix_tracker.read_maturation_manager
|
||||
if maturation_mgr is None:
|
||||
maturation_mgr = ReadMaturationManager(
|
||||
ReadMaturationConfig(
|
||||
enabled=True,
|
||||
quiesce_turns=self.config.read_maturation_quiesce_turns,
|
||||
max_hold_turns=self.config.read_maturation_max_hold_turns,
|
||||
min_size_bytes=self.config.read_maturation_min_size_bytes,
|
||||
),
|
||||
compression_store=get_compression_store(),
|
||||
)
|
||||
prefix_tracker.read_maturation_manager = maturation_mgr
|
||||
maturation = maturation_mgr.apply(
|
||||
optimized_messages,
|
||||
frozen_message_count=frozen_message_count,
|
||||
)
|
||||
if maturation.replacements_applied or maturation.holding_msg_indices:
|
||||
optimized_messages = relocate_cache_breakpoint(
|
||||
maturation.messages,
|
||||
maturation.holding_msg_indices,
|
||||
)
|
||||
optimized_tokens = tokenizer.count_messages(optimized_messages)
|
||||
tokens_saved = max(0, original_tokens - optimized_tokens)
|
||||
if maturation.newly_matured:
|
||||
transforms_applied.append(f"read_maturation:{maturation.newly_matured}")
|
||||
logger.debug(
|
||||
f"[{request_id}] read_maturation: "
|
||||
f"holding={len(maturation.holding_msg_indices)} "
|
||||
f"matured={maturation.newly_matured} "
|
||||
f"replayed={maturation.replacements_applied} "
|
||||
f"bytes_saved={maturation.bytes_saved}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"[{request_id}] read maturation failed: {e}")
|
||||
|
||||
# Hook: post_compress — let hooks observe compression results
|
||||
if self.config.hooks and tokens_saved > 0:
|
||||
from headroom.hooks import CompressEvent
|
||||
|
|
|
|||
|
|
@ -201,6 +201,19 @@ class ProxyConfig:
|
|||
# Read lifecycle management
|
||||
read_lifecycle: bool = True
|
||||
|
||||
# Mechanism B: activity-based read maturation (hold fresh Reads out of
|
||||
# the provider prefix cache; compress once their file quiesces).
|
||||
# Experimental — default off. CLI: --read-maturation;
|
||||
# env: HEADROOM_READ_MATURATION=1
|
||||
read_maturation: bool = False
|
||||
# Read-maturation tuning (only meaningful when read_maturation=True).
|
||||
# Defaults mirror ReadMaturationConfig. CLI: --read-maturation-quiesce-turns,
|
||||
# --read-maturation-max-hold-turns, --read-maturation-min-size-bytes;
|
||||
# env: HEADROOM_READ_MATURATION_QUIESCE_TURNS / _MAX_HOLD_TURNS / _MIN_SIZE_BYTES.
|
||||
read_maturation_quiesce_turns: int = 5
|
||||
read_maturation_max_hold_turns: int = 25
|
||||
read_maturation_min_size_bytes: int = 2048
|
||||
|
||||
# Deprecated compatibility argument. ContentRouter is always active in
|
||||
# the Python proxy; accepting this avoids breaking old config constructors
|
||||
# while keeping it out of runtime state.
|
||||
|
|
|
|||
|
|
@ -3785,6 +3785,12 @@ def _proxy_config_from_env() -> ProxyConfig:
|
|||
http2=_get_env_bool("HEADROOM_HTTP2", True),
|
||||
periodic_toin_stats_enabled=_get_env_bool("HEADROOM_PERIODIC_TOIN_STATS", True),
|
||||
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
|
||||
read_maturation=_get_env_bool("HEADROOM_READ_MATURATION", False),
|
||||
read_maturation_quiesce_turns=_get_env_int("HEADROOM_READ_MATURATION_QUIESCE_TURNS", 5),
|
||||
read_maturation_max_hold_turns=_get_env_int("HEADROOM_READ_MATURATION_MAX_HOLD_TURNS", 25),
|
||||
read_maturation_min_size_bytes=_get_env_int(
|
||||
"HEADROOM_READ_MATURATION_MIN_SIZE_BYTES", 2048
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4339,6 +4345,12 @@ if __name__ == "__main__":
|
|||
max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", args.max_keepalive),
|
||||
keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", args.keepalive_expiry),
|
||||
http2=not args.no_http2 and _get_env_bool("HEADROOM_HTTP2", True),
|
||||
read_maturation=_get_env_bool("HEADROOM_READ_MATURATION", False),
|
||||
read_maturation_quiesce_turns=_get_env_int("HEADROOM_READ_MATURATION_QUIESCE_TURNS", 5),
|
||||
read_maturation_max_hold_turns=_get_env_int("HEADROOM_READ_MATURATION_MAX_HOLD_TURNS", 25),
|
||||
read_maturation_min_size_bytes=_get_env_int(
|
||||
"HEADROOM_READ_MATURATION_MIN_SIZE_BYTES", 2048
|
||||
),
|
||||
tool_profiles=tool_profiles if tool_profiles else None,
|
||||
exclude_tools=exclude_tools if exclude_tools else None,
|
||||
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
|
||||
|
|
|
|||
378
headroom/transforms/read_maturation.py
Normal file
378
headroom/transforms/read_maturation.py
Normal file
|
|
@ -0,0 +1,378 @@
|
|||
"""Mechanism B: hold-back Read maturation — compress before cache entry.
|
||||
|
||||
The prefix cache bills you for everything *after* the first changed byte,
|
||||
so mutating an already-cached Read is ruinously expensive — but bytes that
|
||||
have never been cache-written have no cache entry to bust. This module
|
||||
exploits the one safe window: a fresh Read is deliberately held *out* of
|
||||
the provider cache (the trailing cache breakpoint is relocated to just
|
||||
before it) while its file is active. The model sees the verbatim content
|
||||
the whole time it is working with the file. Once the file has been quiet
|
||||
for `quiesce_turns`, the content is replaced with a CCR-backed marker —
|
||||
and only that final, small form ever enters the cache.
|
||||
|
||||
Timeline for a Read of file F (quiesce_turns=5):
|
||||
|
||||
turn T: model reads F — verbatim, NOT cached
|
||||
T+1..T+k: model edits / re-reads F — read stays verbatim and
|
||||
uncached (every touch resets the quiet clock)
|
||||
T+k+5: F has been quiet 5 turns → read matures into a marker;
|
||||
the breakpoint returns to the tail; the marker form is
|
||||
cache-written once
|
||||
later turns: marker form read from cache at the provider discount
|
||||
|
||||
Why activity-based instead of a fixed hold window: the audit-reads
|
||||
simulation over real traffic showed touch gaps are fat-tailed (next-touch
|
||||
p50 = 4 turns, p90 = 81) — no fixed window covers the tail. The quiesce
|
||||
rule covers the activity cluster, `max_hold_turns` bounds the hold cost
|
||||
for pathologically busy files, and the tail self-heals through the
|
||||
model's *observed* habit of re-reading ranges from disk: 95% of re-reads
|
||||
in real traffic are partial-range reads made while the full text was
|
||||
still in context. The recovery path is the model's existing behavior.
|
||||
|
||||
Two invariants:
|
||||
|
||||
1. **No cached byte is ever mutated.** The verbatim form is never
|
||||
cache-written, so maturation invalidates nothing.
|
||||
2. **Replay is deterministic.** Once matured, the same marker is applied
|
||||
on every subsequent request (state is session-scoped), so the cached
|
||||
prefix stays byte-stable for the rest of the session.
|
||||
|
||||
Recovery contract (same as read_lifecycle): the full original is stored
|
||||
in the CCR compression store, the marker carries the retrieval hash and
|
||||
the file path, and the file itself remains on disk — a confused model
|
||||
re-reads at the cost of one tool call.
|
||||
|
||||
State (matured markers only — holding is derived from the conversation
|
||||
itself, so it survives state loss) lives with the session's prefix
|
||||
tracker. Per-process, like all session state: multi-worker deployments
|
||||
need sticky sessions (existing constraint).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from ..config import ReadMaturationConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Tool names whose results are eligible for maturation.
|
||||
_READ_TOOLS = frozenset({"Read", "read"})
|
||||
# Tool names that count as file activity (reset the quiet clock).
|
||||
_TOUCH_TOOLS = frozenset(
|
||||
{"Read", "read", "Edit", "edit", "Write", "write", "MultiEdit", "NotebookEdit"}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MaturedRead:
|
||||
"""Replayed replacement for a matured Read."""
|
||||
|
||||
marker: str
|
||||
ccr_hash: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Activity:
|
||||
"""Per-request scan of tool activity, in assistant-turn units."""
|
||||
|
||||
# tool_use_id -> (file_path, assistant turn of the Read tool_use)
|
||||
read_calls: dict[str, tuple[str, int]] = field(default_factory=dict)
|
||||
# file_path -> assistant turn of its most recent touch (read or edit)
|
||||
file_last_touch: dict[str, int] = field(default_factory=dict)
|
||||
# Total assistant messages in the conversation ("now").
|
||||
assistant_count: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class MaturationResult:
|
||||
"""Output of one per-request maturation pass."""
|
||||
|
||||
messages: list[dict[str, Any]]
|
||||
# Message indices that contain still-holding Reads (must stay out of
|
||||
# the provider cache this request — feed to relocate_cache_breakpoint).
|
||||
holding_msg_indices: list[int] = field(default_factory=list)
|
||||
holding_reads: int = 0
|
||||
newly_matured: int = 0
|
||||
replacements_applied: int = 0
|
||||
bytes_saved: int = 0
|
||||
|
||||
|
||||
class ReadMaturationManager:
|
||||
"""Per-session Read maturation state machine.
|
||||
|
||||
Construct once per session (or hold in a session-scoped container)
|
||||
and call :meth:`apply` on every request, after read_lifecycle and
|
||||
before breakpoint placement.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: ReadMaturationConfig,
|
||||
compression_store: Any | None = None,
|
||||
):
|
||||
self.config = config
|
||||
self.store = compression_store
|
||||
self._matured: dict[str, MaturedRead] = {}
|
||||
|
||||
# ─── Per-request entry point ────────────────────────────────────────
|
||||
|
||||
def apply(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
frozen_message_count: int = 0,
|
||||
) -> MaturationResult:
|
||||
"""Hold active Reads, mature quiet ones, replay matured markers.
|
||||
|
||||
Args:
|
||||
messages: Conversation messages (Anthropic content-block or
|
||||
OpenAI role="tool" formats).
|
||||
frozen_message_count: Provider-cached message count. Reads
|
||||
inside the frozen prefix were cache-written verbatim
|
||||
before this mechanism saw them (e.g. it was just
|
||||
enabled, or state was lost) — they are never touched.
|
||||
"""
|
||||
result = MaturationResult(messages=messages)
|
||||
if not self.config.enabled:
|
||||
return result
|
||||
|
||||
activity = self._scan_activity(messages)
|
||||
out: list[dict[str, Any]] = []
|
||||
any_changed = False
|
||||
|
||||
for i, msg in enumerate(messages):
|
||||
if i < frozen_message_count:
|
||||
out.append(msg)
|
||||
continue
|
||||
new_msg, msg_holding = self._process_message(msg, activity, result)
|
||||
out.append(new_msg)
|
||||
if new_msg is not msg:
|
||||
any_changed = True
|
||||
if msg_holding:
|
||||
result.holding_msg_indices.append(i)
|
||||
|
||||
if any_changed:
|
||||
result.messages = out
|
||||
return result
|
||||
|
||||
# ─── Internals ──────────────────────────────────────────────────────
|
||||
|
||||
def _scan_activity(self, messages: list[dict[str, Any]]) -> _Activity:
|
||||
"""One pass over assistant messages: read calls, per-file last
|
||||
touch, and the current assistant-turn count."""
|
||||
act = _Activity()
|
||||
for msg in messages:
|
||||
if msg.get("role") != "assistant":
|
||||
continue
|
||||
act.assistant_count += 1
|
||||
turn = act.assistant_count
|
||||
|
||||
for tc in msg.get("tool_calls", []) or []:
|
||||
if not isinstance(tc, dict):
|
||||
continue
|
||||
func = tc.get("function", {})
|
||||
name = func.get("name", "")
|
||||
if name not in _TOUCH_TOOLS:
|
||||
continue
|
||||
try:
|
||||
args = json.loads(func.get("arguments", "{}"))
|
||||
except (ValueError, TypeError):
|
||||
args = {}
|
||||
fp = args.get("file_path") or args.get("path") or ""
|
||||
if fp:
|
||||
act.file_last_touch[fp] = turn
|
||||
if name in _READ_TOOLS:
|
||||
act.read_calls[tc.get("id", "")] = (fp, turn)
|
||||
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
for b in content:
|
||||
if not (isinstance(b, dict) and b.get("type") == "tool_use"):
|
||||
continue
|
||||
name = b.get("name", "")
|
||||
if name not in _TOUCH_TOOLS:
|
||||
continue
|
||||
inp = b.get("input") or {}
|
||||
fp = inp.get("file_path") or inp.get("path") or ""
|
||||
if fp:
|
||||
act.file_last_touch[fp] = turn
|
||||
if name in _READ_TOOLS:
|
||||
act.read_calls[b.get("id", "")] = (fp, turn)
|
||||
return act
|
||||
|
||||
def _process_message(
|
||||
self,
|
||||
msg: dict[str, Any],
|
||||
activity: _Activity,
|
||||
result: MaturationResult,
|
||||
) -> tuple[dict[str, Any], bool]:
|
||||
"""Returns (possibly-replaced message, message_still_holding)."""
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
# OpenAI format: whole message is one tool result.
|
||||
if role == "tool":
|
||||
tc_id = msg.get("tool_call_id", "")
|
||||
if tc_id in activity.read_calls and isinstance(content, str):
|
||||
new_content, holding = self._handle_read(tc_id, content, activity, result)
|
||||
if new_content is not None:
|
||||
return {**msg, "content": new_content}, holding
|
||||
return msg, holding
|
||||
return msg, False
|
||||
|
||||
# Anthropic format: tool_result blocks inside a user message.
|
||||
# NOTE: blocks carrying a client cache_control are NOT skipped —
|
||||
# Claude Code parks its tail breakpoint on the newest content
|
||||
# block, which right after a Read is the Read's tool_result
|
||||
# itself. Under this mechanism the proxy owns breakpoint
|
||||
# placement: relocate_cache_breakpoint() strips/moves breakpoints
|
||||
# in the held region after this pass.
|
||||
if isinstance(content, list):
|
||||
new_blocks: list[Any] = []
|
||||
changed = False
|
||||
holding_any = False
|
||||
for b in content:
|
||||
if (
|
||||
isinstance(b, dict)
|
||||
and b.get("type") == "tool_result"
|
||||
and b.get("tool_use_id", "") in activity.read_calls
|
||||
and isinstance(b.get("content"), str)
|
||||
):
|
||||
tc_id = b["tool_use_id"]
|
||||
new_content, holding = self._handle_read(tc_id, b["content"], activity, result)
|
||||
holding_any = holding_any or holding
|
||||
if new_content is not None:
|
||||
new_blocks.append({**b, "content": new_content})
|
||||
changed = True
|
||||
continue
|
||||
new_blocks.append(b)
|
||||
if changed:
|
||||
return {**msg, "content": new_blocks}, holding_any
|
||||
return msg, holding_any
|
||||
|
||||
return msg, False
|
||||
|
||||
def _handle_read(
|
||||
self,
|
||||
tc_id: str,
|
||||
content: str,
|
||||
activity: _Activity,
|
||||
result: MaturationResult,
|
||||
) -> tuple[str | None, bool]:
|
||||
"""Returns (replacement_content | None, still_holding)."""
|
||||
matured = self._matured.get(tc_id)
|
||||
|
||||
# Matured earlier: replay the recorded marker deterministically.
|
||||
if matured is not None:
|
||||
if content == matured.marker:
|
||||
return None, False
|
||||
result.replacements_applied += 1
|
||||
result.bytes_saved += max(0, len(content) - len(matured.marker))
|
||||
return matured.marker, False
|
||||
|
||||
size = len(content.encode("utf-8", errors="replace"))
|
||||
if size < self.config.min_size_bytes:
|
||||
return None, False
|
||||
# Lifecycle markers (stale/superseded) are already compact — and
|
||||
# read_lifecycle runs first, so respect its replacement.
|
||||
if "Retrieve original: hash=" in content or "Retrieve more: hash=" in content:
|
||||
return None, False
|
||||
|
||||
file_path, read_turn = activity.read_calls[tc_id]
|
||||
last_touch = activity.file_last_touch.get(file_path, read_turn)
|
||||
quiet_turns = activity.assistant_count - last_touch
|
||||
held_turns = activity.assistant_count - read_turn
|
||||
|
||||
if quiet_turns < self.config.quiesce_turns and held_turns < self.config.max_hold_turns:
|
||||
result.holding_reads += 1
|
||||
return None, True # file still active — keep verbatim, uncached
|
||||
|
||||
# File quiesced (or hold cap hit): mature.
|
||||
ccr_hash = hashlib.sha256(content.encode("utf-8", errors="replace")).hexdigest()[:24]
|
||||
if self.store is not None:
|
||||
try:
|
||||
ccr_hash = self.store.store(
|
||||
original=content,
|
||||
compressed="",
|
||||
tool_name="Read",
|
||||
tool_call_id=tc_id,
|
||||
compression_strategy="read_maturation",
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 - storage failure must not break the request
|
||||
logger.warning("read_maturation: CCR store failed for %s: %s", tc_id, e)
|
||||
|
||||
file_display = file_path or "unknown"
|
||||
# NOTE: "Retrieve original: hash=" is load-bearing (marker-
|
||||
# preserving regex + ContentRouter compression pinning).
|
||||
marker = (
|
||||
f"[Read of {file_display} compressed after use — re-read the file "
|
||||
f"if needed. Retrieve original: hash={ccr_hash}]"
|
||||
)
|
||||
self._matured[tc_id] = MaturedRead(marker=marker, ccr_hash=ccr_hash)
|
||||
result.newly_matured += 1
|
||||
result.replacements_applied += 1
|
||||
result.bytes_saved += max(0, len(content) - len(marker))
|
||||
return marker, False
|
||||
|
||||
|
||||
def relocate_cache_breakpoint(
|
||||
messages: list[dict[str, Any]],
|
||||
holding_msg_indices: list[int],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Park the trailing message-level cache breakpoint before held Reads.
|
||||
|
||||
Strips ``cache_control`` from every block at or after the earliest
|
||||
holding message, and places one ephemeral breakpoint on the last
|
||||
block of the latest *eligible* message before it — so the provider
|
||||
caches everything up to (not including) the held Reads. System- and
|
||||
tools-level breakpoints are untouched (they live outside messages).
|
||||
|
||||
Total breakpoints never increase: at most one is added after one or
|
||||
more are removed. Returns the original list unchanged when there is
|
||||
nothing to do.
|
||||
"""
|
||||
if not holding_msg_indices:
|
||||
return messages
|
||||
|
||||
earliest = min(holding_msg_indices)
|
||||
out: list[dict[str, Any]] = list(messages)
|
||||
stripped_any = False
|
||||
|
||||
# 1. Strip breakpoints from the held region [earliest:].
|
||||
for i in range(earliest, len(out)):
|
||||
msg = out[i]
|
||||
content = msg.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
if any(isinstance(b, dict) and "cache_control" in b for b in content):
|
||||
out[i] = {
|
||||
**msg,
|
||||
"content": [
|
||||
{k: v for k, v in b.items() if k != "cache_control"}
|
||||
if isinstance(b, dict)
|
||||
else b
|
||||
for b in content
|
||||
],
|
||||
}
|
||||
stripped_any = True
|
||||
|
||||
if not stripped_any:
|
||||
# No client breakpoint in the held region — nothing was going to
|
||||
# cache the held Reads this request; leave placement alone.
|
||||
return out
|
||||
|
||||
# 2. Re-anchor: ephemeral breakpoint on the last block of the latest
|
||||
# block-style message before the held region.
|
||||
for i in range(earliest - 1, -1, -1):
|
||||
content = out[i].get("content")
|
||||
if isinstance(content, list) and content and isinstance(content[-1], dict):
|
||||
new_content = list(content)
|
||||
new_content[-1] = {**new_content[-1], "cache_control": {"type": "ephemeral"}}
|
||||
out[i] = {**out[i], "content": new_content}
|
||||
break
|
||||
|
||||
return out
|
||||
181
tests/test_live/test_live_maturation.py
Normal file
181
tests/test_live/test_live_maturation.py
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
"""Live validation of Mechanism B's no-bust invariant against the real API.
|
||||
|
||||
The design's central claim, tested empirically:
|
||||
|
||||
1. Request A holds a fresh Read out of the cache (breakpoint relocated
|
||||
to just before it) → the provider's cache_creation must NOT include
|
||||
the read content.
|
||||
2. Request B (one turn later, read matured into a marker, breakpoint
|
||||
back at the tail) → the provider must report a cache READ covering
|
||||
request A's cached prefix — proving the prefix survived the read's
|
||||
replacement, i.e. nothing was busted.
|
||||
|
||||
If assertion 2 fails, breakpoint relocation breaks prefix matching and
|
||||
the mechanism needs redesign before it is enabled anywhere.
|
||||
|
||||
Skipped without ANTHROPIC_API_KEY. Costs ~15K haiku tokens per run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from headroom.config import ReadMaturationConfig
|
||||
from headroom.transforms.read_maturation import (
|
||||
ReadMaturationManager,
|
||||
relocate_cache_breakpoint,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not os.environ.get("ANTHROPIC_API_KEY"),
|
||||
reason="ANTHROPIC_API_KEY not set",
|
||||
)
|
||||
|
||||
MODEL = "claude-haiku-4-5-20251001"
|
||||
API_URL = "https://api.anthropic.com/v1/messages"
|
||||
|
||||
# System pad: must clear the model's minimum cacheable prefix (haiku:
|
||||
# 2048 tokens) on its own, so request A caches system+early messages.
|
||||
SYSTEM_PAD = (
|
||||
"You are a coding assistant. Policy clause %d: always be precise and "
|
||||
"verify against the source before answering. " * 400
|
||||
) % tuple(range(400))
|
||||
|
||||
# Read content: big enough to dominate the message tokens (~4K tokens),
|
||||
# so its presence/absence in cache numbers is unambiguous.
|
||||
FILE_CONTENT = "".join(
|
||||
f" {i}\tdef func_{i}(): return {i} # padding comment line {i}\n" for i in range(700)
|
||||
)
|
||||
|
||||
READ_TOOL = {
|
||||
"name": "Read",
|
||||
"description": "Read a file",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"file_path": {"type": "string"}},
|
||||
"required": ["file_path"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def call(messages: list[dict]) -> dict:
|
||||
resp = httpx.post(
|
||||
API_URL,
|
||||
json={
|
||||
"model": MODEL,
|
||||
"max_tokens": 50,
|
||||
"system": [
|
||||
{"type": "text", "text": SYSTEM_PAD, "cache_control": {"type": "ephemeral"}}
|
||||
],
|
||||
"tools": [READ_TOOL],
|
||||
"messages": messages,
|
||||
},
|
||||
headers={
|
||||
"x-api-key": os.environ["ANTHROPIC_API_KEY"],
|
||||
"anthropic-version": "2023-06-01",
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
assert resp.status_code == 200, f"{resp.status_code}: {resp.text[:500]}"
|
||||
return resp.json()["usage"]
|
||||
|
||||
|
||||
def conv_base() -> list[dict]:
|
||||
return [
|
||||
{"role": "user", "content": [{"type": "text", "text": "Read /src/pad.py please"}]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_r1",
|
||||
"name": "Read",
|
||||
"input": {"file_path": "/src/pad.py"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_r1",
|
||||
"content": FILE_CONTENT,
|
||||
# Claude Code-style tail breakpoint on the newest block.
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class TestNoBustInvariantLive:
|
||||
def test_hold_then_mature_preserves_cache(self):
|
||||
mgr = ReadMaturationManager(ReadMaturationConfig(enabled=True, quiesce_turns=1))
|
||||
|
||||
# ── Request A: fresh read → held, breakpoint relocated before it.
|
||||
msgs_a = conv_base()
|
||||
res_a = mgr.apply(msgs_a)
|
||||
assert res_a.holding_msg_indices == [2], "fixture must trigger holding"
|
||||
fwd_a = relocate_cache_breakpoint(res_a.messages, res_a.holding_msg_indices)
|
||||
assert "cache_control" not in fwd_a[2]["content"][0]
|
||||
|
||||
usage_a = call(fwd_a)
|
||||
created_a = usage_a.get("cache_creation_input_tokens", 0)
|
||||
input_a = usage_a.get("input_tokens", 0)
|
||||
# The held read (~4K tokens of input) must NOT be in the cache
|
||||
# write. created_a covers system pad + first two messages only.
|
||||
assert created_a > 0, f"prefix did not cache at all: {usage_a}"
|
||||
assert input_a > 3000, f"read content missing from input: {usage_a}"
|
||||
total_a = created_a + input_a + usage_a.get("cache_read_input_tokens", 0)
|
||||
assert created_a < total_a * 0.7, (
|
||||
f"cache write covered the held read — hold failed: {usage_a}"
|
||||
)
|
||||
|
||||
# ── Request B: one assistant turn later, file quiet → matured.
|
||||
msgs_b = [
|
||||
*conv_base(),
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "Read it."}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Thanks. Reply with the single word: done",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
# The client breakpoint moved to the new tail; the old read block
|
||||
# no longer carries one.
|
||||
del msgs_b[2]["content"][0]["cache_control"]
|
||||
|
||||
res_b = mgr.apply(msgs_b)
|
||||
assert res_b.newly_matured == 1, "read must mature after quiesce"
|
||||
fwd_b = relocate_cache_breakpoint(res_b.messages, res_b.holding_msg_indices)
|
||||
marker = fwd_b[2]["content"][0]["content"]
|
||||
assert "Retrieve original: hash=" in marker
|
||||
|
||||
usage_b = call(fwd_b)
|
||||
read_b = usage_b.get("cache_read_input_tokens", 0)
|
||||
|
||||
# THE invariant: request A's cached prefix must still be valid —
|
||||
# the matured read sat outside it, so replacing it busts nothing.
|
||||
assert read_b >= created_a * 0.9, (
|
||||
f"NO-BUST INVARIANT FAILED: request B read {read_b} cached tokens "
|
||||
f"but request A created {created_a} — breakpoint relocation broke "
|
||||
f"prefix matching. A={usage_a} B={usage_b}"
|
||||
)
|
||||
# And the matured form is small: B's uncached input should be far
|
||||
# below the read size (marker + two short turns, not 4K tokens).
|
||||
assert usage_b.get("input_tokens", 0) < 2500, (
|
||||
f"matured request still carried heavy uncached input: {usage_b}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
353
tests/test_read_maturation.py
Normal file
353
tests/test_read_maturation.py
Normal file
|
|
@ -0,0 +1,353 @@
|
|||
"""Tests for Mechanism B: activity-based hold-back Read maturation.
|
||||
|
||||
The invariants under test, beyond decision behavior:
|
||||
1. No cached byte is ever mutated — frozen-prefix content and content
|
||||
carrying a client cache_control breakpoint are untouched.
|
||||
2. Replay is deterministic — once matured, the same marker is applied on
|
||||
every subsequent request, byte-identical.
|
||||
3. Holding is derived from the conversation (file activity), so the
|
||||
decision survives state loss; only matured markers are stateful.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.config import ReadMaturationConfig
|
||||
from headroom.transforms.read_maturation import (
|
||||
ReadMaturationManager,
|
||||
relocate_cache_breakpoint,
|
||||
)
|
||||
|
||||
CONTENT = " 1\tdef foo():\n 2\t return 42\n" * 60 # > 2048B
|
||||
SMALL = " 1\tok\n"
|
||||
|
||||
|
||||
def anthropic_read(tc_id: str, file_path: str, content: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": tc_id, "name": "Read", "input": {"file_path": file_path}}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": tc_id, "content": content}],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def anthropic_edit(tc_id: str, file_path: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": tc_id,
|
||||
"name": "Edit",
|
||||
"input": {"file_path": file_path, "old_string": "a", "new_string": "b"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": tc_id, "content": "ok"}],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def openai_read(tc_id: str, file_path: str, content: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tc_id,
|
||||
"function": {"name": "Read", "arguments": f'{{"file_path": "{file_path}"}}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": tc_id, "content": content},
|
||||
]
|
||||
|
||||
|
||||
def quiet(n: int) -> list[dict]:
|
||||
"""n assistant turns with no file activity (advances the quiet clock)."""
|
||||
return [
|
||||
{"role": "assistant", "content": [{"type": "text", "text": f"thinking {i}"}]}
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
|
||||
def base_conv() -> list[dict]:
|
||||
return [{"role": "user", "content": "look"}, *anthropic_read("r1", "/x/foo.py", CONTENT)]
|
||||
|
||||
|
||||
def manager(**overrides) -> ReadMaturationManager:
|
||||
cfg = ReadMaturationConfig(enabled=True, **overrides)
|
||||
return ReadMaturationManager(cfg)
|
||||
|
||||
|
||||
def read_content(res, idx=2):
|
||||
return res.messages[idx]["content"][0]["content"]
|
||||
|
||||
|
||||
class TestActivityDecision:
|
||||
def test_disabled_is_noop(self):
|
||||
m = ReadMaturationManager(ReadMaturationConfig(enabled=False))
|
||||
res = m.apply(base_conv())
|
||||
assert res.messages == base_conv()
|
||||
assert res.holding_msg_indices == []
|
||||
|
||||
def test_fresh_read_holds_verbatim(self):
|
||||
res = manager().apply(base_conv())
|
||||
assert read_content(res) == CONTENT
|
||||
assert res.holding_msg_indices == [2]
|
||||
assert res.holding_reads == 1
|
||||
assert res.newly_matured == 0
|
||||
|
||||
def test_holds_while_file_quiet_below_quiesce(self):
|
||||
msgs = [*base_conv(), *quiet(4)] # quiet = 4 < 5
|
||||
res = manager(quiesce_turns=5).apply(msgs)
|
||||
assert res.holding_msg_indices == [2]
|
||||
assert read_content(res) == CONTENT
|
||||
|
||||
def test_matures_after_quiesce(self):
|
||||
msgs = [*base_conv(), *quiet(5)] # quiet = 5 >= 5
|
||||
res = manager(quiesce_turns=5).apply(msgs)
|
||||
assert res.newly_matured == 1
|
||||
assert res.holding_msg_indices == []
|
||||
marker = read_content(res)
|
||||
assert "compressed after use" in marker
|
||||
assert "/x/foo.py" in marker
|
||||
assert "Retrieve original: hash=" in marker
|
||||
assert res.bytes_saved > 0
|
||||
|
||||
def test_file_activity_resets_quiet_clock(self):
|
||||
# read, 4 quiet, edit same file, 4 quiet: quiet=4 < 5 → still held
|
||||
msgs = [*base_conv(), *quiet(4), *anthropic_edit("e1", "/x/foo.py"), *quiet(4)]
|
||||
res = manager(quiesce_turns=5).apply(msgs)
|
||||
assert res.holding_msg_indices == [2]
|
||||
# one more quiet turn → file quiesced → matures
|
||||
res = manager(quiesce_turns=5).apply([*msgs, *quiet(1)])
|
||||
assert res.newly_matured == 1
|
||||
|
||||
def test_activity_on_other_file_does_not_reset(self):
|
||||
msgs = [*base_conv(), *quiet(3), *anthropic_edit("e1", "/x/OTHER.py"), *quiet(1)]
|
||||
# foo.py quiet for 5 assistant turns (3 quiet + edit-turn + 1 quiet)
|
||||
res = manager(quiesce_turns=5).apply(msgs)
|
||||
assert res.newly_matured == 1
|
||||
|
||||
def test_max_hold_caps_busy_files(self):
|
||||
# File touched every turn — never quiesces — but the hold cap fires.
|
||||
msgs = base_conv()
|
||||
for i in range(6):
|
||||
msgs += anthropic_edit(f"e{i}", "/x/foo.py")
|
||||
res = manager(quiesce_turns=100, max_hold_turns=6).apply(msgs)
|
||||
assert res.newly_matured == 1
|
||||
assert res.holding_msg_indices == []
|
||||
|
||||
def test_replay_is_deterministic_and_stateful(self):
|
||||
m = manager(quiesce_turns=5)
|
||||
matured_msgs = [*base_conv(), *quiet(5)]
|
||||
a = read_content(m.apply(matured_msgs))
|
||||
# Replay applies even when the conversation grows and the file is
|
||||
# touched again later (matured is final).
|
||||
later = [*matured_msgs, *anthropic_edit("e1", "/x/foo.py")]
|
||||
b = read_content(m.apply(later))
|
||||
c = read_content(m.apply([*later, *quiet(3)]))
|
||||
assert a == b == c
|
||||
|
||||
def test_small_reads_ignored(self):
|
||||
msgs = [
|
||||
{"role": "user", "content": "look"},
|
||||
*anthropic_read("r1", "/x/a.py", SMALL),
|
||||
*quiet(10),
|
||||
]
|
||||
res = manager().apply(msgs)
|
||||
assert res.holding_msg_indices == []
|
||||
assert read_content(res) == SMALL
|
||||
|
||||
def test_frozen_prefix_untouched(self):
|
||||
msgs = [*base_conv(), *quiet(10)]
|
||||
m = manager(quiesce_turns=5)
|
||||
res = m.apply(msgs, frozen_message_count=len(msgs))
|
||||
assert res.holding_msg_indices == []
|
||||
assert read_content(res) == CONTENT
|
||||
|
||||
def test_respects_lifecycle_markers(self):
|
||||
marker = "[Read content stale: /x/foo.py ... Retrieve original: hash=abc123]" + " " * 2048
|
||||
msgs = [
|
||||
{"role": "user", "content": "look"},
|
||||
*anthropic_read("r1", "/x/foo.py", marker),
|
||||
*quiet(10),
|
||||
]
|
||||
res = manager().apply(msgs)
|
||||
assert res.holding_msg_indices == []
|
||||
assert res.newly_matured == 0
|
||||
|
||||
def test_client_breakpoint_on_fresh_read_is_held_and_relocated(self):
|
||||
"""Claude Code parks its tail breakpoint on the newest block —
|
||||
right after a Read, that's the Read result itself. The read must
|
||||
still be held (verbatim) and relocation must move the breakpoint
|
||||
off it, otherwise the verbatim form gets cache-written."""
|
||||
msgs = base_conv()
|
||||
msgs[2]["content"][0]["cache_control"] = {"type": "ephemeral"}
|
||||
res = manager().apply(msgs)
|
||||
assert res.holding_msg_indices == [2]
|
||||
assert res.messages[2]["content"][0]["content"] == CONTENT
|
||||
|
||||
out = relocate_cache_breakpoint(res.messages, res.holding_msg_indices)
|
||||
# Breakpoint stripped from the held read, re-anchored before it.
|
||||
assert "cache_control" not in out[2]["content"][0]
|
||||
assert out[1]["content"][-1].get("cache_control") == {"type": "ephemeral"}
|
||||
|
||||
def test_openai_format(self):
|
||||
msgs = [{"role": "user", "content": "look"}, *openai_read("r1", "/x/foo.py", CONTENT)]
|
||||
m = manager(quiesce_turns=5)
|
||||
res = m.apply(msgs)
|
||||
assert res.holding_msg_indices == [2]
|
||||
res = m.apply([*msgs, *quiet(5)])
|
||||
assert "compressed after use" in res.messages[2]["content"]
|
||||
|
||||
def test_files_mature_independently(self):
|
||||
# foo.py quiet for ages; bar.py just read → foo matures, bar holds.
|
||||
msgs = [
|
||||
*base_conv(),
|
||||
*quiet(6),
|
||||
*anthropic_read("r2", "/x/bar.py", CONTENT),
|
||||
]
|
||||
res = manager(quiesce_turns=5).apply(msgs)
|
||||
assert res.newly_matured == 1 # foo.py
|
||||
assert "compressed after use" in read_content(res, 2)
|
||||
assert res.holding_msg_indices == [10] # bar.py result message
|
||||
assert read_content(res, 10) == CONTENT
|
||||
|
||||
def test_decision_survives_state_loss(self):
|
||||
# A fresh manager (proxy restart) makes the same holding/matured
|
||||
# decision because holding is derived from the conversation.
|
||||
msgs = [*base_conv(), *quiet(5)]
|
||||
first = read_content(manager(quiesce_turns=5).apply(msgs))
|
||||
second = read_content(manager(quiesce_turns=5).apply(msgs))
|
||||
assert "compressed after use" in first and "compressed after use" in second
|
||||
# Markers differ only if CCR hashing differed — it must not.
|
||||
assert first == second
|
||||
|
||||
|
||||
class TestCcrIntegration:
|
||||
def test_original_stored_and_retrievable(self):
|
||||
from headroom.cache.backends.memory import InMemoryBackend
|
||||
from headroom.cache.compression_store import CompressionStore
|
||||
|
||||
store = CompressionStore(backend=InMemoryBackend())
|
||||
m = ReadMaturationManager(
|
||||
ReadMaturationConfig(enabled=True, quiesce_turns=5), compression_store=store
|
||||
)
|
||||
res = m.apply([*base_conv(), *quiet(5)])
|
||||
|
||||
marker = read_content(res)
|
||||
ccr_hash = marker.split("hash=")[1].rstrip("]")
|
||||
entry = store.retrieve(ccr_hash)
|
||||
assert entry is not None
|
||||
assert entry.original_content == CONTENT
|
||||
assert entry.compression_strategy == "read_maturation"
|
||||
|
||||
|
||||
class TestProxyWiring:
|
||||
def test_proxy_config_flag_default_off(self):
|
||||
from headroom.proxy.models import ProxyConfig
|
||||
|
||||
assert ProxyConfig().read_maturation is False
|
||||
assert ProxyConfig(read_maturation=True).read_maturation is True
|
||||
|
||||
def test_session_state_rides_on_prefix_tracker(self):
|
||||
"""The handler's session-state flow: a manager attached to the
|
||||
tracker carries matured markers across requests."""
|
||||
from headroom.cache.prefix_tracker import PrefixCacheTracker
|
||||
from headroom.config import ReadMaturationConfig
|
||||
from headroom.transforms.read_maturation import (
|
||||
ReadMaturationManager,
|
||||
relocate_cache_breakpoint,
|
||||
)
|
||||
|
||||
tracker = PrefixCacheTracker("anthropic")
|
||||
assert tracker.read_maturation_manager is None
|
||||
|
||||
# Request 1: fresh read — held; breakpoint relocated.
|
||||
tracker.read_maturation_manager = ReadMaturationManager(
|
||||
ReadMaturationConfig(enabled=True, quiesce_turns=5)
|
||||
)
|
||||
msgs = base_conv()
|
||||
msgs[2]["content"][0] = {
|
||||
**msgs[2]["content"][0],
|
||||
}
|
||||
res = tracker.read_maturation_manager.apply(msgs)
|
||||
assert res.holding_msg_indices == [2]
|
||||
out = relocate_cache_breakpoint(res.messages, res.holding_msg_indices)
|
||||
assert len(out) == len(msgs)
|
||||
|
||||
# Request N (file quiet): same manager matures and replays.
|
||||
later = [*base_conv(), *quiet(5)]
|
||||
res = tracker.read_maturation_manager.apply(later)
|
||||
assert res.newly_matured == 1
|
||||
replay = tracker.read_maturation_manager.apply(later)
|
||||
assert replay.replacements_applied == 1
|
||||
assert replay.newly_matured == 0
|
||||
|
||||
|
||||
class TestBreakpointRelocation:
|
||||
def _msgs_with_tail_breakpoint(self) -> list[dict]:
|
||||
msgs = [
|
||||
{"role": "user", "content": [{"type": "text", "text": "earlier turn"}]},
|
||||
*anthropic_read("r1", "/x/foo.py", CONTENT),
|
||||
]
|
||||
msgs[-1]["content"][-1] = {
|
||||
**msgs[-1]["content"][-1],
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
return msgs
|
||||
|
||||
@staticmethod
|
||||
def _breakpoint_indices(msgs: list[dict]) -> list[int]:
|
||||
return [
|
||||
i
|
||||
for i, m in enumerate(msgs)
|
||||
if isinstance(m.get("content"), list)
|
||||
and any(isinstance(b, dict) and "cache_control" in b for b in m["content"])
|
||||
]
|
||||
|
||||
def test_noop_without_holds(self):
|
||||
msgs = self._msgs_with_tail_breakpoint()
|
||||
assert relocate_cache_breakpoint(msgs, []) is msgs
|
||||
|
||||
def test_relocates_before_held_read(self):
|
||||
msgs = self._msgs_with_tail_breakpoint()
|
||||
out = relocate_cache_breakpoint(msgs, [2])
|
||||
|
||||
# Held region [2:] carries no breakpoint; re-anchored on the
|
||||
# latest eligible message before it (index 1 — the assistant
|
||||
# tool_use message), so everything up to but excluding the held
|
||||
# Read still gets cached.
|
||||
assert self._breakpoint_indices(out) == [1]
|
||||
assert out[1]["content"][-1]["cache_control"] == {"type": "ephemeral"}
|
||||
assert len(self._breakpoint_indices(out)) <= len(self._breakpoint_indices(msgs))
|
||||
|
||||
def test_noop_when_no_breakpoint_in_held_region(self):
|
||||
msgs = [
|
||||
{"role": "user", "content": [{"type": "text", "text": "x"}]},
|
||||
*anthropic_read("r1", "/x/foo.py", CONTENT),
|
||||
]
|
||||
out = relocate_cache_breakpoint(msgs, [2])
|
||||
assert self._breakpoint_indices(out) == []
|
||||
|
||||
def test_originals_not_mutated(self):
|
||||
msgs = self._msgs_with_tail_breakpoint()
|
||||
before = [str(m) for m in msgs]
|
||||
relocate_cache_breakpoint(msgs, [2])
|
||||
assert [str(m) for m in msgs] == before
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
372
tests/test_read_maturation_handler_nobust.py
Normal file
372
tests/test_read_maturation_handler_nobust.py
Normal file
|
|
@ -0,0 +1,372 @@
|
|||
"""Integration: Mechanism B (read maturation) no-bust invariant, through the
|
||||
REAL Anthropic handler, across a multi-turn session.
|
||||
|
||||
The design's central claim is that the verbatim Read is held *out* of the
|
||||
provider prefix cache until it matures, so "no cached byte is ever mutated."
|
||||
The unit tests in ``test_read_maturation.py`` call the manager in isolation
|
||||
with ``frozen_message_count=0``; the live test in ``test_live/`` does a
|
||||
2-request hold->mature with no intermediate turns. Neither exercises the
|
||||
realistic path where a held Read sits across several turns while the prefix
|
||||
tracker advances ``frozen_message_count`` from the provider's reported cache
|
||||
usage.
|
||||
|
||||
This is a regression test for that path. It drives the real handler with a
|
||||
mocked upstream that echoes the cache usage Anthropic would report (caching
|
||||
everything up to the breakpoint the handler chose, system blocks included),
|
||||
so the prefix tracker advances exactly as in production. It then asserts,
|
||||
directly on the FORWARDED bytes:
|
||||
|
||||
1. no-bust: the verbatim Read is never forwarded inside the cached prefix
|
||||
(at or before the last cache_control breakpoint) — if it were, maturing it
|
||||
later would mutate a cached byte and bust the prefix;
|
||||
2. the mechanism actually engages: the Read is held verbatim (out of cache)
|
||||
while the file is active, then matures into a CCR marker once it quiesces,
|
||||
in that order.
|
||||
|
||||
Note on cache-state isolation: the CCR store is persistent (SQLite at
|
||||
~/.headroom/ccr_store.db by default) and shared across processes, so stale
|
||||
entries from prior runs can perturb maturation timing. Run against a clean
|
||||
store for deterministic results.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
MODEL = "claude-haiku-4-5-20251001"
|
||||
SYSTEM = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are a coding assistant. Be terse. " * 200,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
READ_TOOL = {
|
||||
"name": "Read",
|
||||
"description": "Read a file",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"file_path": {"type": "string"}},
|
||||
"required": ["file_path"],
|
||||
},
|
||||
}
|
||||
READ_ID = "toolu_r1"
|
||||
# Big enough to dominate message tokens and clear the maturation min-size gate.
|
||||
BIG = "".join(f" {i}\tdef f_{i}(): return {i} # line {i}\n" for i in range(700))
|
||||
|
||||
|
||||
def _read_pair(tail: bool) -> list[dict]:
|
||||
tr = {"type": "tool_result", "tool_use_id": READ_ID, "content": BIG}
|
||||
if tail:
|
||||
tr["cache_control"] = {"type": "ephemeral"}
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": READ_ID,
|
||||
"name": "Read",
|
||||
"input": {"file_path": "/x/foo.py"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [tr]},
|
||||
]
|
||||
|
||||
|
||||
def _quiet_pair(i: int, tail: bool) -> list[dict]:
|
||||
u = {"type": "text", "text": f"Unrelated question {i}: what is {i}+{i}?"}
|
||||
if tail:
|
||||
u["cache_control"] = {"type": "ephemeral"}
|
||||
return [
|
||||
{"role": "assistant", "content": [{"type": "text", "text": str(2 * i)}]},
|
||||
{"role": "user", "content": [u]},
|
||||
]
|
||||
|
||||
|
||||
def _convo(nquiet: int) -> list[dict]:
|
||||
"""Read of /x/foo.py followed by ``nquiet`` turns that never touch it.
|
||||
The Claude-Code-style tail breakpoint rides the newest user block."""
|
||||
msgs: list[dict] = [{"role": "user", "content": [{"type": "text", "text": "Read /x/foo.py"}]}]
|
||||
msgs += _read_pair(tail=(nquiet == 0))
|
||||
for i in range(1, nquiet + 1):
|
||||
msgs += _quiet_pair(i, tail=(i == nquiet))
|
||||
return msgs
|
||||
|
||||
|
||||
def _breakpoint_index(messages: list[dict]) -> int:
|
||||
"""Index of the last message carrying a cache_control block (-1 if none).
|
||||
Anthropic caches everything up to AND INCLUDING this message."""
|
||||
bp = -1
|
||||
for i, m in enumerate(messages):
|
||||
c = m.get("content")
|
||||
if isinstance(c, list) and any(isinstance(b, dict) and "cache_control" in b for b in c):
|
||||
bp = i
|
||||
return bp
|
||||
|
||||
|
||||
def _read_result_content(message: dict) -> str | None:
|
||||
c = message.get("content")
|
||||
if isinstance(c, list):
|
||||
for b in c:
|
||||
if (
|
||||
isinstance(b, dict)
|
||||
and b.get("type") == "tool_result"
|
||||
and b.get("tool_use_id") == READ_ID
|
||||
):
|
||||
return b.get("content")
|
||||
return None
|
||||
|
||||
|
||||
def _est_tokens(message: dict) -> int:
|
||||
return max(1, len(str(message.get("content", ""))) // 4)
|
||||
|
||||
|
||||
def test_verbatim_read_never_cache_written_before_maturation(monkeypatch):
|
||||
# Isolate the CCR store: it is persistent (SQLite) and shared across
|
||||
# processes by default, so stale entries from other runs would perturb
|
||||
# maturation timing and make this test non-deterministic. The in-memory
|
||||
# backend gives a pristine store per test.
|
||||
from headroom.cache.compression_store import reset_compression_store
|
||||
|
||||
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
|
||||
reset_compression_store()
|
||||
|
||||
# Match the real proxy: cache machinery ON (the prefix tracker + compression
|
||||
# cache are what maturation's hold/frozen-count logic depends on). Disabling
|
||||
# them masks the behavior under test.
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
read_maturation=True,
|
||||
mode="token",
|
||||
cache_enabled=True,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
)
|
||||
app = create_app(config)
|
||||
forwarded: list[list[dict]] = []
|
||||
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
original_retry = proxy._retry_request
|
||||
|
||||
async def _mock_upstream(method, url, headers, body, stream=False, **kwargs):
|
||||
msgs = body.get("messages", []) or []
|
||||
forwarded.append(copy.deepcopy(msgs))
|
||||
# Simulate Anthropic honestly caching up to the handler's breakpoint
|
||||
# (system blocks are cached too), so the prefix tracker advances
|
||||
# frozen_message_count as in prod.
|
||||
bp = _breakpoint_index(msgs)
|
||||
sys_tokens = sum(
|
||||
max(1, len(str(b.get("text", ""))) // 4)
|
||||
for b in (body.get("system") or [])
|
||||
if isinstance(b, dict)
|
||||
)
|
||||
cached = sys_tokens + (sum(_est_tokens(m) for m in msgs[: bp + 1]) if bp >= 0 else 0)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_x",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"usage": {
|
||||
"input_tokens": 20,
|
||||
"output_tokens": 2,
|
||||
"cache_read_input_tokens": cached,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
proxy._retry_request = _mock_upstream
|
||||
try:
|
||||
for n in range(0, 7):
|
||||
r = client.post(
|
||||
"/v1/messages",
|
||||
headers={
|
||||
"x-api-key": "test-key",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"x-headroom-session-id": "nobust-1",
|
||||
"content-type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": MODEL,
|
||||
"max_tokens": 20,
|
||||
"system": SYSTEM,
|
||||
"tools": [READ_TOOL],
|
||||
"messages": _convo(n),
|
||||
},
|
||||
)
|
||||
assert r.status_code == 200, f"turn {n}: {r.text[:300]}"
|
||||
finally:
|
||||
proxy._retry_request = original_retry
|
||||
|
||||
assert forwarded, "no requests were forwarded"
|
||||
|
||||
# Per-turn classification of the Read's forwarded form.
|
||||
held_verbatim = [] # turns where the verbatim Read is OUTSIDE the cache prefix (correct hold)
|
||||
cached_verbatim = [] # turns where the verbatim Read is INSIDE the cache prefix (bust risk)
|
||||
matured = [] # turns where the Read has become a CCR marker
|
||||
for turn, msgs in enumerate(forwarded):
|
||||
bp = _breakpoint_index(msgs)
|
||||
for i, m in enumerate(msgs):
|
||||
content = _read_result_content(m)
|
||||
if content is None:
|
||||
continue
|
||||
if content == BIG:
|
||||
(cached_verbatim if i <= bp else held_verbatim).append(turn)
|
||||
elif "Retrieve original: hash=" in content:
|
||||
matured.append(turn)
|
||||
|
||||
# INVARIANT 1 (no-bust): the verbatim Read must never be forwarded inside
|
||||
# the cached prefix. If it is, maturing it later mutates a cached byte.
|
||||
assert not cached_verbatim, (
|
||||
"no-bust invariant violated: verbatim Read was cache-written before "
|
||||
f"maturation on turn(s) {cached_verbatim}. Maturing it later busts the cache."
|
||||
)
|
||||
|
||||
# INVARIANT 2 (mechanism actually engages): the Read is held verbatim while
|
||||
# the file is active, then matures once it quiesces. Guards against a
|
||||
# vacuous pass where maturation silently no-ops.
|
||||
assert held_verbatim, "expected the fresh Read to be held verbatim out of cache on early turns"
|
||||
assert matured, "expected the Read to mature into a CCR marker after quiescing"
|
||||
# The matured marker only appears AFTER the verbatim hold (ordering).
|
||||
assert min(matured) > max(held_verbatim), (
|
||||
f"maturation must follow the hold: held={held_verbatim} matured={matured}"
|
||||
)
|
||||
|
||||
|
||||
def _drive_session(config, n_turns: int, session_id: str) -> list[list[dict]]:
|
||||
"""Drive ``n_turns`` cumulative turns through the real handler with a mocked
|
||||
upstream; return the forwarded message arrays per turn."""
|
||||
app = create_app(config)
|
||||
forwarded: list[list[dict]] = []
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
original_retry = proxy._retry_request
|
||||
|
||||
async def _mock_upstream(method, url, headers, body, stream=False, **kwargs):
|
||||
msgs = body.get("messages", []) or []
|
||||
forwarded.append(copy.deepcopy(msgs))
|
||||
bp = _breakpoint_index(msgs)
|
||||
sys_tokens = sum(
|
||||
max(1, len(str(b.get("text", ""))) // 4)
|
||||
for b in (body.get("system") or [])
|
||||
if isinstance(b, dict)
|
||||
)
|
||||
cached = sys_tokens + (sum(_est_tokens(m) for m in msgs[: bp + 1]) if bp >= 0 else 0)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_x",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"usage": {
|
||||
"input_tokens": 20,
|
||||
"output_tokens": 2,
|
||||
"cache_read_input_tokens": cached,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
proxy._retry_request = _mock_upstream
|
||||
try:
|
||||
for n in range(n_turns):
|
||||
r = client.post(
|
||||
"/v1/messages",
|
||||
headers={
|
||||
"x-api-key": "test-key",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"x-headroom-session-id": session_id,
|
||||
"content-type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": MODEL,
|
||||
"max_tokens": 20,
|
||||
"system": SYSTEM,
|
||||
"tools": [READ_TOOL],
|
||||
"messages": _convo(n),
|
||||
},
|
||||
)
|
||||
assert r.status_code == 200, f"turn {n}: {r.text[:300]}"
|
||||
finally:
|
||||
proxy._retry_request = original_retry
|
||||
return forwarded
|
||||
|
||||
|
||||
def _first_matured_turn(forwarded: list[list[dict]]) -> int | None:
|
||||
"""The first turn index whose forwarded Read is a CCR marker."""
|
||||
for turn, msgs in enumerate(forwarded):
|
||||
for m in msgs:
|
||||
content = _read_result_content(m)
|
||||
if content and "Retrieve original: hash=" in content:
|
||||
return turn
|
||||
return None
|
||||
|
||||
|
||||
def test_quiesce_turns_config_is_honored(monkeypatch):
|
||||
"""`quiesce_turns` must be runtime-configurable end-to-end: a fresh Read of
|
||||
/x/foo.py matures `quiesce_turns` quiet turns after it appears (the convo
|
||||
builds one quiet assistant turn per step, and the Read sits at assistant
|
||||
turn 1). With quiesce_turns=2 it must mature at turn 2 — not the built-in
|
||||
default of 5. Currently the handler hardcodes ReadMaturationConfig(enabled=
|
||||
True), ignoring the configured value, so this fails (matures at 5)."""
|
||||
from headroom.cache.compression_store import reset_compression_store
|
||||
|
||||
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
|
||||
reset_compression_store()
|
||||
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
read_maturation=True,
|
||||
read_maturation_quiesce_turns=2,
|
||||
mode="token",
|
||||
cache_enabled=True,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
)
|
||||
forwarded = _drive_session(config, n_turns=4, session_id="quiesce-cfg-1")
|
||||
first = _first_matured_turn(forwarded)
|
||||
assert first == 2, (
|
||||
f"expected the Read to mature at turn 2 with quiesce_turns=2, "
|
||||
f"but first matured at turn {first} (handler ignored the configured value)"
|
||||
)
|
||||
|
||||
|
||||
def test_read_maturation_knobs_from_env(monkeypatch):
|
||||
"""Operators must be able to tune maturation via env vars (the pilot
|
||||
playbook says 'pick quiesce_turns')."""
|
||||
from headroom.proxy.server import _MULTI_WORKER_CONFIG_ENV, _proxy_config_from_env
|
||||
|
||||
# _proxy_config_from_env short-circuits on a prebuilt multi-worker JSON
|
||||
# config and ignores the HEADROOM_* vars entirely. Clear it so this test
|
||||
# actually exercises the env-var parsing path it claims to (and isn't
|
||||
# poisoned by a leaked HEADROOM_PROXY_CONFIG_JSON from another test).
|
||||
monkeypatch.delenv(_MULTI_WORKER_CONFIG_ENV, raising=False)
|
||||
|
||||
monkeypatch.setenv("HEADROOM_READ_MATURATION", "1")
|
||||
monkeypatch.setenv("HEADROOM_READ_MATURATION_QUIESCE_TURNS", "3")
|
||||
monkeypatch.setenv("HEADROOM_READ_MATURATION_MAX_HOLD_TURNS", "10")
|
||||
monkeypatch.setenv("HEADROOM_READ_MATURATION_MIN_SIZE_BYTES", "4096")
|
||||
|
||||
cfg = _proxy_config_from_env()
|
||||
|
||||
assert cfg.read_maturation is True
|
||||
assert cfg.read_maturation_quiesce_turns == 3
|
||||
assert cfg.read_maturation_max_hold_turns == 10
|
||||
assert cfg.read_maturation_min_size_bytes == 4096
|
||||
Loading…
Add table
Add a link
Reference in a new issue