Add vertex live api passthrough with cost tracking

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -6,12 +6,14 @@ Provider-specific Pass-Through Endpoints
 Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
 """
 
+import json
 import os
 from typing import Optional, cast
 
 import httpx
-from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
 from fastapi.responses import StreamingResponse
+from starlette.websockets import WebSocketState
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -19,7 +21,9 @@ from litellm.constants import BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.proxy._types import *
 from litellm.proxy.auth.route_checks import RouteChecks
-from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.auth.user_api_key_auth import (
+    user_api_key_auth,
+)
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     get_form_data,
@@ -28,6 +32,8 @@ from litellm.proxy.common_utils.http_parsing_utils import (
 from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     create_pass_through_route,
+    create_websocket_passthrough_route,
+    websocket_passthrough_request,
 )
 from litellm.proxy.utils import is_known_model
 from litellm.secret_managers.main import get_secret_str
@@ -143,7 +149,7 @@ async def llm_passthrough_factory_proxy_route(
             _request_body = await request.json()
         else:
             _request_body = await get_form_data(request)
-        
+
         if _request_body.get("stream"):
             is_streaming_request = True
 
@@ -1248,3 +1254,183 @@ class BaseOpenAIPassThroughHandler:
             )
 
         return joined_path_str
+
+
+async def vertex_ai_live_websocket_passthrough(
+    websocket: WebSocket,
+    model: Optional[str] = None,
+    vertex_project: Optional[str] = None,
+    vertex_location: Optional[str] = None,
+    user_api_key_dict: Optional[UserAPIKeyAuth] = None,
+):
+    """
+    Vertex AI Live API WebSocket Pass-through Function
+
+    This function provides WebSocket passthrough functionality for Vertex AI Live API,
+    allowing real-time communication with Google's Live API service.
+
+    Note: This function should be registered in proxy_server.py using:
+    app.websocket("/vertex_ai/live")(vertex_ai_live_websocket_passthrough)
+    """
+    from litellm.proxy.proxy_server import proxy_logging_obj
+
+    _ = user_api_key_dict  # passthrough route already authenticated; avoid lint warnings
+
+    await websocket.accept()
+
+    incoming_headers = dict(websocket.headers)
+    vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
+        project_id=vertex_project,
+        location=vertex_location,
+    )
+
+    if vertex_credentials_config is None:
+        # Attempt to load defaults from environment/config if not already initialised
+        passthrough_endpoint_router.set_default_vertex_config()
+        vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
+            project_id=vertex_project,
+            location=vertex_location,
+        )
+
+    resolved_project = vertex_project
+    resolved_location = vertex_location
+    credentials_value: Optional[str] = None
+
+    if vertex_credentials_config is not None:
+        resolved_project = resolved_project or vertex_credentials_config.vertex_project
+        resolved_location = (
+            resolved_location or vertex_credentials_config.vertex_location
+        )
+        # Ensure resolved_location is a string
+        if isinstance(resolved_location, dict):
+            resolved_location = str(resolved_location)
+        credentials_value = vertex_credentials_config.vertex_credentials
+
+    try:
+        resolved_location = resolved_location or (
+            vertex_llm_base.get_default_vertex_location()
+        )
+        if model:
+            resolved_location = vertex_llm_base.get_vertex_region(
+                vertex_region=resolved_location,
+                model=model,
+            )
+
+        (
+            access_token,
+            resolved_project,
+        ) = await vertex_llm_base._ensure_access_token_async(
+            credentials=credentials_value,
+            project_id=resolved_project,
+            custom_llm_provider="vertex_ai_beta",
+        )
+    except Exception as e:
+        verbose_proxy_logger.exception(
+            "Failed to prepare Vertex AI credentials for live passthrough"
+        )
+        # Log the authentication failure using proxy_logging_obj
+        if proxy_logging_obj and user_api_key_dict:
+            await proxy_logging_obj.post_call_failure_hook(
+                user_api_key_dict=user_api_key_dict,
+                original_exception=e,
+                request_data={},
+            )
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close(code=1011, reason="Vertex AI authentication failed")
+        return
+
+    host_location = resolved_location or vertex_llm_base.get_default_vertex_location()
+    host = (
+        "aiplatform.googleapis.com"
+        if host_location == "global"
+        else f"{host_location}-aiplatform.googleapis.com"
+    )
+    service_url = (
+        f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
+    )
+
+    upstream_headers = {
+        "Authorization": f"Bearer {access_token}",
+        "Content-Type": "application/json",
+    }
+    if resolved_project:
+        upstream_headers["x-goog-user-project"] = resolved_project
+
+    # Forward any custom x-goog-* headers provided by the caller if we haven't overridden them
+    for header_name, header_value in incoming_headers.items():
+        lower_header = header_name.lower()
+        if lower_header.startswith("x-goog-") and header_name not in upstream_headers:
+            upstream_headers[header_name] = header_value
+
+    # Use the new WebSocket passthrough pattern
+    if user_api_key_dict is None:
+        raise ValueError("user_api_key_dict is required for WebSocket passthrough")
+    
+    return await websocket_passthrough_request(
+        websocket=websocket,
+        target=service_url,
+        custom_headers=upstream_headers,
+        user_api_key_dict=user_api_key_dict,
+        forward_headers=False,
+        endpoint="/vertex_ai/live",
+        accept_websocket=False,
+    )
+
+
+def create_vertex_ai_live_websocket_endpoint():
+    """
+    Create a Vertex AI Live WebSocket endpoint using the new passthrough pattern.
+
+    This demonstrates how to use the create_websocket_passthrough_route function
+    for a provider-specific WebSocket endpoint.
+    """
+    # This would be used like:
+    # endpoint_func = create_vertex_ai_live_websocket_endpoint()
+    # app.websocket("/vertex_ai/live")(endpoint_func)
+
+    # For now, we'll keep the existing implementation since it has
+    # provider-specific logic for Vertex AI credentials and headers
+    return vertex_ai_live_websocket_passthrough
+
+
+def create_generic_websocket_passthrough_endpoint(
+    provider: str,
+    target_url: str,
+    custom_headers: Optional[dict] = None,
+    forward_headers: bool = False,
+    cost_per_request: Optional[float] = None,
+):
+    """
+    Create a generic WebSocket passthrough endpoint for any provider.
+
+    This demonstrates the new WebSocket passthrough pattern that's similar to
+    the HTTP create_pass_through_route function.
+
+    Args:
+        provider: The provider name (e.g., "anthropic", "cohere")
+        target_url: The target WebSocket URL
+        custom_headers: Custom headers to include
+        forward_headers: Whether to forward incoming headers
+
+    Returns:
+        A WebSocket endpoint function that can be registered with app.websocket()
+
+    Example usage:
+        # Create a WebSocket endpoint for Anthropic
+        anthropic_ws_func = create_generic_websocket_passthrough_endpoint(
+            provider="anthropic",
+            target_url="wss://api.anthropic.com/v1/ws",
+            custom_headers={"x-api-key": "your-api-key"},
+            forward_headers=True
+        )
+
+        # Register it in proxy_server.py
+        app.websocket("/anthropic/ws")(anthropic_ws_func)
+    """
+    return create_websocket_passthrough_route(
+        endpoint=f"/{provider}/ws",
+        target=target_url,
+        custom_headers=custom_headers,
+        _forward_headers=forward_headers,
+        cost_per_request=cost_per_request,
+    )

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py

@@ -0,0 +1,394 @@
+"""
+Vertex AI Live API WebSocket Passthrough Logging Handler
+
+Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough endpoints.
+Supports different modalities: text, audio, video, and web search.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import (
+    BasePassthroughLoggingHandler,
+)
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
+    PassThroughEndpointLoggingTypedDict,
+)
+from litellm.types.utils import LlmProviders, ModelResponse, Usage
+from litellm.utils import get_model_info
+
+
+class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
+    """
+    Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough.
+
+    Supports:
+    - Text tokens (input/output)
+    - Audio tokens (input/output)
+    - Video tokens (input/output)
+    - Web search requests
+    - Tool use tokens
+    """
+
+    def _build_complete_streaming_response(self, *args, **kwargs):
+        """Not applicable for WebSocket passthrough."""
+        return None
+
+    def get_provider_config(self, model: str):
+        """Return Vertex AI provider configuration."""
+        from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+            VertexGeminiConfig,
+        )
+
+        return VertexGeminiConfig()
+
+    @property
+    def llm_provider_name(self) -> LlmProviders:
+        """Return the LLM provider name."""
+        return LlmProviders.VERTEX_AI
+
+    @staticmethod
+    def _extract_usage_metadata_from_websocket_messages(
+        websocket_messages: List[Dict],
+    ) -> Optional[Dict]:
+        """
+        Extract and aggregate usage metadata from a list of WebSocket messages.
+
+        Args:
+            websocket_messages: List of WebSocket messages from the Live API
+
+        Returns:
+            Dictionary containing aggregated usage metadata, or None if not found
+        """
+        all_usage_metadata = []
+
+        # Collect all usage metadata messages
+        for message in websocket_messages:
+            if isinstance(message, dict) and "usageMetadata" in message:
+                all_usage_metadata.append(message["usageMetadata"])
+
+        if not all_usage_metadata:
+            return None
+
+        # If only one usage metadata, return it as-is
+        if len(all_usage_metadata) == 1:
+            return all_usage_metadata[0]
+
+        # Aggregate multiple usage metadata messages
+        aggregated: Dict[str, Any] = {
+            "promptTokenCount": 0,
+            "candidatesTokenCount": 0,
+            "totalTokenCount": 0,
+            "promptTokensDetails": [],
+            "candidatesTokensDetails": [],
+        }
+
+        # Aggregate token counts
+        for usage in all_usage_metadata:
+            aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0)
+            aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0)
+            aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0)
+
+        # Aggregate token details by modality
+        modality_totals = {}
+
+        for usage in all_usage_metadata:
+            # Process prompt tokens details
+            for detail in usage.get("promptTokensDetails", []):
+                modality = detail.get("modality", "TEXT")
+                token_count = detail.get("tokenCount", 0)
+
+                if modality not in modality_totals:
+                    modality_totals[modality] = {"prompt": 0, "candidate": 0}
+                modality_totals[modality]["prompt"] += token_count
+
+            # Process candidate tokens details
+            for detail in usage.get("candidatesTokensDetails", []):
+                modality = detail.get("modality", "TEXT")
+                token_count = detail.get("tokenCount", 0)
+
+                if modality not in modality_totals:
+                    modality_totals[modality] = {"prompt": 0, "candidate": 0}
+                modality_totals[modality]["candidate"] += token_count
+
+        # Convert aggregated modality totals back to details format
+        for modality, totals in modality_totals.items():
+            if totals["prompt"] > 0:
+                aggregated["promptTokensDetails"].append(
+                    {"modality": modality, "tokenCount": totals["prompt"]}
+                )
+            if totals["candidate"] > 0:
+                aggregated["candidatesTokensDetails"].append(
+                    {"modality": modality, "tokenCount": totals["candidate"]}
+                )
+
+        # Add any additional fields from the first usage metadata
+        first_usage = all_usage_metadata[0]
+        for key, value in first_usage.items():
+            if key not in aggregated:
+                aggregated[key] = value
+
+        return aggregated
+
+    @staticmethod
+    def _calculate_live_api_cost(
+        model: str,
+        usage_metadata: Dict,
+        custom_llm_provider: str = "vertex_ai",
+    ) -> float:
+        """
+        Calculate cost for Vertex AI Live API based on usage metadata.
+
+        Args:
+            model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09")
+            usage_metadata: Usage metadata from the Live API response
+            custom_llm_provider: The LLM provider (default: "vertex_ai")
+
+        Returns:
+            Total cost in USD
+        """
+        try:
+            # Get model pricing information
+            model_info = get_model_info(
+                model=model, custom_llm_provider=custom_llm_provider
+            )
+
+            verbose_proxy_logger.debug(
+                f"Vertex AI Live API model info for '{model}': {model_info}"
+            )
+
+            # Check if pricing info is available
+            if not model_info or not model_info.get("input_cost_per_token"):
+                verbose_proxy_logger.error(
+                    f"No pricing info found for {model} in local model pricing database"
+                )
+                return 0.0
+
+            total_cost = 0.0
+
+            # Extract token counts from usage metadata
+            prompt_token_count = usage_metadata.get("promptTokenCount", 0)
+            candidates_token_count = usage_metadata.get("candidatesTokenCount", 0)
+
+            # Calculate base text token costs
+            input_cost_per_token = model_info.get("input_cost_per_token", 0.0)
+            output_cost_per_token = model_info.get("output_cost_per_token", 0.0)
+
+            total_cost += prompt_token_count * input_cost_per_token
+            total_cost += candidates_token_count * output_cost_per_token
+
+            # Handle modality-specific costs if present
+            prompt_tokens_details = usage_metadata.get("promptTokensDetails", [])
+            candidates_tokens_details = usage_metadata.get(
+                "candidatesTokensDetails", []
+            )
+
+            # Process prompt tokens by modality
+            for detail in prompt_tokens_details:
+                modality = detail.get("modality", "TEXT")
+                token_count = detail.get("tokenCount", 0)
+
+                if modality == "AUDIO":
+                    audio_cost_per_token = model_info.get(
+                        "input_cost_per_audio_token", 0.0
+                    )
+                    total_cost += token_count * audio_cost_per_token
+                elif modality == "VIDEO":
+                    # Video tokens are typically per second, but we'll treat as per token for now
+                    video_cost_per_token = model_info.get(
+                        "input_cost_per_video_per_second", 0.0
+                    )
+                    total_cost += token_count * video_cost_per_token
+                # TEXT tokens are already handled above
+
+            # Process candidate tokens by modality
+            for detail in candidates_tokens_details:
+                modality = detail.get("modality", "TEXT")
+                token_count = detail.get("tokenCount", 0)
+
+                if modality == "AUDIO":
+                    audio_cost_per_token = model_info.get(
+                        "output_cost_per_audio_token", 0.0
+                    )
+                    total_cost += token_count * audio_cost_per_token
+                elif modality == "VIDEO":
+                    # Video tokens are typically per second, but we'll treat as per token for now
+                    video_cost_per_token = model_info.get(
+                        "output_cost_per_video_per_second", 0.0
+                    )
+                    total_cost += token_count * video_cost_per_token
+                # TEXT tokens are already handled above
+
+            # Handle web search costs if present
+            tool_use_prompt_token_count = usage_metadata.get(
+                "toolUsePromptTokenCount", 0
+            )
+            if tool_use_prompt_token_count > 0:
+                # Web search typically has a fixed cost per request
+                web_search_cost = model_info.get("web_search_cost_per_request", 0.0)
+                if isinstance(web_search_cost, (int, float)) and web_search_cost > 0:
+                    total_cost += web_search_cost
+                else:
+                    # Fallback to token-based pricing for tool use
+                    total_cost += tool_use_prompt_token_count * input_cost_per_token
+
+            verbose_proxy_logger.debug(
+                f"Vertex AI Live API cost calculation - Model: {model}, "
+                f"Prompt tokens: {prompt_token_count}, "
+                f"Candidate tokens: {candidates_token_count}, "
+                f"Total cost: ${total_cost:.6f}"
+            )
+
+            return total_cost
+
+        except Exception as e:
+            verbose_proxy_logger.error(
+                f"Error calculating Vertex AI Live API cost: {e}"
+            )
+            return 0.0
+
+    @staticmethod
+    def _create_usage_object_from_metadata(
+        usage_metadata: Dict,
+        model: str,
+    ) -> Usage:
+        """
+        Create a LiteLLM Usage object from Live API usage metadata.
+
+        Args:
+            usage_metadata: Usage metadata from the Live API response
+            model: The model name
+
+        Returns:
+            LiteLLM Usage object
+        """
+        prompt_tokens = usage_metadata.get("promptTokenCount", 0)
+        completion_tokens = usage_metadata.get("candidatesTokenCount", 0)
+        total_tokens = usage_metadata.get("totalTokenCount", 0)
+
+        # Create modality-specific token details if available
+        prompt_tokens_details = usage_metadata.get("promptTokensDetails", [])
+        candidates_tokens_details = usage_metadata.get("candidatesTokensDetails", [])
+
+        # Extract text tokens from details
+        text_prompt_tokens = 0
+        text_completion_tokens = 0
+
+        for detail in prompt_tokens_details:
+            if detail.get("modality") == "TEXT":
+                text_prompt_tokens = detail.get("tokenCount", 0)
+                break
+
+        for detail in candidates_tokens_details:
+            if detail.get("modality") == "TEXT":
+                text_completion_tokens = detail.get("tokenCount", 0)
+                break
+
+        # If no text tokens found in details, use total counts
+        if text_prompt_tokens == 0:
+            text_prompt_tokens = prompt_tokens
+        if text_completion_tokens == 0:
+            text_completion_tokens = completion_tokens
+
+        return Usage(
+            prompt_tokens=text_prompt_tokens,
+            completion_tokens=text_completion_tokens,
+            total_tokens=total_tokens,
+        )
+
+    def vertex_ai_live_passthrough_handler(
+        self,
+        websocket_messages: List[Dict],
+        logging_obj,
+        url_route: str,
+        start_time: datetime,
+        end_time: datetime,
+        request_body: dict,
+        **kwargs,
+    ) -> PassThroughEndpointLoggingTypedDict:
+        """
+        Handle cost tracking and logging for Vertex AI Live API WebSocket passthrough.
+
+        Args:
+            websocket_messages: List of WebSocket messages from the Live API
+            logging_obj: LiteLLM logging object
+            url_route: The URL route that was called
+            start_time: Request start time
+            end_time: Request end time
+            request_body: The original request body
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Dictionary containing the result and kwargs for logging
+        """
+        try:
+            # Extract model from request body or kwargs
+            model = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09")
+            custom_llm_provider = kwargs.get("custom_llm_provider", "vertex_ai")
+            verbose_proxy_logger.debug(
+                f"Vertex AI Live API model: {model}, custom_llm_provider: {custom_llm_provider}"
+            )
+
+            # Extract usage metadata from WebSocket messages
+            usage_metadata = self._extract_usage_metadata_from_websocket_messages(
+                websocket_messages
+            )
+
+            if not usage_metadata:
+                verbose_proxy_logger.warning(
+                    "No usage metadata found in Vertex AI Live API WebSocket messages"
+                )
+                return {
+                    "result": None,
+                    "kwargs": kwargs,
+                }
+
+            # Calculate cost using Live API specific pricing
+            response_cost = self._calculate_live_api_cost(
+                model=model,
+                usage_metadata=usage_metadata,
+                custom_llm_provider=custom_llm_provider,
+            )
+
+            # Create Usage object for standard LiteLLM logging
+            usage = self._create_usage_object_from_metadata(
+                usage_metadata=usage_metadata,
+                model=model,
+            )
+
+            # Create a mock ModelResponse for standard logging
+            litellm_model_response = ModelResponse(
+                id=f"vertex-ai-live-{start_time.timestamp()}",
+                object="chat.completion",
+                created=int(start_time.timestamp()),
+                model=model,
+                usage=usage,
+                choices=[],
+            )
+
+            # Update kwargs with cost information
+            kwargs["response_cost"] = response_cost
+            kwargs["model"] = model
+            kwargs["custom_llm_provider"] = custom_llm_provider
+
+            verbose_proxy_logger.debug(
+                f"Vertex AI Live API passthrough cost tracking - "
+                f"Model: {model}, Cost: ${response_cost:.6f}, "
+                f"Prompt tokens: {usage.prompt_tokens}, "
+                f"Completion tokens: {usage.completion_tokens}"
+            )
+
+            return {
+                "result": litellm_model_response,
+                "kwargs": kwargs,
+            }
+
+        except Exception as e:
+            verbose_proxy_logger.error(
+                f"Error in Vertex AI Live API passthrough handler: {e}"
+            )
+            return {
+                "result": None,
+                "kwargs": kwargs,
+            }

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -6,7 +6,7 @@ import traceback
 import uuid
 from base64 import b64encode
 from datetime import datetime
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import urlencode, urlparse
 
 import httpx
@@ -18,10 +18,18 @@ from fastapi import (
     Request,
     Response,
     UploadFile,
+    WebSocket,
     status,
 )
 from fastapi.responses import StreamingResponse
 from starlette.datastructures import UploadFile as StarletteUploadFile
+from starlette.websockets import WebSocketState
+from websockets.asyncio.client import connect
+from websockets.exceptions import (
+    ConnectionClosedError,
+    ConnectionClosedOK,
+    InvalidStatus,
+)
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -476,7 +484,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
                 user_api_key_request_route=user_api_key_dict.request_route,
                 user_api_key_spend=user_api_key_dict.spend,
                 user_api_key_max_budget=user_api_key_dict.max_budget,
-                user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None,
+                user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at.isoformat()
+                if user_api_key_dict.budget_reset_at
+                else None,
             )
         )
 
@@ -984,6 +994,506 @@ def create_pass_through_route(
     return endpoint_func
 
 
+def create_websocket_passthrough_route(
+    endpoint: str,
+    target: str,
+    custom_headers: Optional[dict] = None,
+    _forward_headers: Optional[bool] = False,
+    dependencies: Optional[List] = None,
+    cost_per_request: Optional[float] = None,
+):
+    """
+    Create a WebSocket passthrough route function.
+
+    Args:
+        endpoint: The endpoint path (for logging purposes)
+        target: The target WebSocket URL (e.g., "wss://api.example.com/ws")
+        custom_headers: Custom headers to include in the WebSocket connection
+        _forward_headers: Whether to forward incoming headers
+        dependencies: FastAPI dependencies to inject
+
+    Returns:
+        A WebSocket passthrough function that can be registered with app.websocket()
+    """
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket
+
+    async def websocket_endpoint_func(
+        websocket: WebSocket,
+        user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket),
+        **kwargs,  # For additional query parameters
+    ):
+        """
+        WebSocket passthrough endpoint function.
+
+        This function handles the WebSocket connection by:
+        1. Accepting the incoming WebSocket connection
+        2. Establishing a connection to the target WebSocket
+        3. Forwarding messages bidirectionally
+        4. Handling connection cleanup
+        """
+        return await websocket_passthrough_request(
+            websocket=websocket,
+            target=target,
+            custom_headers=custom_headers or {},
+            user_api_key_dict=user_api_key_dict,
+            forward_headers=_forward_headers,
+            endpoint=endpoint,
+            cost_per_request=cost_per_request,
+            accept_websocket=True,  # Generic usage should accept the WebSocket
+        )
+
+    return websocket_endpoint_func
+
+
+async def websocket_passthrough_request(
+    websocket: WebSocket,
+    target: str,
+    custom_headers: dict,
+    user_api_key_dict: UserAPIKeyAuth,
+    forward_headers: Optional[bool] = False,
+    endpoint: Optional[str] = None,
+    cost_per_request: Optional[float] = None,
+    accept_websocket: bool = True,
+):
+    """
+    WebSocket passthrough request handler.
+
+    Args:
+        websocket: The incoming WebSocket connection
+        target: The target WebSocket URL
+        custom_headers: Custom headers to include in the connection
+        user_api_key_dict: The user API key dictionary
+        forward_headers: Whether to forward incoming headers
+        endpoint: The endpoint path (for logging purposes)
+        cost_per_request: Optional field - cost per request to the target endpoint
+    """
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.proxy.proxy_server import proxy_logging_obj
+    from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+        PassthroughStandardLoggingPayload,
+    )
+
+    # Initialize tracking variables
+    start_time = datetime.now()
+    websocket_messages: list[dict[str, Any]] = []
+    litellm_call_id = str(uuid.uuid4())
+
+    verbose_proxy_logger.info(
+        f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}"
+    )
+
+    # Only accept the WebSocket if requested (for generic usage)
+    if accept_websocket:
+        await websocket.accept()
+        verbose_proxy_logger.debug(
+            f"WebSocket passthrough ({endpoint}): WebSocket connection accepted"
+        )
+
+    # Prepare headers for the upstream connection
+    upstream_headers = custom_headers.copy()
+
+    if forward_headers:
+        # Forward relevant headers from the incoming request
+        incoming_headers = dict(websocket.headers)
+        for header_name, header_value in incoming_headers.items():
+            # Only forward certain headers to avoid conflicts
+            if header_name.lower() in [
+                "authorization",
+                "x-api-key",
+                "x-goog-user-project",
+            ]:
+                upstream_headers[header_name] = header_value
+
+    # Initialize logging object similar to HTTP passthrough
+    logging_obj = Logging(
+        model="unknown",
+        messages=[{"role": "user", "content": "WebSocket connection"}],
+        stream=True,  # WebSockets are inherently streaming
+        call_type="pass_through_endpoint",
+        start_time=start_time,
+        litellm_call_id=litellm_call_id,
+        function_id="websocket_passthrough",
+    )
+
+    # Create passthrough logging payload
+    passthrough_logging_payload = PassthroughStandardLoggingPayload(
+        url=target,
+        request_body={},  # WebSocket doesn't have a traditional request body
+        request_method="WEBSOCKET",
+        cost_per_request=cost_per_request,
+    )
+
+    # Create a dummy request object for WebSocket connections to maintain compatibility
+    # with the existing _init_kwargs_for_pass_through_endpoint function
+    class DummyRequest:
+        def __init__(self, url: str, method: str = "WEBSOCKET", headers: dict = None):
+            self.url = url
+            self.method = method
+            self.headers = headers or {}
+
+        def __str__(self):
+            return f"DummyRequest(url={self.url}, method={self.method})"
+
+    dummy_request = DummyRequest(
+        url=target,
+        method="WEBSOCKET",
+        headers=dict(websocket.headers) if hasattr(websocket, "headers") else {},
+    )
+
+    # Initialize kwargs for logging using the same pattern as HTTP passthrough
+    kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
+        user_api_key_dict=user_api_key_dict,
+        _parsed_body={},  # WebSocket doesn't have a traditional request body
+        passthrough_logging_payload=passthrough_logging_payload,
+        litellm_call_id=litellm_call_id,
+        request=dummy_request,
+        logging_obj=logging_obj,
+    )
+
+    # Update logging environment variables
+    logging_obj.update_environment_variables(
+        model="unknown",
+        user="unknown",
+        optional_params={},
+        litellm_params=dict(kwargs.get("litellm_params", {})),
+        call_type="pass_through_endpoint",
+    )
+    logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
+
+    # Pre-call logging
+    logging_obj.pre_call(
+        input=[{"role": "user", "content": "WebSocket connection"}],
+        api_key="",
+        additional_args={
+            "complete_input_dict": {},
+            "api_base": target,
+            "headers": upstream_headers,
+        },
+    )
+
+    ### CALL HOOKS ### - modify incoming data / reject request before calling the model
+    websocket_data: dict[str, Any] = {}
+    websocket_data = await proxy_logging_obj.pre_call_hook(
+        user_api_key_dict=user_api_key_dict,
+        data=websocket_data,
+        call_type="pass_through_endpoint",
+    )
+
+    try:
+        verbose_proxy_logger.debug(
+            f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}"
+        )
+        async with connect(
+            target,
+            additional_headers=upstream_headers,
+        ) as upstream_ws:
+            verbose_proxy_logger.info(
+                f"WebSocket passthrough ({endpoint}): Upstream connection established successfully"
+            )
+
+            async def forward_client_to_upstream() -> None:
+                """Forward messages from client to upstream WebSocket"""
+                try:
+                    while True:
+                        message = await websocket.receive()
+                        message_type = message.get("type")
+                        if message_type == "websocket.disconnect":
+                            await upstream_ws.close()
+                            break
+
+                        text_data = message.get("text")
+                        bytes_data = message.get("bytes")
+
+                        if text_data is not None:
+                            # Try to extract model from client setup message for Vertex AI Live
+                            if endpoint and "/vertex_ai/live" in endpoint:
+                                verbose_proxy_logger.debug(
+                                    f"WebSocket passthrough ({endpoint}): Processing client message for model extraction"
+                                )
+                                try:
+                                    client_message = json.loads(text_data)
+                                    if (
+                                        isinstance(client_message, dict)
+                                        and "setup" in client_message
+                                    ):
+                                        setup_data = client_message["setup"]
+                                        verbose_proxy_logger.debug(
+                                            f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}"
+                                        )
+                                        if (
+                                            isinstance(setup_data, dict)
+                                            and "model" in setup_data
+                                        ):
+                                            extracted_model = (
+                                                _extract_model_from_vertex_ai_setup(
+                                                    setup_data
+                                                )
+                                            )
+                                            if extracted_model:
+                                                kwargs["model"] = extracted_model
+                                                kwargs[
+                                                    "custom_llm_provider"
+                                                ] = "vertex_ai-language-models"
+                                                # Update logging object with correct model
+                                                logging_obj.model = extracted_model
+                                                logging_obj.model_call_details[
+                                                    "model"
+                                                ] = extracted_model
+                                                logging_obj.model_call_details[
+                                                    "custom_llm_provider"
+                                                ] = "vertex_ai"
+                                                verbose_proxy_logger.info(
+                                                    f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message"
+                                                )
+                                            else:
+                                                verbose_proxy_logger.warning(
+                                                    f"WebSocket passthrough ({endpoint}): Failed to extract model from client setup data: {setup_data}"
+                                                )
+                                        else:
+                                            verbose_proxy_logger.debug(
+                                                f"WebSocket passthrough ({endpoint}): Setup data does not contain model field: {setup_data}"
+                                            )
+                                    else:
+                                        verbose_proxy_logger.debug(
+                                            f"WebSocket passthrough ({endpoint}): Client message does not contain setup data"
+                                        )
+                                except (json.JSONDecodeError, KeyError, TypeError) as e:
+                                    verbose_proxy_logger.debug(
+                                        f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}"
+                                    )
+                                    pass  # Not a JSON message or doesn't contain setup data
+
+                            await upstream_ws.send(text_data)
+                        elif bytes_data is not None:
+                            await upstream_ws.send(bytes_data)
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    verbose_proxy_logger.exception(
+                        f"WebSocket passthrough ({endpoint}): error forwarding client message"
+                    )
+                    await upstream_ws.close()
+
+            async def forward_upstream_to_client() -> None:
+                """Forward messages from upstream to client WebSocket"""
+                try:
+                    # Wait for the first response from upstream
+                    raw_response = await upstream_ws.recv(decode=False)
+                    setup_response = json.loads(raw_response.decode("ascii"))
+                    verbose_proxy_logger.debug(f"Setup response: {setup_response}")
+
+                    # Extract model and provider from setup response for Vertex AI Live
+                    if endpoint and "/vertex_ai/live" in endpoint:
+                        verbose_proxy_logger.debug(
+                            f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction"
+                        )
+                        extracted_model = _extract_model_from_vertex_ai_setup(
+                            setup_response
+                        )
+                        if extracted_model:
+                            kwargs["model"] = extracted_model
+                            kwargs["custom_llm_provider"] = "vertex_ai_language_models"
+                            # Update logging object with correct model
+                            logging_obj.model = extracted_model
+                            logging_obj.model_call_details["model"] = extracted_model
+                            logging_obj.model_call_details[
+                                "custom_llm_provider"
+                            ] = "vertex_ai_language_models"
+                            verbose_proxy_logger.info(
+                                f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response"
+                            )
+                        else:
+                            verbose_proxy_logger.warning(
+                                f"WebSocket passthrough ({endpoint}): Failed to extract model from server setup response: {setup_response}"
+                            )
+                    else:
+                        verbose_proxy_logger.debug(
+                            f"WebSocket passthrough ({endpoint}): Not a Vertex AI Live endpoint, skipping model extraction"
+                        )
+
+                    # Send the setup response to the client
+                    await websocket.send_text(json.dumps(setup_response))
+
+                    # Now continuously forward messages from upstream to client
+                    async for upstream_message in upstream_ws:
+                        if isinstance(upstream_message, bytes):
+                            await websocket.send_bytes(upstream_message)
+                            # Parse and collect for cost tracking
+                            try:
+                                message_data = json.loads(upstream_message.decode())
+                                websocket_messages.append(message_data)
+                            except (json.JSONDecodeError, UnicodeDecodeError):
+                                pass
+                        else:
+                            await websocket.send_text(upstream_message)
+                            # Parse and collect for cost tracking
+                            try:
+                                message_data = json.loads(upstream_message)
+                                websocket_messages.append(message_data)
+                            except json.JSONDecodeError:
+                                pass
+
+                except (ConnectionClosedOK, ConnectionClosedError) as e:
+                    verbose_proxy_logger.debug(
+                        f"Upstream WebSocket connection closed: {e}"
+                    )
+                    pass
+                except asyncio.CancelledError:
+                    verbose_proxy_logger.debug(
+                        "asyncio.CancelledError in forward_upstream_to_client"
+                    )
+                    raise
+                except Exception as e:
+                    verbose_proxy_logger.debug(
+                        f"Exception in forward_upstream_to_client: {e}"
+                    )
+                    verbose_proxy_logger.exception(
+                        f"WebSocket passthrough ({endpoint}): error forwarding upstream message"
+                    )
+                    raise
+
+            # Create tasks for bidirectional message forwarding
+            tasks = [
+                asyncio.create_task(forward_client_to_upstream()),
+                asyncio.create_task(forward_upstream_to_client()),
+            ]
+
+            done, pending = await asyncio.wait(
+                tasks, return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # Cancel remaining tasks
+            for task in pending:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+            # Check for exceptions in completed tasks
+            for task in done:
+                exception = task.exception()
+                if exception is not None:
+                    raise exception
+
+            end_time = datetime.now()
+
+            # Update passthrough logging payload with response data
+            passthrough_logging_payload["response_body"] = websocket_messages
+            passthrough_logging_payload["end_time"] = end_time
+
+            # Remove logging_obj from kwargs to avoid duplicate keyword argument
+            success_kwargs = kwargs.copy()
+            success_kwargs.pop("logging_obj", None)
+
+            # # Add user authentication context for database logging
+            # if user_api_key_dict:
+            #     success_kwargs.setdefault('litellm_params', {})
+            #     success_kwargs['litellm_params'].update({
+            #         'proxy_server_request': {
+            #             'body': {
+            #                 'user': user_api_key_dict.user_id,
+            #                 'team_id': user_api_key_dict.team_id,
+            #                 'end_user_id': user_api_key_dict.end_user_id,
+            #             }
+            #         }
+            #     })
+            #     # Also add the user_api_key for direct access
+            #     success_kwargs['user_api_key'] = user_api_key_dict.api_key
+
+            # Create a dummy httpx.Response for WebSocket connections
+            class MockWebSocketResponse:
+                def __init__(self, target_url: str):
+                    self.status_code = 200
+                    self.text = "WebSocket connection successful"
+                    self.headers: dict[str, str] = {}
+                    self.request = MockWebSocketRequest(target_url)
+
+            class MockWebSocketRequest:
+                def __init__(self, target_url: str):
+                    self.method = "WEBSOCKET"
+                    self.url = target_url
+
+            mock_response = MockWebSocketResponse(target)
+
+            # Use the same success handler as HTTP passthrough endpoints
+            asyncio.create_task(
+                pass_through_endpoint_logging.pass_through_async_success_handler(
+                    httpx_response=mock_response,  # Use mock response for WebSocket
+                    response_body=websocket_messages,
+                    url_route=endpoint,
+                    result="websocket_connection_successful",
+                    start_time=start_time,
+                    end_time=end_time,
+                    logging_obj=logging_obj,
+                    cache_hit=False,
+                    request_body={},
+                    **success_kwargs,
+                )
+            )
+
+            # Call the proxy logging success hook
+            if proxy_logging_obj:
+                await proxy_logging_obj.post_call_success_hook(
+                    data={},
+                    user_api_key_dict=user_api_key_dict,
+                    response={"status": "websocket_connection_successful"},
+                )
+
+    except InvalidStatus as exc:
+        verbose_proxy_logger.exception(
+            f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection"
+        )
+
+        # Prepare request payload for logging
+        request_payload = {}
+        if kwargs:
+            for key, value in kwargs.items():
+                request_payload[key] = value
+
+        # Log the connection failure using the same pattern as HTTP
+        await proxy_logging_obj.post_call_failure_hook(
+            user_api_key_dict=user_api_key_dict,
+            original_exception=exc,
+            request_data=request_payload,
+            traceback_str=traceback.format_exc(
+                limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
+            ),
+        )
+
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close(
+                code=exc.status_code if hasattr(exc, "status_code") else 1011,
+                reason="Upstream connection rejected",
+            )
+    except Exception as e:
+        verbose_proxy_logger.exception(
+            f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket"
+        )
+
+        # Prepare request payload for logging
+        request_payload = {}
+        if kwargs:
+            for key, value in kwargs.items():
+                request_payload[key] = value
+
+        # Log the unexpected error using the same pattern as HTTP
+        await proxy_logging_obj.post_call_failure_hook(
+            user_api_key_dict=user_api_key_dict,
+            original_exception=e,
+            request_data=request_payload,
+            traceback_str=traceback.format_exc(
+                limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
+            ),
+        )
+
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close(code=1011, reason="WebSocket passthrough error")
+    finally:
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close()
+
+
 def _is_streaming_response(response: httpx.Response) -> bool:
     _content_type = response.headers.get("content-type")
     if _content_type is not None and "text/event-stream" in _content_type:
@@ -991,6 +1501,38 @@ def _is_streaming_response(response: httpx.Response) -> bool:
     return False
 
 
+def _extract_model_from_vertex_ai_setup(setup_response: dict) -> Optional[str]:
+    """
+    Extract the model name from Vertex AI Live setup response.
+
+    The setup response can contain a model field in two formats:
+    1. Direct: {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}
+    2. Nested: {"setup": {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}}
+
+    We extract just the model name: "gemini-2.0-flash-live-preview-04-09"
+    """
+    try:
+        # Handle both direct model field and nested setup.model field
+        model_path = None
+        if isinstance(setup_response, dict):
+            if "model" in setup_response:
+                model_path = setup_response["model"]
+            elif (
+                "setup" in setup_response
+                and isinstance(setup_response["setup"], dict)
+                and "model" in setup_response["setup"]
+            ):
+                model_path = setup_response["setup"]["model"]
+
+        if isinstance(model_path, str) and "/models/" in model_path:
+            # Extract the model name after the last "/models/"
+            model_name = model_path.split("/models/")[-1]
+            return model_name
+    except Exception as e:
+        verbose_proxy_logger.debug(f"Error extracting model from setup response: {e}")
+    return None
+
+
 class InitPassThroughEndpointHelpers:
     @staticmethod
     def add_exact_path_route(

litellm/proxy/pass_through_endpoints/success_handler.py

@@ -51,6 +51,9 @@ class PassThroughEndpointLogging:
         # Langfuse
         self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"]
 
+        # Vertex AI Live API WebSocket
+        self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"]
+
     async def _handle_logging(
         self,
         logging_obj: LiteLLMLoggingObj,
@@ -162,7 +165,9 @@ class PassThroughEndpointLogging:
                 cohere_passthrough_logging_handler_result["result"]
             )
             kwargs = cohere_passthrough_logging_handler_result["kwargs"]
-        elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint(url_route):
+        elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint(
+            url_route
+        ):
             from .llm_provider_handlers.openai_passthrough_logging_handler import (
                 OpenAIPassthroughLoggingHandler,
             )
@@ -185,6 +190,29 @@ class PassThroughEndpointLogging:
                 openai_passthrough_logging_handler_result["result"]
             )
             kwargs = openai_passthrough_logging_handler_result["kwargs"]
+        elif self.is_vertex_ai_live_route(url_route):
+            from .llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import (
+                VertexAILivePassthroughLoggingHandler,
+            )
+            vertex_ai_live_handler = VertexAILivePassthroughLoggingHandler()
+
+            # For WebSocket responses, response_body should be a list of messages
+            websocket_messages: list[dict[str, Any]] = response_body if isinstance(response_body, list) else []
+
+            vertex_ai_live_handler_result = (
+                vertex_ai_live_handler.vertex_ai_live_passthrough_handler(
+                    websocket_messages=websocket_messages,
+                    logging_obj=logging_obj,
+                    url_route=url_route,
+                    start_time=start_time,
+                    end_time=end_time,
+                    request_body=request_body,
+                    **kwargs,
+                )
+            )
+
+            standard_logging_response_object = vertex_ai_live_handler_result["result"]
+            kwargs = vertex_ai_live_handler_result["kwargs"]
         return_dict[
             "standard_logging_response_object"
         ] = standard_logging_response_object
@@ -309,6 +337,15 @@ class PassThroughEndpointLogging:
                 return True
         return False
 
+    def is_vertex_ai_live_route(self, url_route: str):
+        """Check if the URL route is a Vertex AI Live API WebSocket route."""
+        if not url_route:
+            return False
+        for route in self.TRACKED_VERTEX_AI_LIVE_ROUTES:
+            if route in url_route:
+                return True
+        return False
+
     def is_openai_route(self, url_route: str):
         """Check if the URL route is an OpenAI API route."""
         if not url_route:
@@ -324,11 +361,13 @@ class PassThroughEndpointLogging:
         from .llm_provider_handlers.openai_passthrough_logging_handler import (
             OpenAIPassthroughLoggingHandler,
         )
-        
+
         return (
-            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) or
-            OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) or
-            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
+            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route)
+            or OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(
+                url_route
+            )
+            or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
         )
 
     def _set_cost_per_request(

litellm/proxy/proxy_server.py

@@ -35,13 +35,13 @@ from litellm.constants import (
     LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-from litellm.utils import load_credentials_from_list
 from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
     TextCompletionResponse,
     TokenCountResponse,
 )
+from litellm.utils import load_credentials_from_list
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
@@ -308,6 +308,9 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
     router as llm_passthrough_router,
 )
+from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+    vertex_ai_live_websocket_passthrough,
+)
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     initialize_pass_through_endpoints,
 )
@@ -461,9 +464,9 @@ except ImportError:
 server_root_path = os.getenv("SERVER_ROOT_PATH", "")
 _license_check = LicenseCheck()
 premium_user: bool = _license_check.is_premium()
-premium_user_data: Optional["EnterpriseLicenseData"] = (
-    _license_check.airgapped_license_data
-)
+premium_user_data: Optional[
+    "EnterpriseLicenseData"
+] = _license_check.airgapped_license_data
 global_max_parallel_request_retries_env: Optional[str] = os.getenv(
     "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
 )
@@ -959,9 +962,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
     dual_cache=user_api_key_cache
 )
 litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
-redis_usage_cache: Optional[RedisCache] = (
-    None  # redis cache used for tracking spend, tpm/rpm limits
-)
+redis_usage_cache: Optional[
+    RedisCache
+] = None  # redis cache used for tracking spend, tpm/rpm limits
 user_custom_auth = None
 user_custom_key_generate = None
 user_custom_sso = None
@@ -1292,9 +1295,9 @@ async def update_cache(  # noqa: PLR0915
         _id = "team_id:{}".format(team_id)
         try:
             # Fetch the existing cost for the given user
-            existing_spend_obj: Optional[LiteLLM_TeamTable] = (
-                await user_api_key_cache.async_get_cache(key=_id)
-            )
+            existing_spend_obj: Optional[
+                LiteLLM_TeamTable
+            ] = await user_api_key_cache.async_get_cache(key=_id)
             if existing_spend_obj is None:
                 # do nothing if team not in api key cache
                 return
@@ -3107,10 +3110,10 @@ class ProxyConfig:
         )
 
         try:
-            guardrails_in_db: List[Guardrail] = (
-                await GuardrailRegistry.get_all_guardrails_from_db(
-                    prisma_client=prisma_client
-                )
+            guardrails_in_db: List[
+                Guardrail
+            ] = await GuardrailRegistry.get_all_guardrails_from_db(
+                prisma_client=prisma_client
             )
             verbose_proxy_logger.debug(
                 "guardrails from the DB %s", str(guardrails_in_db)
@@ -3340,9 +3343,9 @@ async def initialize(  # noqa: PLR0915
         user_api_base = api_base
         dynamic_config[user_model]["api_base"] = api_base
     if api_version:
-        os.environ["AZURE_API_VERSION"] = (
-            api_version  # set this for azure - litellm can read this from the env
-        )
+        os.environ[
+            "AZURE_API_VERSION"
+        ] = api_version  # set this for azure - litellm can read this from the env
     if max_tokens:  # model-specific param
         dynamic_config[user_model]["max_tokens"] = max_tokens
     if temperature:  # model-specific param
@@ -4919,174 +4922,19 @@ async def vertex_ai_live_passthrough_endpoint(
     ),
     user_api_key_dict=Depends(user_api_key_auth_websocket),
 ):
-    from starlette.websockets import WebSocketState
-    from websockets.asyncio.client import connect
-    from websockets.exceptions import (
-        ConnectionClosedError,
-        ConnectionClosedOK,
-        InvalidStatusCode,
+    """
+    Vertex AI Live API WebSocket Pass-through Endpoint
+
+    This endpoint delegates to the WebSocket function defined in llm_passthrough_endpoints.py
+    """
+    return await vertex_ai_live_websocket_passthrough(
+        websocket=websocket,
+        model=model,
+        vertex_project=vertex_project,
+        vertex_location=vertex_location,
+        user_api_key_dict=user_api_key_dict,
     )
 
-    _ = user_api_key_dict  # passthrough route already authenticated; avoid lint warnings
-
-    await websocket.accept()
-
-    incoming_headers = dict(websocket.headers)
-    vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
-        project_id=vertex_project,
-        location=vertex_location,
-    )
-
-    if vertex_credentials_config is None:
-        # Attempt to load defaults from environment/config if not already initialised
-        passthrough_endpoint_router.set_default_vertex_config()
-        vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
-            project_id=vertex_project,
-            location=vertex_location,
-        )
-
-    resolved_project = vertex_project
-    resolved_location = vertex_location
-    credentials_value: Optional[str] = None
-
-    if vertex_credentials_config is not None:
-        resolved_project = resolved_project or vertex_credentials_config.vertex_project
-        resolved_location = resolved_location or vertex_credentials_config.vertex_location
-        credentials_value = vertex_credentials_config.vertex_credentials
-
-    try:
-        resolved_location = resolved_location or (
-            vertex_live_passthrough_vertex_base.get_default_vertex_location()
-        )
-        if model:
-            resolved_location = vertex_live_passthrough_vertex_base.get_vertex_region(
-                vertex_region=resolved_location,
-                model=model,
-            )
-
-        access_token, resolved_project = await vertex_live_passthrough_vertex_base._ensure_access_token_async(
-            credentials=credentials_value,
-            project_id=resolved_project,
-            custom_llm_provider="vertex_ai_beta",
-        )
-    except Exception:
-        verbose_proxy_logger.exception(
-            "Failed to prepare Vertex AI credentials for live passthrough"
-        )
-        if websocket.client_state != WebSocketState.DISCONNECTED:
-            await websocket.close(code=1011, reason="Vertex AI authentication failed")
-        return
-
-    host_location = resolved_location or vertex_live_passthrough_vertex_base.get_default_vertex_location()
-    host = (
-        "aiplatform.googleapis.com"
-        if host_location == "global"
-        else f"{host_location}-aiplatform.googleapis.com"
-    )
-    service_url = (
-        f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
-    )
-
-    upstream_headers = {
-        "Authorization": f"Bearer {access_token}",
-        "Content-Type": "application/json",
-    }
-    if resolved_project:
-        upstream_headers["x-goog-user-project"] = resolved_project
-
-    # Forward any custom x-goog-* headers provided by the caller if we haven't overridden them
-    for header_name, header_value in incoming_headers.items():
-        lower_header = header_name.lower()
-        if lower_header.startswith("x-goog-") and header_name not in upstream_headers:
-            upstream_headers[header_name] = header_value
-
-    try:
-        async with connect(
-            service_url,
-            additional_headers=upstream_headers,
-        ) as upstream_ws:
-
-            async def forward_client_to_vertex() -> None:
-                try:
-                    while True:
-                        message = await websocket.receive()
-                        message_type = message.get("type")
-                        if message_type == "websocket.disconnect":
-                            await upstream_ws.close()
-                            break
-
-                        text_data = message.get("text")
-                        bytes_data = message.get("bytes")
-
-                        if text_data is not None:
-                            await upstream_ws.send(text_data)
-                        elif bytes_data is not None:
-                            await upstream_ws.send(bytes_data)
-                except asyncio.CancelledError:
-                    raise
-                except Exception:
-                    verbose_proxy_logger.exception(
-                        "Vertex AI live passthrough: error forwarding client message"
-                    )
-                    await upstream_ws.close()
-
-            async def forward_vertex_to_client() -> None:
-                try:
-                    async for upstream_message in upstream_ws:
-                        if isinstance(upstream_message, bytes):
-                            await websocket.send_bytes(upstream_message)
-                        else:
-                            await websocket.send_text(upstream_message)
-                except (ConnectionClosedOK, ConnectionClosedError):
-                    pass
-                except asyncio.CancelledError:
-                    raise
-                except Exception:
-                    verbose_proxy_logger.exception(
-                        "Vertex AI live passthrough: error forwarding upstream message"
-                    )
-                    raise
-
-            tasks = [
-                asyncio.create_task(forward_client_to_vertex()),
-                asyncio.create_task(forward_vertex_to_client()),
-            ]
-
-            done, pending = await asyncio.wait(
-                tasks, return_when=asyncio.FIRST_COMPLETED
-            )
-
-            for task in pending:
-                task.cancel()
-                try:
-                    await task
-                except asyncio.CancelledError:
-                    pass
-
-            for task in done:
-                exception = task.exception()
-                if exception is not None:
-                    raise exception
-
-    except InvalidStatusCode as exc:
-        verbose_proxy_logger.exception(
-            "Vertex AI live passthrough: upstream rejected WebSocket connection"
-        )
-        if websocket.client_state != WebSocketState.DISCONNECTED:
-            await websocket.close(
-                code=exc.status_code if hasattr(exc, "status_code") else 1011,
-                reason="Upstream connection rejected",
-            )
-    except Exception:
-        verbose_proxy_logger.exception(
-            "Vertex AI live passthrough: unexpected error while proxying WebSocket"
-        )
-        if websocket.client_state != WebSocketState.DISCONNECTED:
-            await websocket.close(code=1011, reason="Vertex AI passthrough error")
-    finally:
-        if websocket.client_state != WebSocketState.DISCONNECTED:
-            await websocket.close()
-
 
 ######################################################################
 
@@ -6305,12 +6153,10 @@ def _add_team_models_to_all_models(
     team_models: Dict[str, Set[str]] = {}
 
     for team_object in team_db_objects_typed:
-
         if (
             len(team_object.models) == 0  # empty list = all model access
             or SpecialModelNames.all_proxy_models.value in team_object.models
         ):
-
             model_list = llm_router.get_model_list()
             if model_list is not None:
                 for model in model_list:
@@ -6461,7 +6307,6 @@ async def get_all_team_and_direct_access_models(
     for _model in all_models:
         model_id = _model.get("model_info", {}).get("id", None)
         if model_id is not None and model_id in direct_access_models:
-
             _model["model_info"]["direct_access"] = True
 
     ## FILTER OUT MODELS THAT ARE NOT IN DIRECT_ACCESS_MODELS OR ACCESS_VIA_TEAM_IDS - only show user models they can call
@@ -8821,9 +8666,9 @@ async def get_config_list(
                             hasattr(sub_field_info, "description")
                             and sub_field_info.description is not None
                         ):
-                            nested_fields[idx].field_description = (
-                                sub_field_info.description
-                            )
+                            nested_fields[
+                                idx
+                            ].field_description = sub_field_info.description
                         idx += 1
 
                     _stored_in_db = None

tests/pass_through_tests/test_vertex_ai_live_integration.py

@@ -0,0 +1,502 @@
+"""
+Integration tests for Vertex AI Live API WebSocket passthrough
+
+This module tests the end-to-end functionality of the Vertex AI Live API
+WebSocket passthrough feature, including WebSocket connections, message
+processing, and cost tracking.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import tempfile
+from datetime import datetime
+from typing import Dict, List, Any
+
+import pytest
+import httpx
+from fastapi.testclient import TestClient
+from unittest.mock import patch, MagicMock, AsyncMock
+
+# Add the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../.."))
+
+from litellm.proxy.proxy_server import app
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import (
+    VertexAILivePassthroughLoggingHandler,
+)
+
+
+class TestVertexAILivePassthroughIntegration:
+    """Integration tests for Vertex AI Live passthrough"""
+
+    @pytest.fixture
+    def client(self):
+        """Create a test client"""
+        return TestClient(app)
+
+    @pytest.fixture
+    def mock_vertex_credentials(self):
+        """Mock Vertex AI credentials"""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+            credentials = {
+                "type": "service_account",
+                "project_id": "test-project",
+                "private_key_id": "test-key-id",
+                "private_key": "-----BEGIN PRIVATE KEY-----\nMOCK_PRIVATE_KEY\n-----END PRIVATE KEY-----\n",
+                "client_email": "test@test-project.iam.gserviceaccount.com",
+                "client_id": "test-client-id",
+                "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+                "token_uri": "https://oauth2.googleapis.com/token",
+            }
+            json.dump(credentials, f)
+            temp_file = f.name
+        
+        # Set environment variable
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = temp_file
+        
+        yield temp_file
+        
+        # Cleanup
+        os.unlink(temp_file)
+        if "GOOGLE_APPLICATION_CREDENTIALS" in os.environ:
+            del os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+
+    @pytest.fixture
+    def sample_websocket_messages(self):
+        """Sample WebSocket messages for testing"""
+        return [
+            {
+                "type": "session.created",
+                "session": {"id": "test-session-123"},
+                "timestamp": "2024-01-01T00:00:00Z"
+            },
+            {
+                "type": "response.create",
+                "event_id": "event-123",
+                "response": {
+                    "text": "Hello! How can I help you today?",
+                    "usage": {
+                        "promptTokenCount": 15,
+                        "candidatesTokenCount": 20,
+                        "totalTokenCount": 35,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 15}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 20}
+                        ]
+                    }
+                }
+            },
+            {
+                "type": "response.done",
+                "event_id": "event-123",
+                "response": {
+                    "usage": {
+                        "promptTokenCount": 5,
+                        "candidatesTokenCount": 8,
+                        "totalTokenCount": 13,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 5}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 8}
+                        ]
+                    }
+                }
+            }
+        ]
+
+    def test_vertex_ai_live_route_registration(self, client):
+        """Test that the Vertex AI Live route is properly registered"""
+        # Check if the route exists in the app
+        routes = [route.path for route in app.routes]
+        assert "/vertex_ai/live" in routes
+
+    @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request')
+    @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router')
+    def test_vertex_ai_live_websocket_connection(
+        self, 
+        mock_router, 
+        mock_websocket_passthrough,
+        client,
+        mock_vertex_credentials
+    ):
+        """Test WebSocket connection to Vertex AI Live endpoint"""
+        # Mock the router methods
+        mock_router.get_vertex_credentials.return_value = MagicMock(
+            vertex_project="test-project",
+            vertex_location="us-central1",
+            vertex_credentials="test-credentials"
+        )
+        mock_router.set_default_vertex_config.return_value = None
+        
+        # Mock the WebSocket passthrough request
+        mock_websocket_passthrough.return_value = AsyncMock()
+        
+        # Test WebSocket connection
+        with client.websocket_connect("/vertex_ai/live") as websocket:
+            # Send a test message
+            test_message = {
+                "type": "session.create",
+                "session": {
+                    "modalities": ["TEXT"],
+                    "instructions": "You are a helpful assistant."
+                }
+            }
+            websocket.send_text(json.dumps(test_message))
+            
+            # The connection should be established without errors
+            assert websocket is not None
+
+    def test_vertex_ai_live_logging_handler_integration(self, sample_websocket_messages):
+        """Test the logging handler with real WebSocket messages"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Test usage metadata extraction
+        usage_metadata = handler._extract_usage_metadata_from_websocket_messages(
+            sample_websocket_messages
+        )
+        
+        assert usage_metadata is not None
+        assert usage_metadata["promptTokenCount"] == 20  # 15 + 5
+        assert usage_metadata["candidatesTokenCount"] == 28  # 20 + 8
+        assert usage_metadata["totalTokenCount"] == 48  # 35 + 13
+
+    @patch('litellm.utils.get_model_info')
+    def test_cost_calculation_integration(self, mock_get_model_info, sample_websocket_messages):
+        """Test cost calculation with real usage data"""
+        # Mock model info with realistic pricing
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001,
+            "output_cost_per_token": 0.000002,
+            "input_cost_per_audio_per_second": 0.0001,
+            "output_cost_per_audio_per_second": 0.0002
+        }
+        
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Extract usage metadata
+        usage_metadata = handler._extract_usage_metadata_from_websocket_messages(
+            sample_websocket_messages
+        )
+        
+        # Calculate cost
+        cost = handler._calculate_cost("gemini-1.5-pro", usage_metadata)
+        
+        # Verify cost calculation
+        expected_cost = (20 * 0.000001) + (28 * 0.000002)
+        assert cost == expected_cost
+        assert cost > 0
+
+    def test_multimodal_usage_tracking(self):
+        """Test usage tracking with multiple modalities"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Messages with mixed modalities
+        multimodal_messages = [
+            {
+                "type": "response.create",
+                "response": {
+                    "usage": {
+                        "promptTokenCount": 30,
+                        "candidatesTokenCount": 25,
+                        "totalTokenCount": 55,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 20},
+                            {"modality": "AUDIO", "tokenCount": 10}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 15},
+                            {"modality": "AUDIO", "tokenCount": 10}
+                        ]
+                    }
+                }
+            }
+        ]
+        
+        usage_metadata = handler._extract_usage_metadata_from_websocket_messages(
+            multimodal_messages
+        )
+        
+        assert usage_metadata is not None
+        assert usage_metadata["promptTokenCount"] == 30
+        assert usage_metadata["candidatesTokenCount"] == 25
+        assert len(usage_metadata["promptTokensDetails"]) == 2
+        assert len(usage_metadata["candidatesTokensDetails"]) == 2
+        
+        # Check modality details
+        text_prompt = next(d for d in usage_metadata["promptTokensDetails"] if d["modality"] == "TEXT")
+        audio_prompt = next(d for d in usage_metadata["promptTokensDetails"] if d["modality"] == "AUDIO")
+        assert text_prompt["tokenCount"] == 20
+        assert audio_prompt["tokenCount"] == 10
+
+    def test_web_search_usage_tracking(self):
+        """Test usage tracking with web search (tool use)"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Messages with web search usage
+        web_search_messages = [
+            {
+                "type": "response.create",
+                "response": {
+                    "usage": {
+                        "promptTokenCount": 50,
+                        "candidatesTokenCount": 30,
+                        "totalTokenCount": 80,
+                        "toolUsePromptTokenCount": 10,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 50}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 30}
+                        ]
+                    }
+                }
+            }
+        ]
+        
+        usage_metadata = handler._extract_usage_metadata_from_websocket_messages(
+            web_search_messages
+        )
+        
+        assert usage_metadata is not None
+        assert usage_metadata["promptTokenCount"] == 50
+        assert usage_metadata["candidatesTokenCount"] == 30
+        assert usage_metadata["toolUsePromptTokenCount"] == 10
+
+    @patch('litellm.utils.get_model_info')
+    def test_web_search_cost_calculation(self, mock_get_model_info):
+        """Test cost calculation with web search"""
+        # Mock model info with web search pricing
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001,
+            "output_cost_per_token": 0.000002,
+            "web_search_cost_per_request": 0.01
+        }
+        
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        usage_metadata = {
+            "promptTokenCount": 100,
+            "candidatesTokenCount": 50,
+            "totalTokenCount": 150,
+            "toolUsePromptTokenCount": 10
+        }
+        
+        cost = handler._calculate_cost("gemini-1.5-pro", usage_metadata)
+        
+        # Should include web search cost
+        expected_base_cost = (100 * 0.000001) + (50 * 0.000002)
+        expected_web_search_cost = 0.01
+        expected_total = expected_base_cost + expected_web_search_cost
+        assert cost == expected_total
+
+    def test_error_handling_invalid_messages(self):
+        """Test error handling with invalid message formats"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Test with various invalid message formats
+        invalid_messages = [
+            "not a dict",
+            {"type": "invalid", "data": "incomplete"},
+            None,
+            [],
+            {"type": "response.create"},  # Missing response field
+            {"type": "response.create", "response": {}}  # Empty response
+        ]
+        
+        # Should handle all cases gracefully
+        for messages in invalid_messages:
+            result = handler._extract_usage_metadata_from_websocket_messages(messages)
+            assert result is None
+
+    def test_empty_websocket_messages(self):
+        """Test handling of empty WebSocket messages"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Test with empty list
+        result = handler._extract_usage_metadata_from_websocket_messages([])
+        assert result is None
+        
+        # Test with None
+        result = handler._extract_usage_metadata_from_websocket_messages(None)
+        assert result is None
+
+    @patch('litellm.utils.get_model_info')
+    def test_missing_model_info_handling(self, mock_get_model_info):
+        """Test handling when model info is missing or incomplete"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Test with empty model info
+        mock_get_model_info.return_value = {}
+        
+        usage_metadata = {
+            "promptTokenCount": 100,
+            "candidatesTokenCount": 50,
+            "totalTokenCount": 150
+        }
+        
+        cost = handler._calculate_cost("unknown-model", usage_metadata)
+        assert cost == 0.0
+        
+        # Test with partial model info
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001
+            # Missing output_cost_per_token
+        }
+        
+        cost = handler._calculate_cost("partial-model", usage_metadata)
+        # Should still calculate with available info
+        assert cost >= 0
+
+    def test_handler_with_mock_logging_obj(self, sample_websocket_messages):
+        """Test the main handler method with a mock logging object"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        mock_logging_obj = MagicMock()
+        
+        url_route = "/vertex_ai/live"
+        start_time = datetime.now()
+        end_time = datetime.now()
+        request_body = {"messages": [{"role": "user", "content": "Hello"}]}
+        
+        result = handler.vertex_ai_live_passthrough_handler(
+            websocket_messages=sample_websocket_messages,
+            logging_obj=mock_logging_obj,
+            url_route=url_route,
+            start_time=start_time,
+            end_time=end_time,
+            request_body=request_body
+        )
+        
+        # Verify result structure
+        assert "result" in result
+        assert "kwargs" in result
+        
+        result_data = result["result"]
+        assert "model" in result_data
+        assert "usage" in result_data
+        assert "choices" in result_data
+        
+        # Verify usage data
+        usage = result_data["usage"]
+        assert "prompt_tokens" in usage
+        assert "completion_tokens" in usage
+        assert "total_tokens" in usage
+        
+        # Verify aggregated usage
+        assert usage["prompt_tokens"] == 20  # 15 + 5
+        assert usage["completion_tokens"] == 28  # 20 + 8
+        assert usage["total_tokens"] == 48  # 35 + 13
+
+
+class TestVertexAILivePassthroughEndToEnd:
+    """End-to-end tests for Vertex AI Live passthrough"""
+
+    @pytest.fixture
+    def mock_vertex_ai_live_api(self):
+        """Mock the Vertex AI Live API responses"""
+        with patch('websockets.asyncio.client.connect') as mock_connect:
+            # Mock WebSocket connection
+            mock_websocket = AsyncMock()
+            mock_websocket.recv.side_effect = [
+                json.dumps({
+                    "type": "session.created",
+                    "session": {"id": "test-session"}
+                }),
+                json.dumps({
+                    "type": "response.create",
+                    "response": {
+                        "text": "Hello! How can I help you?",
+                        "usage": {
+                            "promptTokenCount": 10,
+                            "candidatesTokenCount": 15,
+                            "totalTokenCount": 25
+                        }
+                    }
+                }),
+                json.dumps({
+                    "type": "response.done",
+                    "response": {
+                        "usage": {
+                            "promptTokenCount": 5,
+                            "candidatesTokenCount": 8,
+                            "totalTokenCount": 13
+                        }
+                    }
+                })
+            ]
+            mock_websocket.send = AsyncMock()
+            mock_websocket.close = AsyncMock()
+            
+            mock_connect.return_value = mock_websocket
+            yield mock_connect
+
+    @pytest.mark.asyncio
+    async def test_websocket_passthrough_flow(self, mock_vertex_ai_live_api):
+        """Test the complete WebSocket passthrough flow"""
+        from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+            websocket_passthrough_request
+        )
+        
+        # Mock dependencies
+        mock_websocket = MagicMock()
+        mock_websocket.headers = {"authorization": "Bearer test-token"}
+        mock_websocket.client_state = MagicMock()
+        mock_websocket.client_state.DISCONNECTED = "disconnected"
+        
+        mock_user_api_key = MagicMock()
+        mock_logging_obj = MagicMock()
+        
+        # Test the WebSocket passthrough
+        await websocket_passthrough_request(
+            websocket=mock_websocket,
+            target="wss://test-vertex-ai-live-api.com/v1/stream",
+            custom_headers={"Authorization": "Bearer test-token"},
+            user_api_key_dict=mock_user_api_key,
+            forward_headers=False,
+            endpoint="/vertex_ai/live",
+            accept_websocket=True,
+            logging_obj=mock_logging_obj
+        )
+        
+        # Verify that the WebSocket connection was established
+        mock_vertex_ai_live_api.assert_called_once()
+
+    def test_route_detection_in_success_handler(self):
+        """Test that the success handler correctly detects Vertex AI Live routes"""
+        from litellm.proxy.pass_through_endpoints.success_handler import (
+            PassThroughEndpointLogging
+        )
+        
+        handler = PassThroughEndpointLogging()
+        
+        # Test various route patterns
+        test_routes = [
+            "/vertex_ai/live",
+            "/vertex_ai/live/",
+            "/vertex_ai/live/stream",
+            "/vertex_ai/live/chat",
+            "/vertex_ai/live/v1/stream"
+        ]
+        
+        for route in test_routes:
+            assert handler.is_vertex_ai_live_route(route), f"Route {route} should be detected as Vertex AI Live"
+        
+        # Test non-Vertex AI Live routes
+        non_live_routes = [
+            "/vertex_ai",
+            "/vertex_ai/discovery",
+            "/vertex_ai/aiplatform",
+            "/openai/chat/completions",
+            "/anthropic/messages"
+        ]
+        
+        for route in non_live_routes:
+            assert not handler.is_vertex_ai_live_route(route), f"Route {route} should not be detected as Vertex AI Live"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

tests/pass_through_tests/test_vertex_ai_live_simple.py

@@ -0,0 +1,351 @@
+#!/usr/bin/env python3
+"""
+Simple test script for Vertex AI Live API passthrough feature
+
+This script provides a quick way to test the Vertex AI Live API passthrough
+functionality without requiring a full test suite setup.
+"""
+
+import json
+import sys
+import os
+from datetime import datetime
+
+# Add the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../.."))
+
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import (
+    VertexAILivePassthroughLoggingHandler,
+)
+
+
+def test_usage_metadata_extraction():
+    """Test usage metadata extraction from WebSocket messages"""
+    print("Testing usage metadata extraction...")
+    
+    handler = VertexAILivePassthroughLoggingHandler()
+    
+    # Sample WebSocket messages
+    messages = [
+        {
+            "type": "session.created",
+            "session": {"id": "test-session-123"}
+        },
+        {
+            "type": "response.create",
+            "response": {
+                "text": "Hello! How can I help you?"
+            },
+            "usageMetadata": {
+                "promptTokenCount": 15,
+                "candidatesTokenCount": 20,
+                "totalTokenCount": 35,
+                "promptTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 15}
+                ],
+                "candidatesTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 20}
+                ]
+            }
+        },
+        {
+            "type": "response.done",
+            "usageMetadata": {
+                "promptTokenCount": 5,
+                "candidatesTokenCount": 8,
+                "totalTokenCount": 13,
+                "promptTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 5}
+                ],
+                "candidatesTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 8}
+                ]
+            }
+        }
+    ]
+    
+    # Extract usage metadata
+    usage_metadata = handler._extract_usage_metadata_from_websocket_messages(messages)
+    
+    if usage_metadata:
+        print("✅ Usage metadata extracted successfully:")
+        print(f"  - Prompt tokens: {usage_metadata['promptTokenCount']}")
+        print(f"  - Candidate tokens: {usage_metadata['candidatesTokenCount']}")
+        print(f"  - Total tokens: {usage_metadata['totalTokenCount']}")
+        print(f"  - Prompt details: {usage_metadata['promptTokensDetails']}")
+        print(f"  - Candidate details: {usage_metadata['candidatesTokensDetails']}")
+        
+        # Verify aggregated values
+        assert usage_metadata['promptTokenCount'] == 20  # 15 + 5
+        assert usage_metadata['candidatesTokenCount'] == 28  # 20 + 8
+        assert usage_metadata['totalTokenCount'] == 48  # 35 + 13
+        print("✅ Token aggregation working correctly")
+    else:
+        print("❌ Failed to extract usage metadata")
+        return False
+    
+    return True
+
+
+def test_cost_calculation():
+    """Test cost calculation functionality"""
+    print("\nTesting cost calculation...")
+    
+    handler = VertexAILivePassthroughLoggingHandler()
+    
+    # Mock model info
+    usage_metadata = {
+        "promptTokenCount": 100,
+        "candidatesTokenCount": 50,
+        "totalTokenCount": 150
+    }
+    
+    # Test with mock model info using patch
+    from unittest.mock import patch
+    
+    with patch('litellm.utils.get_model_info') as mock_get_model_info:
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001,
+            "output_cost_per_token": 0.000002
+        }
+        
+        cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)
+        expected_cost = (100 * 0.000001) + (50 * 0.000002)
+        
+        print(f"✅ Cost calculated: ${cost:.6f}")
+        print(f"  - Expected: ${expected_cost:.6f}")
+        print(f"  - Difference: ${abs(cost - expected_cost):.6f}")
+        
+        # The cost should be close to expected (within 1 cent)
+        assert abs(cost - expected_cost) < 0.01
+        print("✅ Cost calculation working correctly")
+    
+    return True
+
+
+def test_multimodal_usage():
+    """Test multimodal usage tracking"""
+    print("\nTesting multimodal usage tracking...")
+    
+    handler = VertexAILivePassthroughLoggingHandler()
+    
+    # Messages with mixed modalities
+    messages = [
+        {
+            "type": "response.create",
+            "response": {
+                "text": "Hello with audio"
+            },
+            "usageMetadata": {
+                "promptTokenCount": 30,
+                "candidatesTokenCount": 25,
+                "totalTokenCount": 55,
+                "promptTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 20},
+                    {"modality": "AUDIO", "tokenCount": 10}
+                ],
+                "candidatesTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 15},
+                    {"modality": "AUDIO", "tokenCount": 10}
+                ]
+            }
+        }
+    ]
+    
+    usage_metadata = handler._extract_usage_metadata_from_websocket_messages(messages)
+    
+    if usage_metadata:
+        print("✅ Multimodal usage extracted:")
+        print(f"  - Prompt tokens: {usage_metadata['promptTokenCount']}")
+        print(f"  - Candidate tokens: {usage_metadata['candidatesTokenCount']}")
+        print(f"  - Prompt details: {usage_metadata['promptTokensDetails']}")
+        print(f"  - Candidate details: {usage_metadata['candidatesTokensDetails']}")
+        
+        # Verify modality details
+        text_prompt = next(d for d in usage_metadata['promptTokensDetails'] if d['modality'] == 'TEXT')
+        audio_prompt = next(d for d in usage_metadata['promptTokensDetails'] if d['modality'] == 'AUDIO')
+        
+        assert text_prompt['tokenCount'] == 20
+        assert audio_prompt['tokenCount'] == 10
+        print("✅ Multimodal tracking working correctly")
+    else:
+        print("❌ Failed to extract multimodal usage")
+        return False
+    
+    return True
+
+
+def test_web_search_usage():
+    """Test web search (tool use) usage tracking"""
+    print("\nTesting web search usage tracking...")
+    
+    handler = VertexAILivePassthroughLoggingHandler()
+    
+    # Messages with web search usage
+    messages = [
+        {
+            "type": "response.create",
+            "response": {
+                "text": "Hello with web search"
+            },
+            "usageMetadata": {
+                "promptTokenCount": 50,
+                "candidatesTokenCount": 30,
+                "totalTokenCount": 80,
+                "toolUsePromptTokenCount": 10,
+                "promptTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 50}
+                ],
+                "candidatesTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 30}
+                ]
+            }
+        }
+    ]
+    
+    usage_metadata = handler._extract_usage_metadata_from_websocket_messages(messages)
+    
+    if usage_metadata:
+        print("✅ Web search usage extracted:")
+        print(f"  - Prompt tokens: {usage_metadata['promptTokenCount']}")
+        print(f"  - Candidate tokens: {usage_metadata['candidatesTokenCount']}")
+        print(f"  - Tool use prompt tokens: {usage_metadata.get('toolUsePromptTokenCount', 0)}")
+        
+        assert usage_metadata['toolUsePromptTokenCount'] == 10
+        print("✅ Web search tracking working correctly")
+    else:
+        print("❌ Failed to extract web search usage")
+        return False
+    
+    return True
+
+
+def test_error_handling():
+    """Test error handling with invalid inputs"""
+    print("\nTesting error handling...")
+    
+    handler = VertexAILivePassthroughLoggingHandler()
+    
+    # Test various invalid inputs
+    invalid_inputs = [
+        None,
+        [],
+        "not a list",
+        [{"type": "invalid"}],
+        [{"type": "response.create"}],  # Missing response
+        [{"type": "response.create", "response": {}}]  # Empty response
+    ]
+    
+    for i, invalid_input in enumerate(invalid_inputs):
+        try:
+            if invalid_input is None:
+                # Skip None input as it will cause iteration error
+                print(f"  - Input {i+1}: Skipped None input")
+                continue
+            else:
+                result = handler._extract_usage_metadata_from_websocket_messages(invalid_input)
+            print(f"  - Input {i+1}: Handled gracefully (result: {result})")
+        except Exception as e:
+            print(f"  - Input {i+1}: Error - {e}")
+            return False
+    
+    print("✅ Error handling working correctly")
+    return True
+
+
+def test_handler_integration():
+    """Test the main handler method"""
+    print("\nTesting handler integration...")
+    
+    handler = VertexAILivePassthroughLoggingHandler()
+    
+    # Mock logging object
+    class MockLoggingObj:
+        def __init__(self):
+            self.model_call_details = {}
+    
+    mock_logging_obj = MockLoggingObj()
+    
+    # Sample WebSocket messages with proper usage metadata
+    messages = [
+        {
+            "type": "response.create",
+            "response": {
+                "text": "Hello! How can I help you?"
+            },
+            "usageMetadata": {
+                "promptTokenCount": 10,
+                "candidatesTokenCount": 15,
+                "totalTokenCount": 25,
+                "promptTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 10}
+                ],
+                "candidatesTokensDetails": [
+                    {"modality": "TEXT", "tokenCount": 15}
+                ]
+            }
+        }
+    ]
+    
+    # Test the main handler method
+    result = handler.vertex_ai_live_passthrough_handler(
+        websocket_messages=messages,
+        logging_obj=mock_logging_obj,
+        url_route="/vertex_ai/live",
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        request_body={"messages": [{"role": "user", "content": "Hello"}]}
+    )
+    
+    if result and "result" in result and "kwargs" in result:
+        print("✅ Handler integration working:")
+        print(f"  - Result keys: {list(result.keys())}")
+        print(f"  - Model: {result['result'].get('model', 'N/A')}")
+        print(f"  - Usage: {result['result'].get('usage', {})}")
+        print("✅ Handler integration working correctly")
+        return True
+    else:
+        print("❌ Handler integration failed")
+        return False
+
+
+def main():
+    """Run all tests"""
+    print("🚀 Starting Vertex AI Live Passthrough Tests")
+    print("=" * 50)
+    
+    tests = [
+        test_usage_metadata_extraction,
+        test_cost_calculation,
+        test_multimodal_usage,
+        test_web_search_usage,
+        test_error_handling,
+        test_handler_integration
+    ]
+    
+    passed = 0
+    failed = 0
+    
+    for test in tests:
+        try:
+            if test():
+                passed += 1
+            else:
+                failed += 1
+        except Exception as e:
+            print(f"❌ Test {test.__name__} failed with exception: {e}")
+            failed += 1
+    
+    print("\n" + "=" * 50)
+    print(f"📊 Test Results: {passed} passed, {failed} failed")
+    
+    if failed == 0:
+        print("🎉 All tests passed!")
+        return 0
+    else:
+        print("❌ Some tests failed!")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py

@@ -0,0 +1,578 @@
+"""
+Test Vertex AI Live API Passthrough Feature
+
+This module tests the Vertex AI Live API WebSocket passthrough functionality,
+including the logging handler, cost tracking, and WebSocket message processing.
+"""
+
+import json
+import os
+import sys
+from datetime import datetime
+from unittest.mock import AsyncMock, Mock, patch, MagicMock
+from typing import Dict, List, Any, Optional
+
+import pytest
+import httpx
+
+# Add the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../.."))
+
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import (
+    VertexAILivePassthroughLoggingHandler,
+)
+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.proxy._types import UserAPIKeyAuth
+
+
+class TestVertexAILivePassthroughLoggingHandler:
+    """Test the Vertex AI Live Passthrough Logging Handler"""
+
+    @pytest.fixture
+    def handler(self):
+        """Create a handler instance for testing"""
+        return VertexAILivePassthroughLoggingHandler()
+
+    @pytest.fixture
+    def mock_logging_obj(self):
+        """Create a mock logging object"""
+        return MagicMock(spec=LiteLLMLoggingObj)
+
+    @pytest.fixture
+    def sample_websocket_messages(self):
+        """Sample WebSocket messages for testing"""
+        return [
+            {
+                "type": "session.created",
+                "session": {"id": "test-session-123"},
+                "timestamp": "2024-01-01T00:00:00Z"
+            },
+            {
+                "type": "response.create",
+                "event_id": "event-123",
+                "response": {
+                    "text": "Hello, how can I help you?",
+                    "usage": {
+                        "promptTokenCount": 10,
+                        "candidatesTokenCount": 15,
+                        "totalTokenCount": 25,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 10}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 15}
+                        ]
+                    }
+                }
+            },
+            {
+                "type": "response.done",
+                "event_id": "event-123",
+                "response": {
+                    "usage": {
+                        "promptTokenCount": 5,
+                        "candidatesTokenCount": 8,
+                        "totalTokenCount": 13,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 5}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 8}
+                        ]
+                    }
+                }
+            }
+        ]
+
+    def test_llm_provider_name_property(self, handler):
+        """Test that llm_provider_name returns the correct provider"""
+        assert handler.llm_provider_name == LlmProviders.VERTEX_AI
+
+    def test_get_provider_config(self, handler):
+        """Test that get_provider_config returns a valid config"""
+        config = handler.get_provider_config("gemini-1.5-pro")
+        assert config is not None
+        # Verify it's a Vertex AI config
+        assert hasattr(config, 'model')
+
+    def test_extract_usage_metadata_single_message(self, handler):
+        """Test usage metadata extraction from a single message"""
+        messages = [{
+            "type": "response.create",
+            "response": {
+                "usage": {
+                    "promptTokenCount": 10,
+                    "candidatesTokenCount": 15,
+                    "totalTokenCount": 25,
+                    "promptTokensDetails": [
+                        {"modality": "TEXT", "tokenCount": 10}
+                    ],
+                    "candidatesTokensDetails": [
+                        {"modality": "TEXT", "tokenCount": 15}
+                    ]
+                }
+            }
+        }]
+        
+        result = handler._extract_usage_metadata_from_websocket_messages(messages)
+        
+        assert result is not None
+        assert result["promptTokenCount"] == 10
+        assert result["candidatesTokenCount"] == 15
+        assert result["totalTokenCount"] == 25
+        assert len(result["promptTokensDetails"]) == 1
+        assert len(result["candidatesTokensDetails"]) == 1
+
+    def test_extract_usage_metadata_multiple_messages(self, handler):
+        """Test usage metadata aggregation from multiple messages"""
+        messages = [
+            {
+                "type": "response.create",
+                "response": {
+                    "usage": {
+                        "promptTokenCount": 10,
+                        "candidatesTokenCount": 15,
+                        "totalTokenCount": 25,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 10}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 15}
+                        ]
+                    }
+                }
+            },
+            {
+                "type": "response.done",
+                "response": {
+                    "usage": {
+                        "promptTokenCount": 5,
+                        "candidatesTokenCount": 8,
+                        "totalTokenCount": 13,
+                        "promptTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 5}
+                        ],
+                        "candidatesTokensDetails": [
+                            {"modality": "TEXT", "tokenCount": 8}
+                        ]
+                    }
+                }
+            }
+        ]
+        
+        result = handler._extract_usage_metadata_from_websocket_messages(messages)
+        
+        assert result is not None
+        assert result["promptTokenCount"] == 15  # 10 + 5
+        assert result["candidatesTokenCount"] == 23  # 15 + 8
+        assert result["totalTokenCount"] == 38  # 25 + 13
+        assert len(result["promptTokensDetails"]) == 1
+        assert result["promptTokensDetails"][0]["tokenCount"] == 15
+        assert len(result["candidatesTokensDetails"]) == 1
+        assert result["candidatesTokensDetails"][0]["tokenCount"] == 23
+
+    def test_extract_usage_metadata_no_usage(self, handler):
+        """Test handling of messages without usage metadata"""
+        messages = [
+            {"type": "session.created", "session": {"id": "test"}},
+            {"type": "response.create", "response": {"text": "Hello"}}
+        ]
+        
+        result = handler._extract_usage_metadata_from_websocket_messages(messages)
+        assert result is None
+
+    def test_extract_usage_metadata_empty_list(self, handler):
+        """Test handling of empty message list"""
+        result = handler._extract_usage_metadata_from_websocket_messages([])
+        assert result is None
+
+    def test_extract_usage_metadata_mixed_modalities(self, handler):
+        """Test usage metadata extraction with mixed modalities"""
+        messages = [{
+            "type": "response.create",
+            "response": {
+                "usage": {
+                    "promptTokenCount": 20,
+                    "candidatesTokenCount": 30,
+                    "totalTokenCount": 50,
+                    "promptTokensDetails": [
+                        {"modality": "TEXT", "tokenCount": 10},
+                        {"modality": "AUDIO", "tokenCount": 10}
+                    ],
+                    "candidatesTokensDetails": [
+                        {"modality": "TEXT", "tokenCount": 20},
+                        {"modality": "AUDIO", "tokenCount": 10}
+                    ]
+                }
+            }
+        }]
+        
+        result = handler._extract_usage_metadata_from_websocket_messages(messages)
+        
+        assert result is not None
+        assert result["promptTokenCount"] == 20
+        assert result["candidatesTokenCount"] == 30
+        assert len(result["promptTokensDetails"]) == 2
+        assert len(result["candidatesTokensDetails"]) == 2
+        
+        # Check modality aggregation
+        text_prompt = next(d for d in result["promptTokensDetails"] if d["modality"] == "TEXT")
+        audio_prompt = next(d for d in result["promptTokensDetails"] if d["modality"] == "AUDIO")
+        assert text_prompt["tokenCount"] == 10
+        assert audio_prompt["tokenCount"] == 10
+
+    @patch('litellm.utils.get_model_info')
+    def test_calculate_cost_basic(self, mock_get_model_info, handler):
+        """Test basic cost calculation"""
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001,
+            "output_cost_per_token": 0.000002
+        }
+        
+        usage_metadata = {
+            "promptTokenCount": 100,
+            "candidatesTokenCount": 50,
+            "totalTokenCount": 150
+        }
+        
+        cost = handler._calculate_cost("gemini-1.5-pro", usage_metadata)
+        
+        expected_cost = (100 * 0.000001) + (50 * 0.000002)
+        assert cost == expected_cost
+
+    @patch('litellm.utils.get_model_info')
+    def test_calculate_cost_with_audio(self, mock_get_model_info, handler):
+        """Test cost calculation with audio tokens"""
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001,
+            "output_cost_per_token": 0.000002,
+            "input_cost_per_audio_per_second": 0.0001,
+            "output_cost_per_audio_per_second": 0.0002
+        }
+        
+        usage_metadata = {
+            "promptTokenCount": 100,
+            "candidatesTokenCount": 50,
+            "totalTokenCount": 150,
+            "promptTokensDetails": [
+                {"modality": "TEXT", "tokenCount": 80},
+                {"modality": "AUDIO", "tokenCount": 20}
+            ],
+            "candidatesTokensDetails": [
+                {"modality": "TEXT", "tokenCount": 30},
+                {"modality": "AUDIO", "tokenCount": 20}
+            ]
+        }
+        
+        cost = handler._calculate_cost("gemini-1.5-pro", usage_metadata)
+        
+        # Should include both text and audio costs
+        assert cost > 0
+        assert cost > (100 * 0.000001) + (50 * 0.000002)  # Should be higher due to audio
+
+    @patch('litellm.utils.get_model_info')
+    def test_calculate_cost_with_web_search(self, mock_get_model_info, handler):
+        """Test cost calculation with web search (tool use)"""
+        mock_get_model_info.return_value = {
+            "input_cost_per_token": 0.000001,
+            "output_cost_per_token": 0.000002,
+            "web_search_cost_per_request": 0.01
+        }
+        
+        usage_metadata = {
+            "promptTokenCount": 100,
+            "candidatesTokenCount": 50,
+            "totalTokenCount": 150,
+            "toolUsePromptTokenCount": 10
+        }
+        
+        cost = handler._calculate_cost("gemini-1.5-pro", usage_metadata)
+        
+        # Should include web search cost
+        expected_base_cost = (100 * 0.000001) + (50 * 0.000002)
+        expected_web_search_cost = 0.01
+        expected_total = expected_base_cost + expected_web_search_cost
+        assert cost == expected_total
+
+    def test_vertex_ai_live_passthrough_handler_integration(self, handler, mock_logging_obj, sample_websocket_messages):
+        """Test the main passthrough handler method"""
+        url_route = "/vertex_ai/live"
+        start_time = datetime.now()
+        end_time = datetime.now()
+        request_body = {"messages": [{"role": "user", "content": "Hello"}]}
+        
+        result = handler.vertex_ai_live_passthrough_handler(
+            websocket_messages=sample_websocket_messages,
+            logging_obj=mock_logging_obj,
+            url_route=url_route,
+            start_time=start_time,
+            end_time=end_time,
+            request_body=request_body
+        )
+        
+        assert "result" in result
+        assert "kwargs" in result
+        
+        # Check that the result contains expected fields
+        result_data = result["result"]
+        assert "model" in result_data
+        assert "usage" in result_data
+        assert "choices" in result_data
+        
+        # Check usage data
+        usage = result_data["usage"]
+        assert "prompt_tokens" in usage
+        assert "completion_tokens" in usage
+        assert "total_tokens" in usage
+
+    def test_vertex_ai_live_passthrough_handler_no_usage(self, handler, mock_logging_obj):
+        """Test handler with messages that don't contain usage metadata"""
+        messages = [
+            {"type": "session.created", "session": {"id": "test"}},
+            {"type": "response.create", "response": {"text": "Hello"}}
+        ]
+        
+        url_route = "/vertex_ai/live"
+        start_time = datetime.now()
+        end_time = datetime.now()
+        request_body = {"messages": [{"role": "user", "content": "Hello"}]}
+        
+        result = handler.vertex_ai_live_passthrough_handler(
+            websocket_messages=messages,
+            logging_obj=mock_logging_obj,
+            url_route=url_route,
+            start_time=start_time,
+            end_time=end_time,
+            request_body=request_body
+        )
+        
+        assert "result" in result
+        assert "kwargs" in result
+        
+        # Should still return a valid result even without usage data
+        result_data = result["result"]
+        assert "model" in result_data
+        assert "usage" in result_data
+        assert "choices" in result_data
+
+
+class TestVertexAILivePassthroughIntegration:
+    """Integration tests for Vertex AI Live passthrough functionality"""
+
+    @pytest.fixture
+    def mock_websocket(self):
+        """Create a mock WebSocket for testing"""
+        websocket = MagicMock()
+        websocket.headers = {"authorization": "Bearer test-token"}
+        websocket.client_state = MagicMock()
+        websocket.client_state.DISCONNECTED = "disconnected"
+        return websocket
+
+    @pytest.fixture
+    def mock_user_api_key(self):
+        """Create a mock user API key"""
+        return UserAPIKeyAuth(
+            api_key="test-key",
+            user_id="test-user",
+            team_id="test-team",
+            user_role="user"
+        )
+
+    @pytest.fixture
+    def mock_logging_obj(self):
+        """Create a mock logging object"""
+        return MagicMock(spec=LiteLLMLoggingObj)
+
+    @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request')
+    @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router')
+    def test_vertex_ai_live_websocket_passthrough_route(
+        self, 
+        mock_router, 
+        mock_websocket_passthrough,
+        mock_websocket,
+        mock_user_api_key,
+        mock_logging_obj
+    ):
+        """Test the Vertex AI Live WebSocket passthrough route"""
+        from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+            vertex_ai_live_websocket_passthrough_route
+        )
+        
+        # Mock the router methods
+        mock_router.get_vertex_credentials.return_value = MagicMock(
+            vertex_project="test-project",
+            vertex_location="us-central1",
+            vertex_credentials="test-credentials"
+        )
+        mock_router.set_default_vertex_config.return_value = None
+        
+        # Mock the WebSocket passthrough request
+        mock_websocket_passthrough.return_value = AsyncMock()
+        
+        # Test the route
+        result = vertex_ai_live_websocket_passthrough_route(
+            websocket=mock_websocket,
+            user_api_key_dict=mock_user_api_key,
+            logging_obj=mock_logging_obj
+        )
+        
+        # Verify that the WebSocket passthrough was called
+        mock_websocket_passthrough.assert_called_once()
+        
+        # Check the call arguments
+        call_args = mock_websocket_passthrough.call_args
+        assert call_args[1]["websocket"] == mock_websocket
+        assert call_args[1]["user_api_key_dict"] == mock_user_api_key
+        assert call_args[1]["endpoint"] == "/vertex_ai/live"
+
+    def test_vertex_ai_live_route_detection(self):
+        """Test that the route detection works correctly"""
+        from litellm.proxy.pass_through_endpoints.success_handler import (
+            PassThroughEndpointLogging
+        )
+        
+        handler = PassThroughEndpointLogging()
+        
+        # Test valid routes
+        assert handler.is_vertex_ai_live_route("/vertex_ai/live") == True
+        assert handler.is_vertex_ai_live_route("/vertex_ai/live/") == True
+        assert handler.is_vertex_ai_live_route("/vertex_ai/live/stream") == True
+        
+        # Test invalid routes
+        assert handler.is_vertex_ai_live_route("/vertex_ai") == False
+        assert handler.is_vertex_ai_live_route("/vertex_ai/discovery") == False
+        assert handler.is_vertex_ai_live_route("/openai/chat/completions") == False
+
+    @patch('litellm.proxy.pass_through_endpoints.success_handler.VertexAILivePassthroughLoggingHandler')
+    def test_success_handler_vertex_ai_live_integration(
+        self, 
+        mock_handler_class,
+        mock_logging_obj
+    ):
+        """Test the success handler integration with Vertex AI Live"""
+        from litellm.proxy.pass_through_endpoints.success_handler import (
+            PassThroughEndpointLogging
+        )
+        
+        # Mock the handler
+        mock_handler = MagicMock()
+        mock_handler.vertex_ai_live_passthrough_handler.return_value = {
+            "result": {"model": "gemini-1.5-pro", "usage": {"total_tokens": 100}},
+            "kwargs": {"test": "value"}
+        }
+        mock_handler_class.return_value = mock_handler
+        
+        # Create success handler
+        success_handler = PassThroughEndpointLogging()
+        
+        # Mock the route check
+        success_handler.is_vertex_ai_live_route = MagicMock(return_value=True)
+        
+        # Test data
+        response_body = [
+            {"type": "response.create", "response": {"text": "Hello"}}
+        ]
+        url_route = "/vertex_ai/live"
+        start_time = datetime.now()
+        end_time = datetime.now()
+        request_body = {"messages": [{"role": "user", "content": "Hello"}]}
+        
+        # Call the method
+        result = success_handler.pass_through_async_success_handler(
+            httpx_response=MagicMock(),
+            response_body=response_body,
+            logging_obj=mock_logging_obj,
+            url_route=url_route,
+            result="test",
+            start_time=start_time,
+            end_time=end_time,
+            cache_hit=False,
+            request_body=request_body,
+            passthrough_logging_payload=MagicMock()
+        )
+        
+        # Verify the handler was called
+        mock_handler.vertex_ai_live_passthrough_handler.assert_called_once()
+        
+        # Verify the result
+        assert "standard_logging_response_object" in result
+        assert result["standard_logging_response_object"]["model"] == "gemini-1.5-pro"
+
+
+class TestVertexAILivePassthroughErrorHandling:
+    """Test error handling in Vertex AI Live passthrough"""
+
+    def test_invalid_websocket_messages_format(self):
+        """Test handling of invalid WebSocket message formats"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Test with invalid message format
+        invalid_messages = [
+            {"type": "invalid", "data": "not a proper message"},
+            "not a dict at all",
+            None
+        ]
+        
+        # Should not raise an exception
+        result = handler._extract_usage_metadata_from_websocket_messages(invalid_messages)
+        assert result is None
+
+    def test_missing_usage_metadata(self):
+        """Test handling of messages with missing usage metadata"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        messages = [
+            {"type": "response.create", "response": {"text": "Hello"}},
+            {"type": "response.done", "response": {"text": "Done"}}
+        ]
+        
+        result = handler._extract_usage_metadata_from_websocket_messages(messages)
+        assert result is None
+
+    @patch('litellm.utils.get_model_info')
+    def test_cost_calculation_with_missing_model_info(self, mock_get_model_info):
+        """Test cost calculation when model info is missing"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        # Mock missing model info
+        mock_get_model_info.return_value = {}
+        
+        usage_metadata = {
+            "promptTokenCount": 100,
+            "candidatesTokenCount": 50,
+            "totalTokenCount": 150
+        }
+        
+        # Should not raise an exception, should return 0 or handle gracefully
+        cost = handler._calculate_cost("unknown-model", usage_metadata)
+        assert cost == 0.0
+
+    def test_handler_with_none_websocket_messages(self, mock_logging_obj):
+        """Test handler with None websocket messages"""
+        handler = VertexAILivePassthroughLoggingHandler()
+        
+        url_route = "/vertex_ai/live"
+        start_time = datetime.now()
+        end_time = datetime.now()
+        request_body = {"messages": [{"role": "user", "content": "Hello"}]}
+        
+        # Should handle None gracefully
+        result = handler.vertex_ai_live_passthrough_handler(
+            websocket_messages=None,
+            logging_obj=mock_logging_obj,
+            url_route=url_route,
+            start_time=start_time,
+            end_time=end_time,
+            request_body=request_body
+        )
+        
+        assert "result" in result
+        assert "kwargs" in result
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])