diff --git a/benchmarks/cache_bust_trace_report.py b/benchmarks/cache_bust_trace_report.py index 80c00b204..339dffceb 100644 --- a/benchmarks/cache_bust_trace_report.py +++ b/benchmarks/cache_bust_trace_report.py @@ -322,7 +322,7 @@ def _render_html(events: list[BustEvent], recent_turns_per_session: int) -> str:

Cache Bust Trace Report

Most recent {recent_turns_per_session} turns per session.

- {''.join(sections)} + {"".join(sections)}
""" @@ -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) diff --git a/benchmarks/cache_validation_bundle.py b/benchmarks/cache_validation_bundle.py index f0c33e24f..c0dfa47a2 100644 --- a/benchmarks/cache_validation_bundle.py +++ b/benchmarks/cache_validation_bundle.py @@ -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: diff --git a/benchmarks/claude_session_branch_compare.py b/benchmarks/claude_session_branch_compare.py index 003b08a37..f9ffa88ce 100644 --- a/benchmarks/claude_session_branch_compare.py +++ b/benchmarks/claude_session_branch_compare.py @@ -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: diff --git a/benchmarks/claude_session_mode_benchmark.py b/benchmarks/claude_session_mode_benchmark.py index 5fa917926..9823aabff 100644 --- a/benchmarks/claude_session_mode_benchmark.py +++ b/benchmarks/claude_session_mode_benchmark.py @@ -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 ) diff --git a/benchmarks/synthetic_long_cache_suite_report.py b/benchmarks/synthetic_long_cache_suite_report.py index 9f52013b5..f76b1e266 100644 --- a/benchmarks/synthetic_long_cache_suite_report.py +++ b/benchmarks/synthetic_long_cache_suite_report.py @@ -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()}, diff --git a/benchmarks/synthetic_token_cache_bust_report.py b/benchmarks/synthetic_token_cache_bust_report.py index c7e8124a9..baec40e42 100644 --- a/benchmarks/synthetic_token_cache_bust_report.py +++ b/benchmarks/synthetic_token_cache_bust_report.py @@ -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"
  • no-cache total cost: {html.escape(winners['no_cache_total_cost'])}
    • " f"
  • window with cache counted: {html.escape(winners['window_with_cache'])}
    • " f"
  • window without cache reads: {html.escape(winners['window_without_cache_reads'])}
    • " - "

    Cache Bust Events

    " - + "".join(event_sections) - + "" + "

    Cache Bust Events

    " + "".join(event_sections) + "" ) html_path.write_text(html_doc, encoding="utf-8") return md_path, json_path, html_path diff --git a/tests/test_proxy_anthropic_cache_stability.py b/tests/test_proxy_anthropic_cache_stability.py index d1eeec97b..72e82c827 100644 --- a/tests/test_proxy_anthropic_cache_stability.py +++ b/tests/test_proxy_anthropic_cache_stability.py @@ -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