fix(bedrock/realtime): keep partial spend on cancelled output task and make the committed-session refusal non-retryable

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

litellm/constants.py

@@ -314,6 +314,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
 # ASGI websocket scope keys the Bedrock realtime bridge uses to carry state across router fallback attempts
 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
 BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
+BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
 
 # SSL/TLS cipher configuration for faster handshakes
 # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones

litellm/llms/bedrock/realtime/handler.py

@@ -10,13 +10,14 @@ import json
 from collections.abc import AsyncIterator, Mapping, MutableMapping
 from dataclasses import dataclass
 from types import MappingProxyType
-from typing import Final, Protocol
+from typing import Final, NoReturn, Protocol
 
 from pydantic import JsonValue, TypeAdapter
 
 import litellm
 from litellm._logging import _redact_string, verbose_proxy_logger
 from litellm.constants import (
+    BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
     BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
     BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
 )
@@ -61,15 +62,6 @@ def _as_bedrock_error(error: BaseException) -> BaseException:
     return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}")
 
 
-class _BedrockForwardingFailed(Exception):
-    """The Bedrock output stream failed after ``logged_events`` were already forwarded to the client."""
-
-    def __init__(self, cause: BaseException, logged_events: tuple[OpenAIRealtimeEvents, ...]) -> None:
-        super().__init__(str(cause))
-        self.cause: Final = cause
-        self.logged_events: Final = logged_events
-
-
 def _json_dict(value: JsonValue) -> dict[str, JsonValue]:
     return value if isinstance(value, dict) else {}
 
@@ -150,14 +142,26 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message:
 def _pending_session_update(scope: Mapping[str, object]) -> str | None:
     """A fallback attempt on the same websocket replays the session.update the failed attempt never acked."""
     if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
+        committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY)
         raise BedrockError(
-            status_code=409,
-            message="Bedrock realtime session already committed to a provider stream; it cannot be replayed",
+            status_code=400,
+            message=(
+                "Bedrock realtime session already committed to a provider stream; it cannot be replayed"
+                + (f". The committed stream failed with: {committed_failure}" if committed_failure else "")
+            ),
         )
     pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY)
     return pending if isinstance(pending, str) else None
 
 
+def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn:
+    error: Final = _as_bedrock_error(failure)
+    verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", _redact_string(str(error)))
+    if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
+        scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = _redact_string(str(error))
+    raise error from failure
+
+
 def _parse_client_message(message: str) -> Mapping[str, JsonValue]:
     try:
         return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
@@ -351,10 +355,7 @@ class BedrockRealtime(BaseAWSLLM):
                 "Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure
             )
             return
-        verbose_proxy_logger.error(
-            "Bedrock Realtime: provider stream failed: %s", _redact_string(str(outcome.provider_failure))
-        )
-        raise _as_bedrock_error(outcome.provider_failure) from outcome.provider_failure
+        _raise_provider_failure(websocket.scope, outcome.provider_failure)
 
     async def _bridge(
         self,
@@ -367,17 +368,13 @@ class BedrockRealtime(BaseAWSLLM):
         initial_message: str | None,
     ) -> _BridgeOutcome:
         """Run both forwarding directions until the client leaves or either side fails."""
+        logged: Final[list[OpenAIRealtimeEvents]] = []  # mutable-ok: events forwarded before a failure are still spend
 
-        async def collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]:
-            logged: Final[list[OpenAIRealtimeEvents]] = []  # mutable-ok: partial events are still logged
-            try:
-                async for event in self._forward_bedrock_to_client(
-                    bedrock_stream, websocket, transformation_config, model, logging_obj, session_state
-                ):
-                    logged.append(event)
-            except Exception as e:
-                raise _BedrockForwardingFailed(e, tuple(logged)) from e
-            return tuple(logged)
+        async def collect_logged_events() -> None:
+            async for event in self._forward_bedrock_to_client(
+                bedrock_stream, websocket, transformation_config, model, logging_obj, session_state
+            ):
+                logged.append(event)
 
         client_task: Final = asyncio.create_task(
             self._forward_client_to_bedrock(
@@ -395,18 +392,12 @@ class BedrockRealtime(BaseAWSLLM):
         client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True)
 
         return _BridgeOutcome(
-            logged_events=(
-                bedrock_outcome.logged_events
-                if isinstance(bedrock_outcome, _BedrockForwardingFailed)
-                else bedrock_outcome
-                if isinstance(bedrock_outcome, tuple)
-                else ()
-            ),
+            logged_events=tuple(logged),
             provider_failure=(
                 client_outcome
                 if isinstance(client_outcome, Exception)
-                else bedrock_outcome.cause
-                if isinstance(bedrock_outcome, _BedrockForwardingFailed)
+                else bedrock_outcome
+                if isinstance(bedrock_outcome, Exception)
                 else None
             ),
             client_disconnected=client_disconnected,

tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py

@@ -168,6 +168,34 @@ class BreakingBedrockReceiver(ScriptedBedrockReceiver):
         return await super().receive()
 
 
+class DrainedThenOpenBedrockReceiver(ScriptedBedrockReceiver):
+    """Delivers its payloads, flags `drained`, then stays open like a live Nova Sonic turn"""
+
+    def __init__(self, payloads):
+        super().__init__(payloads)
+        self.drained = asyncio.Event()
+
+    async def receive(self):
+        if not self._payloads:
+            self.drained.set()
+            await asyncio.Event().wait()
+        return await super().receive()
+
+
+class ResetOnAudioInputStream(FakeInputStream):
+    """Accepts session setup, then the provider resets the input side once the first response was delivered"""
+
+    def __init__(self, drained):
+        super().__init__()
+        self._drained = drained
+
+    async def send(self, event):
+        if "audioInput" in json.loads(event.value.bytes_.decode("utf-8")).get("event", {}):
+            await self._drained.wait()
+            raise RuntimeError("bedrock input stream reset")
+        self.sent.append(event)
+
+
 class ScriptedBedrockStream:
     def __init__(self, payloads, receiver_type=ScriptedBedrockReceiver):
         self.input_stream = FakeInputStream()
@@ -560,10 +588,14 @@ class TestBedrockRealtimeProviderFailurePropagation:
         assert [next(iter(event["event"])) for event in replayed][:2] == ["sessionStart", "promptStart"]
         assert websocket.closed
 
-    @pytest.mark.asyncio
-    async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay(
-        self, stub_aws_sdk_client, monkeypatch
-    ):
+    TEXT_TURN = (
+        json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}),
+        json.dumps({"event": {"textOutput": {"content": "Hi"}}}),
+        json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}),
+    )
+
+    @pytest.fixture
+    def spend_dispatch(self, monkeypatch):
         import litellm.llms.bedrock.realtime.handler as handler_module
 
         dispatched = {}
@@ -577,35 +609,69 @@ class TestBedrockRealtimeProviderFailurePropagation:
                 dispatched["coro"] = coro
 
         monkeypatch.setattr(handler_module, "GLOBAL_LOGGING_WORKER", RecordingLoggingWorker())
+        dispatched["logging_obj"] = RecordingLogging()
+        return dispatched
+
+    @pytest.mark.asyncio
+    async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay(
+        self, stub_aws_sdk_client, spend_dispatch
+    ):
         handler = BedrockRealtime()
         websocket = ConnectedClientWS([self.SESSION_UPDATE])
-        stub_aws_sdk_client["streams"] = [
-            ScriptedBedrockStream(
-                [
-                    json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}),
-                    json.dumps({"event": {"textOutput": {"content": "Hi"}}}),
-                    json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}),
-                ],
-                receiver_type=BreakingBedrockReceiver,
-            )
-        ]
+        stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver)
+        stub_aws_sdk_client["streams"] = [stream]
 
         with pytest.raises(BedrockError) as failure:
             await handler.async_realtime(
-                model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=RecordingLogging(), **self.AWS_PARAMS
+                model="amazon.nova-sonic-v1:0",
+                websocket=websocket,
+                logging_obj=spend_dispatch["logging_obj"],
+                **self.AWS_PARAMS,
             )
 
         assert failure.value.status_code == 424
-        await dispatched["coro"]
-        assert [event["type"] for event in dispatched["events"]] == ["response.done"]
+        await spend_dispatch["coro"]
+        assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
         assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client]
+        flushed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent]
+        assert [next(iter(event["event"])) for event in flushed][-2:] == ["promptEnd", "sessionEnd"]
+        assert stream.input_stream.closed
 
         with pytest.raises(BedrockError) as replay:
             await handler.async_realtime(
-                model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=RecordingLogging(), **self.AWS_PARAMS
+                model="amazon.nova-sonic-v1:0",
+                websocket=websocket,
+                logging_obj=spend_dispatch["logging_obj"],
+                **self.AWS_PARAMS,
             )
 
-        assert replay.value.status_code == 409, "a committed session must not be silently restarted on a fallback"
+        assert replay.value.status_code == 400, "a committed session must not be silently restarted on a fallback"
+        assert not litellm._should_retry(replay.value.status_code), "the router must not retry the replay refusal"
+        assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error"
+
+    @pytest.mark.asyncio
+    async def test_input_side_failure_keeps_spend_for_responses_already_delivered(
+        self, stub_aws_sdk_client, spend_dispatch
+    ):
+        receiver = DrainedThenOpenBedrockReceiver(self.TEXT_TURN)
+        stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=lambda _payloads: receiver)
+        stream.input_stream = ResetOnAudioInputStream(receiver.drained)
+        stub_aws_sdk_client["streams"] = [stream]
+        websocket = ConnectedClientWS(
+            [self.SESSION_UPDATE, json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})]
+        )
+
+        with pytest.raises(RuntimeError, match="bedrock input stream reset"):
+            await BedrockRealtime().async_realtime(
+                model="amazon.nova-sonic-v1:0",
+                websocket=websocket,
+                logging_obj=spend_dispatch["logging_obj"],
+                **self.AWS_PARAMS,
+            )
+
+        assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client]
+        await spend_dispatch["coro"]
+        assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
 
     @pytest.mark.asyncio
     async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client):