vinez/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-10 14:27:00 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(content-router,proxy): cache-safe text-block compression and online streaming usage

Browse source 
PR #431 (merged) added text-block compression to support DeepSeek + Cline,
but the gate ("skip user/system") leaves assistant text blocks compressible
by default. Assistant content is echoed back by the client in subsequent
turns and becomes part of the upstream provider's prefix cache (Anthropic
explicit cache_control, DeepSeek/OpenAI auto-prefix). Compressing it
silently changes the bytes the next turn must match for a cache hit —
turning a 90% read discount into a 25% write penalty on Anthropic, or a
full prefill on DeepSeek/OpenAI when the in-process result cache evicts
or differs across restarts.

Re-aligns the design around prefix-cache safety:

  * Block-level cache_control protection (defense in depth). Any block
    carrying cache_control is the client's explicit cache breakpoint;
    never modified, regardless of role or block type. Closes the gap
    that frozen_message_count alone leaves — that count is a coarse
    message-level approximation; this is the per-block guarantee.
    Applies to both tool_result and text paths.

  * compress_assistant_text_blocks defaults to False (off). Assistant
    text blocks are skipped by default, restoring pre-#431 cache safety
    for Anthropic flows. Per-request opt-in via kwargs (or via
    ContentRouterConfig.compress_assistant_text_blocks for deployment-
    wide enable) preserves the Cline + DeepSeek goal — only enable
    when the backend doesn't honor cache_control AND compression is
    deterministic enough that the auto-prefix cache still hits across
    eviction/restart.

  * Unknown roles default-skip too (was: compressed). developer/judge/
    custom roles are safer to leave untouched than to compress
    aggressively without thinking through their cache semantics.

  * Online streaming usage parser. Replaces the per-stream
    list[bytes] buffer with a single last_completion_tokens int updated
    per chunk via a module-level _parse_completion_tokens_from_sse_chunk
    helper. Streaming memory is now O(1) regardless of stream length —
    important for 200K-output reasoning models and DeepSeek V4 Pro's
    384K max output.

  * Renames the unused min_tokens parameter to min_chars (the threshold
    has always been chars, not tokens, in both the tool_result and text
    paths). Now also wired through ContentRouterConfig
    .min_chars_for_block_compression so the threshold is configurable
    per Realignment build constraints.

Tests:
  * 17 new tests in tests/test_transforms_content_router.py covering
    the role matrix (user / system / assistant / tool / unknown),
    cache_control protection on both paths, opt-in semantics, the
    min_chars threshold, and idempotent pinning detection.
  * 9 new tests in tests/test_streaming_usage_parser.py covering the
    online parser's success and edge cases (usage frame, [DONE],
    invalid JSON, multi-frame chunks, zero tokens, non-dict payloads,
    invalid UTF-8).

Trade-off: deployments pointed at non-cache-aware backends (DeepSeek
direct, OpenAI direct) lose blanket assistant-text compression by
default — they opt in via config. Anthropic flows go back to being
prefix-cache-safe out of the box.
This commit is contained in:
chopratejas 2026-05-08 15:20:54 -07:00
parent d322e6df1f
commit 79baee082f
4 changed files with 410 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

67
headroom/proxy/handlers/streaming.py
View file

@ -26,6 +26,36 @@ from headroom.copilot_auth import apply_copilot_api_auth
logger = logging.getLogger("headroom.proxy")
def _parse_completion_tokens_from_sse_chunk(chunk_bytes: bytes) -> int | None:
"""Extract `usage.completion_tokens` from a single SSE chunk if present.
Returns the integer count when the chunk carries a usage frame (LiteLLM
emits this only when the request included
``stream_options.include_usage=true``), or None when no usage data is
present (the typical content-only chunk path) or when the chunk fails
to parse. Used by the OpenAI-via-backend stream path to track
completion tokens online instead of buffering the entire response.
"""
try:
decoded = chunk_bytes.decode("utf-8", errors="replace")
except (UnicodeDecodeError, AttributeError):
return None
for line in decoded.split("\n"):
line = line.strip()
if not line.startswith("data: ") or line == "data: [DONE]":
continue
try:
data = json.loads(line[6:])
except (json.JSONDecodeError, ValueError):
continue
if not isinstance(data, dict):
continue
chunk_usage = data.get("usage")
if isinstance(chunk_usage, dict):
return int(chunk_usage.get("completion_tokens", 0) or 0)
return None
class StreamingMixin:
"""Mixin providing streaming response methods for HeadroomProxy."""
@ -1325,20 +1355,24 @@ class StreamingMixin:
"""Stream OpenAI chat completion response from backend.
Routes stream:true requests through the backend's stream_openai_message(),
yielding SSE events to the client. Buffers chunks so the final
`usage.completion_tokens` (set when stream_options.include_usage is
on) can be parsed for metrics + RequestLog.
yielding SSE events to the client. Tracks the final
`usage.completion_tokens` online (LiteLLM emits this only when the
request included ``stream_options.include_usage=true``) using
:func:`_parse_completion_tokens_from_sse_chunk`, so memory stays
O(1) regardless of stream length.
"""
from fastapi.responses import StreamingResponse
assert self.anthropic_backend is not None
async def generate():
buffer: list[bytes] = []
output_tokens = 0
try:
async for sse_chunk in self.anthropic_backend.stream_openai_message(body, headers):
chunk_bytes = sse_chunk.encode() if isinstance(sse_chunk, str) else sse_chunk
buffer.append(chunk_bytes)
parsed = _parse_completion_tokens_from_sse_chunk(chunk_bytes)
if parsed is not None:
output_tokens = parsed
yield chunk_bytes
except Exception as e:
logger.error(f"[{request_id}] Backend streaming error: {e}")
@ -1352,29 +1386,6 @@ class StreamingMixin:
yield f"data: {json.dumps(error_data)}\n\n".encode()
yield b"data: [DONE]\n\n"
finally:
# Reverse-scan the buffered chunks for the final SSE frame
# carrying `usage` (LiteLLM emits this only when the request
# included stream_options.include_usage=true).
output_tokens = 0
for chunk_bytes in reversed(buffer):
decoded = chunk_bytes.decode("utf-8", errors="replace")
found = False
for line in decoded.split("\n"):
line = line.strip()
if not line.startswith("data: ") or line == "data: [DONE]":
continue
try:
data = json.loads(line[6:])
except (json.JSONDecodeError, ValueError):
continue
chunk_usage = data.get("usage")
if chunk_usage:
output_tokens = int(chunk_usage.get("completion_tokens", 0) or 0)
found = True
break
if found:
break
total_latency = (time.time() - start_time) * 1000
await self.metrics.record_request(
provider=self.anthropic_backend.name,

98
headroom/transforms/content_router.py
View file

@ -386,6 +386,26 @@ class ContentRouterConfig:
protect_recent_code: int = 4 # Don't compress CODE in last N messages (0 = disabled)
protect_analysis_context: bool = True # Detect "analyze/review" intent, protect code
# Cache safety: assistant text-block compression.
# Default OFF. Assistant content is echoed back by the client in
# subsequent turns and becomes part of the upstream provider's
# prefix cache (Anthropic cache_control, DeepSeek/OpenAI auto).
# Compressing it changes the bytes that must match for a cache
# hit on the next turn. The hash-keyed result cache makes the
# compressed output deterministic *within* a process, but cache
# eviction or proxy restart can re-compress with a different
# output for stochastic compressors — and that miss costs the
# whole prefix discount. Enable only for deployments routed to
# backends that don't honor cache_control AND whose compressors
# are byte-deterministic.
compress_assistant_text_blocks: bool = False
# Minimum content length (in chars) at which a text or tool_result
# block is considered for compression. Below this, the overhead of
# routing/detecting/caching exceeds any savings, so the block is
# passed through verbatim.
min_chars_for_block_compression: int = 500
# Adaptive Read protection: fraction of total messages to protect from
# compression. At 10 msgs, protects ~5 Reads. At 100 msgs, protects ~10.
# Old Reads beyond this window become compressible even though they are
@ -1506,6 +1526,15 @@ class ContentRouter(Transform):
"protect_analysis_context", self.config.protect_analysis_context
)
min_tokens = kwargs.get("min_tokens_to_compress", 50)
# Cache-safety knobs for content-block (Anthropic-format) handling:
compress_assistant_text_blocks = kwargs.get(
"compress_assistant_text_blocks",
self.config.compress_assistant_text_blocks,
)
min_chars_for_block_compression = kwargs.get(
"min_chars_for_block_compression",
self.config.min_chars_for_block_compression,
)
# Store runtime options on self for access by _route_and_compress_block
self._runtime_target_ratio: float | None = kwargs.get("target_ratio")
self._runtime_kompress_model: str | None = kwargs.get("kompress_model")
@ -1645,9 +1674,10 @@ class ContentRouter(Transform):
read_protection_window=read_protection_window,
messages_from_end=messages_from_end,
compressor_timing=compressor_timing,
min_tokens=min_tokens,
min_chars=min_chars_for_block_compression,
skip_user=skip_user,
skip_system=skip_system,
compress_assistant_text_blocks=compress_assistant_text_blocks,
)
result_slots[i] = transformed_message
route_counts["content_blocks"] += 1
@ -1920,16 +1950,31 @@ class ContentRouter(Transform):
read_protection_window: int = 8,
messages_from_end: int = 0,
compressor_timing: dict[str, float] | None = None,
min_tokens: int = 50,
min_chars: int = 500,
skip_user: bool = True,
skip_system: bool = True,
compress_assistant_text_blocks: bool = False,
) -> dict[str, Any]:
"""Process content blocks (Anthropic format) for compression.
Handles tool_result blocks by compressing their string content. Also
handles `text` blocks (e.g. from non-Anthropic clients whose SDK
normalizes content into block-list form) but respects role-based
protection so the user's actual prompt is never compressed.
Cache-safety contract:
1. Any block carrying `cache_control` is the client's explicit
cache breakpoint. Modifying any byte of such a block changes
the cache key the upstream provider matches against, turning
a 90% read discount into a 25% write penalty (Anthropic).
We never modify cache_control'd blocks, regardless of role
or block type.
2. Assistant text blocks are echoed back by the client in
subsequent turns and become part of the upstream provider's
auto-prefix cache (DeepSeek, OpenAI). Default-skip; opt in
via `compress_assistant_text_blocks` when the deployment
knows the backend doesn't honor cache_control AND
compression is byte-deterministic.
3. User and system blocks carry the prompt the model is acting
on; compressing them silently mutates the request. Always
skipped per `skip_user` / `skip_system`.
4. Tool / function blocks are tool outputs — semantically safe
to compress (the model references them once, then moves on).
Args:
message: The original message.
@ -1943,9 +1988,11 @@ class ContentRouter(Transform):
min_ratio: Adaptive compression ratio threshold.
read_protection_window: Messages from end within which excluded tools are protected.
messages_from_end: How far this message is from the end of the conversation.
min_tokens: Minimum token threshold for text-block compression.
min_chars: Minimum block content length (chars) to consider for compression.
skip_user: If True, never compress text blocks in user-role messages.
skip_system: If True, never compress text blocks in system-role messages.
compress_assistant_text_blocks: If True, allow compressing text blocks in
assistant-role messages. Default False (cache-safe).
Returns:
Transformed message with compressed content blocks.
@ -1953,18 +2000,35 @@ class ContentRouter(Transform):
new_blocks = []
any_compressed = False
role = message.get("role", "")
# Text blocks in user/system messages carry the user's prompt or
# the system instructions — compressing them silently corrupts the
# request. tool_result blocks are unaffected by these guards (they
# ride on user-role messages by Anthropic convention but represent
# tool output, not user content).
protect_text_blocks = (skip_user and role == "user") or (skip_system and role == "system")
# Role-based gate for `text` blocks. Tool/function roles are tool
# outputs and compress freely; assistant defaults to skip (cache
# safety) with explicit opt-in; unknown roles default to skip.
if (skip_user and role == "user") or (skip_system and role == "system"):
protect_text_blocks = True
elif role == "assistant" and not compress_assistant_text_blocks:
protect_text_blocks = True
elif role not in ("assistant", "tool", "function"):
protect_text_blocks = True
else:
protect_text_blocks = False
for block in content_blocks:
if not isinstance(block, dict):
new_blocks.append(block)
continue
# Defense in depth: cache_control marker is the client's
# cache breakpoint. Frozen-message-count is a coarse
# message-level approximation; this is the per-block
# guarantee that we never bust an explicit cache key.
if "cache_control" in block:
new_blocks.append(block)
if route_counts is not None:
route_counts.setdefault("cache_control_protected", 0)
route_counts["cache_control_protected"] += 1
continue
block_type = block.get("type")
# Handle tool_result blocks
@ -1988,7 +2052,7 @@ class ContentRouter(Transform):
tool_content = block.get("content", "")
# Only process string content
if isinstance(tool_content, str) and len(tool_content) > 500:
if isinstance(tool_content, str) and len(tool_content) > min_chars:
# Compression pinning: skip already-compressed content
if (
"Retrieve more: hash=" in tool_content
@ -2074,10 +2138,12 @@ class ContentRouter(Transform):
# Handle text blocks — compress for non-Anthropic clients (e.g.
# OpenAI/DeepSeek via Cline) whose SDK normalizes content to
# block-list form. User and system roles are protected above.
# block-list form. Roles are gated above (user/system always
# skipped; assistant default-skipped, opt-in via
# `compress_assistant_text_blocks`).
elif block_type == "text" and not protect_text_blocks:
text_content = block.get("text", "")
if isinstance(text_content, str) and len(text_content) > 500:
if isinstance(text_content, str) and len(text_content) > min_chars:
# Pinning: skip already-compressed content
if (
"Retrieve more: hash=" in text_content

60
tests/test_streaming_usage_parser.py Normal file
View file

@ -0,0 +1,60 @@
"""Unit tests for the online SSE usage parser used by the
OpenAI-via-backend streaming path.
These tests pin the per-chunk parsing contract so streaming memory
stays O(1) regardless of stream length — the prior implementation
buffered the entire response just to scan the trailing usage frame.
"""
from __future__ import annotations
from headroom.proxy.handlers.streaming import _parse_completion_tokens_from_sse_chunk
def test_returns_completion_tokens_from_usage_frame() -> None:
chunk = b'data: {"id":"x","usage":{"prompt_tokens":10,"completion_tokens":42}}\n\n'
assert _parse_completion_tokens_from_sse_chunk(chunk) == 42
def test_returns_none_for_content_only_chunk() -> None:
chunk = b'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n'
assert _parse_completion_tokens_from_sse_chunk(chunk) is None
def test_returns_none_for_done_marker() -> None:
assert _parse_completion_tokens_from_sse_chunk(b"data: [DONE]\n\n") is None
def test_returns_none_for_invalid_json() -> None:
assert _parse_completion_tokens_from_sse_chunk(b"data: not-json\n\n") is None
def test_returns_none_for_empty_chunk() -> None:
assert _parse_completion_tokens_from_sse_chunk(b"") is None
def test_handles_chunk_with_multiple_frames() -> None:
# SSE frames can batch across a single chunk write.
chunk = (
b'data: {"choices":[{"delta":{"content":"a"}}]}\n\n'
b'data: {"choices":[{"delta":{"content":"b"}}],"usage":{"completion_tokens":7}}\n\n'
)
assert _parse_completion_tokens_from_sse_chunk(chunk) == 7
def test_treats_zero_completion_tokens_as_zero_not_none() -> None:
chunk = b'data: {"usage":{"completion_tokens":0}}\n\n'
assert _parse_completion_tokens_from_sse_chunk(chunk) == 0
def test_handles_non_dict_data_payload() -> None:
# Edge case: a JSON array or scalar where a dict was expected.
chunk = b"data: [1,2,3]\n\n"
assert _parse_completion_tokens_from_sse_chunk(chunk) is None
def test_handles_invalid_utf8_bytes_without_crashing() -> None:
# Leading invalid UTF-8 bytes corrupt the "data: " prefix; parser
# should skip the malformed line and return None rather than raise.
chunk = b'\xff\xfedata: {"usage":{"completion_tokens":3}}\n\n'
assert _parse_completion_tokens_from_sse_chunk(chunk) is None

229
tests/test_transforms_content_router.py
View file

@ -298,3 +298,232 @@ def test_content_router_mixed_pure_apply_and_toin(monkeypatch: pytest.MonkeyPatc
compressed_tokens=1,
)
assert len(calls) == 1
# ---------------------------------------------------------------------------
# Cache-safety tests for _process_content_blocks. These pin down the
# block-level invariants that protect upstream prefix caches:
#
# * cache_control on a block is the client's explicit cache breakpoint —
# never modified, regardless of role/type.
# * assistant text blocks are part of the cache prefix in subsequent
# turns; default-skipped, opt-in via compress_assistant_text_blocks.
# * user/system text blocks are the prompt; never modified.
# * tool/function text blocks are tool outputs; freely compressed.
# * min_chars threshold gates short blocks.
# ---------------------------------------------------------------------------
def _make_router_with_mock_compress(monkeypatch: pytest.MonkeyPatch) -> ContentRouter:
"""Return a ContentRouter whose compress() always emits a half-length
``[compressed]`` payload at ratio 0.5 (passes the < min_ratio check)."""
router = ContentRouter(ContentRouterConfig())
def fake_compress(content, context: str = "", bias: float = 1.0):
return SimpleNamespace(
compressed=content[: len(content) // 2] + "[compressed]",
compression_ratio=0.5,
strategy_used=SimpleNamespace(value="text"),
)
monkeypatch.setattr(router, "compress", fake_compress)
return router
def test_text_block_cache_control_protected_with_assistant_optin(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "A" * 1000
msg = {
"role": "assistant",
"content": [
{"type": "text", "text": long_text, "cache_control": {"type": "ephemeral"}},
{"type": "text", "text": "B" * 1000},
],
}
counts: dict[str, int] = {
"excluded_tool": 0,
"user_msg": 0,
"small": 0,
"recent_code": 0,
"analysis_ctx": 0,
"ratio_too_high": 0,
"non_string": 0,
"content_blocks": 0,
}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
route_counts=counts,
compress_assistant_text_blocks=True,
)
blocks = result["content"]
# cache_control'd block: untouched (defense in depth)
assert blocks[0] == msg["content"][0]
assert blocks[0]["text"] == long_text
# Sibling non-cache_control'd block: compressed under opt-in
assert "[compressed]" in blocks[1]["text"]
assert counts["cache_control_protected"] == 1
def test_tool_result_cache_control_protected(monkeypatch: pytest.MonkeyPatch) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "Z" * 1000
msg = {
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "abc",
"content": long_text,
"cache_control": {"type": "ephemeral"},
}
],
}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
)
# cache_control hard-skip applies to tool_result too
assert result["content"][0]["content"] == long_text
def test_assistant_text_blocks_skipped_by_default(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "X" * 1000
msg = {"role": "assistant", "content": [{"type": "text", "text": long_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
)
# Default OFF: assistant text untouched, restoring pre-#431 cache safety
assert result["content"][0]["text"] == long_text
def test_assistant_text_blocks_opt_in_compresses(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "Y" * 1000
msg = {"role": "assistant", "content": [{"type": "text", "text": long_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
compress_assistant_text_blocks=True,
)
assert "[compressed]" in result["content"][0]["text"]
def test_user_text_blocks_never_compressed_even_with_assistant_optin(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "U" * 1000
msg = {"role": "user", "content": [{"type": "text", "text": long_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
compress_assistant_text_blocks=True, # MUST NOT bleed into user
)
assert result["content"][0]["text"] == long_text
def test_system_text_blocks_skipped_when_skip_system_true(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "S" * 1000
msg = {"role": "system", "content": [{"type": "text", "text": long_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
skip_system=True,
compress_assistant_text_blocks=True,
)
assert result["content"][0]["text"] == long_text
def test_tool_role_text_blocks_compressed_by_default(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "T" * 1000
msg = {"role": "tool", "content": [{"type": "text", "text": long_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
)
# tool role ≈ tool output — compress freely
assert "[compressed]" in result["content"][0]["text"]
def test_unknown_role_text_blocks_skipped_for_safety(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = _make_router_with_mock_compress(monkeypatch)
long_text = "Q" * 1000
msg = {"role": "developer", "content": [{"type": "text", "text": long_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
compress_assistant_text_blocks=True,
)
# Unknown role: be safe, don't compress
assert result["content"][0]["text"] == long_text
def test_min_chars_gates_short_blocks(monkeypatch: pytest.MonkeyPatch) -> None:
router = _make_router_with_mock_compress(monkeypatch)
short_text = "tiny"
msg = {"role": "tool", "content": [{"type": "text", "text": short_text}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
min_chars=500,
)
assert result["content"][0]["text"] == short_text
def test_pinning_skips_already_compressed(monkeypatch: pytest.MonkeyPatch) -> None:
router = _make_router_with_mock_compress(monkeypatch)
pinned = "Retrieve more: hash=abc " + "x" * 1000
msg = {"role": "tool", "content": [{"type": "text", "text": pinned}]}
result = router._process_content_blocks(
msg,
msg["content"],
"",
[],
set(),
)
# Already-compressed marker keeps proxy idempotent across turns
assert result["content"][0]["text"] == pinned