fix(proxy/anthropic): don't replay recorded prefix over live history (#3026) (#3052)

## Description

A Claude Code session that reads a large tool result through `headroom
proxy` can fail on turn 2 with Anthropic's `400 prompt_too_long`. The
reporter's controlled comparison completed eight turns through 0.33.0
with 187,986 input tokens, while 0.35.0 failed after five requests with
753,077 input tokens. The local regression uses an actual prior
optimized request to populate tracker state, then a decision-false
bypass turn with Claude-shaped tool-result content. The old
unconditional replay path substitutes the compressed prefix; the
eligibility gate preserves the client's outbound body without claiming a
live provider reproduction.

The Anthropic `/v1/messages` route computes whether a request should be
compressed, but cached-prefix replay currently runs outside that
decision. The replay helper also derives its prefix length from the
original message list and applies that index to the optimized list
without proving the two lists still align. A stale forwarded prefix can
therefore be grafted onto the wrong positions and enlarge later
requests.

This change limits replay to requests whose existing compression
decision permits it and whose pre-upstream backpressure path is
inactive. It also makes `overlay_cached_prefix()` decline misaligned or
inflating candidates while preserving normal append-only replay.

Reported by @itsumonotakumi, whose controlled comparison isolated the
failure from compression, headers, one-request serialization, memory,
code graph, and CCR.

Closes #3026

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Gate Anthropic cached-prefix replay on the existing
`CompressionDecision.should_compress` result and the existing
pre-upstream backpressure state.
- Require positional alignment between optimized and original message
arrays before replay.
- Reject replay candidates that would serialize larger than the current
optimized messages.
- Add focused handler coverage for the decision-false tool-result
regression, bypass and backpressure paths, and outbound optimize-on
preservation.
- Add direct unit coverage for positional mismatch, no-inflation, and
JSON sizing-failure bailouts.
- Update the moved-cache-control and pure-block-append regression
fixtures to keep the no-inflation contract explicit.
- Run the unchanged OpenAI cache-stability preservation proof; no OpenAI
production code was edited.

## Testing

- [x] Unit tests pass (153 focused proxy, helper, cache-control,
block-append, cross-turn, byte-faithful, Anthropic, OpenAI, and
backpressure tests)
- [x] Linting passes (Ruff check and format validation on the seven
changed repository files)
- [ ] Type checking passes (`uv run mypy headroom`)
- [x] New tests added for new functionality when applicable
- [x] Manual testing performed with the in-process proxy and local stub
upstream

### Test Output

```text
python -m pytest tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py tests/test_cache_prefix_overlay.py tests/test_cross_turn_cache_safety.py tests/test_cache_control_move_bust.py tests/test_proxy_byte_faithful_forwarding.py tests/test_proxy_anthropic_cache_stability.py tests/test_anthropic_pre_upstream_backpressure.py -q
python -m pytest tests/test_proxy_openai_cache_stability.py -q
python -m pytest tests/test_issue_2671_block_growth_cache.py::test_pure_append_replays_forwarded_blocks_and_advances_breakpoint -q
153 passed across focused invocations, exit code 0
optimize_off turn2_message_count=3 marker_count=1 outbound_compact_utf8_bytes=2293 client_compact_utf8_bytes=2293
optimize_on turn2_message_count=3 client_message_count=3 marker_count=0 outbound_compact_utf8_bytes=171 client_compact_utf8_bytes=182
python -m ruff check headroom/proxy/handlers/anthropic.py headroom/cache/prefix_tracker.py tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py tests/test_cache_prefix_overlay.py tests/test_cache_control_move_bust.py tests/test_issue_2671_block_growth_cache.py tests/test_proxy_openai_cache_stability.py
All checks passed!, exit code 0

python -m ruff format headroom/proxy/handlers/anthropic.py headroom/cache/prefix_tracker.py tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py tests/test_cache_prefix_overlay.py tests/test_cache_control_move_bust.py tests/test_issue_2671_block_growth_cache.py tests/test_proxy_openai_cache_stability.py --check
7 files already formatted, exit code 0

git diff --check
clean, exit code 0
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12 via `uv`, real Headroom proxy app
with a local stub Anthropic upstream
- Exact command / steps: send an actual optimize-on first request
through the in-process proxy with a deterministic production-pipeline
seam, then send a decision-false bypass turn containing a large
Claude-shaped `tool_result` with moved `cache_control`; separately send
an aligned optimize-on turn with a new suffix
- Observed result: the exact base checkout fails with `AssertionError:
assert 'compressed-tool-result' == 'large-tool-result-marker ...'`; the
guarded path passes with the client marker present once and outbound
compact JSON no larger than the client body. The optimize-on
preservation run records `optimize_on turn2_message_count=3
client_message_count=3 marker_count=0 outbound_compact_utf8_bytes=171
client_compact_utf8_bytes=182`, proving the actual compressed prefix is
outbound before the new suffix without turn-2 growth.
- Not tested: live Claude Code session against api.anthropic.com on this
host

## Runtime Rollout Safety

- Rollout-managed feature(s): none.
- Minimum rollout channel: N/A.
- Stable/default behavior changed: cached-prefix replay now follows the
existing compression and backpressure decision and rejects misaligned or
inflating candidates.
- Kill switch / disable path: no new switch; the existing optimize and
bypass controls remain available.
- Unsafe override required: none.
- Qualification impact: none.
- Rollback path: revert the implementation commit.

## 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
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes

## Additional Notes

`CHANGELOG.md` is not modified because Headroom's release automation
generates it from conventional commits.

This change does not add a context-limit guard or alter compression,
streaming tracker provenance, outbound-body selection, OpenAI behavior,
or provider limits. Local tests prove request-body ownership and replay
bounds. The reporter's live Claude Code completion and Anthropic token
acceptance remain external to this local proof.
This commit is contained in:
Rod Boev
2026-08-17 18:02:18 -04:00
committed by GitHub
parent c502087db7
commit c16be9bbbe
7 changed files with 491 additions and 30 deletions
+49 -7
View File
@@ -449,7 +449,7 @@ def overlay_cached_prefix(
previous_original_messages: list[dict[str, Any]] | None,
previous_forwarded_messages: list[dict[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Replay the previously-forwarded (cached, compressed) prefix byte-identical.
"""Replay a positional, non-inflating cached prefix when it is safe.
Provider-agnostic cache-safety guard for the freeze path. When a message is
"frozen", the compression pipeline may emit the agent's ORIGINAL bytes for
@@ -467,14 +467,23 @@ def overlay_cached_prefix(
``optimized_messages`` unchanged (accept a possible bust rather than forward
wrong content).
This makes freezing byte-identical in BOTH proxy modes, so the only remaining
difference between them is how large a mutable (still-compressible) tail each
leaves — not whether the frozen prefix busts the cache.
The optimized and current-original lists must be positionally aligned, and
compact UTF-8 JSON for the replayed result must not exceed the optimized
candidate. These bounds prefer a cache miss to corrupting or inflating a
client's live history.
"""
prev_orig = previous_original_messages
prev_fwd = previous_forwarded_messages
if not prev_orig or not prev_fwd:
return optimized_messages
if len(optimized_messages) != len(current_original_messages):
logger.debug(
"overlay: optimized/current-original length mismatch (optimized=%d, current=%d) "
"— skipping positional cached-prefix replay",
len(optimized_messages),
len(current_original_messages),
)
return optimized_messages
n = len(prev_orig)
# Positional 1:1 correspondence between prev_orig[i] and prev_fwd[i] holds
# only when last turn forwarded exactly one message per original (the
@@ -534,11 +543,21 @@ def overlay_cached_prefix(
len(current_content) - split,
message_index,
)
return (
replayed = (
list(prev_fwd[:message_index])
+ [merged]
+ list(optimized_messages[message_index + 1 :])
)
replayed_bytes = _compact_json_bytes(replayed)
optimized_bytes = _compact_json_bytes(optimized_messages)
if (
replayed_bytes is None
or optimized_bytes is None
or len(replayed_bytes) > len(optimized_bytes)
):
logger.debug("overlay: block replay inflated compact JSON — skipping")
return optimized_messages
return replayed
# Append-only guard on CONTENT ONLY, message-by-message. Replay the
# previously-forwarded (cached, compressed) bytes for the longest LEADING
# run of messages that is byte-for-byte (content-canonical) identical to
@@ -562,7 +581,7 @@ def overlay_cached_prefix(
# current_original[k] canonicalize-equals prev_orig[k], and prev_fwd[k]
# positionally corresponds to prev_orig[k] (guaranteed by the count check
# above), so no wrong bytes are ever forwarded.
limit = min(n, len(current_original_messages), len(optimized_messages))
limit = min(n, len(current_original_messages))
k = 0
while k < limit and _canonicalize_for_prefix_compare(
current_original_messages[k]
@@ -584,7 +603,30 @@ def overlay_cached_prefix(
)
# Replay the cached (compressed) prefix byte-identical up to the first
# divergence; keep this turn's freshly-produced output for the rest.
return list(prev_fwd[:k]) + list(optimized_messages[k:])
replayed = list(prev_fwd[:k]) + list(optimized_messages[k:])
replayed_bytes = _compact_json_bytes(replayed)
optimized_bytes = _compact_json_bytes(optimized_messages)
if (
replayed_bytes is None
or optimized_bytes is None
or len(replayed_bytes) > len(optimized_bytes)
):
logger.debug("overlay: replay inflated compact JSON — skipping cached-prefix replay")
return optimized_messages
return replayed
def _compact_json_bytes(value: Any) -> bytes | None:
"""Return compact JSON bytes, or ``None`` when sizing cannot be proved."""
try:
return json.dumps(
value,
separators=(",", ":"),
ensure_ascii=False,
default=str,
).encode("utf-8")
except (TypeError, ValueError, OverflowError, UnicodeError):
return None
_STABLE_BOUNDARY_ENV = "HEADROOM_STABLE_BOUNDARY_BREAKPOINT"
+24 -11
View File
@@ -1975,22 +1975,35 @@ class AnthropicHandlerMixin:
overlay_cached_prefix,
)
_overlay_replayed = False
# 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
if _decision.should_compress and not _skip_compression_for_backpressure:
if _cold_recompact_active:
_overlay_replayed = False
else:
_ov = overlay_cached_prefix(
optimized_messages,
original_client_messages,
previous_original_messages,
previous_forwarded_messages,
)
_overlay_replayed = _ov != optimized_messages
if _overlay_replayed:
optimized_messages = _ov
optimized_tokens = tokenizer.count_messages(optimized_messages)
else:
_ov = overlay_cached_prefix(
optimized_messages,
original_client_messages,
previous_original_messages,
previous_forwarded_messages,
replay_skip_reason = (
"pre_upstream_backpressure"
if _skip_compression_for_backpressure
else _decision.passthrough_reason
)
logger.debug(
"[%s] Cached-prefix replay skipped: reason=%s",
request_id,
replay_skip_reason,
)
_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
+4 -7
View File
@@ -61,14 +61,11 @@ def test_marker_move_would_fail_a_raw_dict_guard():
assert _strip_cache_control(CUR_ORIG[:2]) == _strip_cache_control(PREV_ORIG)
def test_overlay_replays_despite_moved_marker():
def test_overlay_skips_inflating_moved_marker_replay():
out = overlay_cached_prefix(OPTIMIZED, CUR_ORIG, PREV_ORIG, PREV_FWD)
# The content-only guard lets the replay happen: the forwarded prefix is now
# byte-identical to what the provider cached (compressed), NOT the freeze's
# original bytes → cache hits instead of busting.
assert out[:2] == PREV_FWD
assert out[:2] != OPTIMIZED[:2]
assert out[2] == OPTIMIZED[2] # compressed tail preserved
# Moving cache_control must not exempt a larger replay candidate from the
# no-inflation bound, even when content-only history alignment succeeds.
assert out == OPTIMIZED
# ── Cross-turn: client moves the marker every turn, provider keys on full bytes ─
+73
View File
@@ -9,6 +9,8 @@ Forwarding original then mismatches the cached prefix and busts the prompt cache
so the cache still hits — in BOTH proxy modes.
"""
import copy
from headroom.cache.prefix_tracker import overlay_cached_prefix
@@ -72,6 +74,77 @@ def test_shorter_current_or_optimized_returns_unchanged():
]
def test_overlay_requires_positional_alignment_with_originals():
optimized = [M("user", "x")]
current = [M("user", "x"), M("assistant", "ok")]
assert overlay_cached_prefix(optimized, current, PREV_ORIG, PREV_FWD) == optimized
optimized = [M("user", "x"), M("assistant", "ok"), M("user", "tail")]
current = [M("user", "x"), M("assistant", "ok")]
previous = [M("user", "x"), M("assistant", "ok")]
forwarded = [M("user", "compressed"), M("assistant", "ok")]
assert overlay_cached_prefix(optimized, current, previous, forwarded) == optimized
def test_overlay_never_inflates_forwarded_payload():
optimized = [M("user", "small"), M("assistant", "ok"), M("user", "tail")]
inflated_forwarded = [M("user", "x" * 1000), M("assistant", "ok")]
previous = [M("user", "small"), M("assistant", "ok")]
current = previous + [M("user", "tail")]
assert overlay_cached_prefix(optimized, current, previous, inflated_forwarded) == optimized
def test_overlay_returns_optimized_when_json_sizing_fails(monkeypatch):
optimized = [M("user", "stable"), M("user", "tail")]
current = [M("user", "stable"), M("user", "tail")]
previous = [M("user", "stable")]
forwarded = [M("user", "compressed")]
monkeypatch.setattr(
"headroom.cache.prefix_tracker.json.dumps",
lambda *args, **kwargs: (_ for _ in ()).throw(TypeError("cannot size")),
)
assert overlay_cached_prefix(optimized, current, previous, forwarded) == optimized
def test_overlay_never_inflates_cache_control_only_replay():
previous = [M("user", "stable"), M("assistant", "ok")]
current = [
M("user", "stable"),
{**M("assistant", "ok"), "cache_control": {"type": "ephemeral"}},
]
optimized = copy.deepcopy(current)
inflated_forwarded = [M("user", "x" * 1000), M("assistant", "ok")]
assert overlay_cached_prefix(optimized, current, previous, inflated_forwarded) == optimized
def test_block_append_overlay_never_inflates_forwarded_payload():
previous = [
{
"role": "user",
"content": [{"type": "text", "text": "stable"}],
}
]
current = [
{
"role": "user",
"content": [
{"type": "text", "text": "stable"},
{"type": "text", "text": "tail"},
],
}
]
optimized = copy.deepcopy(current)
forwarded = [
{
"role": "user",
"content": [{"type": "text", "text": "x" * 1000}],
}
]
assert overlay_cached_prefix(optimized, current, previous, forwarded) == optimized
def test_cache_hit_property_prefix_matches_last_forward():
# The invariant that guarantees a cache hit: forwarded[:n] this turn ==
# forwarded[:n] last turn (== what the provider cached).
+2 -2
View File
@@ -223,14 +223,14 @@ def test_cache_affinity_ignores_only_cache_directive_movement() -> None:
def test_pure_append_replays_forwarded_blocks_and_advances_breakpoint() -> None:
previous_original = _pure_append(30)
previous_forwarded = _message([_text(f"COMPRESSED-{index}") for index in range(30)])
previous_forwarded = _message([_text(f"C-{index}") for index in range(30)])
current = _pure_append(34)
overlaid = overlay_cached_prefix(current, current, previous_original, previous_forwarded)
normalized = normalize_message_cache_control(overlaid, previous_forwarded)
assert [block["text"] for block in normalized[0]["content"][:30]] == [
f"COMPRESSED-{index}" for index in range(30)
f"C-{index}" for index in range(30)
]
assert [block["text"] for block in normalized[0]["content"][30:]] == [
f"block-{index}" for index in range(30, 34)
@@ -0,0 +1,279 @@
"""Regression coverage for Anthropic cached-prefix replay eligibility."""
from __future__ import annotations
import asyncio
import copy
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import httpx
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
MODEL = "claude-sonnet-4-6"
MARKER = "large-tool-result-marker " * 80
def _config(**overrides) -> ProxyConfig:
values = {
"optimize": False,
"cache_enabled": False,
"rate_limit_enabled": False,
"cost_tracking_enabled": False,
"log_requests": False,
"prefix_freeze_enabled": True,
}
values.update(overrides)
return ProxyConfig(**values)
def _response() -> httpx.Response:
return httpx.Response(
200,
json={
"id": "msg_test",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "ok"}],
"model": MODEL,
"usage": {"input_tokens": 10, "output_tokens": 2},
},
)
def _capture_proxy(client: TestClient) -> tuple[list[dict], AsyncMock]:
captured: list[dict] = []
async def capture(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
captured.append(copy.deepcopy(body))
return _response()
retry = AsyncMock(side_effect=capture)
client.app.state.proxy._retry_request = retry
return captured, retry
def _compact_bytes(messages: list[dict]) -> int:
return len(json.dumps(messages, separators=(",", ":"), ensure_ascii=False).encode())
def _pipeline_result(messages: list[dict]) -> SimpleNamespace:
return SimpleNamespace(
messages=messages,
transforms_applied=["test_stub"],
timing={},
tokens_before=1000,
tokens_after=500,
waste_signals=None,
)
def test_no_optimize_second_turn_forwards_client_history_verbatim():
"""A bypass turn owns its Claude-shaped body after an optimized prior turn."""
app = create_app(_config(optimize=True))
tool_result = {
"type": "tool_result",
"tool_use_id": "toolu_01",
"content": MARKER,
"cache_control": {"type": "ephemeral"},
}
first = [{"role": "user", "content": [tool_result]}]
compressed = copy.deepcopy(first)
compressed[0]["content"][0]["content"] = "compressed-tool-result"
app.state.proxy.anthropic_pipeline.apply = Mock(return_value=_pipeline_result(compressed))
with TestClient(app) as client:
captured, _ = _capture_proxy(client)
response = client.post(
"/v1/messages",
json={"model": MODEL, "max_tokens": 16, "messages": first},
)
assert response.status_code == 200, response.text
assistant = {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}
current = [
{
"role": "user",
"content": [{k: v for k, v in tool_result.items() if k != "cache_control"}],
},
assistant,
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_02",
"content": "next tool result",
"cache_control": {"type": "ephemeral"},
}
],
},
]
response = client.post(
"/v1/messages",
headers={"x-headroom-bypass": "true"},
json={"model": MODEL, "max_tokens": 16, "messages": current},
)
assert response.status_code == 200, response.text
outbound = captured[-1]["messages"]
assert outbound[0]["content"][0]["content"] == MARKER
assert "compressed-tool-result" not in json.dumps(outbound)
assert sum(MARKER in json.dumps(message) for message in outbound) == 1
assert _compact_bytes(outbound) <= _compact_bytes(current)
def test_optimize_off_real_exchange_reports_noninflation():
app = create_app(_config(optimize=False))
first = [
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_01",
"content": MARKER,
}
],
}
]
current = first + [
{"role": "assistant", "content": [{"type": "text", "text": "ok"}]},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_02",
"content": "next tool result",
"cache_control": {"type": "ephemeral"},
}
],
},
]
with TestClient(app) as client:
captured, _ = _capture_proxy(client)
assert (
client.post(
"/v1/messages",
json={"model": MODEL, "max_tokens": 16, "messages": first},
).status_code
== 200
)
response = client.post(
"/v1/messages",
json={"model": MODEL, "max_tokens": 16, "messages": current},
)
assert response.status_code == 200, response.text
outbound = captured[-1]["messages"]
assert len(outbound) == 3
assert sum(MARKER in json.dumps(message) for message in outbound) == 1
assert _compact_bytes(outbound) == _compact_bytes(current) == 2293
print(
"optimize_off turn2_message_count=3 marker_count=1 "
"outbound_compact_utf8_bytes=2293 client_compact_utf8_bytes=2293"
)
def test_bypass_header_does_not_invoke_cached_prefix_replay(monkeypatch):
app = create_app(_config(optimize=True))
replay = monkeypatch.setattr
def fail_if_called(*args, **kwargs): # noqa: ANN002, ANN003
raise AssertionError("cached-prefix replay must be bypassed")
replay("headroom.cache.prefix_tracker.overlay_cached_prefix", fail_if_called)
with TestClient(app) as client:
captured, _ = _capture_proxy(client)
messages = [{"role": "user", "content": "bypass"}]
response = client.post(
"/v1/messages",
headers={"x-headroom-bypass": "true"},
json={"model": MODEL, "max_tokens": 16, "messages": messages},
)
assert response.status_code == 200, response.text
assert captured[-1]["messages"] == messages
def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch):
app = create_app(
_config(
optimize=True,
anthropic_pre_upstream_concurrency=1,
anthropic_pre_upstream_acquire_timeout_seconds=0.001,
)
)
def fail_if_called(*args, **kwargs): # noqa: ANN002, ANN003
raise AssertionError("cached-prefix replay must be skipped under backpressure")
monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", fail_if_called)
debug = Mock()
monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.debug", debug)
proxy = app.state.proxy
class _SaturatedSemaphore:
async def acquire(self):
await asyncio.sleep(1)
def release(self):
pass
proxy.anthropic_pre_upstream_sem = _SaturatedSemaphore()
try:
with TestClient(app) as client:
captured, _ = _capture_proxy(client)
messages = [{"role": "user", "content": "backpressure"}]
response = client.post(
"/v1/messages",
json={"model": MODEL, "max_tokens": 16, "messages": messages},
)
finally:
proxy.anthropic_pre_upstream_sem.release()
assert response.status_code == 200, response.text
assert captured[-1]["messages"] == messages
assert any(
call.args and call.args[-1] == "pre_upstream_backpressure" for call in debug.call_args_list
)
def test_optimize_on_aligned_history_preserves_replay():
app = create_app(_config(optimize=True))
first = [{"role": "user", "content": "original prefix"}]
compressed = [{"role": "user", "content": "comp"}]
assistant = {
"role": "assistant",
"content": [{"type": "text", "text": "ok", "cache_control": {"type": "ephemeral"}}],
}
current = first + [assistant, {"role": "user", "content": "new suffix"}]
optimized = compressed + [assistant, {"role": "user", "content": "new suffix"}]
app.state.proxy.anthropic_pipeline.apply = Mock(
side_effect=[_pipeline_result(compressed), _pipeline_result(optimized)]
)
with TestClient(app) as client:
captured, _ = _capture_proxy(client)
response = client.post(
"/v1/messages",
json={"model": MODEL, "max_tokens": 16, "messages": first},
)
assert response.status_code == 200, response.text
response = client.post(
"/v1/messages",
json={"model": MODEL, "max_tokens": 16, "messages": current},
)
assert response.status_code == 200, response.text
outbound = captured[-1]["messages"]
assert outbound[0] == compressed[0]
assert outbound[-1] == optimized[-1]
print(
f"optimize_on turn2_message_count={len(outbound)} client_message_count={len(current)} "
f"marker_count={sum(MARKER in json.dumps(message) for message in outbound)} "
f"outbound_compact_utf8_bytes={_compact_bytes(outbound)} "
f"client_compact_utf8_bytes={_compact_bytes(current)}"
)
+60 -3
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import copy
from types import SimpleNamespace
import httpx
@@ -15,8 +16,15 @@ from headroom.proxy.server import ProxyConfig, create_app
class _FakePrefixTracker:
def __init__(self, frozen_count: int):
def __init__(
self,
frozen_count: int,
previous_original: list[dict] | None = None,
previous_forwarded: list[dict] | None = None,
):
self._frozen_count = frozen_count
self._previous_original = previous_original or []
self._previous_forwarded = previous_forwarded or []
def get_frozen_message_count(self) -> int:
return self._frozen_count
@@ -26,10 +34,10 @@ class _FakePrefixTracker:
# overlay itself is exercised in test_cross_turn_cache_safety.py against the
# real tracker; these stubs just satisfy the handler's overlay call.
def get_last_original_messages(self): # noqa: ANN201
return []
return copy.deepcopy(self._previous_original)
def get_last_forwarded_messages(self): # noqa: ANN201
return []
return copy.deepcopy(self._previous_forwarded)
def update_from_response(self, **kwargs): # noqa: ANN003
return None
@@ -112,6 +120,55 @@ def test_openai_cache_mode_freezes_previous_turns() -> None:
assert captured["frozen_message_count"] == 2
def test_openai_handler_replays_nonempty_cached_prefix() -> None:
captured = {}
previous_original = [{"role": "user", "content": "original prefix"}]
previous_forwarded = [{"role": "user", "content": "comp"}]
fake_tracker = _FakePrefixTracker(0, previous_original, previous_forwarded)
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: (
"stable-session"
)
proxy.session_tracker_store.resolve_tracker = lambda *args, **kwargs: fake_tracker
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
captured["body"] = body
return httpx.Response(
200,
json={
"id": "chatcmpl_overlay",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 20, "completion_tokens": 3, "total_tokens": 23},
},
)
proxy._retry_request = _fake_retry
response = client.post(
"/v1/chat/completions",
headers={"authorization": "Bearer test-key"},
json={
"model": "gpt-4o-mini",
"messages": [
{"role": "user", "content": "original prefix"},
{"role": "user", "content": "new suffix"},
],
},
)
assert response.status_code == 200
assert captured["body"]["messages"] == [
previous_forwarded[0],
{"role": "user", "content": "new suffix"},
]
@pytest.mark.parametrize("tail_role", ["tool", "function"])
def test_openai_cache_mode_keeps_final_tool_observation_mutable(tail_role: str) -> None:
captured = {}