From 9e6f44d6d83e23abadc842ea92b6d3e64d77b063 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 10 Jul 2026 22:47:01 +0200 Subject: [PATCH 1/2] fix websocket workload identity authentication --- src/openai/_client.py | 78 +++++++++++++++---- .../resources/beta/realtime/realtime.py | 40 ++++++++-- .../resources/beta/responses/responses.py | 58 ++++++++------ src/openai/resources/realtime/realtime.py | 70 +++++++++++------ src/openai/resources/responses/responses.py | 58 ++++++++------ tests/test_client.py | 52 +++++++++++++ 6 files changed, 266 insertions(+), 90 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index 66d03b23dd..5ebb4643c1 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -92,6 +92,14 @@ WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER = "workload-identity-auth" +def _is_websocket_unauthorized(exc: Exception) -> bool: + status_code = getattr(exc, "status_code", None) + if status_code is None: + response = getattr(exc, "response", None) + status_code = getattr(response, "status_code", None) + return status_code == 401 + + def _has_header(headers: Headers, header: str) -> bool: header = header.lower() return any(key.lower() == header for key in headers) @@ -445,13 +453,12 @@ def _send_with_auth_retry( retried: bool = False, **kwargs: Unpack[HttpxSendArgs], ) -> httpx.Response: - used_workload_identity_auth = False - - if self._workload_identity_auth is not None: - authorization = request.headers.get("Authorization") - if authorization == f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}": - request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" - used_workload_identity_auth = True + authorization, used_workload_identity_auth = self._resolve_auth_header( + request.headers.get("Authorization"), + workload_identity_placeholder_required=True, + ) + if authorization is not None: + request.headers["Authorization"] = authorization response = super()._send_request(request, stream=stream, **kwargs) if ( @@ -467,6 +474,28 @@ def _send_with_auth_retry( return response + def _resolve_auth_header( + self, authorization: str | None, *, workload_identity_placeholder_required: bool = False + ) -> tuple[str | None, bool]: + if self._workload_identity_auth is None: + return authorization, False + if ( + workload_identity_placeholder_required + and authorization != f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}" + ): + return authorization, False + return f"Bearer {self._workload_identity_auth.get_token()}", True + + def _websocket_auth_headers(self) -> dict[str, str]: + authorization, _ = self._resolve_auth_header(self.auth_headers.get("Authorization")) + return {"Authorization": authorization} if authorization is not None else {} + + def _retry_websocket_auth_headers(self, exc: Exception) -> dict[str, str] | None: + if self._workload_identity_auth is None or not _is_websocket_unauthorized(exc): + return None + self._workload_identity_auth.invalidate_token() + return self._websocket_auth_headers() + @override def _send_request( self, @@ -1041,13 +1070,12 @@ async def _send_with_auth_retry( retried: bool = False, **kwargs: Unpack[HttpxSendArgs], ) -> httpx.Response: - used_workload_identity_auth = False - - if self._workload_identity_auth is not None: - authorization = request.headers.get("Authorization") - if authorization == f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}": - request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" - used_workload_identity_auth = True + authorization, used_workload_identity_auth = await self._resolve_auth_header( + request.headers.get("Authorization"), + workload_identity_placeholder_required=True, + ) + if authorization is not None: + request.headers["Authorization"] = authorization response = await super()._send_request(request, stream=stream, **kwargs) if ( @@ -1063,6 +1091,28 @@ async def _send_with_auth_retry( return response + async def _resolve_auth_header( + self, authorization: str | None, *, workload_identity_placeholder_required: bool = False + ) -> tuple[str | None, bool]: + if self._workload_identity_auth is None: + return authorization, False + if ( + workload_identity_placeholder_required + and authorization != f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}" + ): + return authorization, False + return f"Bearer {await self._workload_identity_auth.get_token_async()}", True + + async def _websocket_auth_headers(self) -> dict[str, str]: + authorization, _ = await self._resolve_auth_header(self.auth_headers.get("Authorization")) + return {"Authorization": authorization} if authorization is not None else {} + + async def _retry_websocket_auth_headers(self, exc: Exception) -> dict[str, str] | None: + if self._workload_identity_auth is None or not _is_websocket_unauthorized(exc): + return None + self._workload_identity_auth.invalidate_token() + return await self._websocket_auth_headers() + @override async def _send_request( self, diff --git a/src/openai/resources/beta/realtime/realtime.py b/src/openai/resources/beta/realtime/realtime.py index 4fa35963b6..df767aa023 100644 --- a/src/openai/resources/beta/realtime/realtime.py +++ b/src/openai/resources/beta/realtime/realtime.py @@ -359,7 +359,7 @@ async def __aenter__(self) -> AsyncRealtimeConnection: extra_query = self.__extra_query await self.__client._refresh_api_key() - auth_headers = self.__client.auth_headers + auth_headers = await self.__client._websocket_auth_headers() if is_async_azure_client(self.__client): url, auth_headers = await self.__client._configure_realtime(self.__model, extra_query) else: @@ -374,8 +374,8 @@ async def __aenter__(self) -> AsyncRealtimeConnection: if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - self.__connection = AsyncRealtimeConnection( - await connect( + try: + websocket = await connect( str(url), user_agent_header=self.__client.user_agent, additional_headers=_merge_mappings( @@ -387,7 +387,19 @@ async def __aenter__(self) -> AsyncRealtimeConnection: ), **self.__websocket_connection_options, ) - ) + except Exception as exc: + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + websocket = await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings( + {**retry_auth_headers, "OpenAI-Beta": "realtime=v1"}, self.__extra_headers + ), + **self.__websocket_connection_options, + ) + self.__connection = AsyncRealtimeConnection(websocket) return self.__connection @@ -542,7 +554,7 @@ def __enter__(self) -> RealtimeConnection: extra_query = self.__extra_query self.__client._refresh_api_key() - auth_headers = self.__client.auth_headers + auth_headers = self.__client._websocket_auth_headers() if is_azure_client(self.__client): url, auth_headers = self.__client._configure_realtime(self.__model, extra_query) else: @@ -557,8 +569,8 @@ def __enter__(self) -> RealtimeConnection: if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - self.__connection = RealtimeConnection( - connect( + try: + websocket = connect( str(url), user_agent_header=self.__client.user_agent, additional_headers=_merge_mappings( @@ -570,7 +582,19 @@ def __enter__(self) -> RealtimeConnection: ), **self.__websocket_connection_options, ) - ) + except Exception as exc: + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + websocket = connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings( + {**retry_auth_headers, "OpenAI-Beta": "realtime=v1"}, self.__extra_headers + ), + **self.__websocket_connection_options, + ) + self.__connection = RealtimeConnection(websocket) return self.__connection diff --git a/src/openai/resources/beta/responses/responses.py b/src/openai/resources/beta/responses/responses.py index 5ac018f744..d71a9897fa 100644 --- a/src/openai/resources/beta/responses/responses.py +++ b/src/openai/resources/beta/responses/responses.py @@ -4381,17 +4381,24 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - return await connect( - str(url), - user_agent_header=self.__client.user_agent, - additional_headers=_merge_mappings( - { - **self.__client.auth_headers, - }, - extra_headers, - ), - **self.__websocket_connection_options, - ) + auth_headers = await self.__client._websocket_auth_headers() + try: + return await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(auth_headers, extra_headers), + **self.__websocket_connection_options, + ) + except Exception as exc: + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + return await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(retry_auth_headers, extra_headers), + **self.__websocket_connection_options, + ) def _prepare_url(self) -> httpx.URL: if self.__client.websocket_base_url is not None: @@ -4826,17 +4833,24 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - return connect( - str(url), - user_agent_header=self.__client.user_agent, - additional_headers=_merge_mappings( - { - **self.__client.auth_headers, - }, - extra_headers, - ), - **self.__websocket_connection_options, - ) + auth_headers = self.__client._websocket_auth_headers() + try: + return connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(auth_headers, extra_headers), + **self.__websocket_connection_options, + ) + except Exception as exc: + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + return connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(retry_auth_headers, extra_headers), + **self.__websocket_connection_options, + ) def _prepare_url(self) -> httpx.URL: if self.__client.websocket_base_url is not None: diff --git a/src/openai/resources/realtime/realtime.py b/src/openai/resources/realtime/realtime.py index e4c5bd8163..568957c3af 100644 --- a/src/openai/resources/realtime/realtime.py +++ b/src/openai/resources/realtime/realtime.py @@ -686,7 +686,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc await self.__client._refresh_api_key() - auth_headers = self.__client.auth_headers + auth_headers = await self.__client._websocket_auth_headers() if self.__call_id is not omit: extra_query = {**extra_query, "call_id": self.__call_id} if is_async_azure_client(self.__client): @@ -707,17 +707,28 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - return await connect( - str(url), - user_agent_header=self.__client.user_agent, - additional_headers=_merge_mappings( - { - **auth_headers, - }, - extra_headers, - ), - **self.__websocket_connection_options, - ) + try: + return await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings( + { + **auth_headers, + }, + extra_headers, + ), + **self.__websocket_connection_options, + ) + except Exception as exc: + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + return await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(retry_auth_headers, extra_headers), + **self.__websocket_connection_options, + ) def _prepare_url(self) -> httpx.URL: if self.__client.websocket_base_url is not None: @@ -1154,7 +1165,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc self.__client._refresh_api_key() - auth_headers = self.__client.auth_headers + auth_headers = self.__client._websocket_auth_headers() if self.__call_id is not omit: extra_query = {**extra_query, "call_id": self.__call_id} if is_azure_client(self.__client): @@ -1175,17 +1186,28 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - return connect( - str(url), - user_agent_header=self.__client.user_agent, - additional_headers=_merge_mappings( - { - **auth_headers, - }, - extra_headers, - ), - **self.__websocket_connection_options, - ) + try: + return connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings( + { + **auth_headers, + }, + extra_headers, + ), + **self.__websocket_connection_options, + ) + except Exception as exc: + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + return connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(retry_auth_headers, extra_headers), + **self.__websocket_connection_options, + ) def _prepare_url(self) -> httpx.URL: if self.__client.websocket_base_url is not None: diff --git a/src/openai/resources/responses/responses.py b/src/openai/resources/responses/responses.py index 8131b57ca2..de920e8194 100644 --- a/src/openai/resources/responses/responses.py +++ b/src/openai/resources/responses/responses.py @@ -4332,17 +4332,24 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - return await connect( - str(url), - user_agent_header=self.__client.user_agent, - additional_headers=_merge_mappings( - { - **self.__client.auth_headers, - }, - extra_headers, - ), - **self.__websocket_connection_options, - ) + auth_headers = await self.__client._websocket_auth_headers() + try: + return await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(auth_headers, extra_headers), + **self.__websocket_connection_options, + ) + except Exception as exc: + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + return await connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(retry_auth_headers, extra_headers), + **self.__websocket_connection_options, + ) def _prepare_url(self) -> httpx.URL: if self.__client.websocket_base_url is not None: @@ -4777,17 +4784,24 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - return connect( - str(url), - user_agent_header=self.__client.user_agent, - additional_headers=_merge_mappings( - { - **self.__client.auth_headers, - }, - extra_headers, - ), - **self.__websocket_connection_options, - ) + auth_headers = self.__client._websocket_auth_headers() + try: + return connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(auth_headers, extra_headers), + **self.__websocket_connection_options, + ) + except Exception as exc: + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + if retry_auth_headers is None: + raise + return connect( + str(url), + user_agent_header=self.__client.user_agent, + additional_headers=_merge_mappings(retry_auth_headers, extra_headers), + **self.__websocket_connection_options, + ) def _prepare_url(self) -> httpx.URL: if self.__client.websocket_base_url is not None: diff --git a/tests/test_client.py b/tests/test_client.py index 2d8955a58e..c761d906f1 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -101,6 +101,14 @@ class Counter: value: int = 0 +class MockWebSocketUnauthorized(Exception): + status_code = 401 + + +class MockWebSocketUnauthorizedWithResponse(Exception): + response = httpx.Response(401) + + def _make_sync_iterator(iterable: Iterable[T], counter: Optional[Counter] = None) -> Iterator[T]: for item in iterable: if counter: @@ -2949,3 +2957,47 @@ def provider() -> str: assert len(calls) == 2 assert provider_call_count == 1 + + +class TestWorkloadIdentityWebSocketAuth: + def test_resolves_token_and_refreshes_after_unauthorized_handshake(self) -> None: + client = OpenAI(workload_identity=workload_identity) + auth = mock.Mock() + auth.get_token.side_effect = ["openai-access-token-1", "openai-access-token-2"] + client._workload_identity_auth = auth + + assert client._websocket_auth_headers() == {"Authorization": "Bearer openai-access-token-1"} + assert client._retry_websocket_auth_headers(MockWebSocketUnauthorized()) == { + "Authorization": "Bearer openai-access-token-2" + } + auth.invalidate_token.assert_called_once_with() + + def test_does_not_refresh_after_non_unauthorized_handshake_error(self) -> None: + client = OpenAI(workload_identity=workload_identity) + auth = mock.Mock() + client._workload_identity_auth = auth + + assert client._retry_websocket_auth_headers(Exception("connection failed")) is None + auth.invalidate_token.assert_not_called() + + +class TestAsyncWorkloadIdentityWebSocketAuth: + async def test_resolves_token_and_refreshes_after_unauthorized_handshake(self) -> None: + client = AsyncOpenAI(workload_identity=workload_identity) + auth = mock.Mock() + auth.get_token_async = mock.AsyncMock(side_effect=["openai-access-token-1", "openai-access-token-2"]) + client._workload_identity_auth = auth + + assert await client._websocket_auth_headers() == {"Authorization": "Bearer openai-access-token-1"} + assert await client._retry_websocket_auth_headers(MockWebSocketUnauthorizedWithResponse()) == { + "Authorization": "Bearer openai-access-token-2" + } + auth.invalidate_token.assert_called_once_with() + + async def test_does_not_refresh_after_non_unauthorized_handshake_error(self) -> None: + client = AsyncOpenAI(workload_identity=workload_identity) + auth = mock.Mock() + client._workload_identity_auth = auth + + assert await client._retry_websocket_auth_headers(Exception("connection failed")) is None + auth.invalidate_token.assert_not_called() From edd0404230de4e09d6560e01f81267bda179a3ba Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 11 Jul 2026 06:07:58 +0200 Subject: [PATCH 2/2] fix websocket authorization overrides --- src/openai/_client.py | 28 +++++++++---- .../resources/beta/realtime/realtime.py | 8 ++-- .../resources/beta/responses/responses.py | 8 ++-- src/openai/resources/realtime/realtime.py | 8 ++-- src/openai/resources/responses/responses.py | 8 ++-- tests/test_client.py | 39 ++++++++++++++++--- 6 files changed, 69 insertions(+), 30 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index 5ebb4643c1..0092b82985 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -486,15 +486,21 @@ def _resolve_auth_header( return authorization, False return f"Bearer {self._workload_identity_auth.get_token()}", True - def _websocket_auth_headers(self) -> dict[str, str]: + def _websocket_auth_headers(self, extra_headers: Headers) -> dict[str, str]: + if _has_header(extra_headers, "Authorization"): + return {} authorization, _ = self._resolve_auth_header(self.auth_headers.get("Authorization")) return {"Authorization": authorization} if authorization is not None else {} - def _retry_websocket_auth_headers(self, exc: Exception) -> dict[str, str] | None: - if self._workload_identity_auth is None or not _is_websocket_unauthorized(exc): + def _retry_websocket_auth_headers(self, exc: Exception, extra_headers: Headers) -> dict[str, str] | None: + if ( + self._workload_identity_auth is None + or _has_header(extra_headers, "Authorization") + or not _is_websocket_unauthorized(exc) + ): return None self._workload_identity_auth.invalidate_token() - return self._websocket_auth_headers() + return self._websocket_auth_headers(extra_headers) @override def _send_request( @@ -1103,15 +1109,21 @@ async def _resolve_auth_header( return authorization, False return f"Bearer {await self._workload_identity_auth.get_token_async()}", True - async def _websocket_auth_headers(self) -> dict[str, str]: + async def _websocket_auth_headers(self, extra_headers: Headers) -> dict[str, str]: + if _has_header(extra_headers, "Authorization"): + return {} authorization, _ = await self._resolve_auth_header(self.auth_headers.get("Authorization")) return {"Authorization": authorization} if authorization is not None else {} - async def _retry_websocket_auth_headers(self, exc: Exception) -> dict[str, str] | None: - if self._workload_identity_auth is None or not _is_websocket_unauthorized(exc): + async def _retry_websocket_auth_headers(self, exc: Exception, extra_headers: Headers) -> dict[str, str] | None: + if ( + self._workload_identity_auth is None + or _has_header(extra_headers, "Authorization") + or not _is_websocket_unauthorized(exc) + ): return None self._workload_identity_auth.invalidate_token() - return await self._websocket_auth_headers() + return await self._websocket_auth_headers(extra_headers) @override async def _send_request( diff --git a/src/openai/resources/beta/realtime/realtime.py b/src/openai/resources/beta/realtime/realtime.py index df767aa023..3344c20c52 100644 --- a/src/openai/resources/beta/realtime/realtime.py +++ b/src/openai/resources/beta/realtime/realtime.py @@ -359,7 +359,7 @@ async def __aenter__(self) -> AsyncRealtimeConnection: extra_query = self.__extra_query await self.__client._refresh_api_key() - auth_headers = await self.__client._websocket_auth_headers() + auth_headers = await self.__client._websocket_auth_headers(self.__extra_headers) if is_async_azure_client(self.__client): url, auth_headers = await self.__client._configure_realtime(self.__model, extra_query) else: @@ -388,7 +388,7 @@ async def __aenter__(self) -> AsyncRealtimeConnection: **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc, self.__extra_headers) if retry_auth_headers is None: raise websocket = await connect( @@ -554,7 +554,7 @@ def __enter__(self) -> RealtimeConnection: extra_query = self.__extra_query self.__client._refresh_api_key() - auth_headers = self.__client._websocket_auth_headers() + auth_headers = self.__client._websocket_auth_headers(self.__extra_headers) if is_azure_client(self.__client): url, auth_headers = self.__client._configure_realtime(self.__model, extra_query) else: @@ -583,7 +583,7 @@ def __enter__(self) -> RealtimeConnection: **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc, self.__extra_headers) if retry_auth_headers is None: raise websocket = connect( diff --git a/src/openai/resources/beta/responses/responses.py b/src/openai/resources/beta/responses/responses.py index d71a9897fa..d2dce8269c 100644 --- a/src/openai/resources/beta/responses/responses.py +++ b/src/openai/resources/beta/responses/responses.py @@ -4381,7 +4381,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - auth_headers = await self.__client._websocket_auth_headers() + auth_headers = await self.__client._websocket_auth_headers(extra_headers) try: return await connect( str(url), @@ -4390,7 +4390,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc, extra_headers) if retry_auth_headers is None: raise return await connect( @@ -4833,7 +4833,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - auth_headers = self.__client._websocket_auth_headers() + auth_headers = self.__client._websocket_auth_headers(extra_headers) try: return connect( str(url), @@ -4842,7 +4842,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc, extra_headers) if retry_auth_headers is None: raise return connect( diff --git a/src/openai/resources/realtime/realtime.py b/src/openai/resources/realtime/realtime.py index 568957c3af..acd3a6ae45 100644 --- a/src/openai/resources/realtime/realtime.py +++ b/src/openai/resources/realtime/realtime.py @@ -686,7 +686,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc await self.__client._refresh_api_key() - auth_headers = await self.__client._websocket_auth_headers() + auth_headers = await self.__client._websocket_auth_headers(extra_headers) if self.__call_id is not omit: extra_query = {**extra_query, "call_id": self.__call_id} if is_async_azure_client(self.__client): @@ -720,7 +720,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc, extra_headers) if retry_auth_headers is None: raise return await connect( @@ -1165,7 +1165,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc self.__client._refresh_api_key() - auth_headers = self.__client._websocket_auth_headers() + auth_headers = self.__client._websocket_auth_headers(extra_headers) if self.__call_id is not omit: extra_query = {**extra_query, "call_id": self.__call_id} if is_azure_client(self.__client): @@ -1199,7 +1199,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc, extra_headers) if retry_auth_headers is None: raise return connect( diff --git a/src/openai/resources/responses/responses.py b/src/openai/resources/responses/responses.py index de920e8194..4dba10ecd6 100644 --- a/src/openai/resources/responses/responses.py +++ b/src/openai/resources/responses/responses.py @@ -4332,7 +4332,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - auth_headers = await self.__client._websocket_auth_headers() + auth_headers = await self.__client._websocket_auth_headers(extra_headers) try: return await connect( str(url), @@ -4341,7 +4341,7 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = await self.__client._retry_websocket_auth_headers(exc, extra_headers) if retry_auth_headers is None: raise return await connect( @@ -4784,7 +4784,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Connection options: %s", self.__websocket_connection_options) - auth_headers = self.__client._websocket_auth_headers() + auth_headers = self.__client._websocket_auth_headers(extra_headers) try: return connect( str(url), @@ -4793,7 +4793,7 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo **self.__websocket_connection_options, ) except Exception as exc: - retry_auth_headers = self.__client._retry_websocket_auth_headers(exc) + retry_auth_headers = self.__client._retry_websocket_auth_headers(exc, extra_headers) if retry_auth_headers is None: raise return connect( diff --git a/tests/test_client.py b/tests/test_client.py index c761d906f1..cb141fb7b6 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -2966,18 +2966,30 @@ def test_resolves_token_and_refreshes_after_unauthorized_handshake(self) -> None auth.get_token.side_effect = ["openai-access-token-1", "openai-access-token-2"] client._workload_identity_auth = auth - assert client._websocket_auth_headers() == {"Authorization": "Bearer openai-access-token-1"} - assert client._retry_websocket_auth_headers(MockWebSocketUnauthorized()) == { + assert client._websocket_auth_headers({}) == {"Authorization": "Bearer openai-access-token-1"} + assert client._retry_websocket_auth_headers(MockWebSocketUnauthorized(), {}) == { "Authorization": "Bearer openai-access-token-2" } auth.invalidate_token.assert_called_once_with() + @pytest.mark.parametrize("authorization", ["Bearer custom", Omit()]) + def test_respects_authorization_override_without_resolving_token(self, authorization: str | Omit) -> None: + client = OpenAI(workload_identity=workload_identity) + auth = mock.Mock() + client._workload_identity_auth = auth + extra_headers = {"authorization": authorization} + + assert client._websocket_auth_headers(extra_headers) == {} + assert client._retry_websocket_auth_headers(MockWebSocketUnauthorized(), extra_headers) is None + auth.get_token.assert_not_called() + auth.invalidate_token.assert_not_called() + def test_does_not_refresh_after_non_unauthorized_handshake_error(self) -> None: client = OpenAI(workload_identity=workload_identity) auth = mock.Mock() client._workload_identity_auth = auth - assert client._retry_websocket_auth_headers(Exception("connection failed")) is None + assert client._retry_websocket_auth_headers(Exception("connection failed"), {}) is None auth.invalidate_token.assert_not_called() @@ -2988,16 +3000,31 @@ async def test_resolves_token_and_refreshes_after_unauthorized_handshake(self) - auth.get_token_async = mock.AsyncMock(side_effect=["openai-access-token-1", "openai-access-token-2"]) client._workload_identity_auth = auth - assert await client._websocket_auth_headers() == {"Authorization": "Bearer openai-access-token-1"} - assert await client._retry_websocket_auth_headers(MockWebSocketUnauthorizedWithResponse()) == { + assert await client._websocket_auth_headers({}) == {"Authorization": "Bearer openai-access-token-1"} + assert await client._retry_websocket_auth_headers(MockWebSocketUnauthorizedWithResponse(), {}) == { "Authorization": "Bearer openai-access-token-2" } auth.invalidate_token.assert_called_once_with() + @pytest.mark.parametrize("authorization", ["Bearer custom", Omit()]) + async def test_respects_authorization_override_without_resolving_token(self, authorization: str | Omit) -> None: + client = AsyncOpenAI(workload_identity=workload_identity) + auth = mock.Mock() + auth.get_token_async = mock.AsyncMock() + client._workload_identity_auth = auth + extra_headers = {"authorization": authorization} + + assert await client._websocket_auth_headers(extra_headers) == {} + assert ( + await client._retry_websocket_auth_headers(MockWebSocketUnauthorizedWithResponse(), extra_headers) is None + ) + auth.get_token_async.assert_not_awaited() + auth.invalidate_token.assert_not_called() + async def test_does_not_refresh_after_non_unauthorized_handshake_error(self) -> None: client = AsyncOpenAI(workload_identity=workload_identity) auth = mock.Mock() client._workload_identity_auth = auth - assert await client._retry_websocket_auth_headers(Exception("connection failed")) is None + assert await client._retry_websocket_auth_headers(Exception("connection failed"), {}) is None auth.invalidate_token.assert_not_called()