fix(model_armor): keep upstream usage fields when masking a raw Anthropic stream

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/guardrails/anthropic_sse.py

@@ -111,6 +111,65 @@ def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]:
     )
 
 
+def _is_text_delta_event(event: str) -> bool:
+    from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
+        AnthropicPassthroughLoggingHandler,
+    )
+
+    event_data: Final = AnthropicPassthroughLoggingHandler._extract_sse_data(event)  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses
+    if event_data is None or event_data.get("type") != "content_block_delta":
+        return False
+    delta: Final = event_data.get("delta")
+    return isinstance(delta, dict) and delta.get("type") == "text_delta"
+
+
+def _text_delta_frame(text: str, index: object) -> bytes:
+    payload: Final = {
+        "type": "content_block_delta",
+        "index": index if isinstance(index, int) else 0,
+        "delta": {"type": "text_delta", "text": text},
+    }
+    return f"event: content_block_delta\ndata: {json.dumps(payload)}\n\n".encode()
+
+
+def _content_block_index(event: str) -> object:
+    from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
+        AnthropicPassthroughLoggingHandler,
+    )
+
+    event_data: Final = AnthropicPassthroughLoggingHandler._extract_sse_data(event)  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses
+    return None if event_data is None else event_data.get("index")
+
+
+def rewrite_anthropic_sse_text(all_chunks: Sequence[object], replacement: str) -> tuple[bytes, ...] | None:
+    """Re-emit the upstream frames with the assistant text replaced by ``replacement``.
+
+    Rebuilding the stream from an assembled ``ModelResponse`` loses everything the assembler does
+    not model, such as the cache, server-tool-use and service-tier usage fields Anthropic sends in
+    ``message_start``. Rewriting the frames in place keeps them. Returns None when the stream holds
+    no text delta to replace.
+    """
+    from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
+        AnthropicPassthroughLoggingHandler,
+    )
+
+    sse_stream: Final = _joined_sse_stream(all_chunks)
+    if sse_stream is None:
+        return None
+    events: Final = AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream)  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses
+    text_positions: Final = tuple(position for position, event in enumerate(events) if _is_text_delta_event(event))
+    if not text_positions:
+        return None
+    dropped: Final = frozenset(text_positions[1:])
+    return tuple(
+        _text_delta_frame(replacement, _content_block_index(event))
+        if position == text_positions[0]
+        else f"{event}\n\n".encode()
+        for position, event in enumerate(events)
+        if position not in dropped
+    )
+
+
 def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:
     from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
         LiteLLMAnthropicMessagesAdapter,

litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py

@@ -842,10 +842,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
         from litellm.main import stream_chunk_builder
         from litellm.proxy.guardrails.anthropic_sse import (
-            anthropic_sse_chunks_from_response,
             anthropic_sse_error_frames,
             assemble_anthropic_sse_stream,
             is_raw_sse_stream,
+            rewrite_anthropic_sse_text,
         )
 
         # Collect all chunks
@@ -853,8 +853,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         async for chunk in response:
             all_chunks.append(chunk)
 
-        # /v1/messages arrives as raw SSE frames and /v1/responses as Responses API event
-        # objects; stream_chunk_builder can assemble neither
         raw_sse: Final = is_raw_sse_stream(all_chunks)
         chat_stream: Final = bool(all_chunks) and all(isinstance(chunk, ModelResponseStream) for chunk in all_chunks)
         assembled_response: Final = (
@@ -920,7 +918,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
 
                             # Return sanitized stream
                             if raw_sse:
-                                for sse_chunk in anthropic_sse_chunks_from_response(assembled_response):
+                                rewritten: Final = rewrite_anthropic_sse_text(all_chunks, sanitized_content)
+                                for sse_chunk in rewritten or anthropic_sse_error_frames(
+                                    f"{self.guardrail_name}: sanitized response could not be re-emitted, blocking it"
+                                ):
                                     yield sse_chunk
                                 return
                             mock_response: Final = MockResponseIterator(model_response=assembled_response)

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py

@@ -625,7 +625,8 @@ async def test_model_armor_streaming_block_yields_sse_error():
 
 _ANTHROPIC_SSE_CHUNKS = (
     b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message",'
-    b'"role":"assistant","model":"claude","content":[],"usage":{"input_tokens":5,"output_tokens":0}}}\n\n',
+    b'"role":"assistant","model":"claude","content":[],"usage":{"input_tokens":5,"output_tokens":0,'
+    b'"cache_read_input_tokens":4,"service_tier":"standard"}}}\n\n',
     b'event: content_block_start\ndata: {"type":"content_block_start","index":0,'
     b'"content_block":{"type":"text","text":""}}\n\n',
     b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,'
@@ -762,6 +763,10 @@ async def test_streaming_hook_masks_raw_anthropic_sse():
     body = b"".join(chunk for chunk in delivered if isinstance(chunk, bytes))
     assert b"[REDACTED]" in body
     assert b"123-45-6789" not in body
+    # the masked stream is a rewrite of the upstream frames, so usage the assembler does not
+    # model (cache counts, service tier) still reaches the client
+    assert b'"cache_read_input_tokens":4' in body
+    assert b'"service_tier":"standard"' in body
 
 
 @pytest.mark.asyncio