From ea0e5773857cc1ca7e126dce6be751d53ea8ceab Mon Sep 17 00:00:00 2001 From: L4XB Date: Thu, 17 Sep 2026 19:19:42 +0200 Subject: [PATCH] refactor(proxy): read the websocket key inline again The key extraction moves back into user_api_key_auth_websocket at its old place, with the lines that read the key unchanged from main. The behaviour change stays: a missing key reaches user_api_key_auth as None, and rejections raise WebSocketException without closing the socket first --- litellm/proxy/auth/user_api_key_auth.py | 45 ++++++++----------------- 1 file changed, 14 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 3b437d7865e..2028fed3159 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -629,36 +629,6 @@ def _apply_budget_limits_to_end_user_params( verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) -def _get_websocket_api_key(websocket: WebSocket) -> str | None: - """Read the API key a WebSocket client presented, or None when it presented none. - - Whether a key is required is decided by ``user_api_key_auth``, which allows a - keyless request when no master key is configured. - """ - authorization: Final = websocket.headers.get("authorization") - if authorization: - if not authorization.startswith("Bearer "): - raise WebSocketException( - code=status.WS_1008_POLICY_VIOLATION, - reason="Invalid Authorization header format", - ) - return authorization[len("Bearer ") :].strip() - - header_key: Final = websocket.headers.get("api-key") - if header_key: - return header_key - - subprotocol_prefix: Final = "openai-insecure-api-key." - return next( - ( - protocol.strip()[len(subprotocol_prefix) :] - for protocol in websocket.headers.get("sec-websocket-protocol", "").split(",") - if protocol.strip().startswith(subprotocol_prefix) - ), - None, - ) - - async def user_api_key_auth_websocket(websocket: WebSocket): # Accept the WebSocket connection @@ -690,7 +660,20 @@ async def user_api_key_auth_websocket(websocket: WebSocket): request.body = return_body - api_key: Final = _get_websocket_api_key(websocket) + authorization: Final = websocket.headers.get("authorization") + if not authorization: + api_key = websocket.headers.get("api-key") + if not api_key: + for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): + protocol = protocol.strip() + if protocol.startswith("openai-insecure-api-key."): + api_key = protocol[len("openai-insecure-api-key.") :] + break + else: + if not authorization.startswith("Bearer "): + raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid Authorization header format") + + api_key = authorization[len("Bearer ") :].strip() try: return await user_api_key_auth(