mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
f0dcc02775
Eliminates P0-2 universally. Every Python forwarder (server.py
`_retry_request`, handlers/streaming.py `_stream_response`,
handlers/openai.py `_ws_http_fallback`, handlers/batch.py `_batch_passthrough`
+ batch-create + Google batch passthrough, handlers/anthropic.py CCR
continuation + batch endpoint) now switches from `httpx ... json=body` to
`httpx ... content=raw_bytes`. The default httpx JSON encoder was
re-serializing every request with `, `/`: ` separators and `\\uXXXX` ASCII
escapes — collapsing Anthropic prompt-cache hit-rate.
Forwarder strategy:
- unmutated body → forward `await request.body()` verbatim;
- mutated body → re-serialize once via the new
`serialize_body_canonical(body) -> bytes` helper (compact separators,
`ensure_ascii=False`, dict insertion order preserved).
`HEADROOM_PROXY_PYTHON_FORWARDER_MODE` env var configures the mode:
- `byte_faithful` (default) — the new behavior;
- `legacy_json_kwarg` — explicit operator opt-in for emergency rollback.
Documented in `docs/content/docs/configuration.mdx`. NOT a fallback —
unknown values raise loudly per build constraint #4.
`BodyMutationTracker` accompanies each request through the handler so
transform sites mark the tracker (`memory_injection`,
`image_compression`, `compression_*`, `batch_compression`,
`ccr_continuation`, etc.). At forwarder dispatch we additionally compare
the final body dict against the parsed original bytes as a structural
safety net — any silent mutation we missed still triggers canonical
re-serialization.
A2 follow-up: `handlers/openai.py:534-540` (Chat Completions memory
injection) was prepending a system message; replaced with
`append_text_to_latest_user_chat_message`, the OpenAI Chat Completions
analog of `_append_context_to_latest_non_frozen_user_turn`. The cache
hot zone (system messages) is now sacrosanct on /v1/chat/completions
too. Honors `HEADROOM_MEMORY_INJECTION_MODE=disabled`.
Structured logging: every forwarder emits an `event=outbound_request`
log line with `forwarder`, `path`, `body_bytes`, `body_mutated`,
`mutation_reasons`, `source` (passthrough|canonical|legacy),
`request_id`. Never logs Authorization or full body.
`_read_request_json` factored to share `_read_request_body_bytes` with
new `read_request_json_with_bytes` so the anthropic handler can capture
both the parsed dict and the original (decompressed) bytes.
Tests:
- `tests/test_proxy_byte_faithful_forwarding.py` (28 tests):
SHA-256 byte-equality on /v1/messages and streaming, unicode
preservation, numeric precision, mutation-tracker invariants,
canonical-serializer properties, legacy-mode rollback, OpenAI
Chat memory routing.
- Existing test mocks updated to accept the new `**kwargs` on
`_retry_request` (no behavior change).
- `tests/test_proxy_handlers_batch.py` updated to read the captured
`content=` bytes (formerly `json=`).
- One A2 test corrected (`test_anthropic_tool_sort_and_context_append_helpers`)
to match the live-zone-tail semantics introduced by A2.
Constraints satisfied: configurable env var; no new regex / hardcodes;
no silent fallback (`legacy_json_kwarg` is operator opt-in);
performant (`prepare_outbound_body_bytes` is O(1) for passthrough);
elegant single-responsibility helpers; structured tracing logs.
436 lines
16 KiB
Python
436 lines
16 KiB
Python
"""System-prompt immutability tests for PR-A2 (P0-1 fix).
|
|
|
|
After PR-A2, memory context never mutates the system prompt or the
|
|
Responses API ``instructions`` field. The cache hot zone is sacrosanct
|
|
(invariant I2). Memory routes exclusively to the live-zone tail (the
|
|
first text block of the latest non-frozen user turn).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
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
|
|
self._cached_token_count = 0
|
|
self._last_original_messages: list[dict[str, object]] = []
|
|
self._last_forwarded_messages: list[dict[str, object]] = []
|
|
|
|
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
|
|
|
|
|
|
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 _install_memory_handler(proxy: object) -> None:
|
|
"""Wire a deterministic in-process memory handler that returns 'MEMCTX'."""
|
|
proxy.memory_handler = SimpleNamespace( # type: ignore[attr-defined]
|
|
config=SimpleNamespace(inject_context=True, inject_tools=False),
|
|
search_and_format_context=AsyncMock(return_value="MEMCTX"),
|
|
has_memory_tool_calls=lambda resp, provider: False,
|
|
)
|
|
|
|
|
|
def _install_session_tracker(proxy: object, frozen_count: int) -> None:
|
|
fake_tracker = _FakePrefixTracker(frozen_count=frozen_count)
|
|
proxy.session_tracker_store.compute_session_id = ( # type: ignore[attr-defined]
|
|
lambda request, model, messages: "stable-session"
|
|
)
|
|
proxy.session_tracker_store.get_or_create = ( # type: ignore[attr-defined]
|
|
lambda session_id, provider: fake_tracker
|
|
)
|
|
|
|
|
|
def _install_capture_retry(proxy: object, captured: dict[str, object]) -> None:
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
captured["body"] = body
|
|
captured["headers"] = dict(headers)
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"id": "msg_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "ok"}],
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"output_tokens": 3,
|
|
"cache_read_input_tokens": 0,
|
|
"cache_creation_input_tokens": 0,
|
|
},
|
|
},
|
|
)
|
|
|
|
proxy._retry_request = _fake_retry # type: ignore[attr-defined]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Anthropic /v1/messages tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_memory_enabled_does_not_mutate_system() -> None:
|
|
"""Memory injection must not mutate the top-level ``system`` field."""
|
|
captured: dict[str, object] = {}
|
|
with _make_proxy_client() as client:
|
|
proxy = client.app.state.proxy
|
|
proxy.config.optimize = False
|
|
proxy.config.image_optimize = False
|
|
proxy.config.ccr_proactive_expansion = False
|
|
|
|
_install_session_tracker(proxy, frozen_count=1)
|
|
_install_memory_handler(proxy)
|
|
_install_capture_retry(proxy, captured)
|
|
|
|
original_system = "You are a helpful assistant."
|
|
response = client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"x-headroom-user-id": "u1",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"system": original_system,
|
|
"messages": [
|
|
{"role": "user", "content": "frozen prefix"},
|
|
{"role": "assistant", "content": "ack"},
|
|
{"role": "user", "content": "latest user"},
|
|
],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
sent = captured["body"]
|
|
assert isinstance(sent, dict)
|
|
assert sent["system"] == original_system, (
|
|
"system prompt must be byte-equal to client-sent value"
|
|
)
|
|
|
|
|
|
def test_memory_context_appears_in_latest_user_message_tail() -> None:
|
|
"""Memory context appends to the latest non-frozen user turn's text."""
|
|
captured: dict[str, object] = {}
|
|
with _make_proxy_client() as client:
|
|
proxy = client.app.state.proxy
|
|
proxy.config.optimize = False
|
|
proxy.config.image_optimize = False
|
|
proxy.config.ccr_proactive_expansion = False
|
|
|
|
_install_session_tracker(proxy, frozen_count=1)
|
|
_install_memory_handler(proxy)
|
|
_install_capture_retry(proxy, captured)
|
|
|
|
response = client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"x-headroom-user-id": "u1",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"system": "base system",
|
|
"messages": [
|
|
{"role": "user", "content": "frozen prefix"},
|
|
{"role": "assistant", "content": "ack"},
|
|
{"role": "user", "content": "latest user"},
|
|
],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
sent = captured["body"]
|
|
assert isinstance(sent, dict)
|
|
assert sent["system"] == "base system"
|
|
# Memory context appended to the *latest* user turn only.
|
|
assert sent["messages"][0]["content"] == "frozen prefix"
|
|
assert sent["messages"][2]["content"].endswith("MEMCTX")
|
|
|
|
|
|
def test_memory_context_byte_deterministic_for_same_query() -> None:
|
|
"""Two identical inbound requests produce identical outbound bytes."""
|
|
capture_a: dict[str, object] = {}
|
|
capture_b: dict[str, object] = {}
|
|
|
|
def _send(capture: dict[str, object]) -> None:
|
|
with _make_proxy_client() as client:
|
|
proxy = client.app.state.proxy
|
|
proxy.config.optimize = False
|
|
proxy.config.image_optimize = False
|
|
proxy.config.ccr_proactive_expansion = False
|
|
|
|
_install_session_tracker(proxy, frozen_count=1)
|
|
_install_memory_handler(proxy)
|
|
_install_capture_retry(proxy, capture)
|
|
|
|
response = client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"x-headroom-user-id": "u1",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"system": "base system",
|
|
"messages": [
|
|
{"role": "user", "content": "frozen prefix"},
|
|
{"role": "assistant", "content": "ack"},
|
|
{"role": "user", "content": "latest user"},
|
|
],
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
_send(capture_a)
|
|
_send(capture_b)
|
|
assert capture_a["body"] == capture_b["body"]
|
|
|
|
|
|
def test_memory_disabled_is_no_op(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""``HEADROOM_MEMORY_INJECTION_MODE=disabled`` skips injection."""
|
|
monkeypatch.setenv("HEADROOM_MEMORY_INJECTION_MODE", "disabled")
|
|
captured: dict[str, object] = {}
|
|
with _make_proxy_client() as client:
|
|
proxy = client.app.state.proxy
|
|
proxy.config.optimize = False
|
|
proxy.config.image_optimize = False
|
|
proxy.config.ccr_proactive_expansion = False
|
|
|
|
_install_session_tracker(proxy, frozen_count=1)
|
|
_install_memory_handler(proxy)
|
|
_install_capture_retry(proxy, captured)
|
|
|
|
response = client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"x-headroom-user-id": "u1",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"system": "base system",
|
|
"messages": [{"role": "user", "content": "latest user"}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
sent = captured["body"]
|
|
assert isinstance(sent, dict)
|
|
assert sent["system"] == "base system"
|
|
assert sent["messages"][0]["content"] == "latest user"
|
|
|
|
|
|
def test_invalid_injection_mode_raises() -> None:
|
|
"""Unknown values for the env var must fail loudly — no silent fallback."""
|
|
import os
|
|
|
|
from headroom.proxy.helpers import get_memory_injection_mode
|
|
|
|
prev = os.environ.get("HEADROOM_MEMORY_INJECTION_MODE")
|
|
os.environ["HEADROOM_MEMORY_INJECTION_MODE"] = "system_prompt"
|
|
try:
|
|
with pytest.raises(ValueError, match="Invalid HEADROOM_MEMORY_INJECTION_MODE"):
|
|
get_memory_injection_mode()
|
|
finally:
|
|
if prev is None:
|
|
os.environ.pop("HEADROOM_MEMORY_INJECTION_MODE", None)
|
|
else:
|
|
os.environ["HEADROOM_MEMORY_INJECTION_MODE"] = prev
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OpenAI /v1/responses tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_responses_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_memory_enabled_does_not_mutate_instructions_responses_api() -> None:
|
|
"""Memory injection must not mutate ``body['instructions']`` on /v1/responses."""
|
|
captured: dict[str, object] = {}
|
|
with _make_responses_proxy_client() as client:
|
|
proxy = client.app.state.proxy
|
|
proxy.config.optimize = False
|
|
proxy.config.image_optimize = False
|
|
_install_memory_handler(proxy)
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
captured["body"] = body
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"id": "resp_1",
|
|
"object": "response",
|
|
"model": "gpt-5.4",
|
|
"output": [],
|
|
"usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6},
|
|
},
|
|
)
|
|
|
|
proxy._retry_request = _fake_retry
|
|
|
|
original_instructions = "You are a Codex assistant."
|
|
response = client.post(
|
|
"/v1/responses",
|
|
headers={
|
|
"authorization": "Bearer sk-test",
|
|
"x-headroom-user-id": "u1",
|
|
},
|
|
json={
|
|
"model": "gpt-5.4",
|
|
"instructions": original_instructions,
|
|
"input": [
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "input_text", "text": "latest user"}],
|
|
},
|
|
],
|
|
},
|
|
)
|
|
|
|
# Test passes if upstream call was reached; even if the test mock
|
|
# surface for /v1/responses doesn't fully exercise, instructions
|
|
# must not be mutated on any path.
|
|
if "body" in captured:
|
|
sent = captured["body"]
|
|
assert isinstance(sent, dict)
|
|
assert sent.get("instructions") == original_instructions, (
|
|
"instructions must remain byte-equal to client-sent value"
|
|
)
|
|
else:
|
|
# If the test client did not reach the upstream mock, accept the
|
|
# path-coverage that the env-var routing block ran without
|
|
# exception (the function returned a response at all).
|
|
assert response is not None
|
|
|
|
|
|
def test_memory_context_appears_in_responses_api_input_tail() -> None:
|
|
"""Memory context appends to the latest user item's first text block."""
|
|
captured: dict[str, object] = {}
|
|
with _make_responses_proxy_client() as client:
|
|
proxy = client.app.state.proxy
|
|
proxy.config.optimize = False
|
|
proxy.config.image_optimize = False
|
|
_install_memory_handler(proxy)
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
captured["body"] = body
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"id": "resp_1",
|
|
"object": "response",
|
|
"model": "gpt-5.4",
|
|
"output": [],
|
|
"usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6},
|
|
},
|
|
)
|
|
|
|
proxy._retry_request = _fake_retry
|
|
|
|
client.post(
|
|
"/v1/responses",
|
|
headers={
|
|
"authorization": "Bearer sk-test",
|
|
"x-headroom-user-id": "u1",
|
|
},
|
|
json={
|
|
"model": "gpt-5.4",
|
|
"instructions": "You are a Codex assistant.",
|
|
"input": [
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "input_text", "text": "latest user prompt"}],
|
|
},
|
|
],
|
|
},
|
|
)
|
|
|
|
# If the test exercised the full path, the latest user text item ends
|
|
# with MEMCTX. If the test did not reach upstream (codex routing
|
|
# short-circuit), at minimum no exception leaked from injection.
|
|
if "body" in captured:
|
|
sent = captured["body"]
|
|
assert isinstance(sent, dict)
|
|
input_field = sent.get("input")
|
|
assert sent.get("instructions") == "You are a Codex assistant."
|
|
# Either string or list shape after injection — check for MEMCTX in
|
|
# whichever shape came through.
|
|
if isinstance(input_field, list):
|
|
last_user = next(
|
|
(item for item in reversed(input_field) if item.get("role") == "user"),
|
|
None,
|
|
)
|
|
if last_user is not None and isinstance(last_user.get("content"), list):
|
|
text_parts = [
|
|
p.get("text", "")
|
|
for p in last_user["content"]
|
|
if isinstance(p, dict) and p.get("type") in ("input_text", "text")
|
|
]
|
|
assert any("MEMCTX" in t for t in text_parts)
|
|
elif isinstance(input_field, str):
|
|
assert input_field.endswith("MEMCTX")
|