diff --git a/litellm/constants.py b/litellm/constants.py index 7d86080086a..a1a737153fa 100644 --- a/litellm/constants.py +++ b/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 diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 126a7ae5c3e..130dec20a4b 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/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, diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 085ffeb5129..2eadaee9e1a 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/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):