From 22e7c7a5338a96009607c612804fe90c0dad9a1a Mon Sep 17 00:00:00 2001
From: Marty Sullivan <marty@martysullivan.com>
Date: Mon, 7 Sep 2026 06:25:07 -0400
Subject: [PATCH] test(vertex-live): type the session helpers in the Live
 passthrough tests

The four session helpers this PR added were unannotated. Typing them needs a
name for the (text, audio) pair each turn carries, so _LiveTurn is a TypedDict
rather than a Mapping union that would leave sum() over a prompt pair
ill-typed, and AUDIO_SESSION is declared with it. The message list reuses
list[dict[str, object]], the annotation the passthrough already uses where it
collects those messages
---
 .../test_vertex_ai_live_passthrough.py        | 31 +++++++++++++++----
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py
index bfff52de673..b4d0a6c06e5 100644
--- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py
+++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py
@@ -6,12 +6,14 @@ including the logging handler, cost tracking, and WebSocket message processing.
 """
 
 import json
+from collections.abc import Sequence
 from datetime import datetime
 from unittest.mock import AsyncMock, Mock, patch, MagicMock
 from typing import Dict, List, Any, Optional
 
 import pytest
 import httpx
+from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 # Add the parent directory to the system path
 
@@ -22,10 +24,16 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
     PassThroughEndpointLogging,
 )
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.types.utils import LlmProviders
+from litellm.types.utils import LlmProviders, Usage
 from litellm.proxy._types import UserAPIKeyAuth
 
 
+class _LiveTurn(TypedDict):
+    prompt: ReadOnly[tuple[int, int]]
+    candidates: ReadOnly[tuple[int, int]]
+    candidate_audio_token_count_missing: NotRequired[ReadOnly[bool]]
+
+
 class TestVertexAILivePassthroughLoggingHandler:
     """Test the Vertex AI Live Passthrough Logging Handler"""
 
@@ -257,7 +265,7 @@ class TestVertexAILivePassthroughLoggingHandler:
     # window, so the prompt side repeats the accumulated audio while the candidates side reports
     # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is
     # the shape Live really emits at the end of a spoken answer.
-    AUDIO_SESSION = (
+    AUDIO_SESSION: tuple[_LiveTurn, ...] = (
         {"prompt": (14, 122), "candidates": (8, 20)},
         {"prompt": (21, 182), "candidates": (5, 50)},
         {"prompt": (24, 203), "candidates": (13, 27)},
@@ -265,7 +273,7 @@ class TestVertexAILivePassthroughLoggingHandler:
     )
 
     @staticmethod
-    def _live_messages(turns):
+    def _live_messages(turns: Sequence[_LiveTurn]) -> list[dict[str, object]]:
         """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits."""
         return [{"type": "session.created", "session": {"id": "s"}}] + [
             {
@@ -292,7 +300,12 @@ class TestVertexAILivePassthroughLoggingHandler:
         ]
 
     @staticmethod
-    def _session_usage(handler, mock_logging_obj, messages, model):
+    def _session_usage(
+        handler: VertexAILivePassthroughLoggingHandler,
+        mock_logging_obj: MagicMock,
+        messages: list[dict[str, object]],
+        model: str,
+    ) -> Usage:
         result = handler.vertex_ai_live_passthrough_handler(
             websocket_messages=messages,
             logging_obj=mock_logging_obj,
@@ -306,7 +319,13 @@ class TestVertexAILivePassthroughLoggingHandler:
         return result["result"].usage
 
     @classmethod
-    def _session_cost(cls, handler, mock_logging_obj, messages, model):
+    def _session_cost(
+        cls,
+        handler: VertexAILivePassthroughLoggingHandler,
+        mock_logging_obj: MagicMock,
+        messages: list[dict[str, object]],
+        model: str,
+    ) -> float:
         from litellm.cost_calculator import completion_cost
         from litellm.types.utils import ModelResponse
 
@@ -321,7 +340,7 @@ class TestVertexAILivePassthroughLoggingHandler:
         )
 
     @classmethod
-    def _expected_session_cost(cls, turns):
+    def _expected_session_cost(cls, turns: Sequence[_LiveTurn]) -> float:
         from litellm.utils import get_model_info
 
         info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai")