fix(bedrock): dispatch success handlers for realtime sessions so spend is logged

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

litellm/llms/bedrock/realtime/handler.py

@@ -13,6 +13,8 @@ from pydantic import TypeAdapter
 
 from litellm._logging import _redact_string, verbose_proxy_logger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.types.llms.openai import OpenAIRealtimeEvents
 
 from ..base_aws_llm import BaseAWSLLM
 from ..common_utils import BedrockError
@@ -154,6 +156,7 @@ class BedrockRealtime(BaseAWSLLM):
                 )
             )
 
+            logged_events: Final[list[OpenAIRealtimeEvents]] = []  # mutable-ok: events accumulate across stream loop iterations for spend logging
             bedrock_to_client_task: Final = asyncio.create_task(
                 self._forward_bedrock_to_client(
                     bedrock_stream,
@@ -162,6 +165,7 @@ class BedrockRealtime(BaseAWSLLM):
                     model,
                     logging_obj,
                     session_state,
+                    logged_events,
                 )
             )
 
@@ -172,6 +176,11 @@ class BedrockRealtime(BaseAWSLLM):
                 return_exceptions=True,
             )
 
+            if logged_events:
+                GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
+                    logging_obj.dispatch_success_handlers(logged_events, prefer_async_handlers=True)
+                )
+
         except Exception as e:
             verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e)
             try:
@@ -252,6 +261,7 @@ class BedrockRealtime(BaseAWSLLM):
         model: str,
         logging_obj: LiteLLMLogging,
         session_state: dict,
+        logged_events: "list[OpenAIRealtimeEvents] | None" = None,
     ):
         """Forward messages from Bedrock stream to client WebSocket."""
         try:
@@ -304,6 +314,8 @@ class BedrockRealtime(BaseAWSLLM):
                     # Send transformed messages to client
                     openai_messages = transformed_response.get("response", [])
                     for openai_message in openai_messages:
+                        if logged_events is not None and isinstance(openai_message, dict):
+                            logged_events.append(openai_message)
                         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])

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

@@ -104,6 +104,26 @@ class RealtimeClientWS:
         self.closed = True
 
 
+class ScriptedBedrockReceiver:
+    def __init__(self, payloads):
+        self._payloads = list(payloads)
+
+    async def receive(self):
+        if not self._payloads:
+            return None
+        payload = self._payloads.pop(0)
+        return SimpleNamespace(value=SimpleNamespace(bytes_=payload.encode("utf-8")))
+
+
+class ScriptedBedrockStream:
+    def __init__(self, payloads):
+        self.input_stream = FakeInputStream()
+        self._receiver = ScriptedBedrockReceiver(payloads)
+
+    async def await_output(self):
+        return (None, self._receiver)
+
+
 class ImmediatelyEndingBedrockStream:
     def __init__(self):
         self.input_stream = FakeInputStream()
@@ -271,6 +291,34 @@ class TestBedrockRealtimeHandler:
         assert "sessionEnd" in event_names
         assert stream.input_stream.closed
 
+    @pytest.mark.asyncio
+    async def test_forwarded_events_are_collected_for_spend_logging(self):
+        handler = BedrockRealtime()
+        stream = ScriptedBedrockStream(
+            [
+                json.dumps({"event": {"userSpeechStart": {}}}),
+                json.dumps({"event": {"userSpeechEnd": {}}}),
+            ]
+        )
+        client_ws = RealtimeClientWS()
+        logged_events = []
+
+        await handler._forward_bedrock_to_client(
+            stream,
+            client_ws,
+            BedrockRealtimeConfig(),
+            "amazon.nova-sonic-v1:0",
+            FakeLogging(),
+            {},
+            logged_events,
+        )
+
+        assert [event["type"] for event in logged_events] == [
+            "input_audio_buffer.speech_started",
+            "input_audio_buffer.speech_stopped",
+        ]
+        assert client_ws.closed
+
     @pytest.mark.asyncio
     async def test_bedrock_stream_end_closes_client_websocket(self):
         handler = BedrockRealtime()