Skip to content

Fix workload identity authentication for WebSockets #3485

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fallintoplace wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix workload identity authentication for WebSockets #3485

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9e6f44d
fix websocket workload identity authentication
fallintoplace Jul 10, 2026
edd0404
fix websocket authorization overrides
fallintoplace Jul 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 76 additions & 14 deletions src/openai/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 (
Expand All @@ -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"))

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Respect WebSocket Authorization overrides before token exchange

With a workload-identity client this resolves a workload token unconditionally, while the WebSocket managers only merge extra_headers afterward. If a caller supplies extra_headers={"Authorization": ...} or omits that header to override auth for a specific socket, that override would win later, but the workload token exchange still runs first and can fail or block the connection despite the supplied header; HTTP requests avoid this by resolving only when the merged Authorization is the workload placeholder. Consider deferring workload-token resolution until after WebSocket extra headers are considered.

Useful? React with 👍 / 👎.

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,
Expand Down Expand Up @@ -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 (
Expand All @@ -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,
Expand Down
40 changes: 32 additions & 8 deletions src/openai/resources/beta/realtime/realtime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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(
Expand All @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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(
Expand All @@ -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

Expand Down
58 changes: 36 additions & 22 deletions src/openai/resources/beta/responses/responses.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
70 changes: 46 additions & 24 deletions src/openai/resources/realtime/realtime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand All @@ -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:
Expand Down
Loading