fix(responses): bill a WebSocket turn's guardrail cost when the deployment refuses it

A turn the per-turn guardrail passes and the deployment then refuses now writes
its own failed spend row with that turn's guardrail information, on both relays,
the way a guardrail-blocked frame and an HTTP refusal already do

litellm/proxy/response_api_endpoints/endpoints.py

@@ -1858,7 +1858,7 @@ def _responses_ws_frame_guard(
     }
     model: Final = admitted.get("model")
 
-    # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation
+    # blocked and refused frames bill as their own failed requests; only the first turn settles the reservation
     rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None})
 
     async def guard(frame: str) -> GuardedResponsesWsFrame:
@@ -1906,6 +1906,7 @@ def _responses_ws_frame_guard(
             frame=json.dumps(rewritten),
             guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)),
             pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")),
+            bill_refusal=partial(_bill_rejected_frame, rejected_frame_auth, frame_data=guarded_data),
         )
 
     return guard

litellm/responses/streaming_iterator.py

@@ -1962,11 +1962,6 @@ def _turn_rejection_frame(rejection: Exception) -> str:
             return _turn_admission_error_frame(rejection)
 
 
-async def _answer_rejected_turn(websocket: ResponsesClientWebSocket, rejection: Exception) -> None:
-    with contextlib.suppress(Exception):
-        await websocket.send_text(_turn_rejection_frame(rejection))
-
-
 RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
     "response.created",
     "response.completed",
@@ -2039,16 +2034,28 @@ def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], un
 @dataclass(frozen=True, slots=True)
 class GuardedResponsesWsFrame:
     """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails
-    may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the
-    Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply."""
+    may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, the
+    Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply, and how
+    to bill those guardrail records when the deployment then refuses the turn."""
 
     frame: str
     guardrail_information: tuple[Mapping[str, object], ...] = ()
     pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
+    bill_refusal: Callable[[Exception], Awaitable[None]] | None = None
 
 
 ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]]
 
+
+async def _reject_guarded_turn(
+    websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception
+) -> None:
+    with contextlib.suppress(Exception):
+        await websocket.send_text(_turn_rejection_frame(rejection))
+    if guarded.guardrail_information and guarded.bill_refusal is not None:
+        await guarded.bill_refusal(rejection)
+
+
 _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES
 
 
@@ -2310,7 +2317,7 @@ class ResponsesWebSocketStreaming:
         self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None
     ) -> None:
         guarded: Final = await self._guard_frame(message)
-        if guarded is None or not await self._admit_or_reject_turn(message, turn_admission):
+        if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission):
             return
         if guarded.pii_tokens:
             self._remember_pii_tokens(guarded.pii_tokens)
@@ -2708,13 +2715,15 @@ class ResponsesWebSocketStreaming:
             return False
         return True
 
-    async def _admit_or_reject_turn(self, message: str, turn_admission: Callable[[], Awaitable[None]] | None) -> bool:
+    async def _admit_or_reject_turn(
+        self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None
+    ) -> bool:
         rejection: Final = await _admit_frame_turn(turn_admission, message)
         if rejection is None:
             return True
         if self._failure is None:
             self._failure = rejection
-        await _answer_rejected_turn(self.websocket, rejection)
+        await _reject_guarded_turn(self.websocket, guarded, rejection)
         return False
 
     async def _send_turn_refusal(self, refusal: RateLimitError) -> None:
@@ -3290,7 +3299,7 @@ class ManagedResponsesWebSocketHandler:
         if rejection is not None:
             if self._failure is None:
                 self._failure = rejection
-            await _answer_rejected_turn(self.websocket, rejection)
+            await _reject_guarded_turn(self.websocket, guarded, rejection)
             return
         if guarded.pii_tokens:
             self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens}

tests/unit/proxy/response_api_endpoints/test_endpoints.py

@@ -2054,6 +2054,44 @@ class TestResponsesWSFirstFrameModelAuth:
         assert failure["request_data"]["input"] == "hi"
         assert len(failure["request_data"]["litellm_call_id"]) == 36
 
+    @pytest.mark.asyncio
+    async def test_frame_guard_bills_a_passed_frame_the_deployment_refuses_with_its_guardrail_records(self):
+        from litellm.exceptions import RateLimitError
+        from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
+        from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
+
+        async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
+            data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}]
+            return {**data, "input": data["input"] + " [guarded]"}
+
+        proxy_logging_obj = MagicMock()
+        proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
+        spend_logger = MagicMock(spec=ProxyDBLogger)
+        spend_logger.async_post_call_failure_hook = AsyncMock()
+        rate_limiter = MagicMock()
+        rate_limiter.async_post_call_failure_hook = AsyncMock()
+        user_api_key_dict = MagicMock()
+        refusal = RateLimitError(message="crossed budget", llm_provider="", model="")
+        guard = _responses_ws_frame_guard(
+            {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj
+        )
+
+        guarded = await guard(json.dumps({"type": "response.create", "input": "hi"}))
+        assert guarded.bill_refusal is not None
+        with patch.object(litellm, "callbacks", [rate_limiter, spend_logger]):
+            await guarded.bill_refusal(refusal)
+
+        rate_limiter.async_post_call_failure_hook.assert_not_awaited()
+        spend_logger.async_post_call_failure_hook.assert_awaited_once()
+        failure = spend_logger.async_post_call_failure_hook.await_args.kwargs
+        user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None})
+        assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value
+        assert failure["original_exception"] is refusal
+        assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [
+            {"guardrail_cost": 0.5}
+        ]
+        assert failure["request_data"]["input"] == "hi [guarded]"
+
     @pytest.mark.asyncio
     async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self):
         from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard

tests/unit/responses/test_responses_websocket_all_providers.py

@@ -1632,6 +1632,17 @@ def _guard_blocking(word: str):
     return guard
 
 
+def _guard_recording(guardrail_information: tuple[Mapping[str, object], ...], bill_refusal: object):
+    from litellm.responses.streaming_iterator import GuardedResponsesWsFrame
+
+    async def guard(frame: str) -> GuardedResponsesWsFrame:
+        return GuardedResponsesWsFrame(
+            frame=frame, guardrail_information=guardrail_information, bill_refusal=bill_refusal
+        )
+
+    return guard
+
+
 def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]:
     return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"]
 
@@ -1833,6 +1844,90 @@ class TestWebSocketTurnAdmission:
         ]
 
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)]
+    )
+    async def test_native_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses(
+        self, guardrail_information: tuple[Mapping[str, object], ...], billed: bool
+    ):
+        from unittest.mock import AsyncMock, call
+
+        from litellm.exceptions import RateLimitError
+
+        first, second, third = _turn_frames()
+        refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="")
+        bill_refusal: Final = AsyncMock()
+        backend_ws: Final = MagicMock()
+        backend_ws.send = AsyncMock()
+        handler: Final = _make_streaming(
+            websocket=_ClientSocket(second, third),
+            backend_ws=backend_ws,
+            request_data={},
+            first_message=first,
+            frame_guard=_guard_recording(guardrail_information, bill_refusal),
+            turn_admission=AsyncMock(side_effect=[refusal, None]),
+        )
+
+        await handler.client_to_backend()
+
+        assert [json.loads(sent.args[0])["input"] for sent in backend_ws.send.await_args_list] == ["one", "three"]
+        assert bill_refusal.await_args_list == ([call(refusal)] if billed else [])
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)]
+    )
+    async def test_managed_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        guardrail_information: tuple[Mapping[str, object], ...],
+        billed: bool,
+    ):
+        from unittest.mock import AsyncMock, call
+
+        import litellm
+        from litellm.exceptions import RateLimitError
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler
+
+        answered: Final[list[object]] = []
+
+        async def fake_aresponses(*args: object, **kwargs: object):
+            answered.append(kwargs.get("input"))
+
+            async def _empty():
+                return
+                yield
+
+            return _empty()
+
+        monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
+        first, second, third = _turn_frames()
+        refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="")
+        bill_refusal: Final = AsyncMock()
+        handler: Final = ManagedResponsesWebSocketHandler(
+            websocket=_ClientSocket(second, third),
+            model="test-model",
+            logging_obj=Logging(
+                model="test-model",
+                messages=[],
+                stream=True,
+                call_type="aresponses",
+                start_time=0,
+                litellm_call_id="test-id",
+                function_id="test-func",
+            ),
+            first_message=first,
+            frame_guard=_guard_recording(guardrail_information, bill_refusal),
+            turn_admission=AsyncMock(side_effect=[refusal, None]),
+        )
+
+        await handler.run()
+
+        assert answered == ["one", "three"]
+        assert bill_refusal.await_args_list == ([call(refusal)] if billed else [])
+
     @pytest.mark.asyncio
     async def test_native_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving(self):
         from unittest.mock import AsyncMock