diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 39d23abab..f3b4f4e45 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -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, diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index b75761e4f..05eacc8a5 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -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 diff --git a/tests/test_streaming_usage_parser.py b/tests/test_streaming_usage_parser.py new file mode 100644 index 000000000..c93b710a9 --- /dev/null +++ b/tests/test_streaming_usage_parser.py @@ -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 diff --git a/tests/test_transforms_content_router.py b/tests/test_transforms_content_router.py index 0c20cd08a..fc82e20c3 100644 --- a/tests/test_transforms_content_router.py +++ b/tests/test_transforms_content_router.py @@ -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