headroom/tests/test_proxy_openai_cache_stability.py
vscunha 05932d7165
fix(proxy): compress OpenCode tool schemas and embedded JSON (#1535)
## Description

Fixes two remaining OpenCode/OpenAI Chat compression gaps after `main`
incorporated the original savings-profile threading and user
content-block work from this PR.

OpenCode requests can still report very low savings when most input
tokens live in verbose `tools` schemas rather than messages. They can
also route poorly when a short instruction wraps a valid JSON block but
does not satisfy the existing long-prose heuristic.

Closes #1534

## 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 causes existing functionality
to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Compact OpenAI Chat Completions `tools` schemas whenever request
compression is active, reusing the existing OpenAI Responses schema
compactor. The outbound tool invocation shape is preserved while
non-semantic annotations such as `$schema`, `title`, and `examples` are
removed.
- Include the tool-schema token delta in Headroom's savings accounting
and expose `openai:chat:tool_schema_compaction` in the applied
transforms.
- Detect valid JSON blocks surrounded by prose or log text as mixed
content, so short OpenCode instructions route through mixed/SmartCrusher
handling instead of falling through or producing a no-op.
- Adapt the mixed-content change to the new
`headroom.transforms.mixed_content` module introduced on `main` by
#1939.

## Why the Focus Changed

The original headline fix—threading savings-profile kwargs into
`/v1/chat/completions`—is now already present on `main`, as is the user
content-block opt-in behavior. Those duplicate changes were removed
during the merge.

The branch also no longer changes developer/system role protection or
forced-Kompress semantics. It follows `main` for both, so the earlier
instruction-role safety concern is outside the current diff.

The resulting PR is limited to two OpenCode-specific compression gaps
that remain reproducible on current `main`.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check` on changed files)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for the fixed behavior
- [ ] Manual live-upstream testing performed after the latest rebase

### Test Output

```text
59 passed, 1 warning in 83.53s
All checks passed!  # ruff check
4 files already formatted  # ruff format --check
python -m py_compile: passed
git diff --check: passed
```

Focused test coverage includes:

- OpenAI Chat tool-schema compaction, transform reporting, outbound
schema shape, and positive token savings.
- Embedded JSON mixed-content detection, SmartCrusher routing, positive
savings, and preservation of a critical sentinel value.
- Current `main` regressions for savings-profile threading, user content
blocks, turn hooks, and forced-Kompress behavior.

## Real Behavior Proof

- Environment: Linux ARM64, Python 3.13.12, current `main` at `9bacf481`
merged into the branch.
- Exact command / steps: focused pytest run across the OpenAI
cache-stability, content-router, mixed-content, savings-profile,
user-block, turn-hook, and forced-Kompress suites.
- Observed result: 59 tests passed; the chat request test forwarded
compacted tools and reported positive savings, while the embedded-JSON
fixture used mixed routing and preserved `CRITICAL_NEEDLE_42`.
- Not tested: full repository suite and a live external OpenCode request
after the latest merge; those remain for CI/live follow-up.

## Review Readiness

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

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented the non-obvious behavior
- [ ] I have made corresponding documentation changes — N/A; internal
routing behavior only
- [x] My changes generate no new warnings
- [x] I have added tests that prove the fixes are effective
- [x] New and existing focused tests pass locally
- [ ] I have updated the changelog — N/A; release automation handles fix
entries

## Screenshots

N/A — proxy/transform behavior only.

## Additional Notes

- Current diff versus `main`: 4 files, 172 insertions, no role-policy or
forced-Kompress changes.
- The mixed-content conflict was resolved by extending the new isolated
parser module rather than reintroducing parsing code into
`ContentRouter`.
2026-07-15 19:58:42 +00:00

454 lines
16 KiB
Python

"""Regression tests for OpenAI cache-mode stability in proxy mode."""
from __future__ import annotations
from types import SimpleNamespace
import httpx
import pytest
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
class _FakePrefixTracker:
def __init__(self, frozen_count: int):
self._frozen_count = frozen_count
def get_frozen_message_count(self) -> int:
return self._frozen_count
# Empty history → overlay_cached_prefix() is a no-op here, so these tests
# keep asserting the cache-freeze behavior they always have. The cross-turn
# 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 []
def get_last_forwarded_messages(self): # noqa: ANN201
return []
def update_from_response(self, **kwargs): # noqa: ANN003
return None
def _make_proxy_client() -> TestClient:
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
return TestClient(app)
def test_openai_cache_mode_freezes_previous_turns() -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "cache"
fake_tracker = _FakePrefixTracker(frozen_count=0)
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: (
"stable-session"
)
proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
def _fake_apply(**kwargs):
captured["frozen_message_count"] = kwargs.get("frozen_message_count")
return SimpleNamespace(
messages=kwargs["messages"],
transforms_applied=[],
timing={},
tokens_before=60,
tokens_after=60,
waste_signals=None,
)
proxy.openai_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
return httpx.Response(
200,
json={
"id": "chatcmpl_1",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 60, "completion_tokens": 3, "total_tokens": 63},
},
)
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": "turn1"},
{"role": "assistant", "content": "turn1-assistant"},
{"role": "user", "content": "current turn"},
],
},
)
assert response.status_code == 200
assert captured["frozen_message_count"] == 2
@pytest.mark.parametrize("tail_role", ["tool", "function"])
def test_openai_cache_mode_keeps_final_tool_observation_mutable(tail_role: str) -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "cache"
fake_tracker = _FakePrefixTracker(frozen_count=0)
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: (
"stable-session"
)
proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
def _fake_apply(**kwargs):
captured.setdefault("calls", []).append(
{
"frozen_message_count": kwargs.get("frozen_message_count"),
"roles": [msg.get("role") for msg in kwargs["messages"]],
"mode": proxy.config.mode,
}
)
return SimpleNamespace(
messages=kwargs["messages"],
transforms_applied=["test:compress-tail"],
timing={},
tokens_before=120,
tokens_after=80,
waste_signals=None,
)
proxy.openai_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
return httpx.Response(
200,
json={
"id": "chatcmpl_tool_tail",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 80, "completion_tokens": 3, "total_tokens": 83},
},
)
proxy._retry_request = _fake_retry
tail = {
"role": tail_role,
"content": "large command observation " * 200,
}
if tail_role == "tool":
tail["tool_call_id"] = "call_1"
else:
tail["name"] = "bash"
response = client.post(
"/v1/chat/completions",
headers={"authorization": "Bearer test-key"},
json={
"model": "gpt-4o-mini",
"messages": [
{"role": "user", "content": "turn1"},
{"role": "assistant", "content": "run command"},
tail,
],
},
)
assert response.status_code == 200
assert any(call["frozen_message_count"] == 2 for call in captured["calls"]), captured[
"calls"
]
def test_openai_cache_mode_restores_mutated_frozen_prefix() -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "cache"
fake_tracker = _FakePrefixTracker(frozen_count=0)
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: (
"stable-session"
)
proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
original_messages = [
{"role": "user", "content": "turn1"},
{"role": "assistant", "content": "turn1-assistant"},
{"role": "user", "content": "current turn"},
]
def _fake_apply(**kwargs):
mutated = list(kwargs["messages"])
mutated[0] = {**mutated[0], "content": "MUTATED_PREFIX"}
return SimpleNamespace(
messages=mutated,
transforms_applied=["fake:mutated"],
timing={},
tokens_before=70,
tokens_after=65,
waste_signals=None,
)
proxy.openai_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
captured["body"] = body
return httpx.Response(
200,
json={
"id": "chatcmpl_2",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 65, "completion_tokens": 3, "total_tokens": 68},
},
)
proxy._retry_request = _fake_retry
response = client.post(
"/v1/chat/completions",
headers={"authorization": "Bearer test-key"},
json={
"model": "gpt-4o-mini",
"messages": original_messages,
},
)
assert response.status_code == 200
sent_messages = captured["body"]["messages"]
assert sent_messages[0] == original_messages[0]
assert sent_messages[1] == original_messages[1]
# ─── Issue #327 cross-handler regression ────────────────────────────────
#
# The OpenAI handler was never affected by issue #327's content-keyed walker
# bug — it has only ever used `compute_frozen_count` (positional). This test
# locks that property by spying on the OpenAI traffic path and asserting that
# the buggy walker functions (`should_defer_compression`, `mark_stable`) are
# never called from the production handler. If a future refactor accidentally
# adds the same walker to OpenAI, this test fails immediately.
def test_issue_327_openai_handler_does_not_call_walker_functions() -> None:
calls: list[tuple[str, tuple, dict]] = []
class _SpyCompCache:
def apply_cached(self, messages): # noqa: ANN001
calls.append(("apply_cached", (), {}))
return list(messages)
def compute_frozen_count(self, messages): # noqa: ANN001
calls.append(("compute_frozen_count", (), {}))
return 0
def update_from_result(self, originals, compressed): # noqa: ANN001
calls.append(("update_from_result", (), {}))
def mark_stable_from_messages(self, messages, up_to): # noqa: ANN001
calls.append(("mark_stable_from_messages", (up_to,), {}))
# Methods below MUST NOT be called from OpenAI handler.
def should_defer_compression(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003
calls.append(("should_defer_compression", args, kwargs))
return False
def mark_stable(self, content_hash): # noqa: ANN001
calls.append(("mark_stable", (content_hash,), {}))
@staticmethod
def content_hash(content): # noqa: ANN001
return f"H({content[:40] if isinstance(content, str) else 'list'})"
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "token" # token mode is where Anthropic had the bug
fake_tracker = _FakePrefixTracker(frozen_count=0)
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: (
"openai-spy-session"
)
proxy.session_tracker_store.get_or_create = lambda s, p: fake_tracker
proxy._get_compression_cache = lambda s: _SpyCompCache()
def _fake_apply(**kwargs): # noqa: ANN003
return SimpleNamespace(
messages=list(kwargs["messages"]),
transforms_applied=[],
timing={},
tokens_before=60,
tokens_after=60,
waste_signals=None,
)
proxy.openai_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
return httpx.Response(
200,
json={
"id": "cmpl",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 60, "completion_tokens": 3, "total_tokens": 63},
},
)
proxy._retry_request = _fake_retry
# Drive 5 turns so any walker bug would have time to fire repeatedly.
for turn in range(5):
r = client.post(
"/v1/chat/completions",
headers={"authorization": "Bearer test-key"},
json={
"model": "gpt-4o-mini",
"messages": [
{"role": "user", "content": f"turn-{turn}-q"},
{"role": "assistant", "content": f"turn-{turn}-a"},
{"role": "tool", "tool_call_id": "t1", "content": "x" * 600},
{"role": "user", "content": f"continue-{turn}"},
],
},
)
assert r.status_code == 200
method_names = [c[0] for c in calls]
assert "should_defer_compression" not in method_names, (
f"OpenAI handler unexpectedly called should_defer_compression. "
f"Calls observed: {method_names}"
)
assert "mark_stable" not in method_names, (
f"OpenAI handler unexpectedly called mark_stable (the walker side-effect). "
f"Calls observed: {method_names}"
)
# Sanity: the safe positional methods DID fire.
assert "compute_frozen_count" in method_names
assert "apply_cached" in method_names
def test_openai_chat_completions_compacts_tools_when_profile_enabled() -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "token"
proxy.config.savings_profile = "agent-90"
def _fake_apply(**kwargs): # noqa: ANN003
return SimpleNamespace(
messages=kwargs["messages"],
transforms_applied=[],
timing={},
tokens_before=10,
tokens_after=10,
waste_signals=None,
)
proxy.openai_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
captured["body"] = body
return httpx.Response(
200,
json={
"id": "chatcmpl_tools",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 500, "completion_tokens": 3, "total_tokens": 503},
},
)
proxy._retry_request = _fake_retry
verbose_schema_note = "schema annotation repeated for opencode tool definitions " * 50
tools = [
{
"type": "function",
"function": {
"name": "read_file",
"description": "read file helper",
"parameters": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "ReadFileParameters",
"type": "object",
"properties": {
"path": {
"type": "string",
"title": "Path",
"description": verbose_schema_note,
"examples": [verbose_schema_note],
}
},
"required": ["path"],
},
},
}
]
response = client.post(
"/v1/chat/completions",
headers={"authorization": "Bearer test-key"},
json={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "inspect this file"}],
"tools": tools,
},
)
assert response.status_code == 200
assert "openai:chat:tool_schema_compaction" in response.headers["x-headroom-transforms"]
assert int(response.headers["x-headroom-tokens-saved"]) > 0
sent_params = captured["body"]["tools"][0]["function"]["parameters"]
assert "$schema" not in sent_params
assert "title" not in sent_params
assert "title" not in sent_params["properties"]["path"]
assert "examples" not in sent_params["properties"]["path"]