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

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

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