diff --git a/src/openai/_client.py b/src/openai/_client.py index 66d03b23dd..0092b82985 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,34 @@ 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, 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, 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(extra_headers) + @override def _send_request( self, @@ -1041,13 +1076,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 +1097,34 @@ 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, 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, 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(extra_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..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 = self.__client.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: @@ -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, self.__extra_headers) + 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(self.__extra_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, self.__extra_headers) + 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..d2dce8269c 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(extra_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, extra_headers) + 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(extra_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, extra_headers) + 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..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 = self.__client.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): @@ -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, extra_headers) + 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(extra_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, extra_headers) + 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..4dba10ecd6 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(extra_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, extra_headers) + 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(extra_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, extra_headers) + 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..cb141fb7b6 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,74 @@ 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() + + @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 + 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() + + @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 + auth.invalidate_token.assert_not_called()