diff --git a/litellm/constants.py b/litellm/constants.py index 58a889d279a..e8aafda1797 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -314,6 +314,10 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float( # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 +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 # This balances performance with broad compatibility diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index ca2370303f2..7ab5a14dfc2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,13 +7,20 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json -from collections.abc import AsyncIterator, Mapping -from typing import Final, Protocol +from collections.abc import AsyncIterator, Mapping, MutableMapping +from dataclasses import dataclass +from types import MappingProxyType +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, +) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -28,6 +35,32 @@ from .transformation import BedrockRealtimeConfig _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) + +_BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType( + { + "AccessDeniedException": 403, + "ConflictException": 400, + "InternalServerException": 500, + "ModelErrorException": 424, + "ModelNotReadyException": 429, + "ModelStreamErrorException": 424, + "ModelTimeoutException": 408, + "ResourceNotFoundException": 404, + "ServiceQuotaExceededException": 400, + "ServiceUnavailableException": 503, + "ThrottlingException": 429, + "ValidationException": 400, + } +) + + +def _as_bedrock_error(error: BaseException) -> BaseException: + status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__) + if status_code is None: + return error + return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}") + def _json_dict(value: JsonValue) -> dict[str, JsonValue]: return value if isinstance(value, dict) else {} @@ -51,6 +84,8 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool: class RealtimeClientWebSocket(Protocol): """The client-facing websocket surface the realtime bridge talks to.""" + scope: MutableMapping[str, object] # mutable-ok: the ASGI scope is the per-connection state store + async def receive_text(self) -> str: ... async def send_text(self, data: str) -> None: ... @@ -85,6 +120,81 @@ class BedrockBidirectionalStream(Protocol): async def await_output(self) -> tuple[object, BedrockOutputStream]: ... +@dataclass(frozen=True, slots=True) +class _BridgeOutcome: + logged_events: tuple[OpenAIRealtimeEvents, ...] + provider_failure: BaseException | None + client_disconnected: bool + + +async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]: + if initial_message is not None: + yield initial_message + while True: + try: + yield await client_ws.receive_text() + except Exception as e: # noqa: BLE001 # any receive failure means the client is gone + verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + return + + +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=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)) + except ValueError: + return _EMPTY_JSON_OBJECT + + +async def _ack_session_update( + client_ws: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + logging_obj: LiteLLMLogging | None, + parsed_client_message: Mapping[str, JsonValue], +) -> bool: + """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" + await bedrock_stream.await_output() + client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None) + client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + if logging_obj is None: + return True + requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( + _json_dict(parsed_client_message.get("session")).get("modalities") + ) + try: + await client_ws.send_text( + json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities)) + ) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + return False + return True + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -132,6 +242,8 @@ class BedrockRealtime(BaseAWSLLM): except ImportError: raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime") + pending_session_update: Final = _pending_session_update(websocket.scope) + # Get AWS region if aws_region_name is None: optional_params: Final = { @@ -190,90 +302,105 @@ class BedrockRealtime(BaseAWSLLM): transformation_config: Final = BedrockRealtimeConfig() - try: - # Initialize the bidirectional stream - bedrock_stream: Final = await open_bidirectional_stream() + bedrock_stream: Final = await open_bidirectional_stream() - verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") + verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") + if pending_session_update is None: await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") - # Track state for transformation - session_state: Final[RealtimeResponseTransformInput] = { - "current_output_item_id": None, - "current_response_id": None, - "current_conversation_id": None, - "current_delta_chunks": None, - "current_item_chunks": None, - "current_delta_type": None, - "session_configuration_request": None, - } + # Track state for transformation + session_state: Final[RealtimeResponseTransformInput] = { + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } - # Create tasks for bidirectional forwarding - client_to_bedrock_task: Final = asyncio.create_task( - self._forward_client_to_bedrock( - websocket, - bedrock_stream, - transformation_config, - model, - session_state, - logging_obj, + outcome: Final = await self._bridge( + websocket, + bedrock_stream, + transformation_config, + model, + session_state, + logging_obj, + initial_message=pending_session_update, + ) + + logged_events: Final = ( + *outcome.logged_events, + *( + leftover_event + for leftover_event in transformation_config.leftover_usage_done_events() + if _should_log_event(leftover_event) + ), + ) + if logged_events: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + logging_obj.dispatch_success_handlers( + list(logged_events), # mutable-ok: realtime spend logging requires a list result + prefer_async_handlers=True, ) ) - async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]: - return tuple( - [ - event - async for event in self._forward_bedrock_to_client( - bedrock_stream, - websocket, - transformation_config, - model, - logging_obj, - session_state, - ) - ] - ) - - bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events()) - - # Wait for both tasks to complete - await asyncio.gather( - client_to_bedrock_task, - bedrock_to_client_task, - return_exceptions=True, + if outcome.provider_failure is None: + return + if outcome.client_disconnected: + verbose_proxy_logger.debug( + "Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure ) + return + _raise_provider_failure(websocket.scope, outcome.provider_failure) - forwarded_logged_events: Final = ( - bedrock_to_client_task.result() - if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None - else () - ) - logged_events: Final = ( - *forwarded_logged_events, - *( - leftover_event - for leftover_event in transformation_config.leftover_usage_done_events() - if _should_log_event(leftover_event) - ), - ) - if logged_events: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - logging_obj.dispatch_success_handlers( - list(logged_events), # mutable-ok: realtime spend logging requires a list result - prefer_async_handlers=True, - ) - ) + async def _bridge( + self, + websocket: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + session_state: RealtimeResponseTransformInput, + logging_obj: LiteLLMLogging, + 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 - except Exception as e: - verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e) - try: - await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}")) - except Exception: - pass - raise + 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( + websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message + ) + ) + bedrock_task: Final = asyncio.create_task(collect_logged_events()) + + await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION) + client_disconnected: Final = ( + client_task.done() and not client_task.cancelled() and client_task.exception() is None + ) + client_task.cancel() + bedrock_task.cancel() + client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True) + + return _BridgeOutcome( + logged_events=tuple(logged), + provider_failure=( + client_outcome + if isinstance(client_outcome, Exception) + else bedrock_outcome + if isinstance(bedrock_outcome, Exception) + else None + ), + client_disconnected=client_disconnected, + ) async def _forward_client_to_bedrock( self, @@ -283,8 +410,12 @@ class BedrockRealtime(BaseAWSLLM): model: str, session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging | None = None, - ): - """Forward messages from client WebSocket to Bedrock stream.""" + initial_message: str | None = None, + ) -> None: + """Forward messages from client WebSocket to Bedrock stream. + + Returns once the client is gone; provider failures (input stream or readiness) propagate to the caller. + """ from aws_sdk_bedrock_runtime.models import ( BidirectionalInputPayloadPart, InvokeModelWithBidirectionalStreamInputChunk, @@ -299,41 +430,28 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) try: - while True: - # Receive message from client - message = await client_ws.receive_text() + async for message in _client_messages(client_ws, initial_message): verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200]) + parsed_client_message = _parse_client_message(message) + is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" + if is_session_update: + client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = ( + message # rebind-ok: scope outlives the attempt + ) - # Transform OpenAI format to Bedrock format transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, session_configuration_request=session_state.get("session_configuration_request"), ) - - # Send transformed messages to Bedrock for bedrock_message in transformed_messages: await send_to_bedrock(bedrock_message) - if logging_obj is not None: - client_message_type: str | None = None - requested_modalities: list[str] | None = None - with contextlib.suppress(Exception): - parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) - client_message_type = _json_str(parsed_client_message.get("type")) - if client_message_type == "session.update": - requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( - _json_dict(parsed_client_message.get("session")).get("modalities") - ) - if client_message_type == "session.update": - await client_ws.send_text( - json.dumps( - transformation_config.session_updated_event(model, logging_obj, requested_modalities) - ) - ) - - except Exception as e: - verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + if is_session_update and not await _ack_session_update( + client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message + ): + break + finally: for close_message in transformation_config.session_close_messages(): with contextlib.suppress(Exception): await send_to_bedrock(close_message) @@ -349,68 +467,71 @@ class BedrockRealtime(BaseAWSLLM): logging_obj: LiteLLMLogging, session_state: RealtimeResponseTransformInput, ) -> AsyncIterator[OpenAIRealtimeEvents]: - """Forward messages from Bedrock to the client, yielding the ones to record for spend logging.""" - try: - while True: - # Receive from Bedrock - output = await bedrock_stream.await_output() - result = await output[1].receive() + """Forward messages from Bedrock to the client, yielding the ones to record for spend logging. - if result is None: - verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended") - break + Provider failures propagate to the caller; the client websocket is only closed on a normal stream end. + """ - payload_bytes = result.value.bytes_ if result.value else None - if payload_bytes: - bedrock_response = payload_bytes.decode("utf-8") - verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200]) - - # Transform Bedrock format to OpenAI format - realtime_response_transform_input: RealtimeResponseTransformInput = { - "current_output_item_id": session_state.get("current_output_item_id"), - "current_response_id": session_state.get("current_response_id"), - "current_conversation_id": session_state.get("current_conversation_id"), - "current_delta_chunks": session_state.get("current_delta_chunks"), - "current_item_chunks": session_state.get("current_item_chunks"), - "current_delta_type": session_state.get("current_delta_type"), - "session_configuration_request": session_state.get("session_configuration_request"), - } - - transformed_response = transformation_config.transform_realtime_response( - message=bedrock_response, - model=model, - logging_obj=logging_obj, - realtime_response_transform_input=realtime_response_transform_input, - ) - - # Update session state - session_state.update( - { - "current_output_item_id": transformed_response.get("current_output_item_id"), - "current_response_id": transformed_response.get("current_response_id"), - "current_conversation_id": transformed_response.get("current_conversation_id"), - "current_delta_chunks": transformed_response.get("current_delta_chunks"), - "current_item_chunks": transformed_response.get("current_item_chunks"), - "current_delta_type": transformed_response.get("current_delta_type"), - "session_configuration_request": transformed_response.get("session_configuration_request"), - } - ) - - # Send transformed messages to client - response_value = transformed_response["response"] - openai_messages = response_value if isinstance(response_value, list) else (response_value,) - for openai_message in openai_messages: - message_json = json.dumps(openai_message) - await client_ws.send_text(message_json) - verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) - if _should_log_event(openai_message): - yield openai_message - - except Exception as e: - verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True) - finally: - # Close the client WebSocket + async def send_to_client(message_json: str) -> bool: try: - await client_ws.close() - except Exception: - pass + await client_ws.send_text(message_json) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True) + return False + verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) + return True + + output: Final = await bedrock_stream.await_output() + while True: + result = await output[1].receive() + + if result is None: + verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended") + with contextlib.suppress(Exception): + await client_ws.close() + return + + payload_bytes = result.value.bytes_ if result.value else None + if payload_bytes: + bedrock_response = payload_bytes.decode("utf-8") + verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200]) + + # Transform Bedrock format to OpenAI format + realtime_response_transform_input: RealtimeResponseTransformInput = { + "current_output_item_id": session_state.get("current_output_item_id"), + "current_response_id": session_state.get("current_response_id"), + "current_conversation_id": session_state.get("current_conversation_id"), + "current_delta_chunks": session_state.get("current_delta_chunks"), + "current_item_chunks": session_state.get("current_item_chunks"), + "current_delta_type": session_state.get("current_delta_type"), + "session_configuration_request": session_state.get("session_configuration_request"), + } + + transformed_response = transformation_config.transform_realtime_response( + message=bedrock_response, + model=model, + logging_obj=logging_obj, + realtime_response_transform_input=realtime_response_transform_input, + ) + + # Update session state + session_state.update( + { + "current_output_item_id": transformed_response.get("current_output_item_id"), + "current_response_id": transformed_response.get("current_response_id"), + "current_conversation_id": transformed_response.get("current_conversation_id"), + "current_delta_chunks": transformed_response.get("current_delta_chunks"), + "current_item_chunks": transformed_response.get("current_item_chunks"), + "current_delta_type": transformed_response.get("current_delta_type"), + "session_configuration_request": transformed_response.get("session_configuration_request"), + } + ) + + # Send transformed messages to client + response_value = transformed_response["response"] + openai_messages = response_value if isinstance(response_value, list) else (response_value,) + for openai_message in openai_messages: + if not await send_to_client(json.dumps(openai_message)): + return + if _should_log_event(openai_message): + yield openai_message 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 0ea5b7ad4a1..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 @@ -1,3 +1,4 @@ +import asyncio import json import sys import types @@ -51,6 +52,27 @@ class FakeBedrockStream: def __init__(self, input_stream=None): self.input_stream = input_stream if input_stream is not None else FakeInputStream() + async def await_output(self): + return (None, EndedBedrockReceiver()) + + +class ServiceUnavailableException(Exception): + """Named like the modeled AWS SDK error so the handler maps it to HTTP 503""" + + +class ModelStreamErrorException(Exception): + """Named like the modeled AWS SDK error so the handler maps it to HTTP 424""" + + +class UnavailableBedrockStream: + """Lazy duplex stream whose HTTP response only fails once the output is awaited""" + + def __init__(self): + self.input_stream = FakeInputStream() + + async def await_output(self): + raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable") + class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): @@ -61,6 +83,7 @@ class DisconnectingClientWS: def __init__(self, messages): self._messages = list(messages) self.sent_to_client = [] + self.scope = {} async def receive_text(self): if self._messages: @@ -93,6 +116,7 @@ class RealtimeClientWS: def __init__(self): self.closed = False self.sent_to_client = [] + self.scope = {} async def receive_text(self): raise RuntimeError("client disconnected") @@ -104,6 +128,25 @@ class RealtimeClientWS: self.closed = True +class ConnectedClientWS(RealtimeClientWS): + """Client that sends its scripted messages and then stays connected until the server closes it""" + + def __init__(self, messages): + super().__init__() + self._messages = list(messages) + self._closed_event = asyncio.Event() + + async def receive_text(self): + if self._messages: + return self._messages.pop(0) + await self._closed_event.wait() + raise RuntimeError("client disconnected") + + async def close(self, code=None, reason=None): + self.closed = True + self._closed_event.set() + + class ScriptedBedrockReceiver: def __init__(self, payloads): self._payloads = list(payloads) @@ -115,10 +158,48 @@ class ScriptedBedrockReceiver: return SimpleNamespace(value=SimpleNamespace(bytes_=payload.encode("utf-8"))) -class ScriptedBedrockStream: +class BreakingBedrockReceiver(ScriptedBedrockReceiver): + """Delivers its payloads, then the provider stream breaks instead of ending normally""" + + async def receive(self): + if not self._payloads: + await asyncio.sleep(0) + raise ModelStreamErrorException("Nova Sonic stream broke") + 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() - self._receiver = ScriptedBedrockReceiver(payloads) + self._receiver = receiver_type(payloads) async def await_output(self): return (None, self._receiver) @@ -163,6 +244,11 @@ def stub_aws_sdk_client(monkeypatch): async def invoke_model_with_bidirectional_stream(self, operation_input): captured["operation_input"] = operation_input + if captured.get("streams"): + stream = captured["streams"].pop(0) + if isinstance(stream, Exception): + raise stream + return stream return ScriptedBedrockStream(captured.get("scripted_payloads", [])) package = types.ModuleType("aws_sdk_bedrock_runtime") @@ -263,7 +349,8 @@ class TestBedrockRealtimeHandler: [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] ) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + with pytest.raises(RuntimeError, match="bedrock send failed"): + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) assert stream.input_stream.closed @@ -464,6 +551,154 @@ class TestBedrockRealtimeSessionLifecycle: assert client_ws.sent_to_client == [] +class TestBedrockRealtimeProviderFailurePropagation: + """Deferred Nova Sonic failures must escape async_realtime so the router can fall back / cool down (LIT-6484)""" + + SESSION_UPDATE = json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}}) + AWS_PARAMS = {"aws_region_name": "us-east-1", "aws_access_key_id": "k", "aws_secret_access_key": "s"} + + @pytest.mark.asyncio + async def test_readiness_failure_escapes_and_fallback_replays_session_update(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.SESSION_UPDATE]) + healthy_stream = ScriptedBedrockStream([]) + eager_failure = ServiceUnavailableException("fault injected before the stream was returned") + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), eager_failure, healthy_stream] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert failure.value.status_code == 503 + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] + assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close" + + with pytest.raises(ServiceUnavailableException): + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created", "session.updated"] + replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent] + assert [next(iter(event["event"])) for event in replayed][:2] == ["sessionStart", "promptStart"] + assert websocket.closed + + 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 = {} + + class RecordingLogging(FakeLogging): + async def dispatch_success_handlers(self, result=None, prefer_async_handlers=False, **kwargs): + dispatched["events"] = result + + class RecordingLoggingWorker: + def ensure_initialized_and_enqueue(self, coro): + 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]) + 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=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + + assert failure.value.status_code == 424 + 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=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + + 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): + stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert stream.input_stream.closed + + @pytest.mark.asyncio + async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models): + handler = BedrockRealtime() + stream = UnavailableBedrockStream() + client_ws = DisconnectingClientWS([self.SESSION_UPDATE]) + + with pytest.raises(ServiceUnavailableException): + await handler._forward_client_to_bedrock( + client_ws, stream, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", {}, FakeLogging() + ) + + assert client_ws.sent_to_client == [] + assert stream.input_stream.closed + + class TestBedrockRealtimeAwsAuth: """AWS auth params passed via litellm_params must reach the Smithy client config (LIT-3923 regression)"""