From b7d7b93eb9024c2cdec2cd007f3be87974a58e9d Mon Sep 17 00:00:00 2001 From: joereyna Date: Thu, 9 Apr 2026 11:14:34 -0700 Subject: [PATCH] fix(responses-ws): use urllib.parse to append model param, fix test mocking --- litellm/llms/custom_httpx/llm_http_handler.py | 11 +- .../test_responses_websocket_all_providers.py | 117 ++++++++++++++---- 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4e5b7a3410..7a8820a878 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,5 +1,6 @@ import json import ssl +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, Any, @@ -5029,8 +5030,14 @@ class BaseLLMHTTPHandler: ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") # OpenAI's WebSocket responses endpoint requires ?model= in the URL, # matching the Realtime API convention (wss://.../v1/realtime?model=...). - if "?" not in ws_url: - ws_url = f"{ws_url}?model={model}" + # Use urllib.parse so existing query params (e.g. api-version) are preserved. + _parsed = urlparse(ws_url) + _qs = parse_qs(_parsed.query) + if "model" not in _qs: + _qs["model"] = [model] + ws_url = urlunparse( + _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) + ) try: ssl_context = get_shared_realtime_ssl_context() diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 09bde86593..efec841cc4 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -974,32 +974,103 @@ class TestWebSocketChunkTypes: class TestNativeWebSocketUrlConstruction: - """Test that native WebSocket URLs include the model query parameter.""" + """Test that native WebSocket URLs include the model query parameter. - def test_openai_ws_url_includes_model(self): - """ws_url for OpenAI native WebSocket must include ?model= so OpenAI - knows which model to use before the first response.create event.""" - config = OpenAIResponsesAPIConfig() - http_url = config.get_complete_url(api_base=None, litellm_params={}) - base_ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + These tests mock websockets.connect so they exercise the actual URL-building + code inside BaseLLMHTTPHandler.async_responses_websocket rather than + reimplementing the logic themselves. + """ - # get_complete_url should not include query params - assert "?" not in base_ws_url + @pytest.mark.asyncio + async def test_openai_ws_url_includes_model(self): + """Handler must pass ?model= in the URL to the backend WebSocket.""" + from unittest.mock import AsyncMock, MagicMock, patch - # The handler appends ?model= when none is present - model = "gpt-4o-mini" - ws_url = f"{base_ws_url}?model={model}" if "?" not in base_ws_url else base_ws_url - assert ws_url == "wss://api.openai.com/v1/responses?model=gpt-4o-mini" + captured_urls = [] - def test_ws_url_model_not_duplicated_if_query_already_present(self): - """If api_base already has query params, the ?model= should not be appended.""" - http_url = "https://custom.example.com/v1/responses?api-version=2024-05-01" - ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) - # Fix: only append when no query string present - if "?" not in ws_url: - ws_url = f"{ws_url}?model=gpt-4o" + async def __aenter__(self): + raise Exception("stop") - assert "api-version=2024-05-01" in ws_url - assert ws_url.count("?") == 1 - assert "model=gpt-4o" not in ws_url + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = "https://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o-mini", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o-mini"], f"Expected model in URL, got: {captured_urls[0]}" + + @pytest.mark.asyncio + async def test_ws_url_preserves_existing_params_and_adds_model(self): + """When api_base already has query params, model is added alongside them.""" + from unittest.mock import AsyncMock, MagicMock, patch + + captured_urls = [] + + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = ( + "https://custom.example.com/v1/responses?api-version=2024-05-01" + ) + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o"], f"model missing from URL: {captured_urls[0]}" + assert qs.get("api-version") == ["2024-05-01"], f"existing param lost: {captured_urls[0]}"