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

test: align anthropic cache stability fixtures

Browse source 
Sync the Anthropic cache stability test double with the prefix tracker contract used by the handler.

Format the benchmark scripts that were failing ruff format --check in CI.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
JerrettDavis 2026-04-06 23:26:51 -05:00
parent 78163199c4
commit 6f701033a1
7 changed files with 62 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

6
benchmarks/cache_bust_trace_report.py
View file

@ -322,7 +322,7 @@ def _render_html(events: list[BustEvent], recent_turns_per_session: int) -> str:
<div class="shell">
<h1>Cache Bust Trace Report</h1>
<p>Most recent {recent_turns_per_session} turns per session.</p>
{''.join(sections)}
{"".join(sections)}
</div>
</body>
</html>"""
@ -344,7 +344,9 @@ def main() -> int:
json_path = output_dir / "cache_bust_trace.json"
html_path = output_dir / "cache_bust_trace.html"
md_path.write_text(_render_markdown(events, recent_turns_per_session), encoding="utf-8")
json_path.write_text(json.dumps([asdict(event) for event in events], indent=2), encoding="utf-8")
json_path.write_text(
json.dumps([asdict(event) for event in events], indent=2), encoding="utf-8"
)
html_path.write_text(_render_html(events, recent_turns_per_session), encoding="utf-8")
print(md_path)
print(json_path)

36
benchmarks/cache_validation_bundle.py
View file

@ -186,7 +186,9 @@ def _collect_real_processed_events(
stable_prefix_message_count=prior_context_message_count,
)
if rewrite:
prior_forwarded = pending.forwarded if pending is not None else previous_forwarded
prior_forwarded = (
pending.forwarded if pending is not None else previous_forwarded
)
prior_ts = pending.turn.timestamp if pending is not None else previous_timestamp
eligible = bool(
prior_ts is not None
@ -210,23 +212,33 @@ def _collect_real_processed_events(
events.append(
{
"mode": mode,
"session_id": replay.session_id if include_content else _redact_text(replay.session_id, prefix="session"),
"project": replay.decoded_project_path if include_content else _redact_path(replay.decoded_project_path),
"request_id": turn.request_id if include_content else _redact_text(turn.request_id, prefix="request"),
"session_id": replay.session_id
if include_content
else _redact_text(replay.session_id, prefix="session"),
"project": replay.decoded_project_path
if include_content
else _redact_path(replay.decoded_project_path),
"request_id": turn.request_id
if include_content
else _redact_text(turn.request_id, prefix="request"),
"timestamp": turn.timestamp.isoformat(),
"cache_eligible": eligible,
"prefix_preserved": prefix_preserved,
"retroactive_rewrite": retro,
"first_diff_index": first_diff_index,
"original_tail": [
_message_preview(m, max_chars=max_chars) if include_content else {
_message_preview(m, max_chars=max_chars)
if include_content
else {
"role": str(m.get("role")),
"content_excerpt": "[redacted]",
}
for m in conversation[max(0, len(conversation) - 4) :]
],
"forwarded_tail": [
_message_preview(m, max_chars=max_chars) if include_content else {
_message_preview(m, max_chars=max_chars)
if include_content
else {
"role": str(m.get("role")),
"content_excerpt": "[redacted]",
}
@ -248,6 +260,7 @@ def _collect_real_processed_events(
message_token_counts=[tokenizer.count_message(msg) for msg in forwarded],
original_messages=conversation,
)
class Pending:
pass
@ -264,7 +277,9 @@ def _collect_real_processed_events(
return {"events": events}
def _write_processed_event_reports(output_dir: Path, payload: dict[str, Any]) -> tuple[Path, Path, Path]:
def _write_processed_event_reports(
output_dir: Path, payload: dict[str, Any]
) -> tuple[Path, Path, Path]:
out_dir = output_dir / "real_processed"
out_dir.mkdir(parents=True, exist_ok=True)
json_path = out_dir / "real_processed_rewrite_report.json"
@ -524,7 +539,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
parser.add_argument("--recent-turns-per-session", type=int, default=None)
parser.add_argument("--workers", type=int, default=1)
parser.add_argument("--cache-ttl-minutes", type=int, default=real_bench.DEFAULT_CACHE_TTL_MINUTES)
parser.add_argument(
"--cache-ttl-minutes", type=int, default=real_bench.DEFAULT_CACHE_TTL_MINUTES
)
parser.add_argument("--cache-write-multiplier", type=float, default=1.25)
parser.add_argument("--max-sessions", type=int, default=None)
parser.add_argument("--max-real-events-per-mode", type=int, default=8)
@ -604,7 +621,8 @@ def main() -> int:
token_bust.bench._make_proxy = lambda mode: token_bust._FakeProxy()
try:
_, token_bust_summaries = token_bust.simulate_replays(
[token_bust_replay], cache_ttl_minutes=token_bust.TTL_MINUTES if hasattr(token_bust, "TTL_MINUTES") else 5
[token_bust_replay],
cache_ttl_minutes=token_bust.TTL_MINUTES if hasattr(token_bust, "TTL_MINUTES") else 5,
)
token_bust_events = token_bust._build_bust_events(token_bust_replay)
finally:

6
benchmarks/claude_session_branch_compare.py
View file

@ -212,9 +212,9 @@ def _run_branch_benchmark(
workers=args.workers,
)
env = os.environ.copy()
env["PYTHONPATH"] = os.pathsep.join(
[str(worktree_dir), env.get("PYTHONPATH", "")]
).rstrip(os.pathsep)
env["PYTHONPATH"] = os.pathsep.join([str(worktree_dir), env.get("PYTHONPATH", "")]).rstrip(
os.pathsep
)
subprocess.run(command, cwd=worktree_dir, check=True, env=env)
return _load_branch_result(repo_root, ref, label, branch_output_dir)
finally:

4
benchmarks/claude_session_mode_benchmark.py
View file

@ -1183,7 +1183,9 @@ def _simulate_single_replay_mode(
summary.retroactive_rewrite_turns += 1
else:
summary.latest_turn_only_rewrite_turns += 1
prior_forwarded_for_rewrite = pending.forwarded if pending is not None else previous_forwarded
prior_forwarded_for_rewrite = (
pending.forwarded if pending is not None else previous_forwarded
)
prior_timestamp_for_rewrite = (
pending.turn.timestamp if pending is not None else previous_timestamp
)

5
benchmarks/synthetic_long_cache_suite_report.py
View file

@ -230,10 +230,7 @@ def _write_report(
"total_turns": TURNS_PER_SCENARIO * len(per_scenario),
"ttl_minutes": TTL_MINUTES,
"scenarios": {
session_id: {
mode: _summary_payload(summary)
for mode, summary in summaries.items()
}
session_id: {mode: _summary_payload(summary) for mode, summary in summaries.items()}
for session_id, summaries in per_scenario.items()
},
"aggregate": {mode: _summary_payload(summary) for mode, summary in aggregate.items()},

19
benchmarks/synthetic_token_cache_bust_report.py
View file

@ -146,13 +146,10 @@ def _build_bust_events(replay: SessionReplay) -> dict[str, list[dict[str, object
previous_forwarded_messages=previous_forwarded_context,
)
if (
previous_forwarded_request is not None
and _cache_gap_within_ttl(
turn.timestamp,
previous_timestamp,
ttl=bench.timedelta(minutes=ttl_minutes),
)
if previous_forwarded_request is not None and _cache_gap_within_ttl(
turn.timestamp,
previous_timestamp,
ttl=bench.timedelta(minutes=ttl_minutes),
):
prefix_preserved = (
len(forwarded) >= len(previous_forwarded_request)
@ -192,7 +189,9 @@ def _build_bust_events(replay: SessionReplay) -> dict[str, list[dict[str, object
conversation.append(copy.deepcopy(turn.assistant_message))
previous_original = copy.deepcopy(conversation)
previous_forwarded_context = copy.deepcopy(forwarded) + [copy.deepcopy(turn.assistant_message)]
previous_forwarded_context = copy.deepcopy(forwarded) + [
copy.deepcopy(turn.assistant_message)
]
previous_forwarded_request = copy.deepcopy(forwarded)
previous_request_id = turn.request_id
previous_timestamp = turn.timestamp
@ -328,9 +327,7 @@ def _write_report(
f"<li>no-cache total cost: <code>{html.escape(winners['no_cache_total_cost'])}</code></li>"
f"<li>window with cache counted: <code>{html.escape(winners['window_with_cache'])}</code></li>"
f"<li>window without cache reads: <code>{html.escape(winners['window_without_cache_reads'])}</code></li>"
"</ul><h2>Cache Bust Events</h2>"
+ "".join(event_sections)
+ "</body></html>"
"</ul><h2>Cache Bust Events</h2>" + "".join(event_sections) + "</body></html>"
)
html_path.write_text(html_doc, encoding="utf-8")
return md_path, json_path, html_path

16
tests/test_proxy_anthropic_cache_stability.py
View file

@ -19,11 +19,27 @@ from headroom.proxy.server import ProxyConfig, create_app
class _FakePrefixTracker:
def __init__(self, frozen_count: int):
self._frozen_count = frozen_count
self._cached_token_count = 0
self._last_original_messages = []
self._last_forwarded_messages = []
def get_frozen_message_count(self) -> int:
return self._frozen_count
def get_last_original_messages(self): # noqa: ANN201
return self._last_original_messages.copy()
def get_last_forwarded_messages(self): # noqa: ANN201
return self._last_forwarded_messages.copy()
def update_from_response(self, **kwargs): # noqa: ANN003
self._cached_token_count = kwargs.get("cache_read_tokens", 0) + kwargs.get(
"cache_write_tokens", 0
)
self._last_original_messages = kwargs.get(
"original_messages", kwargs.get("messages", [])
).copy()
self._last_forwarded_messages = kwargs.get("messages", []).copy()
return None