fix(responses): preserve explicit websocket API keys (internal copy of #40853) (#46724)

* fix(responses): preserve explicit websocket API keys

Refresh the one-line fix from BerriAI/litellm#34253 against the current default branch

Co-authored-by: puffer <389053104@qq.com>

* test(responses): verify websocket authorization

Cover explicit deployment keys, precedence over global keys, and the existing global fallback at the upstream WebSocket transport boundary

* test(responses): type the authorization capture

Replace the mutable header list with a typed, final Future carrying the upstream authorization value, addressing the review on #40853

* Refactor test_openai_websocket_authorization method

* fix(responses): resolve websocket keys with the deployment's provider

Pass the deployment custom_llm_provider to get_llm_provider so an api_base
host never swaps in another provider's env key, and drop the OpenAI-only
global key fallbacks so non-OpenAI websocket deployments resolve their own
env keys the way the HTTP route does

* fix(responses): let a deployment's own websocket key beat env keys

A deployment with an api_base on a host litellm recognizes, and no declared
provider, got that host's env key ahead of its own api_key. The explicit key
now always wins, matching the HTTP Responses route. Drop the litellm_params
api_key fallback, which could never be set because api_key is a named argument

* fix(responses): keep openai deployments on other providers' known hosts on the HTTP stream

An openai deployment whose api_base is a known OpenAI-compatible provider
host (DeepSeek, Groq, ...) opened a native Responses WebSocket there and got
a 405, because those hosts serve /responses over HTTP only. Route it through
the managed HTTP stream with the deployment's own key instead.

* fix(responses): resolve openai deployments on known provider hosts as that provider

An openai deployment whose api_base is a known provider host (DeepSeek,
Mistral, ...) now resolves to that host's provider on the Responses
WebSocket again, as before this branch, so it keeps that provider's
chat completions bridge, pricing and env key fallback. Its own api_key
still wins over the env key. The known-host test now mocks the upstream
HTTP route and checks the route, the bearer and response.completed.

* test(responses): prove websocket deployment keys and declared providers reach the provider socket

* test(responses): use the reported bare model shape for the websocket deployment key cell

---------

Co-authored-by: Ysh <734645956@qq.com>
Co-authored-by: puffer <389053104@qq.com>
This commit is contained in:
Mateo Wang 2026-10-16 01:28:12 +00:00 • committed by GitHub
parent d66d8ac127
commit b47f7a4712
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 286 additions and 17 deletions

View file

@ -969,7 +969,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = {
"context_management": None,
}
openai_compatible_endpoints: Final[list] = [
openai_compatible_endpoints: Final[list[str]] = [
"api.perplexity.ai",
"api.endpoints.anyscale.com/v1",
"api.deepinfra.com/v1/openai",

View file

@ -61,6 +61,10 @@ def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
return url_path == endpoint_path or url_path.startswith(endpoint_path + "/")
def is_openai_compatible_endpoint(api_base: str) -> bool:
return any(_endpoint_matches_api_base(endpoint, api_base) for endpoint in litellm.openai_compatible_endpoints)
def is_non_openai_azure_model(model: str) -> bool:
try:
model_name: Final = model.split("/", 1)[1]

View file

@ -59,8 +59,8 @@ if TYPE_CHECKING:
else:
ResponseText = str # Fallback for ResponseText import
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.get_llm_provider_logic import is_openai_compatible_endpoint
from litellm.llms.openai.data_residency import infer_openai_data_residency
from litellm.secret_managers.main import get_secret_str
from litellm.types.responses.main import *
from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
from litellm.types.router import GenericLiteLLMParams
@ -2461,6 +2461,13 @@ def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) ->
)
def _websocket_custom_llm_provider(custom_llm_provider: str | None, api_base: str | None) -> str | None:
is_openai_on_known_provider_host: Final = (
custom_llm_provider == "openai" and api_base is not None and is_openai_compatible_endpoint(api_base)
)
return None if is_openai_on_known_provider_host else custom_llm_provider
@client
async def _aresponses_websocket(
model: str,
@ -2491,6 +2498,7 @@ async def _aresponses_websocket(
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=_websocket_custom_llm_provider(litellm_params.custom_llm_provider, api_base),
api_base=api_base,
api_key=api_key,
litellm_params=litellm_params,
@ -2512,22 +2520,18 @@ async def _aresponses_websocket(
)
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
responses_api_provider_config: BaseResponsesAPIConfig | None = None
if _custom_llm_provider is not None:
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
responses_api_provider_config: Final = (
None
if _custom_llm_provider is None
else ProviderConfigManager.get_provider_responses_api_config(
model=resolved_model,
provider=litellm.LlmProviders(_custom_llm_provider),
api_base=resolved_api_base,
)
resolved_api_key: Final = (
dynamic_api_key
or litellm_params.api_key
or litellm.api_key
or litellm.openai_key
or get_secret_str("OPENAI_API_KEY")
)
resolved_api_key: Final = api_key or dynamic_api_key
# Extract params that we're passing explicitly to avoid duplicates in **kwargs
_explicit_keys: Final = {
"user_api_key_dict",

View file

@ -27,14 +27,23 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
PROVIDER_MODEL: Final = "ws-peer-model"
USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}
TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"})
PROXY_OPENAI_KEY: Final = "sk-proxy-environment-openai-key"
PROXY_AZURE_KEY: Final = "proxy-environment-azure-key"
_Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]]
@dataclass(frozen=True, slots=True)
class Credentials:
authorization: str | None
api_key: str | None
@dataclass(frozen=True, slots=True)
class Peer:
url: str
paths: SimpleQueue[str]
credentials: SimpleQueue[Credentials]
frames: SimpleQueue[dict[str, JsonValue]]
@ -75,10 +84,17 @@ def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]:
async def _answer(
connection: ServerConnection,
paths: SimpleQueue[str],
credentials: SimpleQueue[Credentials],
frames: SimpleQueue[dict[str, JsonValue]],
script: _Script,
) -> None:
paths.put(connection.request.path if connection.request is not None else "")
request: Final = connection.request
paths.put(request.path if request is not None else "")
credentials.put(
Credentials(request.headers.get("Authorization"), request.headers.get("api-key"))
if request is not None
else Credentials(None, None)
)
turns: Final = itertools.count(1)
async for raw in connection:
frame: Final = json.loads(raw)
@ -92,12 +108,15 @@ async def _answer(
async def _serve(
tls: ssl.SSLContext,
paths: SimpleQueue[str],
credentials: SimpleQueue[Credentials],
frames: SimpleQueue[dict[str, JsonValue]],
ports: SimpleQueue[int],
stop: asyncio.Event,
script: _Script,
) -> None:
async with serve(lambda connection: _answer(connection, paths, frames, script), "127.0.0.1", 0, ssl=tls) as server:
async with serve(
lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls
) as server:
ports.put(next(iter(server.sockets)).getsockname()[1])
await stop.wait()
@ -107,16 +126,17 @@ def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterato
loop: Final = asyncio.new_event_loop()
stop: Final = asyncio.Event()
paths: Final = SimpleQueue[str]()
credentials: Final = SimpleQueue[Credentials]()
frames: Final = SimpleQueue[dict[str, JsonValue]]()
ports: Final = SimpleQueue[int]()
thread: Final = threading.Thread(
target=loop.run_until_complete,
args=(_serve(server_context(*cert), paths, frames, ports, stop, script),),
args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),),
daemon=True,
)
thread.start()
try:
yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, frames)
yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames)
finally:
loop.call_soon_threadsafe(stop.set)
thread.join(timeout=10)
@ -180,7 +200,11 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
@pytest.fixture(scope="module")
def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]:
with gateway_from_environment() as base:
with owned_proxy(base, tmp_path_factory.mktemp("responses-ws"), {"SSL_CERT_FILE": str(cert[0])}) as proxy:
with owned_proxy(
base,
tmp_path_factory.mktemp("responses-ws"),
{"SSL_CERT_FILE": str(cert[0]), "OPENAI_API_KEY": PROXY_OPENAI_KEY, "AZURE_API_KEY": PROXY_AZURE_KEY},
) as proxy:
yield proxy
@ -253,3 +277,67 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a
error: Final = object_value(object_value(row["metadata"])["error_information"])
assert str(error["error_code"]) == "429", error
assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error
def _one_turn(candidate: Gateway, key: str, model: str) -> tuple[dict[str, JsonValue], ...]:
(events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",)))
return events
def test_a_bare_model_declared_as_openai_sends_its_deployment_key_to_the_provider_socket(
candidate: Gateway, cert: tuple[Path, Path]
) -> None:
deployment_key: Final = f"sk-deployment-{uuid.uuid4().hex}"
with responses_peer(cert) as peer, candidate.scenario() as scenario:
model: Final = scenario.model(
model=PROVIDER_MODEL, custom_llm_provider="openai", api_base=peer.url, api_key=deployment_key
)
events: Final = _one_turn(candidate, scenario.key(models=[model]), model)
assert _completed(events)["status"] == "completed"
assert peer.credentials.get_nowait().authorization == f"Bearer {deployment_key}"
def test_an_openai_deployment_without_its_own_key_uses_the_proxy_openai_key(
candidate: Gateway, cert: tuple[Path, Path]
) -> None:
with responses_peer(cert) as peer, candidate.scenario() as scenario:
model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url, api_key=None)
events: Final = _one_turn(candidate, scenario.key(models=[model]), model)
assert _completed(events)["status"] == "completed"
assert peer.credentials.get_nowait().authorization == f"Bearer {PROXY_OPENAI_KEY}"
def test_an_azure_deployment_without_its_own_key_gets_the_azure_key_and_never_the_openai_key(
candidate: Gateway, cert: tuple[Path, Path]
) -> None:
with responses_peer(cert) as peer, candidate.scenario() as scenario:
model: Final = scenario.model(
model=f"azure/{PROVIDER_MODEL}", api_base=peer.url.removesuffix("/v1"), api_key=None
)
events: Final = _one_turn(candidate, scenario.key(models=[model]), model)
assert _completed(events)["status"] == "completed"
credentials: Final = peer.credentials.get_nowait()
assert credentials.api_key == PROXY_AZURE_KEY
assert PROXY_OPENAI_KEY not in (credentials.authorization or "")
def test_a_declared_openai_provider_wins_over_another_providers_model_prefix(
candidate: Gateway, cert: tuple[Path, Path]
) -> None:
gateway_model: Final = f"anthropic/{PROVIDER_MODEL}"
with responses_peer(cert) as peer, candidate.scenario() as scenario:
model: Final = scenario.model(model=gateway_model, custom_llm_provider="openai", api_base=peer.url)
key: Final = scenario.key(models=[model])
events: Final = _one_turn(candidate, key, model)
assert _completed(events)["status"] == "completed"
(forwarded,) = _drain(peer.frames)
assert forwarded["model"] == gateway_model
(row,) = eventually(
lambda: read_rows(
'SELECT custom_llm_provider FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
(hashlib.sha256(key.encode()).hexdigest(),),
),
lambda rows: len(rows) >= 1,
seconds=70,
)
assert row["custom_llm_provider"] == "openai", row

View file

@ -8,6 +8,7 @@ Tests that:
"""
import json
from collections.abc import Mapping
from typing import Final
from unittest.mock import MagicMock
@ -2805,6 +2806,178 @@ class TestWebSocketChunkTypes:
assert "Complete reasoning" in serialized
class TestResponsesWebSocketCredentials:
@staticmethod
async def _upstream_handshake_headers(**call_kwargs: object) -> Mapping[str, str]:
from asyncio import Future, get_running_loop
from types import TracebackType
from unittest.mock import AsyncMock, patch
from litellm.responses import main as responses_main
captured_headers: Final[Future[Mapping[str, str]]] = get_running_loop().create_future()
class FakeConnect:
def __init__(
self,
_url: str,
*,
additional_headers: Mapping[str, str],
**_kwargs: object,
) -> None:
captured_headers.set_result(dict(additional_headers))
async def __aenter__(self) -> None:
raise RuntimeError("stop at upstream handshake")
async def __aexit__(
self,
_exc_type: type[BaseException] | None,
_exc: BaseException | None,
_tb: TracebackType | None,
) -> None:
return None
websocket: Final = MagicMock()
websocket.close = AsyncMock()
with patch("websockets.connect", FakeConnect):
await responses_main._aresponses_websocket.__wrapped__(
websocket=websocket,
litellm_logging_obj=MagicMock(),
**call_kwargs,
)
assert captured_headers.done(), "the request never reached the native upstream websocket"
return captured_headers.result()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"api_base,explicit_key,global_key,expected_key",
[
("https://upstream.invalid/v1", "deployment-key", None, "deployment-key"),
("https://upstream.invalid/v1", "deployment-key", "global-key", "deployment-key"),
("https://upstream.invalid/v1", None, "global-key", "global-key"),
],
)
async def test_openai_websocket_authorization(
self,
monkeypatch: pytest.MonkeyPatch,
api_base: str,
explicit_key: str | None,
global_key: str | None,
expected_key: str,
) -> None:
import litellm
monkeypatch.setattr(litellm, "api_key", global_key)
monkeypatch.setattr(litellm, "openai_key", None)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
monkeypatch.setenv("LITELLM_RUST", "false")
headers: Final = await self._upstream_handshake_headers(
model="gpt-6-astra",
custom_llm_provider="openai",
api_key=explicit_key,
api_base=api_base,
)
assert headers["Authorization"] == f"Bearer {expected_key}"
@pytest.mark.asyncio
async def test_azure_websocket_without_its_own_key_never_gets_the_openai_key(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
import litellm
monkeypatch.setattr(litellm, "api_key", None)
monkeypatch.setattr(litellm, "openai_key", None)
monkeypatch.setattr(litellm, "azure_key", None)
monkeypatch.setenv("OPENAI_API_KEY", "openai-global-key")
monkeypatch.setenv("AZURE_API_KEY", "azure-env-key")
monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False)
monkeypatch.setenv("LITELLM_RUST", "false")
headers: Final = await self._upstream_handshake_headers(
model="gpt-6-astra",
custom_llm_provider="azure",
api_base="https://resource.openai.azure.com",
)
assert headers["api-key"] == "azure-env-key"
assert "openai-global-key" not in " ".join(headers.values())
@pytest.mark.asyncio
@pytest.mark.parametrize(
"custom_llm_provider,api_base,deployment_key,expected_key",
[
(None, "https://api.deepseek.com/v1", "deployment-key", "deployment-key"),
("openai", "https://api.deepseek.com/v1", "deployment-key", "deployment-key"),
("openai", "https://api.deepseek.com/v1", None, "deepseek-env-key"),
("openai", "https://api.mistral.ai/v1", "deployment-key", "deployment-key"),
],
)
async def test_deployment_on_a_known_provider_host_completes_through_that_provider(
self,
monkeypatch: pytest.MonkeyPatch,
custom_llm_provider: str | None,
api_base: str,
deployment_key: str | None,
expected_key: str,
) -> None:
from unittest.mock import AsyncMock, patch
import httpx
import respx
import litellm
from litellm.responses import main as responses_main
chunks: Final = (
{"choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}, "finish_reason": None}]},
{"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
)
stream: Final = "".join(
f"data: {json.dumps({'id': 'chatcmpl-1', 'object': 'chat.completion.chunk', 'created': 1, 'model': 'astra-chat', **chunk})}\n\n"
for chunk in chunks
)
def refuse_native_socket(*_args: object, **_kwargs: object) -> None:
raise ConnectionRefusedError("this host has no Responses WebSocket")
monkeypatch.setattr(litellm, "api_key", None)
monkeypatch.setattr(litellm, "openai_key", None)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "openai-env-key")
monkeypatch.setenv("DEEPSEEK_API_KEY", "deepseek-env-key")
monkeypatch.setenv("MISTRAL_API_KEY", "mistral-env-key")
monkeypatch.setenv("LITELLM_RUST", "false")
websocket: Final = MagicMock()
websocket.send_text = AsyncMock()
websocket.close = AsyncMock()
websocket.receive_text = AsyncMock(side_effect=RuntimeError("client closed"))
with respx.mock(assert_all_called=False) as upstream, patch("websockets.connect", refuse_native_socket):
chat_completions: Final = upstream.post(f"{api_base}/chat/completions").mock(
return_value=httpx.Response(
200, text=stream + "data: [DONE]\n\n", headers={"content-type": "text/event-stream"}
)
)
await responses_main._aresponses_websocket.__wrapped__(
model="astra-chat",
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=deployment_key,
first_message=json.dumps({"type": "response.create", "model": "astra-chat", "input": "Reply OK"}),
websocket=websocket,
litellm_logging_obj=MagicMock(),
)
sent_types: Final = [json.loads(call.args[0]).get("type") for call in websocket.send_text.await_args_list]
assert chat_completions.called, f"the request never reached {api_base}/chat/completions, client saw {sent_types}"
assert chat_completions.calls.last.request.headers["Authorization"] == f"Bearer {expected_key}"
assert "response.completed" in sent_types
class TestNativeWebSocketUrlConstruction:
"""Test that native WebSocket URLs include the model query parameter.