fix: report the blocked LLM response's real token usage (#31217)

When a guardrail blocks a post-call response, the synthetic violation response
reported hard-coded zero usage, discarding the token usage the upstream call
had already consumed.

Fix the root cause rather than re-counting tokens:
- Add an optional `original_response` field to ModifyResponseException.
- The unified guardrail's post-call success hook attaches the blocked LLM
  response to the exception.
- The /v1/messages and OpenAI-format (/v1/chat/completions, /v1/completions)
  block handlers report `original_response.usage` directly. Pre-call blocks
  never invoked the LLM, so usage is zero.

Mock-based tests cover the helper (returns original usage / zero), the success
hook attaching original_response, and the endpoint reporting it end-to-end.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

litellm/exceptions.py

@@ -1165,12 +1165,18 @@ class ModifyResponseException(Exception):
         request_data: Dict[str, Any],
         guardrail_name: Optional[str] = None,
         detection_info: Optional[Dict[str, Any]] = None,
+        original_response: Optional[Any] = None,
     ):
         self.message = message
         self.model = model
         self.request_data = request_data
         self.guardrail_name = guardrail_name
         self.detection_info = detection_info or {}
+        # The LLM response that was blocked (post-call). Carries the real token
+        # usage the upstream call consumed, so the synthetic block response can
+        # report it instead of discarding it. None for pre-call blocks (the LLM
+        # was never invoked).
+        self.original_response = original_response
         super().__init__(message)
 
 

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -19,6 +19,7 @@ from litellm.proxy.common_request_processing import (
     create_response,
 )
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 from litellm.types.utils import TokenCountResponse
 
 router = APIRouter()
@@ -58,6 +59,33 @@ def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
         usage.pop("total_tokens", None)
 
 
+def _blocked_response_usage(original_response: Optional[Any]) -> AnthropicUsage:
+    """
+    Token usage for a synthetic guardrail-blocked response.
+
+    A post-call block replaces the LLM's response with the violation message,
+    but the upstream call already consumed tokens -- report that real usage
+    (carried on ``ModifyResponseException.original_response``) rather than
+    discarding it. Pre-call blocks never invoked the LLM (no original_response),
+    so usage is zero.
+    """
+    usage_obj: Any = None
+    if isinstance(original_response, dict):
+        usage_obj = original_response.get("usage")
+    elif original_response is not None:
+        usage_obj = getattr(original_response, "usage", None)
+
+    def _tokens(key: str) -> int:
+        if isinstance(usage_obj, dict):
+            return int(usage_obj.get(key, 0) or 0)
+        return int(getattr(usage_obj, key, 0) or 0)
+
+    return AnthropicUsage(
+        input_tokens=_tokens("input_tokens"),
+        output_tokens=_tokens("output_tokens"),
+    )
+
+
 @router.post(
     "/v1/messages",
     tags=["[beta] Anthropic `/v1/messages`"],
@@ -134,6 +162,10 @@ async def anthropic_response(
 
         from litellm.types.utils import AnthropicMessagesResponse
 
+        # Report the blocked LLM response's real token usage (carried on the
+        # exception) instead of discarding it; zero for pre-call blocks.
+        _usage = _blocked_response_usage(e.original_response)
+
         _anthropic_response = AnthropicMessagesResponse(
             id=f"msg_{str(uuid.uuid4())}",
             type="message",
@@ -141,7 +173,7 @@ async def anthropic_response(
             content=[{"type": "text", "text": e.message}],
             model=e.model,
             stop_reason="end_turn",
-            usage={"input_tokens": 0, "output_tokens": 0},
+            usage=_usage,
         )
 
         if data.get("stream", None) is not None and data["stream"] is True:

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -197,6 +197,10 @@ class UnifiedLLMGuardrails(CustomLogger):
         )
         from litellm.types.guardrails import GuardrailEventHooks
 
+        # Local import avoids a module-level cyclic import with
+        # litellm.integrations.custom_guardrail.
+        from litellm.integrations.custom_guardrail import ModifyResponseException
+
         guardrail_to_apply: CustomGuardrail = data.pop("guardrail_to_apply", None)
 
         if guardrail_to_apply is None:
@@ -238,13 +242,22 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
 
-        response = await endpoint_translation.process_output_response(
-            response=response,  # type: ignore
-            guardrail_to_apply=guardrail_to_apply,
-            litellm_logging_obj=data.get("litellm_logging_obj"),
-            user_api_key_dict=user_api_key_dict,
-            request_data=data,
-        )
+        try:
+            response = await endpoint_translation.process_output_response(
+                response=response,  # type: ignore
+                guardrail_to_apply=guardrail_to_apply,
+                litellm_logging_obj=data.get("litellm_logging_obj"),
+                user_api_key_dict=user_api_key_dict,
+                request_data=data,
+            )
+        except ModifyResponseException as e:
+            # The guardrail blocked the response. Attach the original LLM
+            # response so the endpoint handler can report its real token usage
+            # instead of discarding it (the block replaces the content, but the
+            # upstream call already consumed those tokens).
+            if e.original_response is None:
+                e.original_response = response
+            raise
         # Add guardrail to applied guardrails header
         add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_to_apply.guardrail_name)
 

litellm/proxy/proxy_server.py

@@ -8361,6 +8361,22 @@ async def model_info(
     )
 
 
+def _blocked_response_usage(original_response: Optional[Any]) -> "litellm.Usage":
+    """
+    Token usage for a synthetic guardrail-blocked response.
+
+    A post-call block replaces the LLM's response with the violation message,
+    but the upstream call already consumed tokens -- report that real usage
+    (carried on ``ModifyResponseException.original_response``) rather than
+    discarding it. Pre-call blocks never invoked the LLM (no original_response),
+    so usage is zero.
+    """
+    usage = getattr(original_response, "usage", None) if original_response is not None else None
+    if isinstance(usage, litellm.Usage):
+        return usage
+    return litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
+
+
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(user_api_key_auth)],
@@ -8467,6 +8483,9 @@ async def chat_completion(
         _chat_response.model = e.model  # type: ignore
         _chat_response.choices[0].message.content = e.message  # type: ignore
         _chat_response.choices[0].finish_reason = "content_filter"  # type: ignore
+        # Report the blocked LLM response's real usage (set before the stream
+        # branch so both paths carry it); zero for pre-call blocks.
+        _chat_response.usage = _blocked_response_usage(e.original_response)  # type: ignore
 
         if data.get("stream", None) is not None and data["stream"] is True:
             _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True)
@@ -8488,8 +8507,6 @@ async def chat_completion(
                 media_type="text/event-stream",
                 status_code=200,  # Return 200 for passthrough mode
             )
-        _usage = litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
-        _chat_response.usage = _usage  # type: ignore
         return _chat_response
     except RejectedRequestError as e:
         _data = e.request_data
@@ -8618,11 +8635,7 @@ async def completion(
             # Set text attribute dynamically for text completion format
             setattr(_text_response.choices[0], "text", e.message)
             _text_response.model = e.model  # type: ignore[assignment]
-            _usage = litellm.Usage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0,
-            )
+            _usage = _blocked_response_usage(e.original_response)
             # Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition)
             setattr(_text_response, "usage", _usage)
             _iterator = litellm.utils.ModelResponseIterator(model_response=_text_response, convert_to_delta=True)
@@ -8647,11 +8660,7 @@ async def completion(
             _response = litellm.TextCompletionResponse()
             _response.choices[0].text = e.message
             _response.model = e.model  # type: ignore
-            _usage = litellm.Usage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0,
-            )
+            _usage = _blocked_response_usage(e.original_response)
             _response.usage = _usage  # type: ignore
             return _response
     except RejectedRequestError as e:

tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py

@@ -58,15 +58,71 @@ class TestAnthropicEndpoints(unittest.TestCase):
         self.assertEqual(result, expected_result)
 
         # Assert safe_dumps was called for dictionary objects
-        mock_safe_dumps.assert_any_call(
-            {"type": "message_start", "message": {"id": "msg_123"}}
+        mock_safe_dumps.assert_any_call({"type": "message_start", "message": {"id": "msg_123"}})
+        mock_safe_dumps.assert_any_call({"type": "content_block_delta", "delta": {"text": "more data"}})
+        assert mock_safe_dumps.call_count == 2  # Called twice, once for each dict object
+
+
+class TestBlockedResponseUsage:
+    """Blocked responses report the blocked LLM response's real usage."""
+
+    def test_uses_original_response_usage(self):
+        from litellm.proxy.anthropic_endpoints.endpoints import _blocked_response_usage
+
+        # original_response is the AnthropicMessagesResponse the LLM produced
+        # before the guardrail blocked it; its usage is real.
+        original = {"usage": {"input_tokens": 31, "output_tokens": 9}}
+        assert _blocked_response_usage(original) == {
+            "input_tokens": 31,
+            "output_tokens": 9,
+        }
+
+    def test_zero_usage_when_no_original_response(self):
+        from litellm.proxy.anthropic_endpoints.endpoints import _blocked_response_usage
+
+        # Pre-call blocks never invoked the LLM -> nothing consumed.
+        assert _blocked_response_usage(None) == {
+            "input_tokens": 0,
+            "output_tokens": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_blocked_endpoint_response_carries_original_usage(self):
+        """The /v1/messages block handler reports the blocked response's real
+        usage, carried on ModifyResponseException.original_response."""
+        from unittest.mock import AsyncMock, MagicMock
+
+        import litellm.proxy.anthropic_endpoints.endpoints as ep
+        import litellm.proxy.proxy_server as proxy_server
+        from litellm.integrations.custom_guardrail import ModifyResponseException
+
+        exc = ModifyResponseException(
+            message="blocked by guardrail",
+            model="claude-3-5-sonnet-20240620",
+            request_data={"messages": [{"role": "user", "content": "hi"}]},
+            guardrail_name="rubrik",
+            original_response={"usage": {"input_tokens": 12, "output_tokens": 5}},
         )
-        mock_safe_dumps.assert_any_call(
-            {"type": "content_block_delta", "delta": {"text": "more data"}}
-        )
-        assert (
-            mock_safe_dumps.call_count == 2
-        )  # Called twice, once for each dict object
+
+        with (
+            patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})),
+            patch.object(
+                ep.ProxyBaseLLMRequestProcessing,
+                "base_process_llm_request",
+                new=AsyncMock(side_effect=exc),
+            ),
+            patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
+        ):
+            mock_logging.post_call_failure_hook = AsyncMock()
+            response = await ep.anthropic_response(
+                fastapi_response=MagicMock(),
+                request=MagicMock(),
+                user_api_key_dict=MagicMock(),
+            )
+
+        assert response["content"][0]["text"] == "blocked by guardrail"
+        assert response["usage"] == {"input_tokens": 12, "output_tokens": 5}
+        mock_logging.post_call_failure_hook.assert_awaited_once()
 
 
 class TestEventLoggingBatchEndpoint:
@@ -159,9 +215,7 @@ class TestStripTotalTokens(unittest.TestCase):
 
         # SimpleNamespace mimics the .usage attribute access pattern; the
         # helper's contract: if .usage is dict-shaped, strip total_tokens.
-        response = SimpleNamespace(
-            usage={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}
-        )
+        response = SimpleNamespace(usage={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150})
         _strip_total_tokens_from_anthropic_response(response)
         assert "total_tokens" not in response.usage
         assert response.usage == {"input_tokens": 100, "output_tokens": 50}

tests/test_litellm/proxy/test_blocked_response_usage.py

@@ -0,0 +1,84 @@
+"""
+Token usage on synthetic guardrail-blocked responses for the OpenAI-format
+proxy endpoints (/v1/chat/completions and /v1/completions).
+
+A post-call block replaces the LLM response with the violation message, but the
+upstream call already consumed tokens. `_blocked_response_usage` reports that
+real usage (carried on `ModifyResponseException.original_response`) rather than
+zero; a pre-call block never invoked the LLM, so usage is zero.
+"""
+
+import pytest
+
+import litellm
+from litellm.proxy.proxy_server import _blocked_response_usage
+
+
+def test_uses_original_response_usage():
+    resp = litellm.ModelResponse()
+    resp.usage = litellm.Usage(prompt_tokens=42, completion_tokens=7, total_tokens=49)
+
+    usage = _blocked_response_usage(resp)
+
+    assert usage.prompt_tokens == 42
+    assert usage.completion_tokens == 7
+    assert usage.total_tokens == 49
+
+
+def test_zero_usage_when_no_original_response():
+    usage = _blocked_response_usage(None)
+
+    assert usage.prompt_tokens == 0
+    assert usage.completion_tokens == 0
+    assert usage.total_tokens == 0
+
+
+@pytest.mark.asyncio
+async def test_success_hook_attaches_original_response_on_block():
+    """The unified guardrail's post-call success hook must attach the blocked
+    LLM response to ModifyResponseException so its real usage isn't discarded."""
+    from unittest.mock import AsyncMock, MagicMock, patch
+
+    import litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail as ug
+    from litellm.integrations.custom_guardrail import ModifyResponseException
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.utils import CallTypes
+
+    response = litellm.ModelResponse()
+    response.usage = litellm.Usage(prompt_tokens=15, completion_tokens=3, total_tokens=18)
+
+    guardrail = MagicMock()
+    guardrail.should_run_guardrail.return_value = True
+    guardrail.guardrail_name = "rubrik"
+
+    # The translation layer raises a block without pre-setting original_response.
+    translation = MagicMock()
+    translation.process_output_response = AsyncMock(
+        side_effect=ModifyResponseException(
+            message="blocked",
+            model="gpt-4o",
+            request_data={},
+            guardrail_name="rubrik",
+        )
+    )
+
+    unified = ug.UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions")
+    data = {"guardrail_to_apply": guardrail, "model": "gpt-4o"}
+
+    # Inject our translation for the inferred call type (the module global is
+    # cached across tests, so patch it directly rather than the loader).
+    with patch.object(
+        ug,
+        "endpoint_guardrail_translation_mappings",
+        {
+            CallTypes.acompletion: lambda: translation,
+            CallTypes.completion: lambda: translation,
+        },
+    ):
+        with pytest.raises(ModifyResponseException) as excinfo:
+            await unified.async_post_call_success_hook(
+                data=data, user_api_key_dict=user_api_key_dict, response=response
+            )
+
+    assert excinfo.value.original_response is response