From d83f5cb65829085f37d886ab2ff00b20cad7c1ba Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 01:18:48 -0700 Subject: [PATCH 01/11] feat(auth): support X.509 workload identity federation --- README.md | 64 +- examples/x509_workload_identity.py | 41 ++ examples/x509_workload_identity_async.py | 42 ++ src/openai/_client.py | 50 +- src/openai/auth/__init__.py | 6 + src/openai/auth/_workload.py | 60 +- src/openai/auth/_x509.py | 250 +++++++ tests/test_client.py | 4 +- tests/test_x509_workload_identity.py | 842 +++++++++++++++++++++++ 9 files changed, 1334 insertions(+), 25 deletions(-) create mode 100644 examples/x509_workload_identity.py create mode 100644 examples/x509_workload_identity_async.py create mode 100644 src/openai/auth/_x509.py create mode 100644 tests/test_x509_workload_identity.py diff --git a/README.md b/README.md index 93f3071371..52564cdbbc 100644 --- a/README.md +++ b/README.md @@ -167,6 +167,56 @@ client = OpenAI( ) ``` +#### X.509 workload identity (mutual TLS) + +For X.509 workload identity federation, configure the client certificate and +server trust on an HTTPX2 client, then pass only the identity-provider and +service-account IDs to the SDK: + +```python +import os +import ssl + +from openai import OpenAI, DefaultHttpx2Client +from openai.auth import x509_workload_identity + +tls_context = ssl.create_default_context( + cafile=os.getenv("OPENAI_MTLS_CA_BUNDLE"), +) +tls_context.load_cert_chain( + certfile=os.environ["OPENAI_MTLS_CERTIFICATE_CHAIN"], + keyfile=os.environ["OPENAI_MTLS_PRIVATE_KEY"], + password=os.getenv("OPENAI_MTLS_PRIVATE_KEY_PASSWORD"), +) + +client = OpenAI( + workload_identity=x509_workload_identity( + identity_provider_id=os.environ["OPENAI_IDENTITY_PROVIDER_ID"], + service_account_id=os.environ["OPENAI_SERVICE_ACCOUNT_ID"], + # refresh_buffer_seconds=120.0, + ), + http_client=DefaultHttpx2Client( + verify=tls_context, + follow_redirects=False, + ), +) +``` + +X.509 mode defaults to `https://mtls.api.openai.com/v1` when neither `base_url` +nor `OPENAI_BASE_URL` is set. The same configured HTTP client presents its +certificate to the fixed mTLS token-exchange endpoint and to the API. Tokens +are exchanged lazily, cached, and refreshed automatically. Certificate files, +private keys, passwords, server trust, proxies, and rotation remain application +and transport concerns. + +For asynchronous requests, use `AsyncOpenAI` with +`DefaultAsyncHttpx2Client`. See the complete [sync rollout-toggle +example](examples/x509_workload_identity.py) and [async rollout-toggle +example](examples/x509_workload_identity_async.py), which select API-key or +X.509 authentication with the application-owned `OPENAI_AUTH_MODE` +environment variable. X.509 workload identity currently supports HTTP APIs; +Realtime and WebSockets are not included. + ### Vision With an image URL: @@ -950,9 +1000,11 @@ client = AsyncOpenAI( See the complete [sync HTTPX2](examples/mtls_httpx2.py) and [async HTTPX2](examples/mtls_httpx2_async.py) examples. -The certificate-bearing HTTP client is transport-wide. Dedicate it to the -selected mTLS origin; do not reuse it for other services or pass it through -`with_options()` with a different `base_url`. If redirects are required, add an +The certificate-bearing HTTP client is transport-wide. For API-key mTLS, +dedicate it to the selected API origin; X.509 workload identity also uses the +fixed OpenAI mTLS token-exchange origin. Do not reuse the client for unrelated +services or pass it through `with_options()` with a different `base_url`. +If redirects are required for API-key mTLS, add an HTTPX2 request hook that rejects requests whose scheme, host, or port differs from the configured mTLS origin before enabling `follow_redirects`. @@ -969,9 +1021,9 @@ For certificate rotation, build a new `SSLContext`, HTTP client, and `OpenAI` or client after its in-flight requests finish. Do not assume existing TLS connections will renegotiate. -This recipe applies to ordinary API-key HTTP traffic. It does not implement -certificate-only X.509 workload identity, token exchange, or Realtime WebSocket -mTLS. +This recipe applies to ordinary API-key HTTP traffic. For certificate-backed +token exchange, use the X.509 workload identity configuration described above. +Realtime WebSocket mTLS is not included. ### Managing HTTP resources diff --git a/examples/x509_workload_identity.py b/examples/x509_workload_identity.py new file mode 100644 index 0000000000..abf2d38697 --- /dev/null +++ b/examples/x509_workload_identity.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import os +import ssl + +from openai import OpenAI, DefaultHttpx2Client +from openai.auth import x509_workload_identity + + +def create_client() -> OpenAI: + mode = os.getenv("OPENAI_AUTH_MODE", "api_key") + if mode == "api_key": + return OpenAI() + if mode != "x509": + raise ValueError("OPENAI_AUTH_MODE must be 'api_key' or 'x509'") + + tls_context = ssl.create_default_context(cafile=os.getenv("OPENAI_MTLS_CA_BUNDLE")) + tls_context.load_cert_chain( + certfile=os.environ["OPENAI_MTLS_CERTIFICATE_CHAIN"], + keyfile=os.environ["OPENAI_MTLS_PRIVATE_KEY"], + password=os.getenv("OPENAI_MTLS_PRIVATE_KEY_PASSWORD"), + ) + + return OpenAI( + workload_identity=x509_workload_identity( + identity_provider_id=os.environ["OPENAI_IDENTITY_PROVIDER_ID"], + service_account_id=os.environ["OPENAI_SERVICE_ACCOUNT_ID"], + ), + base_url=os.getenv("OPENAI_BASE_URL"), + http_client=DefaultHttpx2Client(verify=tls_context, follow_redirects=False), + ) + + +def main() -> None: + with create_client() as client: + response = client.responses.create(model="gpt-5.5", input="Hello!") + print(response.output_text) + + +if __name__ == "__main__": + main() diff --git a/examples/x509_workload_identity_async.py b/examples/x509_workload_identity_async.py new file mode 100644 index 0000000000..da5529a4d8 --- /dev/null +++ b/examples/x509_workload_identity_async.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import os +import ssl +import asyncio + +from openai import AsyncOpenAI, DefaultAsyncHttpx2Client +from openai.auth import x509_workload_identity + + +def create_client() -> AsyncOpenAI: + mode = os.getenv("OPENAI_AUTH_MODE", "api_key") + if mode == "api_key": + return AsyncOpenAI() + if mode != "x509": + raise ValueError("OPENAI_AUTH_MODE must be 'api_key' or 'x509'") + + tls_context = ssl.create_default_context(cafile=os.getenv("OPENAI_MTLS_CA_BUNDLE")) + tls_context.load_cert_chain( + certfile=os.environ["OPENAI_MTLS_CERTIFICATE_CHAIN"], + keyfile=os.environ["OPENAI_MTLS_PRIVATE_KEY"], + password=os.getenv("OPENAI_MTLS_PRIVATE_KEY_PASSWORD"), + ) + + return AsyncOpenAI( + workload_identity=x509_workload_identity( + identity_provider_id=os.environ["OPENAI_IDENTITY_PROVIDER_ID"], + service_account_id=os.environ["OPENAI_SERVICE_ACCOUNT_ID"], + ), + base_url=os.getenv("OPENAI_BASE_URL"), + http_client=DefaultAsyncHttpx2Client(verify=tls_context, follow_redirects=False), + ) + + +async def main() -> None: + async with create_client() as client: + response = await client.responses.create(model="gpt-5.5", input="Hello!") + print(response.output_text) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/openai/_client.py b/src/openai/_client.py index 5f980c8cb6..b72a7b696e 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -34,6 +34,12 @@ from ._version import __version__ from ._provider import _Provider, _provider_name, _ProviderRuntime, _configure_provider from ._streaming import Stream as Stream, AsyncStream as AsyncStream +from .auth._x509 import ( + MTLS_API_BASE_URL, + SyncX509WorkloadIdentityAuth, + AsyncX509WorkloadIdentityAuth, + is_x509_workload_identity, +) from ._exceptions import OpenAIError, APIStatusError from ._base_client import ( DEFAULT_MAX_RETRIES, @@ -246,12 +252,13 @@ def __init__( self.websocket_base_url = websocket_base_url + x509_identity = workload_identity if is_x509_workload_identity(workload_identity) else None if provider_runtime is not None: base_url = provider_runtime.base_url elif base_url is None: base_url = os.environ.get("OPENAI_BASE_URL") if base_url is None: - base_url = f"https://api.openai.com/v1" + base_url = MTLS_API_BASE_URL if x509_identity is not None else "https://api.openai.com/v1" custom_headers_env = os.environ.get("OPENAI_CUSTOM_HEADERS") if provider_runtime is None else None if custom_headers_env is not None: @@ -274,10 +281,15 @@ def __init__( ) if workload_identity is not None: - self._workload_identity_auth = WorkloadIdentityAuth( - workload_identity=workload_identity, - _use_httpx2=is_httpx2_sync_client(self._client), - ) + if x509_identity is not None: + self._workload_identity_auth = SyncX509WorkloadIdentityAuth( + workload_identity=x509_identity, http_client=self._client, max_retries=max_retries + ) + else: + self._workload_identity_auth = WorkloadIdentityAuth( + workload_identity=workload_identity, + _use_httpx2=is_httpx2_sync_client(self._client), + ) self._default_stream_cls = Stream @@ -459,12 +471,16 @@ def _send_with_auth_retry( **kwargs: Unpack[HttpxSendArgs], ) -> httpx2.Response: used_workload_identity_auth = False + request_is_replayable = 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 + request_is_replayable = self._workload_identity_auth._can_retry_request(request) + if self._workload_identity_auth._follow_redirects is not None: + kwargs["follow_redirects"] = self._workload_identity_auth._follow_redirects response = super()._send_request(request, stream=stream, **kwargs) if ( @@ -472,8 +488,10 @@ def _send_with_auth_retry( and self._workload_identity_auth is not None and used_workload_identity_auth and not retried + and request_is_replayable ): response.close() + self._workload_identity_auth._prepare_retry_request(request) self._workload_identity_auth.invalidate_token() request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @@ -852,12 +870,13 @@ def __init__( self.websocket_base_url = websocket_base_url + x509_identity = workload_identity if is_x509_workload_identity(workload_identity) else None if provider_runtime is not None: base_url = provider_runtime.base_url elif base_url is None: base_url = os.environ.get("OPENAI_BASE_URL") if base_url is None: - base_url = f"https://api.openai.com/v1" + base_url = MTLS_API_BASE_URL if x509_identity is not None else "https://api.openai.com/v1" custom_headers_env = os.environ.get("OPENAI_CUSTOM_HEADERS") if provider_runtime is None else None if custom_headers_env is not None: @@ -880,10 +899,15 @@ def __init__( ) if workload_identity is not None: - self._workload_identity_auth = WorkloadIdentityAuth( - workload_identity=workload_identity, - _use_httpx2=is_httpx2_async_client(self._client), - ) + if x509_identity is not None: + self._workload_identity_auth = AsyncX509WorkloadIdentityAuth( + workload_identity=x509_identity, http_client=self._client, max_retries=max_retries + ) + else: + self._workload_identity_auth = WorkloadIdentityAuth( + workload_identity=workload_identity, + _use_httpx2=is_httpx2_async_client(self._client), + ) self._default_stream_cls = AsyncStream @@ -1065,12 +1089,16 @@ async def _send_with_auth_retry( **kwargs: Unpack[HttpxSendArgs], ) -> httpx2.Response: used_workload_identity_auth = False + request_is_replayable = 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 + request_is_replayable = self._workload_identity_auth._can_retry_request(request) + if self._workload_identity_auth._follow_redirects is not None: + kwargs["follow_redirects"] = self._workload_identity_auth._follow_redirects response = await super()._send_request(request, stream=stream, **kwargs) if ( @@ -1078,8 +1106,10 @@ async def _send_with_auth_retry( and self._workload_identity_auth is not None and used_workload_identity_auth and not retried + and request_is_replayable ): await response.aclose() + self._workload_identity_auth._prepare_retry_request(request) self._workload_identity_auth.invalidate_token() request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) diff --git a/src/openai/auth/__init__.py b/src/openai/auth/__init__.py index 367aa86b72..fbe72c8378 100644 --- a/src/openai/auth/__init__.py +++ b/src/openai/auth/__init__.py @@ -4,7 +4,10 @@ WorkloadIdentity as WorkloadIdentity, SubjectTokenProvider as SubjectTokenProvider, WorkloadIdentityAuth as WorkloadIdentityAuth, + X509WorkloadIdentity as X509WorkloadIdentity, + SubjectTokenWorkloadIdentity as SubjectTokenWorkloadIdentity, gcp_id_token_provider as gcp_id_token_provider, + x509_workload_identity as x509_workload_identity, k8s_service_account_token_provider as k8s_service_account_token_provider, azure_managed_identity_token_provider as azure_managed_identity_token_provider, ) @@ -12,7 +15,10 @@ __all__ = [ "SubjectTokenProvider", "WorkloadIdentity", + "SubjectTokenWorkloadIdentity", + "X509WorkloadIdentity", "WorkloadIdentityAuth", + "x509_workload_identity", "k8s_service_account_token_provider", "azure_managed_identity_token_provider", "gcp_id_token_provider", diff --git a/src/openai/auth/_workload.py b/src/openai/auth/_workload.py index 73b110b445..4f3d6ab67a 100644 --- a/src/openai/auth/_workload.py +++ b/src/openai/auth/_workload.py @@ -4,7 +4,7 @@ import threading from typing import Any, Callable, TypedDict, cast from pathlib import Path -from typing_extensions import Literal, NotRequired +from typing_extensions import Literal, TypeAlias, NotRequired import httpx2 @@ -27,7 +27,7 @@ class SubjectTokenProvider(TypedDict): get_token: Callable[[], str] -class WorkloadIdentity(TypedDict): +class SubjectTokenWorkloadIdentity(TypedDict): """Identity provider resource id in WIFAPI.""" identity_provider_id: str @@ -42,6 +42,35 @@ class WorkloadIdentity(TypedDict): refresh_buffer_seconds: NotRequired[float] +class X509WorkloadIdentity(TypedDict): + """Authenticate with the client certificate configured on the HTTP transport.""" + + type: Literal["x509"] + identity_provider_id: str + service_account_id: str + refresh_buffer_seconds: NotRequired[float] + + +WorkloadIdentity: TypeAlias = SubjectTokenWorkloadIdentity | X509WorkloadIdentity + + +def x509_workload_identity( + *, + identity_provider_id: str, + service_account_id: str, + refresh_buffer_seconds: float | None = None, +) -> X509WorkloadIdentity: + """Configure X.509 workload identity without handling certificate material.""" + identity: X509WorkloadIdentity = { + "type": "x509", + "identity_provider_id": identity_provider_id, + "service_account_id": service_account_id, + } + if refresh_buffer_seconds is not None: + identity["refresh_buffer_seconds"] = refresh_buffer_seconds + return identity + + def k8s_service_account_token_provider( token_file_path: str | Path = "/var/run/secrets/kubernetes.io/serviceaccount/token", ) -> SubjectTokenProvider: @@ -182,6 +211,7 @@ def __init__( self.workload_identity = workload_identity self.token_exchange_url = token_exchange_url self._use_httpx2 = _use_httpx2 + self._follow_redirects: bool | None = None self._cached_token: str | None = None self._cached_token_expires_at_monotonic: float | None = None @@ -230,6 +260,9 @@ def invalidate_token(self) -> None: def _perform_refresh(self) -> None: token_data = self._fetch_token_from_exchange() + self._store_token(token_data) + + def _store_token(self, token_data: dict[str, Any]) -> None: now = time.monotonic() expires_in = token_data["expires_in"] @@ -241,7 +274,8 @@ def _perform_refresh(self) -> None: def _fetch_token_from_exchange(self) -> dict[str, Any]: subject_token = self._get_subject_token() - token_type = self.workload_identity["provider"]["token_type"] + identity = cast(SubjectTokenWorkloadIdentity, self.workload_identity) + token_type = identity["provider"]["token_type"] subject_token_type = SUBJECT_TOKEN_TYPES.get(token_type) if subject_token_type is None: raise OpenAIError( @@ -282,16 +316,19 @@ def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]: expires_in = body.get("expires_in") if not isinstance(access_token, str) or not access_token: raise OpenAIError("Token exchange response did not include a valid access_token") - if not isinstance(expires_in, (int, float)): - raise OpenAIError("Token exchange response did not include a valid expires_in") - return {"access_token": access_token, "expires_in": float(expires_in)} + return {"access_token": access_token, "expires_in": self._validate_expires_in(expires_in)} raise OpenAIError( f"Token exchange failed with status {response.status_code}", ) + def _validate_expires_in(self, expires_in: object) -> float: + if not isinstance(expires_in, (int, float)): + raise OpenAIError("Token exchange response did not include a valid expires_in") + return float(expires_in) + def _get_subject_token(self) -> str: - provider = self.workload_identity["provider"] + provider = cast(SubjectTokenWorkloadIdentity, self.workload_identity)["provider"] subject_token = provider["get_token"]() if not subject_token: raise OpenAIError("The workload identity provider returned an empty subject token") @@ -314,3 +351,12 @@ def _refresh_delay_seconds(self, expires_in: float) -> float: configured_buffer = self.workload_identity.get("refresh_buffer_seconds", DEFAULT_REFRESH_BUFFER_SECONDS) effective_buffer = min(configured_buffer, expires_in / 2) return max(expires_in - effective_buffer, 0.0) + + def _can_retry_request(self, request: httpx2.Request) -> bool: + """Preserve the established subject-token request retry behavior.""" + del request + return True + + def _prepare_retry_request(self, request: httpx2.Request) -> None: + """Preserve the established subject-token request retry behavior.""" + del request diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py new file mode 100644 index 0000000000..3c82c6ee81 --- /dev/null +++ b/src/openai/auth/_x509.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import math +import time +import email.utils +from typing import Any, NoReturn, cast +from typing_extensions import TypeGuard, override + +import anyio +import httpx2 + +from .._utils import is_dict +from .._httpx2 import timeout_exceptions, _loaded_legacy_httpx +from ._workload import ( + TOKEN_EXCHANGE_GRANT_TYPE, + WorkloadIdentity, + WorkloadIdentityAuth, + X509WorkloadIdentity, +) +from .._constants import MAX_RETRY_DELAY, INITIAL_RETRY_DELAY, MAX_RETRY_AFTER_DELAY +from .._exceptions import OAuthError, OpenAIError, APITimeoutError, APIConnectionError + +MTLS_API_BASE_URL = "https://mtls.api.openai.com/v1" +_X509_TOKEN_EXCHANGE_URL = "https://mtls.auth.openai.com/oauth/token" +_X509_SUBJECT_TOKEN_TYPE = "urn:openai:params:oauth:token-type:x509" +_MAX_EXCHANGE_RETRIES = 2 +_REPLAY_POSITION_EXTENSION = "openai_x509_replay_position" +_ALLOWED_IDENTITY_FIELDS = {"type", "identity_provider_id", "service_account_id", "refresh_buffer_seconds"} + + +def is_x509_workload_identity(identity: WorkloadIdentity | None) -> TypeGuard[X509WorkloadIdentity]: + return identity is not None and identity.get("type") == "x509" + + +def _validate_identity(identity: X509WorkloadIdentity) -> None: + if "provider" in identity or "client_id" in identity: + raise OpenAIError("X.509 workload identity does not accept a subject-token provider or client ID") + + if set(identity) - _ALLOWED_IDENTITY_FIELDS: + raise OpenAIError("X.509 workload identity accepts only identity IDs and an optional refresh buffer") + + if not identity["identity_provider_id"] or not identity["service_account_id"]: + raise OpenAIError("X.509 workload identity requires identity-provider and service-account IDs") + + refresh_buffer = cast(object, identity.get("refresh_buffer_seconds")) + if refresh_buffer is not None and ( + isinstance(refresh_buffer, bool) + or not isinstance(refresh_buffer, (int, float)) + or not math.isfinite(refresh_buffer) + or refresh_buffer < 0 + ): + raise OpenAIError("X.509 workload identity requires a finite, non-negative refresh buffer") + + +def _exchange_payload(identity: X509WorkloadIdentity) -> dict[str, str]: + return { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "subject_token_type": _X509_SUBJECT_TOKEN_TYPE, + "identity_provider_id": identity["identity_provider_id"], + "service_account_id": identity["service_account_id"], + } + + +def _retry_delay(response: httpx2.Response | None, attempt: int) -> float | None: + if response is not None: + if response.status_code not in (408, 409, 429) and response.status_code < 500: + return None + + retry_after = response.headers.get("retry-after") + if retry_after is not None: + try: + delay = float(retry_after) + except ValueError: + try: + parsed = email.utils.parsedate_tz(retry_after) + delay = float(email.utils.mktime_tz(parsed) - time.time()) if parsed is not None else -1 + except (OverflowError, OSError, ValueError): + delay = -1 + + if math.isfinite(delay) and 0 <= delay <= MAX_RETRY_AFTER_DELAY: + return delay + if delay > MAX_RETRY_AFTER_DELAY: + return None + + return float(min(INITIAL_RETRY_DELAY * 2**attempt, MAX_RETRY_DELAY)) + + +def _is_replayable_request(request: httpx2.Request) -> bool: + if isinstance(getattr(request, "_content", None), bytes): + return True + + stream = request.stream + fields = getattr(stream, "fields", None) + if isinstance(fields, list): + for field in cast(list[object], fields): + file = getattr(field, "file", None) + if file is None or isinstance(file, (str, bytes)): + continue + seekable = getattr(file, "seekable", None) + if not callable(seekable) or not seekable(): + return False + return True + + source = getattr(stream, "_stream", stream) + seekable = getattr(source, "seekable", None) + seek = getattr(source, "seek", None) + tell = getattr(source, "tell", None) + if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell): + return False + request.extensions[_REPLAY_POSITION_EXTENSION] = tell() + return True + + +def _transport_errors() -> tuple[type[Exception], ...]: + legacy_httpx = _loaded_legacy_httpx() + if legacy_httpx is None: + return (httpx2.TransportError,) + legacy_transport_error = cast(type[Exception], getattr(legacy_httpx, "TransportError", httpx2.TransportError)) + return (httpx2.TransportError, legacy_transport_error) + + +def _raise_transport_error(error: Exception) -> NoReturn: + request = cast(httpx2.Request | None, getattr(error, "request", None)) + if request is None: + raise OpenAIError("X.509 token exchange connection failed") from error + if isinstance(error, timeout_exceptions()): + raise APITimeoutError(request=request) from error + raise APIConnectionError(request=request) from error + + +class _X509WorkloadIdentityAuth(WorkloadIdentityAuth): + def __init__(self, *, workload_identity: X509WorkloadIdentity, max_retries: int) -> None: + _validate_identity(workload_identity) + super().__init__(workload_identity=workload_identity, token_exchange_url=_X509_TOKEN_EXCHANGE_URL) + self._max_exchange_retries = min(max(max_retries, 0), _MAX_EXCHANGE_RETRIES) + self._follow_redirects = False + + @override + def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]: + if response.status_code not in (400, 401, 403): + return super()._handle_token_response(response) + + try: + response_body = response.json() if response.content else None + except ValueError: + response_body = None + + oauth_error = response_body.get("error") if is_dict(response_body) else None + safe_body = {"error": oauth_error} if isinstance(oauth_error, str) else None + raise OAuthError(response=response, body=safe_body) + + @override + def _validate_expires_in(self, expires_in: object) -> float: + if isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)): + raise OpenAIError("X.509 token exchange response did not include a positive, finite expires_in") + if not math.isfinite(expires_in) or expires_in <= 0: + raise OpenAIError("X.509 token exchange response did not include a positive, finite expires_in") + return float(expires_in) + + @override + def _can_retry_request(self, request: httpx2.Request) -> bool: + return _is_replayable_request(request) + + @override + def _prepare_retry_request(self, request: httpx2.Request) -> None: + position = request.extensions.get(_REPLAY_POSITION_EXTENSION) + if not isinstance(position, int): + return + source = getattr(request.stream, "_stream", request.stream) + seek = getattr(source, "seek", None) + if callable(seek): + seek(position) + + +class SyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth): + def __init__( + self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.Client, max_retries: int + ) -> None: + super().__init__(workload_identity=workload_identity, max_retries=max_retries) + self._http_client = http_client + + @override + def _fetch_token_from_exchange(self) -> dict[str, Any]: + identity = cast(X509WorkloadIdentity, self.workload_identity) + for attempt in range(self._max_exchange_retries + 1): + try: + response = self._http_client.post( + _X509_TOKEN_EXCHANGE_URL, + json=_exchange_payload(identity), + timeout=10.0, + follow_redirects=False, + ) + except _transport_errors() as error: + if attempt >= self._max_exchange_retries: + _raise_transport_error(error) + delay = _retry_delay(None, attempt) + else: + delay = _retry_delay(response, attempt) + if attempt >= self._max_exchange_retries or delay is None: + return self._handle_token_response(response) + + if delay is not None: + time.sleep(delay) + + raise AssertionError("X.509 token exchange retry loop exhausted unexpectedly") + + +class AsyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth): + def __init__( + self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.AsyncClient, max_retries: int + ) -> None: + super().__init__(workload_identity=workload_identity, max_retries=max_retries) + self._http_client = http_client + self._async_lock = anyio.Lock() + + @override + async def get_token_async(self) -> str: + async with self._async_lock: + with self._lock: + if not self._token_unusable() and not self._needs_refresh(): + return cast(str, self._cached_token) + + token_data = await self._fetch_token_from_exchange_async() + self._store_token(token_data) + with self._lock: + return cast(str, self._cached_token) + + async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: + identity = cast(X509WorkloadIdentity, self.workload_identity) + for attempt in range(self._max_exchange_retries + 1): + try: + response = await self._http_client.post( + _X509_TOKEN_EXCHANGE_URL, + json=_exchange_payload(identity), + timeout=10.0, + follow_redirects=False, + ) + except _transport_errors() as error: + if attempt >= self._max_exchange_retries: + _raise_transport_error(error) + delay = _retry_delay(None, attempt) + else: + delay = _retry_delay(response, attempt) + if attempt >= self._max_exchange_retries or delay is None: + return self._handle_token_response(response) + + if delay is not None: + await anyio.sleep(delay) + + raise AssertionError("X.509 token exchange retry loop exhausted unexpectedly") diff --git a/tests/test_client.py b/tests/test_client.py index 79a66098fa..d82c39e616 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -17,7 +17,7 @@ from pydantic import ValidationError from openai import OpenAI, AsyncOpenAI, OpenAIError, APIResponseValidationError -from openai.auth import WorkloadIdentity +from openai.auth import SubjectTokenWorkloadIdentity from tests.respx2 import MockRouter from openai._types import Omit from openai._utils import asyncify @@ -42,7 +42,7 @@ base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") api_key = "My API Key" admin_api_key = "My Admin API Key" -workload_identity: WorkloadIdentity = { +workload_identity: SubjectTokenWorkloadIdentity = { "identity_provider_id": "provider_123", "service_account_id": "service_account_123", "provider": { diff --git a/tests/test_x509_workload_identity.py b/tests/test_x509_workload_identity.py new file mode 100644 index 0000000000..9521b807bc --- /dev/null +++ b/tests/test_x509_workload_identity.py @@ -0,0 +1,842 @@ +from __future__ import annotations + +import io +import os +import ssl +import json +import runpy +import asyncio +import inspect +import logging +import importlib +import threading +from typing import Any, Callable, Iterable, Iterator, AsyncIterator, cast +from pathlib import Path +from typing_extensions import override +from concurrent.futures import ThreadPoolExecutor + +import anyio +import httpx2 +import pytest + +import openai.auth._x509 as x509_auth +from openai import OpenAI, OAuthError, AsyncOpenAI, OpenAIError, APIStatusError, APITimeoutError, APIConnectionError +from openai.auth import X509WorkloadIdentity, x509_workload_identity +from openai._client import WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER + +_TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" +_API_URL = "https://mtls.api.openai.com/v1/models" +_MTLS_FIXTURES = Path(__file__).parent / "fixtures" / "mtls" + + +def _identity(**kwargs: Any) -> X509WorkloadIdentity: + return x509_workload_identity( + identity_provider_id="idp_123", + service_account_id="svc_acct_123", + **kwargs, + ) + + +def _token_response( + request: httpx2.Request, *, token: str = "access-token", expires_in: object = 3600 +) -> httpx2.Response: + return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": expires_in}) + + +def _models_response(request: httpx2.Request, status_code: int = 200) -> httpx2.Response: + return httpx2.Response(status_code, request=request, json={"object": "list", "data": []}) + + +def test_x509_helper_returns_only_typed_identity_configuration() -> None: + assert _identity() == { + "type": "x509", + "identity_provider_id": "idp_123", + "service_account_id": "svc_acct_123", + } + assert _identity(refresh_buffer_seconds=45.0).get("refresh_buffer_seconds") == 45.0 + assert "token_exchange_url" not in inspect.signature(x509_workload_identity).parameters + + +@pytest.mark.parametrize("invalid_key", ["provider", "client_id"]) +def test_x509_rejects_subject_token_configuration(invalid_key: str) -> None: + identity = cast(X509WorkloadIdentity, {**_identity(), invalid_key: "not-allowed"}) + with pytest.raises(OpenAIError, match="does not accept a subject-token provider or client ID"): + OpenAI(workload_identity=identity) + + +@pytest.mark.parametrize( + "invalid_key", + ["certificate", "certificate_chain", "private_key", "password", "subject_token", "token_exchange_url"], +) +def test_x509_rejects_certificate_and_token_material_without_leaking_it(invalid_key: str) -> None: + secret = "certificate-or-token-material-never-visible" + identity = cast(X509WorkloadIdentity, {**_identity(), invalid_key: secret}) + with pytest.raises(OpenAIError, match="only identity IDs and an optional refresh buffer") as error: + OpenAI(workload_identity=identity) + assert secret not in str(error.value) + + +@pytest.mark.parametrize("refresh_buffer", [-1.0, float("inf"), float("nan"), True]) +def test_x509_rejects_invalid_refresh_buffer(refresh_buffer: float) -> None: + with pytest.raises(OpenAIError, match="finite, non-negative refresh buffer"): + OpenAI(workload_identity=_identity(refresh_buffer_seconds=refresh_buffer)) + + +def test_sync_x509_uses_one_transport_pinned_endpoint_and_cached_token() -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return _token_response(request) + return _models_response(request) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + client = OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + assert requests == [] + assert str(client.base_url) == "https://mtls.api.openai.com/v1/" + + assert client.models.list().object == "list" + assert client.models.list().object == "list" + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL, _API_URL] + assert json.loads(requests[0].content) == { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:openai:params:oauth:token-type:x509", + "identity_provider_id": "idp_123", + "service_account_id": "svc_acct_123", + } + assert "subject_token" not in json.loads(requests[0].content) + assert [request.headers["authorization"] for request in requests[1:]] == ["Bearer access-token"] * 2 + assert not http_client.is_closed + client.close() + + +def test_sync_x509_does_not_mutate_or_independently_close_caller_transport() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return _token_response(request) if str(request.url) == _TOKEN_URL else _models_response(request) + + tls_context = ssl.create_default_context(cafile=_MTLS_FIXTURES / "root.pem") + tls_context.load_cert_chain(_MTLS_FIXTURES / "client-chain.pem", _MTLS_FIXTURES / "client.key") + transport = httpx2.MockTransport(handler) + http_client = httpx2.Client(transport=transport, verify=tls_context, follow_redirects=True, trust_env=False) + initial_timeout = http_client.timeout + initial_headers = dict(http_client.headers) + + client = OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + assert client.models.list().object == "list" + assert http_client._transport is transport + assert http_client.follow_redirects is True + assert http_client.timeout == initial_timeout + assert dict(http_client.headers) == initial_headers + assert not http_client.is_closed + auth = client._workload_identity_auth + assert auth is not None + assert all(not isinstance(value, ssl.SSLContext) for value in vars(auth).values()) + + client.close() + assert http_client.is_closed + + +async def test_async_x509_uses_one_transport_without_threaded_exchange(monkeypatch: pytest.MonkeyPatch) -> None: + requests: list[httpx2.Request] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return _token_response(request) + return _models_response(request) + + async def reject_thread(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("X.509 exchange must not use a worker thread") + + monkeypatch.setattr("openai.auth._workload.to_thread", reject_thread) + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert requests == [] + assert str(client.base_url) == "https://mtls.api.openai.com/v1/" + assert (await client.models.list()).object == "list" + assert (await client.models.list()).object == "list" + assert not http_client.is_closed + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL, _API_URL] + assert "subject_token" not in json.loads(requests[0].content) + + +async def test_async_x509_does_not_mutate_or_independently_close_caller_transport() -> None: + async def handler(request: httpx2.Request) -> httpx2.Response: + return _token_response(request) if str(request.url) == _TOKEN_URL else _models_response(request) + + tls_context = ssl.create_default_context(cafile=_MTLS_FIXTURES / "root.pem") + tls_context.load_cert_chain(_MTLS_FIXTURES / "client-chain.pem", _MTLS_FIXTURES / "client.key") + transport = httpx2.MockTransport(handler) + http_client = httpx2.AsyncClient(transport=transport, verify=tls_context, follow_redirects=True, trust_env=False) + initial_timeout = http_client.timeout + initial_headers = dict(http_client.headers) + + client = AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + assert (await client.models.list()).object == "list" + assert http_client._transport is transport + assert http_client.follow_redirects is True + assert http_client.timeout == initial_timeout + assert dict(http_client.headers) == initial_headers + assert not http_client.is_closed + auth = client._workload_identity_auth + assert auth is not None + assert all(not isinstance(value, ssl.SSLContext) for value in vars(auth).values()) + + await client.close() + assert http_client.is_closed + + +@pytest.mark.parametrize("base_url", ["https://custom.example/v1", "https://eu.api.openai.com/v1"]) +def test_x509_preserves_explicit_base_url(base_url: str) -> None: + client = OpenAI(workload_identity=_identity(), base_url=base_url) + assert str(client.base_url) == f"{base_url}/" + client.close() + + +def test_x509_preserves_environment_base_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://custom.example/v1") + client = OpenAI(workload_identity=_identity()) + assert str(client.base_url) == "https://custom.example/v1/" + client.close() + + +def test_api_key_clients_keep_ordinary_api_endpoint() -> None: + with OpenAI(api_key="ordinary-api-key") as client: + assert str(client.base_url) == "https://api.openai.com/v1/" + + +@pytest.mark.parametrize("expires_in", [0, -1, True, "3600", None]) +def test_x509_rejects_nonpositive_or_nonnumeric_expiration(expires_in: object) -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return _token_response(request, expires_in=expires_in) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="expires_in"): + client.models.list() + + +def test_x509_short_lived_token_clamps_refresh_buffer_to_half_ttl() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return _token_response(request, expires_in=8) if str(request.url) == _TOKEN_URL else _models_response(request) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(refresh_buffer_seconds=1200), http_client=http_client) as client: + client.models.list() + auth = client._workload_identity_auth + assert auth is not None + assert auth._cached_token_expires_at_monotonic is not None + assert auth._cached_token_refresh_at_monotonic is not None + assert abs(auth._cached_token_expires_at_monotonic - auth._cached_token_refresh_at_monotonic - 4.0) < 0.001 + + +@pytest.mark.parametrize("deadline", ["_cached_token_expires_at_monotonic", "_cached_token_refresh_at_monotonic"]) +def test_sync_x509_refreshes_expired_or_proactively_stale_token(deadline: str) -> None: + token_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal token_count + if str(request.url) == _TOKEN_URL: + token_count += 1 + return _token_response(request, token=f"token-{token_count}") + return _models_response(request) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + client.models.list() + auth = client._workload_identity_auth + assert auth is not None + setattr(auth, deadline, 0.0) + client.models.list() + + assert token_count == 2 + + +@pytest.mark.parametrize("deadline", ["_cached_token_expires_at_monotonic", "_cached_token_refresh_at_monotonic"]) +async def test_async_x509_refreshes_expired_or_proactively_stale_token(deadline: str) -> None: + token_count = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal token_count + if str(request.url) == _TOKEN_URL: + token_count += 1 + return _token_response(request, token=f"token-{token_count}") + return _models_response(request) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + await client.models.list() + auth = client._workload_identity_auth + assert auth is not None + setattr(auth, deadline, 0.0) + await client.models.list() + + assert token_count == 2 + + +def test_sync_x509_concurrent_requests_share_one_exchange() -> None: + exchange_calls = 0 + lock = threading.Lock() + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + with lock: + exchange_calls += 1 + threading.Event().wait(0.03) + return _token_response(request) + return _models_response(request) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + + def list_models(_: int) -> str: + return client.models.list().object + + with ThreadPoolExecutor(max_workers=8) as executor: + assert list(executor.map(list_models, range(8))) == ["list"] * 8 + + assert exchange_calls == 1 + + +async def test_async_x509_concurrent_requests_share_one_exchange() -> None: + exchange_calls = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + await anyio.sleep(0.03) + return _token_response(request) + return _models_response(request) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + responses = await asyncio.gather(*(client.models.list() for _ in range(8))) + + assert [response.object for response in responses] == ["list"] * 8 + assert exchange_calls == 1 + + +async def test_async_x509_cancelled_waiter_does_not_cancel_shared_refresh() -> None: + exchange_started = anyio.Event() + finish_exchange = anyio.Event() + exchange_calls = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + exchange_started.set() + await finish_exchange.wait() + return _token_response(request) + return _models_response(request) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + + async def list_models() -> str: + return (await client.models.list()).object + + owner = asyncio.create_task(list_models()) + await exchange_started.wait() + waiter = asyncio.create_task(list_models()) + await anyio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + finish_exchange.set() + assert await owner == "list" + assert (await client.models.list()).object == "list" + + assert exchange_calls == 1 + + +async def test_async_x509_cancelled_refresh_owner_releases_waiters() -> None: + exchange_started = anyio.Event() + exchange_calls = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + if exchange_calls == 1: + exchange_started.set() + await anyio.Event().wait() + return _token_response(request) + return _models_response(request) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + + async def list_models() -> str: + return (await client.models.list()).object + + owner = asyncio.create_task(list_models()) + await exchange_started.wait() + waiter = asyncio.create_task(list_models()) + await anyio.sleep(0) + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + assert await waiter == "list" + + assert exchange_calls == 2 + + +@pytest.mark.parametrize("status_code", [408, 409, 429, 500, 503]) +def test_sync_x509_retries_transient_exchange_and_honors_retry_after( + status_code: int, monkeypatch: pytest.MonkeyPatch +) -> None: + exchange_calls = 0 + delays: list[float] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) != _TOKEN_URL: + return _models_response(request) + exchange_calls += 1 + if exchange_calls == 1: + return httpx2.Response(status_code, request=request, headers={"retry-after": "0.25"}) + return _token_response(request) + + monkeypatch.setattr(x509_auth.time, "sleep", delays.append) + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=2) as client: + assert client.models.list().object == "list" + + assert exchange_calls == 2 + assert delays == [0.25] + + +async def test_async_x509_retries_transient_exchange() -> None: + exchange_calls = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) != _TOKEN_URL: + return _models_response(request) + exchange_calls += 1 + if exchange_calls == 1: + return httpx2.Response(429, request=request, headers={"retry-after": "0"}) + return _token_response(request) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=2) as client: + assert (await client.models.list()).object == "list" + + assert exchange_calls == 2 + + +async def test_async_x509_honors_retry_after_without_blocking(monkeypatch: pytest.MonkeyPatch) -> None: + exchange_calls = 0 + delays: list[float] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) != _TOKEN_URL: + return _models_response(request) + exchange_calls += 1 + if exchange_calls == 1: + return httpx2.Response(429, request=request, headers={"retry-after": "0.25"}) + return _token_response(request) + + async def record_sleep(delay: float) -> None: + delays.append(delay) + + monkeypatch.setattr(x509_auth.anyio, "sleep", record_sleep) + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=2) as client: + assert (await client.models.list()).object == "list" + + assert delays == [0.25] + + +def test_x509_exchange_retries_are_bounded(monkeypatch: pytest.MonkeyPatch) -> None: + exchange_calls = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + exchange_calls += 1 + return httpx2.Response(503, request=request) + + def skip_sleep(_: float) -> None: + return None + + monkeypatch.setattr(x509_auth.time, "sleep", skip_sleep) + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=20) as client: + with pytest.raises(OpenAIError, match="status 503"): + client.models.list() + + assert exchange_calls == 3 + + +@pytest.mark.parametrize( + ("failure", "expected_error"), + [(httpx2.ConnectError, APIConnectionError), (httpx2.ReadTimeout, APITimeoutError)], +) +def test_x509_connection_retries_are_bounded_without_outer_api_retries( + failure: type[httpx2.TransportError], expected_error: type[APIConnectionError], monkeypatch: pytest.MonkeyPatch +) -> None: + exchange_calls = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + exchange_calls += 1 + raise failure("exchange unavailable", request=request) + + def skip_sleep(_: float) -> None: + return None + + monkeypatch.setattr(x509_auth.time, "sleep", skip_sleep) + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=20) as client: + with pytest.raises(expected_error) as error: + client.models.list() + + assert isinstance(error.value.__cause__, failure) + assert exchange_calls == 3 + + +async def test_async_x509_exchange_connection_error_preserves_original_cause() -> None: + async def handler(request: httpx2.Request) -> httpx2.Response: + raise httpx2.ConnectError("exchange unavailable", request=request) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIConnectionError) as error: + await client.models.list() + + assert isinstance(error.value.__cause__, httpx2.ConnectError) + + +@pytest.mark.parametrize("status_code", [400, 401, 403]) +def test_x509_does_not_retry_oauth_failures(status_code: int) -> None: + exchange_calls = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + exchange_calls += 1 + return httpx2.Response(status_code, request=request, json={"error": "invalid_grant"}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=2) as client: + with pytest.raises(OAuthError): + client.models.list() + + assert exchange_calls == 1 + + +@pytest.mark.parametrize("status_code", [400, 401, 403]) +def test_x509_oauth_errors_redact_untrusted_descriptions(status_code: int, caplog: pytest.LogCaptureFixture) -> None: + token_secret = "bearer-token-never-visible" + certificate_secret = "certificate-subject-never-visible" + + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response( + status_code, + request=request, + json={ + "error": "invalid_grant", + "error_description": f"{token_secret} {certificate_secret}", + "access_token": token_secret, + }, + ) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with caplog.at_level(logging.DEBUG): + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=2) as client: + with pytest.raises(OAuthError) as error: + client.models.list() + + assert error.value.error == "invalid_grant" + assert error.value.body == {"error": "invalid_grant"} + assert token_secret not in str(error.value) + assert certificate_secret not in str(error.value) + assert token_secret not in caplog.text + assert certificate_secret not in caplog.text + + +def test_x509_success_logs_and_auth_repr_do_not_leak_tokens(caplog: pytest.LogCaptureFixture) -> None: + token_secret = "bearer-token-never-visible" + + def handler(request: httpx2.Request) -> httpx2.Response: + return ( + _token_response(request, token=token_secret) + if str(request.url) == _TOKEN_URL + else _models_response(request) + ) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with caplog.at_level(logging.DEBUG): + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + assert token_secret not in repr(client._workload_identity_auth) + + assert token_secret not in caplog.text + + +def test_x509_never_falls_back_to_environment_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_API_KEY", "api-key-never-used") + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return httpx2.Response(401, request=request, json={"error": "invalid_grant"}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OAuthError): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +def test_x509_refuses_exchange_redirects_even_when_transport_follows_them() -> None: + urls: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + urls.append(str(request.url)) + return httpx2.Response(302, request=request, headers={"location": "https://other.example/oauth/token"}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), follow_redirects=True, trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="status 302"): + client.models.list() + + assert urls == [_TOKEN_URL] + + +def test_x509_refuses_api_redirects_even_when_transport_follows_them() -> None: + urls: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + urls.append(str(request.url)) + if str(request.url) == _TOKEN_URL: + return _token_response(request) + return httpx2.Response(302, request=request, headers={"location": "https://other.example/v1/models"}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), follow_redirects=True, trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError): + client.models.list() + + assert urls == [_TOKEN_URL, _API_URL] + + +def test_sync_x509_retries_replayable_401_request_once() -> None: + requests: list[httpx2.Request] = [] + api_authorizations: list[str] = [] + token_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal token_count + requests.append(request) + if str(request.url) == _TOKEN_URL: + token_count += 1 + return _token_response(request, token=f"token-{token_count}") + api_authorizations.append(request.headers["authorization"]) + return _models_response(request, 401 if token_count == 1 else 200) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL, _TOKEN_URL, _API_URL] + assert api_authorizations == ["Bearer token-1", "Bearer token-2"] + + +async def test_async_x509_retries_replayable_401_request_once() -> None: + api_authorizations: list[str] = [] + token_count = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal token_count + if str(request.url) == _TOKEN_URL: + token_count += 1 + return _token_response(request, token=f"token-{token_count}") + api_authorizations.append(request.headers["authorization"]) + return _models_response(request, 401 if token_count == 1 else 200) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert api_authorizations == ["Bearer token-1", "Bearer token-2"] + + +@pytest.mark.parametrize("seekable", [True, False]) +def test_sync_x509_401_retries_only_replayable_uploads(seekable: bool) -> None: + requests: list[httpx2.Request] = [] + + class Upload(io.BytesIO): + @override + def seekable(self) -> bool: + return seekable + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return _token_response(request, token=f"token-{len(requests)}") + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError): + client.files.create(file=("document.txt", Upload(b"contents")), purpose="assistants") + + assert len(requests) == (4 if seekable else 2) + + +def test_sync_x509_rewinds_seekable_request_stream_to_its_original_position() -> None: + request_bodies: list[bytes] = [] + token_count = 0 + + class StreamingTransport(httpx2.BaseTransport): + @override + def handle_request(self, request: httpx2.Request) -> httpx2.Response: + nonlocal token_count + if str(request.url) == _TOKEN_URL: + token_count += 1 + return _token_response(request, token=f"token-{token_count}") + + request_bodies.append(b"".join(cast(Iterable[bytes], request.stream))) + return _models_response(request, 401 if token_count == 1 else 200) + + stream = io.BytesIO(b"prefix-body") + stream.seek(len(b"prefix-")) + request = httpx2.Request( + "POST", + _API_URL, + content=stream, + headers={"authorization": f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}"}, + ) + http_client = httpx2.Client(transport=StreamingTransport(), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + response = client._send_request(request, stream=False) + + assert response.status_code == 200 + assert request_bodies == [b"body", b"body"] + + +@pytest.mark.parametrize("seekable", [True, False]) +async def test_async_x509_401_retries_only_replayable_uploads(seekable: bool) -> None: + requests: list[httpx2.Request] = [] + + class Upload(io.BytesIO): + @override + def seekable(self) -> bool: + return seekable + + async def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return _token_response(request, token=f"token-{len(requests)}") + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError): + await client.files.create(file=("document.txt", Upload(b"contents")), purpose="assistants") + + assert len(requests) == (4 if seekable else 2) + + +def test_x509_does_not_retry_one_shot_sync_request_stream() -> None: + def chunks() -> Iterator[bytes]: + yield b"body" + + request = httpx2.Request("POST", _API_URL, content=chunks()) + assert not x509_auth._is_replayable_request(request) + + +async def test_x509_does_not_retry_one_shot_async_request_stream() -> None: + async def chunks() -> AsyncIterator[bytes]: + yield b"body" + + request = httpx2.Request("POST", _API_URL, content=chunks()) + assert not x509_auth._is_replayable_request(request) + + +@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") +def test_sync_x509_reuses_legacy_httpx_transport() -> None: + httpx = cast(Any, importlib.import_module("httpx")) + requests: list[Any] = [] + + def handler(request: Any) -> Any: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return httpx.Response(200, request=request, json={"access_token": "legacy-token", "expires_in": 3600}) + return httpx.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx.Client(transport=httpx.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + + +@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") +async def test_async_x509_reuses_legacy_httpx_transport() -> None: + httpx = cast(Any, importlib.import_module("httpx")) + requests: list[Any] = [] + + async def handler(request: Any) -> Any: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return httpx.Response(200, request=request, json={"access_token": "legacy-token", "expires_in": 3600}) + return httpx.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + + +def test_x509_copy_reuses_effective_transport_and_preserves_identity() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return _token_response(request) if str(request.url) == _TOKEN_URL else _models_response(request) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + client = OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + copied = client.with_options() + assert copied._client is http_client + assert copied.workload_identity == client.workload_identity + assert copied.models.list().object == "list" + client.close() + + +@pytest.mark.parametrize("example", ["x509_workload_identity.py", "x509_workload_identity_async.py"]) +@pytest.mark.parametrize("mode", ["api_key", "x509"]) +def test_x509_rollout_examples_construct_clients_without_network( + example: str, mode: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_AUTH_MODE", mode) + monkeypatch.setenv("OPENAI_API_KEY", "example-api-key") + monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_example") + monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "svc_acct_example") + monkeypatch.setenv("OPENAI_MTLS_CA_BUNDLE", str(_MTLS_FIXTURES / "root.pem")) + monkeypatch.setenv("OPENAI_MTLS_CERTIFICATE_CHAIN", str(_MTLS_FIXTURES / "client-chain.pem")) + monkeypatch.setenv("OPENAI_MTLS_PRIVATE_KEY", str(_MTLS_FIXTURES / "client.key")) + + namespace = runpy.run_path(str(Path(__file__).parent.parent / "examples" / example)) + create_client = cast(Callable[[], OpenAI | AsyncOpenAI], namespace["create_client"]) + client = create_client() + if mode == "x509": + assert client.workload_identity == { + "type": "x509", + "identity_provider_id": "idp_example", + "service_account_id": "svc_acct_example", + } + assert client._client.follow_redirects is False + else: + assert client.api_key == "example-api-key" + + if isinstance(client, OpenAI): + client.close() + else: + anyio.run(client.close) From c712ee36482a50b2a01fa02b6605ab47623d0436 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 09:57:03 -0700 Subject: [PATCH 02/11] fix(auth): address X.509 review feedback and compatibility --- src/openai/_client.py | 94 +++++----- src/openai/auth/_workload.py | 30 +++- src/openai/auth/_x509.py | 8 +- tests/test_auth.py | 16 +- tests/test_x509_workload_identity.py | 8 +- ...test_x509_workload_identity_regressions.py | 160 ++++++++++++++++++ 6 files changed, 261 insertions(+), 55 deletions(-) create mode 100644 tests/test_x509_workload_identity_regressions.py diff --git a/src/openai/_client.py b/src/openai/_client.py index b72a7b696e..81554f0424 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -10,7 +10,7 @@ from . import _exceptions from ._qs import Querystring -from .auth import WorkloadIdentity, WorkloadIdentityAuth +from .auth import WorkloadIdentity, WorkloadIdentityAuth, X509WorkloadIdentity from ._types import ( Omit, Headers, @@ -115,7 +115,7 @@ class OpenAI(SyncAPIClient): # client options api_key: str admin_api_key: str | None - workload_identity: WorkloadIdentity | None + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None organization: str | None project: str | None webhook_secret: str | None @@ -136,7 +136,7 @@ def __init__( *, api_key: str | Callable[[], str] | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, organization: str | None = None, project: str | None = None, webhook_secret: str | None = None, @@ -252,7 +252,14 @@ def __init__( self.websocket_base_url = websocket_base_url - x509_identity = workload_identity if is_x509_workload_identity(workload_identity) else None + if is_x509_workload_identity(workload_identity): + x509_identity = workload_identity + subject_token_identity = None + else: + x509_identity = None + subject_token_identity = ( + workload_identity if workload_identity is not None and "provider" in workload_identity else None + ) if provider_runtime is not None: base_url = provider_runtime.base_url elif base_url is None: @@ -280,16 +287,15 @@ def __init__( _strict_response_validation=_strict_response_validation, ) - if workload_identity is not None: - if x509_identity is not None: - self._workload_identity_auth = SyncX509WorkloadIdentityAuth( - workload_identity=x509_identity, http_client=self._client, max_retries=max_retries - ) - else: - self._workload_identity_auth = WorkloadIdentityAuth( - workload_identity=workload_identity, - _use_httpx2=is_httpx2_sync_client(self._client), - ) + if x509_identity is not None: + self._workload_identity_auth = SyncX509WorkloadIdentityAuth( + workload_identity=x509_identity, http_client=self._client, max_retries=max_retries + ) + elif subject_token_identity is not None: + self._workload_identity_auth = WorkloadIdentityAuth( + workload_identity=subject_token_identity, + _use_httpx2=is_httpx2_sync_client(self._client), + ) self._default_stream_cls = Stream @@ -470,29 +476,29 @@ def _send_with_auth_retry( retried: bool = False, **kwargs: Unpack[HttpxSendArgs], ) -> httpx2.Response: - used_workload_identity_auth = False + used_access_token: str | None = None request_is_replayable = False if self._workload_identity_auth is not None: + if self._workload_identity_auth._follow_redirects is not None: + kwargs["follow_redirects"] = self._workload_identity_auth._follow_redirects 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 + used_access_token = self._workload_identity_auth.get_token() + request.headers["Authorization"] = f"Bearer {used_access_token}" request_is_replayable = self._workload_identity_auth._can_retry_request(request) - if self._workload_identity_auth._follow_redirects is not None: - kwargs["follow_redirects"] = self._workload_identity_auth._follow_redirects response = super()._send_request(request, stream=stream, **kwargs) if ( response.status_code == 401 and self._workload_identity_auth is not None - and used_workload_identity_auth + and used_access_token is not None and not retried and request_is_replayable ): response.close() self._workload_identity_auth._prepare_retry_request(request) - self._workload_identity_auth.invalidate_token() + self._workload_identity_auth.invalidate_token(used_access_token) request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @@ -733,7 +739,7 @@ class AsyncOpenAI(AsyncAPIClient): # client options api_key: str admin_api_key: str | None - workload_identity: WorkloadIdentity | None + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None organization: str | None project: str | None webhook_secret: str | None @@ -754,7 +760,7 @@ def __init__( *, api_key: str | Callable[[], Awaitable[str]] | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, organization: str | None = None, project: str | None = None, webhook_secret: str | None = None, @@ -870,7 +876,14 @@ def __init__( self.websocket_base_url = websocket_base_url - x509_identity = workload_identity if is_x509_workload_identity(workload_identity) else None + if is_x509_workload_identity(workload_identity): + x509_identity = workload_identity + subject_token_identity = None + else: + x509_identity = None + subject_token_identity = ( + workload_identity if workload_identity is not None and "provider" in workload_identity else None + ) if provider_runtime is not None: base_url = provider_runtime.base_url elif base_url is None: @@ -898,16 +911,15 @@ def __init__( _strict_response_validation=_strict_response_validation, ) - if workload_identity is not None: - if x509_identity is not None: - self._workload_identity_auth = AsyncX509WorkloadIdentityAuth( - workload_identity=x509_identity, http_client=self._client, max_retries=max_retries - ) - else: - self._workload_identity_auth = WorkloadIdentityAuth( - workload_identity=workload_identity, - _use_httpx2=is_httpx2_async_client(self._client), - ) + if x509_identity is not None: + self._workload_identity_auth = AsyncX509WorkloadIdentityAuth( + workload_identity=x509_identity, http_client=self._client, max_retries=max_retries + ) + elif subject_token_identity is not None: + self._workload_identity_auth = WorkloadIdentityAuth( + workload_identity=subject_token_identity, + _use_httpx2=is_httpx2_async_client(self._client), + ) self._default_stream_cls = AsyncStream @@ -1088,29 +1100,29 @@ async def _send_with_auth_retry( retried: bool = False, **kwargs: Unpack[HttpxSendArgs], ) -> httpx2.Response: - used_workload_identity_auth = False + used_access_token: str | None = None request_is_replayable = False if self._workload_identity_auth is not None: + if self._workload_identity_auth._follow_redirects is not None: + kwargs["follow_redirects"] = self._workload_identity_auth._follow_redirects 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 + used_access_token = await self._workload_identity_auth.get_token_async() + request.headers["Authorization"] = f"Bearer {used_access_token}" request_is_replayable = self._workload_identity_auth._can_retry_request(request) - if self._workload_identity_auth._follow_redirects is not None: - kwargs["follow_redirects"] = self._workload_identity_auth._follow_redirects response = await super()._send_request(request, stream=stream, **kwargs) if ( response.status_code == 401 and self._workload_identity_auth is not None - and used_workload_identity_auth + and used_access_token is not None and not retried and request_is_replayable ): await response.aclose() self._workload_identity_auth._prepare_retry_request(request) - self._workload_identity_auth.invalidate_token() + self._workload_identity_auth.invalidate_token(used_access_token) request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) diff --git a/src/openai/auth/_workload.py b/src/openai/auth/_workload.py index 4f3d6ab67a..39bb65644b 100644 --- a/src/openai/auth/_workload.py +++ b/src/openai/auth/_workload.py @@ -8,6 +8,7 @@ import httpx2 +from .._utils import is_dict from .._httpx2 import DefaultHttpx2Client, _loaded_legacy_httpx from .._exceptions import OAuthError, OpenAIError, SubjectTokenProviderError from .._utils._sync import to_thread @@ -27,7 +28,7 @@ class SubjectTokenProvider(TypedDict): get_token: Callable[[], str] -class SubjectTokenWorkloadIdentity(TypedDict): +class WorkloadIdentity(TypedDict): """Identity provider resource id in WIFAPI.""" identity_provider_id: str @@ -42,6 +43,9 @@ class SubjectTokenWorkloadIdentity(TypedDict): refresh_buffer_seconds: NotRequired[float] +SubjectTokenWorkloadIdentity: TypeAlias = WorkloadIdentity + + class X509WorkloadIdentity(TypedDict): """Authenticate with the client certificate configured on the HTTP transport.""" @@ -51,9 +55,6 @@ class X509WorkloadIdentity(TypedDict): refresh_buffer_seconds: NotRequired[float] -WorkloadIdentity: TypeAlias = SubjectTokenWorkloadIdentity | X509WorkloadIdentity - - def x509_workload_identity( *, identity_provider_id: str, @@ -208,9 +209,22 @@ def __init__( token_exchange_url: str = DEFAULT_TOKEN_EXCHANGE_URL, _use_httpx2: bool = True, ): + self._initialize_token_cache( + workload_identity=workload_identity, + token_exchange_url=token_exchange_url, + use_httpx2=_use_httpx2, + ) + + def _initialize_token_cache( + self, + *, + workload_identity: WorkloadIdentity | X509WorkloadIdentity, + token_exchange_url: str, + use_httpx2: bool = True, + ) -> None: self.workload_identity = workload_identity self.token_exchange_url = token_exchange_url - self._use_httpx2 = _use_httpx2 + self._use_httpx2 = use_httpx2 self._follow_redirects: bool | None = None self._cached_token: str | None = None @@ -252,8 +266,10 @@ def get_token(self) -> str: async def get_token_async(self) -> str: return await to_thread(self.get_token) - def invalidate_token(self) -> None: + def invalidate_token(self, token: str | None = None) -> None: with self._lock: + if token is not None and self._cached_token != token: + return self._cached_token = None self._cached_token_expires_at_monotonic = None self._cached_token_refresh_at_monotonic = None @@ -312,6 +328,8 @@ def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]: if response.is_success: if body is None: raise OpenAIError("Token exchange succeeded but response body was empty") + if not is_dict(body): + raise OpenAIError("Token exchange succeeded but response body was not a JSON object") access_token = body.get("access_token") expires_in = body.get("expires_in") if not isinstance(access_token, str) or not access_token: diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 3c82c6ee81..60e005ac08 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -4,7 +4,7 @@ import time import email.utils from typing import Any, NoReturn, cast -from typing_extensions import TypeGuard, override +from typing_extensions import TypeIs, override import anyio import httpx2 @@ -28,7 +28,9 @@ _ALLOWED_IDENTITY_FIELDS = {"type", "identity_provider_id", "service_account_id", "refresh_buffer_seconds"} -def is_x509_workload_identity(identity: WorkloadIdentity | None) -> TypeGuard[X509WorkloadIdentity]: +def is_x509_workload_identity( + identity: WorkloadIdentity | X509WorkloadIdentity | None, +) -> TypeIs[X509WorkloadIdentity]: return identity is not None and identity.get("type") == "x509" @@ -131,7 +133,7 @@ def _raise_transport_error(error: Exception) -> NoReturn: class _X509WorkloadIdentityAuth(WorkloadIdentityAuth): def __init__(self, *, workload_identity: X509WorkloadIdentity, max_retries: int) -> None: _validate_identity(workload_identity) - super().__init__(workload_identity=workload_identity, token_exchange_url=_X509_TOKEN_EXCHANGE_URL) + self._initialize_token_cache(workload_identity=workload_identity, token_exchange_url=_X509_TOKEN_EXCHANGE_URL) self._max_exchange_retries = min(max(max_retries, 0), _MAX_EXCHANGE_RETRIES) self._follow_redirects = False diff --git a/tests/test_auth.py b/tests/test_auth.py index f9b1719cb5..fc717973e6 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -1,5 +1,5 @@ import json -from typing import cast +from typing import cast, get_type_hints from pathlib import Path import httpx2 @@ -8,6 +8,7 @@ from tests import respx2 from openai import OpenAI, OAuthError +from openai.auth import WorkloadIdentity, WorkloadIdentityAuth, SubjectTokenWorkloadIdentity from tests.respx2.models import Call from openai.auth._workload import ( gcp_id_token_provider, @@ -16,6 +17,19 @@ ) +def test_workload_identity_preserves_callable_typed_dict_api() -> None: + identity = WorkloadIdentity( + identity_provider_id="idp_existing", + service_account_id="svc_acct_existing", + provider={"token_type": "jwt", "get_token": lambda: "subject-token"}, + ) + + assert identity["provider"]["get_token"]() == "subject-token" + assert SubjectTokenWorkloadIdentity is WorkloadIdentity + assert get_type_hints(WorkloadIdentityAuth.__init__)["workload_identity"] is WorkloadIdentity + assert WorkloadIdentityAuth(workload_identity=identity).workload_identity is identity + + @respx2.mock def test_basic_auth(): respx2.post("https://auth.openai.com/oauth/token").mock( diff --git a/tests/test_x509_workload_identity.py b/tests/test_x509_workload_identity.py index 9521b807bc..97ecde50ad 100644 --- a/tests/test_x509_workload_identity.py +++ b/tests/test_x509_workload_identity.py @@ -345,8 +345,8 @@ async def list_models() -> str: waiter = asyncio.create_task(list_models()) await anyio.sleep(0) waiter.cancel() - with pytest.raises(asyncio.CancelledError): - await waiter + await asyncio.gather(waiter, return_exceptions=True) + assert waiter.cancelled() finish_exchange.set() assert await owner == "list" assert (await client.models.list()).object == "list" @@ -379,8 +379,8 @@ async def list_models() -> str: waiter = asyncio.create_task(list_models()) await anyio.sleep(0) owner.cancel() - with pytest.raises(asyncio.CancelledError): - await owner + await asyncio.gather(owner, return_exceptions=True) + assert owner.cancelled() assert await waiter == "list" assert exchange_calls == 2 diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py new file mode 100644 index 0000000000..ebd26aeee7 --- /dev/null +++ b/tests/test_x509_workload_identity_regressions.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor + +import anyio +import httpx2 +import pytest + +from openai import OpenAI, AsyncOpenAI, OpenAIError +from openai.auth import X509WorkloadIdentity, x509_workload_identity + +_TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" +_API_URL = "https://mtls.api.openai.com/v1/models" + + +def _identity() -> X509WorkloadIdentity: + return x509_workload_identity(identity_provider_id="idp_123", service_account_id="svc_acct_123") + + +@pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) +def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: + urls: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + urls.append(str(request.url)) + return httpx2.Response(302, request=request, headers={"location": "https://other.example/v1/models"}) + + headers = {} if authorization is None else {"Authorization": authorization} + request = httpx2.Request("GET", _API_URL, headers=headers) + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), follow_redirects=True, trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + response = client._send_request(request, stream=False) + + assert response.status_code == 302 + assert urls == [_API_URL] + + +@pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) +async def test_async_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: + urls: list[str] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + urls.append(str(request.url)) + return httpx2.Response(302, request=request, headers={"location": "https://other.example/v1/models"}) + + headers = {} if authorization is None else {"Authorization": authorization} + request = httpx2.Request("GET", _API_URL, headers=headers) + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), follow_redirects=True, trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + response = await client._send_request(request, stream=False) + + assert response.status_code == 302 + assert urls == [_API_URL] + + +@pytest.mark.parametrize("response_body", [[], "not-an-object", 42, True]) +def test_sync_x509_rejects_non_object_token_responses(response_body: object) -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, request=request, json=response_body) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="response body was not a JSON object"): + client.models.list() + + +@pytest.mark.parametrize("response_body", [[], "not-an-object", 42, True]) +async def test_async_x509_rejects_non_object_token_responses(response_body: object) -> None: + async def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, request=request, json=response_body) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="response body was not a JSON object"): + await client.models.list() + + +def test_sync_x509_concurrent_stale_401_responses_share_one_replacement_token() -> None: + exchange_calls = 0 + stale_requests = 0 + state_lock = threading.Lock() + both_stale_requests = threading.Barrier(2) + replacement_issued = threading.Event() + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls, stale_requests + if str(request.url) == _TOKEN_URL: + with state_lock: + exchange_calls += 1 + token_number = exchange_calls + if token_number == 2: + replacement_issued.set() + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{token_number}", "expires_in": 3600} + ) + + authorization = request.headers["Authorization"] + if authorization == "Bearer token-1": + with state_lock: + stale_requests += 1 + request_number = stale_requests + both_stale_requests.wait(timeout=5) + if request_number == 2: + assert replacement_issued.wait(timeout=5) + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + + assert authorization == "Bearer token-2" + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + + def list_models(_: int) -> str: + return client.models.list().object + + with ThreadPoolExecutor(max_workers=2) as executor: + responses = list(executor.map(list_models, range(2))) + + assert responses == ["list", "list"] + assert exchange_calls == 2 + + +async def test_async_x509_concurrent_stale_401_responses_share_one_replacement_token() -> None: + exchange_calls = 0 + stale_requests = 0 + both_stale_requests = anyio.Event() + replacement_issued = anyio.Event() + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls, stale_requests + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + if exchange_calls == 2: + replacement_issued.set() + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + authorization = request.headers["Authorization"] + if authorization == "Bearer token-1": + stale_requests += 1 + request_number = stale_requests + if stale_requests == 2: + both_stale_requests.set() + await both_stale_requests.wait() + if request_number == 2: + await replacement_issued.wait() + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + + assert authorization == "Bearer token-2" + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + responses = await asyncio.gather(client.models.list(), client.models.list()) + + assert [response.object for response in responses] == ["list", "list"] + assert exchange_calls == 2 From 55e1460da30441baeb376b30d0ee7d3f5e0e7e21 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 10:06:20 -0700 Subject: [PATCH 03/11] refactor(auth): share typed workload-identity initialization --- src/openai/_client.py | 4 +- src/openai/auth/_workload.py | 113 ++++++++++-------- src/openai/auth/_x509.py | 12 +- ...test_x509_workload_identity_regressions.py | 17 +++ 4 files changed, 84 insertions(+), 62 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index 81554f0424..be47dda37a 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -119,7 +119,7 @@ class OpenAI(SyncAPIClient): organization: str | None project: str | None webhook_secret: str | None - _workload_identity_auth: WorkloadIdentityAuth | None + _workload_identity_auth: WorkloadIdentityAuth | SyncX509WorkloadIdentityAuth | None _provider: _Provider | None _provider_runtime: _ProviderRuntime | None @@ -743,7 +743,7 @@ class AsyncOpenAI(AsyncAPIClient): organization: str | None project: str | None webhook_secret: str | None - _workload_identity_auth: WorkloadIdentityAuth | None + _workload_identity_auth: WorkloadIdentityAuth | AsyncX509WorkloadIdentityAuth | None _provider: _Provider | None _provider_runtime: _ProviderRuntime | None diff --git a/src/openai/auth/_workload.py b/src/openai/auth/_workload.py index 39bb65644b..4f9797f092 100644 --- a/src/openai/auth/_workload.py +++ b/src/openai/auth/_workload.py @@ -2,9 +2,9 @@ import time import threading -from typing import Any, Callable, TypedDict, cast +from typing import Any, Generic, TypeVar, Callable, TypedDict, cast from pathlib import Path -from typing_extensions import Literal, TypeAlias, NotRequired +from typing_extensions import Literal, TypeAlias, NotRequired, override import httpx2 @@ -55,6 +55,9 @@ class X509WorkloadIdentity(TypedDict): refresh_buffer_seconds: NotRequired[float] +_WorkloadIdentityT = TypeVar("_WorkloadIdentityT", WorkloadIdentity, X509WorkloadIdentity) + + def x509_workload_identity( *, identity_provider_id: str, @@ -201,30 +204,17 @@ def get_token() -> str: return {"token_type": "id", "get_token": get_token} -class WorkloadIdentityAuth: +class _WorkloadIdentityAuth(Generic[_WorkloadIdentityT]): def __init__( self, *, - workload_identity: WorkloadIdentity, + workload_identity: _WorkloadIdentityT, token_exchange_url: str = DEFAULT_TOKEN_EXCHANGE_URL, _use_httpx2: bool = True, - ): - self._initialize_token_cache( - workload_identity=workload_identity, - token_exchange_url=token_exchange_url, - use_httpx2=_use_httpx2, - ) - - def _initialize_token_cache( - self, - *, - workload_identity: WorkloadIdentity | X509WorkloadIdentity, - token_exchange_url: str, - use_httpx2: bool = True, ) -> None: - self.workload_identity = workload_identity + self.workload_identity: _WorkloadIdentityT = workload_identity self.token_exchange_url = token_exchange_url - self._use_httpx2 = use_httpx2 + self._use_httpx2 = _use_httpx2 self._follow_redirects: bool | None = None self._cached_token: str | None = None @@ -288,33 +278,7 @@ def _store_token(self, token_data: dict[str, Any]) -> None: self._cached_token_refresh_at_monotonic = now + self._refresh_delay_seconds(expires_in) def _fetch_token_from_exchange(self) -> dict[str, Any]: - subject_token = self._get_subject_token() - - identity = cast(SubjectTokenWorkloadIdentity, self.workload_identity) - token_type = identity["provider"]["token_type"] - subject_token_type = SUBJECT_TOKEN_TYPES.get(token_type) - if subject_token_type is None: - raise OpenAIError( - f"Unsupported token type: {token_type!r}. Supported types: {', '.join(SUBJECT_TOKEN_TYPES.keys())}" - ) - - legacy_httpx = _loaded_legacy_httpx() if not self._use_httpx2 else None - exchange_client = ( - legacy_httpx.Client() if legacy_httpx is not None else DefaultHttpx2Client(follow_redirects=False) - ) - with exchange_client as client: - response = client.post( - self.token_exchange_url, - json={ - "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, - "subject_token": subject_token, - "subject_token_type": subject_token_type, - "identity_provider_id": self.workload_identity["identity_provider_id"], - "service_account_id": self.workload_identity["service_account_id"], - }, - timeout=10.0, - ) - return self._handle_token_response(response) + raise NotImplementedError("Workload identity authentication must implement token exchange") def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]: try: @@ -345,13 +309,6 @@ def _validate_expires_in(self, expires_in: object) -> float: raise OpenAIError("Token exchange response did not include a valid expires_in") return float(expires_in) - def _get_subject_token(self) -> str: - provider = cast(SubjectTokenWorkloadIdentity, self.workload_identity)["provider"] - subject_token = provider["get_token"]() - if not subject_token: - raise OpenAIError("The workload identity provider returned an empty subject token") - return subject_token - def _token_unusable(self) -> bool: return self._cached_token is None or self._token_expired() @@ -378,3 +335,53 @@ def _can_retry_request(self, request: httpx2.Request) -> bool: def _prepare_retry_request(self, request: httpx2.Request) -> None: """Preserve the established subject-token request retry behavior.""" del request + + +class WorkloadIdentityAuth(_WorkloadIdentityAuth[WorkloadIdentity]): + def __init__( + self, + *, + workload_identity: WorkloadIdentity, + token_exchange_url: str = DEFAULT_TOKEN_EXCHANGE_URL, + _use_httpx2: bool = True, + ) -> None: + super().__init__( + workload_identity=workload_identity, + token_exchange_url=token_exchange_url, + _use_httpx2=_use_httpx2, + ) + + @override + def _fetch_token_from_exchange(self) -> dict[str, Any]: + subject_token = self._get_subject_token() + token_type = self.workload_identity["provider"]["token_type"] + subject_token_type = SUBJECT_TOKEN_TYPES.get(token_type) + if subject_token_type is None: + raise OpenAIError( + f"Unsupported token type: {token_type!r}. Supported types: {', '.join(SUBJECT_TOKEN_TYPES.keys())}" + ) + + legacy_httpx = _loaded_legacy_httpx() if not self._use_httpx2 else None + exchange_client = ( + legacy_httpx.Client() if legacy_httpx is not None else DefaultHttpx2Client(follow_redirects=False) + ) + with exchange_client as client: + response = client.post( + self.token_exchange_url, + json={ + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "subject_token": subject_token, + "subject_token_type": subject_token_type, + "identity_provider_id": self.workload_identity["identity_provider_id"], + "service_account_id": self.workload_identity["service_account_id"], + }, + timeout=10.0, + ) + return self._handle_token_response(response) + + def _get_subject_token(self) -> str: + provider = self.workload_identity["provider"] + subject_token = provider["get_token"]() + if not subject_token: + raise OpenAIError("The workload identity provider returned an empty subject token") + return subject_token diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 60e005ac08..cc2315bd9d 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -14,8 +14,8 @@ from ._workload import ( TOKEN_EXCHANGE_GRANT_TYPE, WorkloadIdentity, - WorkloadIdentityAuth, X509WorkloadIdentity, + _WorkloadIdentityAuth, ) from .._constants import MAX_RETRY_DELAY, INITIAL_RETRY_DELAY, MAX_RETRY_AFTER_DELAY from .._exceptions import OAuthError, OpenAIError, APITimeoutError, APIConnectionError @@ -130,10 +130,10 @@ def _raise_transport_error(error: Exception) -> NoReturn: raise APIConnectionError(request=request) from error -class _X509WorkloadIdentityAuth(WorkloadIdentityAuth): +class _X509WorkloadIdentityAuth(_WorkloadIdentityAuth[X509WorkloadIdentity]): def __init__(self, *, workload_identity: X509WorkloadIdentity, max_retries: int) -> None: _validate_identity(workload_identity) - self._initialize_token_cache(workload_identity=workload_identity, token_exchange_url=_X509_TOKEN_EXCHANGE_URL) + super().__init__(workload_identity=workload_identity, token_exchange_url=_X509_TOKEN_EXCHANGE_URL) self._max_exchange_retries = min(max(max_retries, 0), _MAX_EXCHANGE_RETRIES) self._follow_redirects = False @@ -183,12 +183,11 @@ def __init__( @override def _fetch_token_from_exchange(self) -> dict[str, Any]: - identity = cast(X509WorkloadIdentity, self.workload_identity) for attempt in range(self._max_exchange_retries + 1): try: response = self._http_client.post( _X509_TOKEN_EXCHANGE_URL, - json=_exchange_payload(identity), + json=_exchange_payload(self.workload_identity), timeout=10.0, follow_redirects=False, ) @@ -228,12 +227,11 @@ async def get_token_async(self) -> str: return cast(str, self._cached_token) async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: - identity = cast(X509WorkloadIdentity, self.workload_identity) for attempt in range(self._max_exchange_retries + 1): try: response = await self._http_client.post( _X509_TOKEN_EXCHANGE_URL, - json=_exchange_payload(identity), + json=_exchange_payload(self.workload_identity), timeout=10.0, follow_redirects=False, ) diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index ebd26aeee7..45ed8f36b1 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -10,6 +10,7 @@ from openai import OpenAI, AsyncOpenAI, OpenAIError from openai.auth import X509WorkloadIdentity, x509_workload_identity +from openai.auth._workload import _WorkloadIdentityAuth _TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" _API_URL = "https://mtls.api.openai.com/v1/models" @@ -19,6 +20,22 @@ def _identity() -> X509WorkloadIdentity: return x509_workload_identity(identity_provider_id="idp_123", service_account_id="svc_acct_123") +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_x509_auth_initializes_its_typed_common_superclass(client_type: type[OpenAI] | type[AsyncOpenAI]) -> None: + client = client_type(workload_identity=_identity()) + auth = client._workload_identity_auth + + assert isinstance(auth, _WorkloadIdentityAuth) + assert auth.workload_identity == _identity() + assert auth._cached_token is None + assert auth._follow_redirects is False + + if isinstance(client, OpenAI): + client.close() + else: + anyio.run(client.close) + + @pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: urls: list[str] = [] From 5df57ae901c240557af6c066130b42f32c75b43b Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 10:19:00 -0700 Subject: [PATCH 04/11] fix(auth): preserve X.509 client copies and invalidate failed uploads --- src/openai/_client.py | 50 ++++---- src/openai/lib/azure.py | 6 +- src/openai/lib/bedrock.py | 6 +- ...test_x509_workload_identity_regressions.py | 116 +++++++++++++++++- 4 files changed, 143 insertions(+), 35 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index be47dda37a..5ce6a47021 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -489,20 +489,17 @@ def _send_with_auth_retry( request_is_replayable = self._workload_identity_auth._can_retry_request(request) response = super()._send_request(request, stream=stream, **kwargs) - if ( - response.status_code == 401 - and self._workload_identity_auth is not None - and used_access_token is not None - and not retried - and request_is_replayable - ): - response.close() - self._workload_identity_auth._prepare_retry_request(request) - self._workload_identity_auth.invalidate_token(used_access_token) - request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" - return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) + if response.status_code != 401 or self._workload_identity_auth is None or used_access_token is None: + return response - return response + self._workload_identity_auth.invalidate_token(used_access_token) + if retried or not request_is_replayable: + return response + + response.close() + self._workload_identity_auth._prepare_retry_request(request) + request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" + return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @override def _send_request( @@ -612,7 +609,7 @@ def copy( *, api_key: str | Callable[[], str] | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, provider: _Provider | None | NotGiven = not_given, organization: str | None = None, project: str | None = None, @@ -1113,20 +1110,17 @@ async def _send_with_auth_retry( request_is_replayable = self._workload_identity_auth._can_retry_request(request) response = await super()._send_request(request, stream=stream, **kwargs) - if ( - response.status_code == 401 - and self._workload_identity_auth is not None - and used_access_token is not None - and not retried - and request_is_replayable - ): - await response.aclose() - self._workload_identity_auth._prepare_retry_request(request) - self._workload_identity_auth.invalidate_token(used_access_token) - request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" - return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) + if response.status_code != 401 or self._workload_identity_auth is None or used_access_token is None: + return response - return response + self._workload_identity_auth.invalidate_token(used_access_token) + if retried or not request_is_replayable: + return response + + await response.aclose() + self._workload_identity_auth._prepare_retry_request(request) + request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" + return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @override async def _send_request( @@ -1247,7 +1241,7 @@ def copy( *, api_key: str | Callable[[], Awaitable[str]] | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, provider: _Provider | None | NotGiven = not_given, organization: str | None = None, project: str | None = None, diff --git a/src/openai/lib/azure.py b/src/openai/lib/azure.py index 5e0b6f5203..2a1be84d6f 100644 --- a/src/openai/lib/azure.py +++ b/src/openai/lib/azure.py @@ -7,7 +7,7 @@ import httpx2 -from ..auth import WorkloadIdentity +from ..auth import WorkloadIdentity, X509WorkloadIdentity from .._types import NOT_GIVEN, Omit, Query, Headers, Timeout, NotGiven from .._utils import is_given, is_mapping from .._client import OpenAI, AsyncOpenAI @@ -284,7 +284,7 @@ def copy( *, api_key: str | Callable[[], str] | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, provider: _Provider | None | NotGiven = NOT_GIVEN, organization: str | None = None, project: str | None = None, @@ -608,7 +608,7 @@ def copy( *, api_key: str | Callable[[], Awaitable[str]] | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, provider: _Provider | None | NotGiven = NOT_GIVEN, organization: str | None = None, project: str | None = None, diff --git a/src/openai/lib/bedrock.py b/src/openai/lib/bedrock.py index 25a4a56453..327f7c7ec8 100644 --- a/src/openai/lib/bedrock.py +++ b/src/openai/lib/bedrock.py @@ -10,7 +10,7 @@ import httpx2 -from ..auth import WorkloadIdentity +from ..auth import WorkloadIdentity, X509WorkloadIdentity from .._types import NOT_GIVEN, Timeout, NotGiven from .._utils import is_given from .._client import OpenAI, AsyncOpenAI @@ -513,7 +513,7 @@ def copy( *, api_key: str | BedrockTokenProvider | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, provider: _Provider | None | NotGiven = NOT_GIVEN, bedrock_token_provider: BedrockTokenProvider | None = None, aws_region: str | None = None, @@ -749,7 +749,7 @@ def copy( *, api_key: str | AsyncBedrockTokenProvider | None = None, admin_api_key: str | None = None, - workload_identity: WorkloadIdentity | None = None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None = None, provider: _Provider | None | NotGiven = NOT_GIVEN, bedrock_token_provider: AsyncBedrockTokenProvider | None = None, aws_region: str | None = None, diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index 45ed8f36b1..1c5cc8292c 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -1,14 +1,17 @@ from __future__ import annotations +import io import asyncio import threading +from typing import get_args, get_type_hints +from typing_extensions import override from concurrent.futures import ThreadPoolExecutor import anyio import httpx2 import pytest -from openai import OpenAI, AsyncOpenAI, OpenAIError +from openai import OpenAI, AsyncOpenAI, OpenAIError, APIStatusError from openai.auth import X509WorkloadIdentity, x509_workload_identity from openai.auth._workload import _WorkloadIdentityAuth @@ -36,6 +39,52 @@ def test_x509_auth_initializes_its_typed_common_superclass(client_type: type[Ope anyio.run(client.close) +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_copy_signatures_accept_typed_x509_workload_identities(client_type: type[OpenAI] | type[AsyncOpenAI]) -> None: + assert X509WorkloadIdentity in get_args(get_type_hints(client_type.copy)["workload_identity"]) + assert X509WorkloadIdentity in get_args(get_type_hints(client_type.with_options)["workload_identity"]) + + +@pytest.mark.parametrize("method", ["copy", "with_options"]) +def test_sync_copy_accepts_an_explicit_x509_identity(method: str) -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "copied-token", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + replacement = x509_workload_identity(identity_provider_id="idp_replacement", service_account_id="svc_replacement") + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + copied = ( + client.copy(workload_identity=replacement) + if method == "copy" + else client.with_options(workload_identity=replacement) + ) + assert copied.workload_identity == replacement + assert copied._client is http_client + assert copied.models.list().object == "list" + + +@pytest.mark.parametrize("method", ["copy", "with_options"]) +async def test_async_copy_accepts_an_explicit_x509_identity(method: str) -> None: + async def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "copied-token", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + replacement = x509_workload_identity(identity_provider_id="idp_replacement", service_account_id="svc_replacement") + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + copied = ( + client.copy(workload_identity=replacement) + if method == "copy" + else client.with_options(workload_identity=replacement) + ) + assert copied.workload_identity == replacement + assert copied._client is http_client + assert (await copied.models.list()).object == "list" + + @pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: urls: list[str] = [] @@ -94,6 +143,71 @@ async def handler(request: httpx2.Request) -> httpx2.Response: await client.models.list() +def test_sync_x509_invalidates_rejected_token_without_replaying_one_shot_upload() -> None: + exchange_calls = 0 + api_authorizations: list[str] = [] + + class OneShotUpload(io.BytesIO): + @override + def seekable(self) -> bool: + return False + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + api_authorizations.append(request.headers["Authorization"]) + if exchange_calls == 1: + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + return httpx2.Response(200, request=request, json={"id": "file_123", "object": "file"}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError): + client.files.create(file=("first.txt", OneShotUpload(b"first")), purpose="assistants") + assert client.files.create(file=("second.txt", OneShotUpload(b"second")), purpose="assistants").id == "file_123" + + assert exchange_calls == 2 + assert api_authorizations == ["Bearer token-1", "Bearer token-2"] + + +async def test_async_x509_invalidates_rejected_token_without_replaying_one_shot_upload() -> None: + exchange_calls = 0 + api_authorizations: list[str] = [] + + class OneShotUpload(io.BytesIO): + @override + def seekable(self) -> bool: + return False + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + api_authorizations.append(request.headers["Authorization"]) + if exchange_calls == 1: + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + return httpx2.Response(200, request=request, json={"id": "file_123", "object": "file"}) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError): + await client.files.create(file=("first.txt", OneShotUpload(b"first")), purpose="assistants") + uploaded = await client.files.create(file=("second.txt", OneShotUpload(b"second")), purpose="assistants") + + assert uploaded.id == "file_123" + assert exchange_calls == 2 + assert api_authorizations == ["Bearer token-1", "Bearer token-2"] + + def test_sync_x509_concurrent_stale_401_responses_share_one_replacement_token() -> None: exchange_calls = 0 stale_requests = 0 From 2252fbe80f140e7943f28d59b5038caea01b707d Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 10:38:17 -0700 Subject: [PATCH 05/11] fix(auth): harden X.509 exchange and rejected-token handling --- src/openai/_client.py | 26 ++-- src/openai/auth/_x509.py | 4 +- src/openai/lib/azure.py | 5 + tests/lib/test_azure.py | 13 ++ ...test_x509_workload_identity_regressions.py | 123 +++++++++++++++++- 5 files changed, 159 insertions(+), 12 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index 5ce6a47021..f0a5df23ec 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -255,11 +255,14 @@ def __init__( if is_x509_workload_identity(workload_identity): x509_identity = workload_identity subject_token_identity = None - else: + elif workload_identity is None: x509_identity = None - subject_token_identity = ( - workload_identity if workload_identity is not None and "provider" in workload_identity else None - ) + subject_token_identity = None + elif "provider" in workload_identity: + x509_identity = None + subject_token_identity = workload_identity + else: + raise OpenAIError("Invalid `workload_identity` configuration: expected an X.509 or subject-token identity") if provider_runtime is not None: base_url = provider_runtime.base_url elif base_url is None: @@ -498,7 +501,7 @@ def _send_with_auth_retry( response.close() self._workload_identity_auth._prepare_retry_request(request) - request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" + request.headers["Authorization"] = f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}" return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @override @@ -876,11 +879,14 @@ def __init__( if is_x509_workload_identity(workload_identity): x509_identity = workload_identity subject_token_identity = None - else: + elif workload_identity is None: x509_identity = None - subject_token_identity = ( - workload_identity if workload_identity is not None and "provider" in workload_identity else None - ) + subject_token_identity = None + elif "provider" in workload_identity: + x509_identity = None + subject_token_identity = workload_identity + else: + raise OpenAIError("Invalid `workload_identity` configuration: expected an X.509 or subject-token identity") if provider_runtime is not None: base_url = provider_runtime.base_url elif base_url is None: @@ -1119,7 +1125,7 @@ async def _send_with_auth_retry( await response.aclose() self._workload_identity_auth._prepare_retry_request(request) - request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" + request.headers["Authorization"] = f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}" return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @override diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index cc2315bd9d..20dec7dbba 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -41,7 +41,7 @@ def _validate_identity(identity: X509WorkloadIdentity) -> None: if set(identity) - _ALLOWED_IDENTITY_FIELDS: raise OpenAIError("X.509 workload identity accepts only identity IDs and an optional refresh buffer") - if not identity["identity_provider_id"] or not identity["service_account_id"]: + if not identity.get("identity_provider_id") or not identity.get("service_account_id"): raise OpenAIError("X.509 workload identity requires identity-provider and service-account IDs") refresh_buffer = cast(object, identity.get("refresh_buffer_seconds")) @@ -188,6 +188,7 @@ def _fetch_token_from_exchange(self) -> dict[str, Any]: response = self._http_client.post( _X509_TOKEN_EXCHANGE_URL, json=_exchange_payload(self.workload_identity), + auth=lambda request: request, timeout=10.0, follow_redirects=False, ) @@ -232,6 +233,7 @@ async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: response = await self._http_client.post( _X509_TOKEN_EXCHANGE_URL, json=_exchange_payload(self.workload_identity), + auth=lambda request: request, timeout=10.0, follow_redirects=False, ) diff --git a/src/openai/lib/azure.py b/src/openai/lib/azure.py index 2a1be84d6f..b0ddbf4e5d 100644 --- a/src/openai/lib/azure.py +++ b/src/openai/lib/azure.py @@ -16,6 +16,7 @@ from .._models import SecurityOptions, FinalRequestOptions from .._provider import _Provider from .._streaming import Stream, AsyncStream +from ..auth._x509 import is_x509_workload_identity from .._exceptions import OpenAIError from .._base_client import DEFAULT_MAX_RETRIES, BaseClient @@ -309,6 +310,8 @@ def copy( """ if not isinstance(provider, NotGiven): raise OpenAIError("Configure `provider` on `OpenAI`, not on `AzureOpenAI.with_options()`.") + if is_x509_workload_identity(workload_identity): + raise OpenAIError("X.509 workload identity is not supported by Azure clients") return super().copy( api_key=api_key, @@ -633,6 +636,8 @@ def copy( """ if not isinstance(provider, NotGiven): raise OpenAIError("Configure `provider` on `AsyncOpenAI`, not on `AsyncAzureOpenAI.with_options()`.") + if is_x509_workload_identity(workload_identity): + raise OpenAIError("X.509 workload identity is not supported by Azure clients") return super().copy( api_key=api_key, diff --git a/tests/lib/test_azure.py b/tests/lib/test_azure.py index 691a911144..f7540facd4 100644 --- a/tests/lib/test_azure.py +++ b/tests/lib/test_azure.py @@ -8,6 +8,7 @@ import pytest from openai import OpenAIError +from openai.auth import x509_workload_identity from tests.utils import update_env from tests.respx2 import MockRouter from openai._types import Omit @@ -79,6 +80,18 @@ def test_client_copying_override_options(client: Client) -> None: assert copied._custom_query == {"api-version": "2022-05-01"} +@pytest.mark.parametrize("client", [sync_client, async_client]) +@pytest.mark.parametrize("method", ["copy", "with_options"]) +def test_client_copying_rejects_x509_workload_identity(client: Client, method: Literal["copy", "with_options"]) -> None: + identity = x509_workload_identity(identity_provider_id="idp_123", service_account_id="svc_acct_123") + + with pytest.raises(OpenAIError, match="X.509 workload identity is not supported by Azure clients"): + if method == "copy": + client.copy(workload_identity=identity) + else: + client.with_options(workload_identity=identity) + + def test_enforce_credentials_false_sync() -> None: with update_env(AZURE_OPENAI_API_KEY=Omit(), AZURE_OPENAI_AD_TOKEN=Omit()): AzureOpenAI( diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index 1c5cc8292c..0573b4dae1 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -3,7 +3,7 @@ import io import asyncio import threading -from typing import get_args, get_type_hints +from typing import cast, get_args, get_type_hints from typing_extensions import override from concurrent.futures import ThreadPoolExecutor @@ -45,6 +45,33 @@ def test_copy_signatures_accept_typed_x509_workload_identities(client_type: type assert X509WorkloadIdentity in get_args(get_type_hints(client_type.with_options)["workload_identity"]) +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize( + "identity", + [ + {"identity_provider_id": "idp_123", "service_account_id": "svc_acct_123"}, + {"type": "x590", "identity_provider_id": "idp_123", "service_account_id": "svc_acct_123"}, + {}, + ], +) +def test_client_rejects_unrecognized_workload_identity_shapes( + client_type: type[OpenAI] | type[AsyncOpenAI], identity: dict[str, str] +) -> None: + with pytest.raises(OpenAIError, match="Invalid `workload_identity` configuration"): + client_type(workload_identity=cast(X509WorkloadIdentity, identity)) + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize("missing_key", ["identity_provider_id", "service_account_id"]) +def test_client_rejects_x509_identity_missing_a_required_id( + client_type: type[OpenAI] | type[AsyncOpenAI], missing_key: str +) -> None: + identity = {key: value for key, value in _identity().items() if key != missing_key} + + with pytest.raises(OpenAIError, match="requires identity-provider and service-account IDs"): + client_type(workload_identity=cast(X509WorkloadIdentity, identity)) + + @pytest.mark.parametrize("method", ["copy", "with_options"]) def test_sync_copy_accepts_an_explicit_x509_identity(method: str) -> None: def handler(request: httpx2.Request) -> httpx2.Response: @@ -121,6 +148,46 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert urls == [_API_URL] +def test_sync_x509_exchange_does_not_inherit_caller_http_auth() -> None: + exchange_authorizations: list[str | None] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + exchange_authorizations.append(request.headers.get("Authorization")) + return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.Client( + transport=httpx2.MockTransport(handler), + auth=httpx2.BasicAuth("caller", "private-api-credential"), + trust_env=False, + ) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert exchange_authorizations == [None] + + +async def test_async_x509_exchange_does_not_inherit_caller_http_auth() -> None: + exchange_authorizations: list[str | None] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + exchange_authorizations.append(request.headers.get("Authorization")) + return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.AsyncClient( + transport=httpx2.MockTransport(handler), + auth=httpx2.BasicAuth("caller", "private-api-credential"), + trust_env=False, + ) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert exchange_authorizations == [None] + + @pytest.mark.parametrize("response_body", [[], "not-an-object", 42, True]) def test_sync_x509_rejects_non_object_token_responses(response_body: object) -> None: def handler(request: httpx2.Request) -> httpx2.Response: @@ -208,6 +275,60 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert api_authorizations == ["Bearer token-1", "Bearer token-2"] +def test_sync_x509_invalidates_a_rejected_replay_token() -> None: + exchange_calls = 0 + api_authorizations: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + api_authorizations.append(request.headers["Authorization"]) + if exchange_calls < 3: + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError, match="401"): + client.models.list() + assert client.models.list().object == "list" + + assert exchange_calls == 3 + assert api_authorizations == ["Bearer token-1", "Bearer token-2", "Bearer token-3"] + + +async def test_async_x509_invalidates_a_rejected_replay_token() -> None: + exchange_calls = 0 + api_authorizations: list[str] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + api_authorizations.append(request.headers["Authorization"]) + if exchange_calls < 3: + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(APIStatusError, match="401"): + await client.models.list() + assert (await client.models.list()).object == "list" + + assert exchange_calls == 3 + assert api_authorizations == ["Bearer token-1", "Bearer token-2", "Bearer token-3"] + + def test_sync_x509_concurrent_stale_401_responses_share_one_replacement_token() -> None: exchange_calls = 0 stale_requests = 0 From 2e426fa5355f10454f43bca0cd47e12b5c346a24 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 10:49:46 -0700 Subject: [PATCH 06/11] fix(auth): preserve X.509 transport and replay invariants --- src/openai/_client.py | 28 +++- src/openai/auth/_x509.py | 19 ++- ...test_x509_workload_identity_regressions.py | 131 ++++++++++++++++++ 3 files changed, 171 insertions(+), 7 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index f0a5df23ec..b770d2df1b 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -596,7 +596,7 @@ def _prepare_request(self, request: httpx2.Request) -> None: @override def _custom_auth(self, security: SecurityOptions) -> httpx2.Auth | None: - if self._provider_runtime is not None: + if self._provider_runtime is not None or isinstance(self._workload_identity_auth, SyncX509WorkloadIdentityAuth): return httpx2.Auth() return super()._custom_auth(security) @@ -674,11 +674,18 @@ def copy( "base_url": base_url, } else: + inherited_base_url = ( + None + if is_x509_workload_identity(workload_identity) and str(self.base_url) == "https://api.openai.com/v1/" + else self.base_url + ) auth_options = { - "api_key": api_key or self._api_key_provider or self.api_key, + "api_key": api_key + if workload_identity is not None + else api_key or self._api_key_provider or self.api_key, "admin_api_key": admin_api_key or self.admin_api_key, "workload_identity": workload_identity or self.workload_identity, - "base_url": base_url or self.base_url, + "base_url": base_url or inherited_base_url, } return self.__class__( @@ -1231,7 +1238,9 @@ async def _prepare_request(self, request: httpx2.Request) -> None: @property @override def custom_auth(self) -> httpx2.Auth | None: - if self._provider_runtime is not None: + if self._provider_runtime is not None or isinstance( + self._workload_identity_auth, AsyncX509WorkloadIdentityAuth + ): return httpx2.Auth() return super().custom_auth @@ -1308,11 +1317,18 @@ def copy( "base_url": base_url, } else: + inherited_base_url = ( + None + if is_x509_workload_identity(workload_identity) and str(self.base_url) == "https://api.openai.com/v1/" + else self.base_url + ) auth_options = { - "api_key": api_key or self._api_key_provider or self.api_key, + "api_key": api_key + if workload_identity is not None + else api_key or self._api_key_provider or self.api_key, "admin_api_key": admin_api_key or self.admin_api_key, "workload_identity": workload_identity or self.workload_identity, - "base_url": base_url or self.base_url, + "base_url": base_url or inherited_base_url, } return self.__class__( diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 20dec7dbba..90e38d4711 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -25,6 +25,7 @@ _X509_SUBJECT_TOKEN_TYPE = "urn:openai:params:oauth:token-type:x509" _MAX_EXCHANGE_RETRIES = 2 _REPLAY_POSITION_EXTENSION = "openai_x509_replay_position" +_REPLAY_FILE_POSITIONS_EXTENSION = "openai_x509_replay_file_positions" _ALLOWED_IDENTITY_FIELDS = {"type", "identity_provider_id", "service_account_id", "refresh_buffer_seconds"} @@ -94,13 +95,21 @@ def _is_replayable_request(request: httpx2.Request) -> bool: stream = request.stream fields = getattr(stream, "fields", None) if isinstance(fields, list): + file_positions: list[tuple[object, int]] = [] for field in cast(list[object], fields): file = getattr(field, "file", None) if file is None or isinstance(file, (str, bytes)): continue seekable = getattr(file, "seekable", None) - if not callable(seekable) or not seekable(): + seek = getattr(file, "seek", None) + tell = getattr(file, "tell", None) + if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell): return False + position = tell() + if not isinstance(position, int): + return False + file_positions.append((file, position)) + request.extensions[_REPLAY_FILE_POSITIONS_EXTENSION] = file_positions return True source = getattr(stream, "_stream", stream) @@ -165,6 +174,14 @@ def _can_retry_request(self, request: httpx2.Request) -> bool: @override def _prepare_retry_request(self, request: httpx2.Request) -> None: + file_positions = request.extensions.get(_REPLAY_FILE_POSITIONS_EXTENSION) + if isinstance(file_positions, list): + for file, file_position in cast(list[tuple[object, int]], file_positions): + seek = getattr(file, "seek", None) + if callable(seek): + seek(file_position) + return + position = request.extensions.get(_REPLAY_POSITION_EXTENSION) if not isinstance(position, int): return diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index 0573b4dae1..13487fe418 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -13,6 +13,7 @@ from openai import OpenAI, AsyncOpenAI, OpenAIError, APIStatusError from openai.auth import X509WorkloadIdentity, x509_workload_identity +from openai._client import WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER from openai.auth._workload import _WorkloadIdentityAuth _TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" @@ -112,6 +113,62 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert (await copied.models.list()).object == "list" +@pytest.mark.parametrize("method", ["copy", "with_options"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +def test_sync_copy_can_switch_from_api_key_to_x509_identity(method: str, base_url: str | None) -> None: + api_authorizations: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "switched-token", "expires_in": 3600}) + api_authorizations.append(request.headers["Authorization"]) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(api_key="original-api-key", base_url=base_url, http_client=http_client, max_retries=0) as client: + copied = ( + client.copy(workload_identity=_identity()) + if method == "copy" + else client.with_options(workload_identity=_identity()) + ) + assert client.api_key == "original-api-key" + assert copied.api_key == WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER + assert str(copied.base_url) == f"{base_url or 'https://mtls.api.openai.com/v1'}/" + assert copied._client is http_client + assert copied.models.list().object == "list" + + assert api_authorizations == ["Bearer switched-token"] + + +@pytest.mark.parametrize("method", ["copy", "with_options"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +async def test_async_copy_can_switch_from_api_key_to_x509_identity(method: str, base_url: str | None) -> None: + api_authorizations: list[str] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "switched-token", "expires_in": 3600}) + api_authorizations.append(request.headers["Authorization"]) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI( + api_key="original-api-key", base_url=base_url, http_client=http_client, max_retries=0 + ) as client: + copied = ( + client.copy(workload_identity=_identity()) + if method == "copy" + else client.with_options(workload_identity=_identity()) + ) + assert client.api_key == "original-api-key" + assert copied.api_key == WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER + assert str(copied.base_url) == f"{base_url or 'https://mtls.api.openai.com/v1'}/" + assert copied._client is http_client + assert (await copied.models.list()).object == "list" + + assert api_authorizations == ["Bearer switched-token"] + + @pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: urls: list[str] = [] @@ -150,11 +207,13 @@ async def handler(request: httpx2.Request) -> httpx2.Response: def test_sync_x509_exchange_does_not_inherit_caller_http_auth() -> None: exchange_authorizations: list[str | None] = [] + api_authorizations: list[str | None] = [] def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_authorizations.append(request.headers.get("Authorization")) return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + api_authorizations.append(request.headers.get("Authorization")) return httpx2.Response(200, request=request, json={"object": "list", "data": []}) http_client = httpx2.Client( @@ -166,15 +225,18 @@ def handler(request: httpx2.Request) -> httpx2.Response: assert client.models.list().object == "list" assert exchange_authorizations == [None] + assert api_authorizations == ["Bearer safe-token"] async def test_async_x509_exchange_does_not_inherit_caller_http_auth() -> None: exchange_authorizations: list[str | None] = [] + api_authorizations: list[str | None] = [] async def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_authorizations.append(request.headers.get("Authorization")) return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + api_authorizations.append(request.headers.get("Authorization")) return httpx2.Response(200, request=request, json={"object": "list", "data": []}) http_client = httpx2.AsyncClient( @@ -186,6 +248,7 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert (await client.models.list()).object == "list" assert exchange_authorizations == [None] + assert api_authorizations == ["Bearer safe-token"] @pytest.mark.parametrize("response_body", [[], "not-an-object", 42, True]) @@ -275,6 +338,74 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert api_authorizations == ["Bearer token-1", "Bearer token-2"] +def test_sync_x509_rewinds_seekable_multipart_upload_to_its_original_position() -> None: + upload = io.BytesIO(b"prefix-upload-payload") + upload.seek(len(b"prefix-")) + initial_position = upload.tell() + positions: list[int] = [] + bodies: list[bytes] = [] + exchange_calls = 0 + + class ConsumingTransport(httpx2.BaseTransport): + @override + def handle_request(self, request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + positions.append(upload.tell()) + bodies.append(request.read()) + if exchange_calls == 1: + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + return httpx2.Response(200, request=request, json={"id": "file_123", "object": "file"}) + + http_client = httpx2.Client(transport=ConsumingTransport(), trust_env=False) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + uploaded = client.files.create(file=("upload.txt", upload), purpose="assistants") + + assert uploaded.id == "file_123" + assert positions == [initial_position, initial_position] + assert len(bodies) == 2 and bodies[0] == bodies[1] + assert b"upload-payload" in bodies[0] + + +async def test_async_x509_rewinds_seekable_multipart_upload_to_its_original_position() -> None: + upload = io.BytesIO(b"prefix-upload-payload") + upload.seek(len(b"prefix-")) + initial_position = upload.tell() + positions: list[int] = [] + bodies: list[bytes] = [] + exchange_calls = 0 + + class ConsumingTransport(httpx2.AsyncBaseTransport): + @override + async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_calls + if str(request.url) == _TOKEN_URL: + exchange_calls += 1 + return httpx2.Response( + 200, request=request, json={"access_token": f"token-{exchange_calls}", "expires_in": 3600} + ) + + positions.append(upload.tell()) + bodies.append(await request.aread()) + if exchange_calls == 1: + return httpx2.Response(401, request=request, json={"error": {"message": "unauthorized"}}) + return httpx2.Response(200, request=request, json={"id": "file_123", "object": "file"}) + + http_client = httpx2.AsyncClient(transport=ConsumingTransport(), trust_env=False) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + uploaded = await client.files.create(file=("upload.txt", upload), purpose="assistants") + + assert uploaded.id == "file_123" + assert positions == [initial_position, initial_position] + assert len(bodies) == 2 and bodies[0] == bodies[1] + assert b"upload-payload" in bodies[0] + + def test_sync_x509_invalidates_a_rejected_replay_token() -> None: exchange_calls = 0 api_authorizations: list[str] = [] From 9512899eb98aab917ec9cf4c59342c77eb45c394 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 11:00:00 -0700 Subject: [PATCH 07/11] fix(auth): support bidirectional X.509 client copies --- src/openai/_client.py | 24 +++++--- ...test_x509_workload_identity_regressions.py | 56 +++++++++++++++++++ 2 files changed, 72 insertions(+), 8 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index b770d2df1b..d5c0c33480 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -674,17 +674,21 @@ def copy( "base_url": base_url, } else: + next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity + if api_key is not None and workload_identity is None: + next_workload_identity = None + current_x509 = is_x509_workload_identity(self.workload_identity) + next_x509 = is_x509_workload_identity(next_workload_identity) + current_default_base_url = f"{MTLS_API_BASE_URL}/" if current_x509 else "https://api.openai.com/v1/" inherited_base_url = ( - None - if is_x509_workload_identity(workload_identity) and str(self.base_url) == "https://api.openai.com/v1/" - else self.base_url + None if current_x509 != next_x509 and str(self.base_url) == current_default_base_url else self.base_url ) auth_options = { "api_key": api_key if workload_identity is not None else api_key or self._api_key_provider or self.api_key, "admin_api_key": admin_api_key or self.admin_api_key, - "workload_identity": workload_identity or self.workload_identity, + "workload_identity": next_workload_identity, "base_url": base_url or inherited_base_url, } @@ -1317,17 +1321,21 @@ def copy( "base_url": base_url, } else: + next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity + if api_key is not None and workload_identity is None: + next_workload_identity = None + current_x509 = is_x509_workload_identity(self.workload_identity) + next_x509 = is_x509_workload_identity(next_workload_identity) + current_default_base_url = f"{MTLS_API_BASE_URL}/" if current_x509 else "https://api.openai.com/v1/" inherited_base_url = ( - None - if is_x509_workload_identity(workload_identity) and str(self.base_url) == "https://api.openai.com/v1/" - else self.base_url + None if current_x509 != next_x509 and str(self.base_url) == current_default_base_url else self.base_url ) auth_options = { "api_key": api_key if workload_identity is not None else api_key or self._api_key_provider or self.api_key, "admin_api_key": admin_api_key or self.admin_api_key, - "workload_identity": workload_identity or self.workload_identity, + "workload_identity": next_workload_identity, "base_url": base_url or inherited_base_url, } diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index 13487fe418..7d79649202 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -169,6 +169,62 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert api_authorizations == ["Bearer switched-token"] +@pytest.mark.parametrize("method", ["copy", "with_options"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +def test_sync_copy_can_switch_from_x509_identity_to_api_key(method: str, base_url: str | None) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False) + with OpenAI(workload_identity=_identity(), base_url=base_url, http_client=http_client, max_retries=0) as client: + copied = ( + client.copy(api_key="replacement-api-key") + if method == "copy" + else client.with_options(api_key="replacement-api-key") + ) + assert client.workload_identity == _identity() + assert copied.workload_identity is None + assert copied.api_key == "replacement-api-key" + assert str(copied.base_url) == f"{base_url or 'https://api.openai.com/v1'}/" + assert copied._client is http_client + assert copied.models.list().object == "list" + + assert len(requests) == 1 + assert requests[0].headers["Authorization"] == "Bearer replacement-api-key" + + +@pytest.mark.parametrize("method", ["copy", "with_options"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +async def test_async_copy_can_switch_from_x509_identity_to_api_key(method: str, base_url: str | None) -> None: + requests: list[httpx2.Request] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False) + async with AsyncOpenAI( + workload_identity=_identity(), base_url=base_url, http_client=http_client, max_retries=0 + ) as client: + copied = ( + client.copy(api_key="replacement-api-key") + if method == "copy" + else client.with_options(api_key="replacement-api-key") + ) + assert client.workload_identity == _identity() + assert copied.workload_identity is None + assert copied.api_key == "replacement-api-key" + assert str(copied.base_url) == f"{base_url or 'https://api.openai.com/v1'}/" + assert copied._client is http_client + assert (await copied.models.list()).object == "list" + + assert len(requests) == 1 + assert requests[0].headers["Authorization"] == "Bearer replacement-api-key" + + @pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: urls: list[str] = [] From 145d3f2a02c993b620fe79b2b8fba4b34e375bfb Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 18:34:24 +0000 Subject: [PATCH 08/11] fix(auth): isolate X.509 exchange authorization --- src/openai/auth/_x509.py | 9 +++- ...test_x509_workload_identity_regressions.py | 46 +++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 90e38d4711..02be50d1a2 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -64,6 +64,11 @@ def _exchange_payload(identity: X509WorkloadIdentity) -> dict[str, str]: } +def _strip_authorization(request: httpx2.Request) -> httpx2.Request: + request.headers.pop("Authorization", None) + return request + + def _retry_delay(response: httpx2.Response | None, attempt: int) -> float | None: if response is not None: if response.status_code not in (408, 409, 429) and response.status_code < 500: @@ -205,7 +210,7 @@ def _fetch_token_from_exchange(self) -> dict[str, Any]: response = self._http_client.post( _X509_TOKEN_EXCHANGE_URL, json=_exchange_payload(self.workload_identity), - auth=lambda request: request, + auth=_strip_authorization, timeout=10.0, follow_redirects=False, ) @@ -250,7 +255,7 @@ async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: response = await self._http_client.post( _X509_TOKEN_EXCHANGE_URL, json=_exchange_payload(self.workload_identity), - auth=lambda request: request, + auth=_strip_authorization, timeout=10.0, follow_redirects=False, ) diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index 7d79649202..73ed0fe54b 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -307,6 +307,52 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert api_authorizations == ["Bearer safe-token"] +def test_sync_x509_exchange_does_not_inherit_caller_authorization_header() -> None: + exchange_authorizations: list[str | None] = [] + api_authorizations: list[str | None] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + exchange_authorizations.append(request.headers.get("Authorization")) + return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + api_authorizations.append(request.headers.get("Authorization")) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.Client( + transport=httpx2.MockTransport(handler), + headers={"Authorization": "Bearer private-api-credential"}, + trust_env=False, + ) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert exchange_authorizations == [None] + assert api_authorizations == ["Bearer safe-token"] + + +async def test_async_x509_exchange_does_not_inherit_caller_authorization_header() -> None: + exchange_authorizations: list[str | None] = [] + api_authorizations: list[str | None] = [] + + async def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + exchange_authorizations.append(request.headers.get("Authorization")) + return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + api_authorizations.append(request.headers.get("Authorization")) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + http_client = httpx2.AsyncClient( + transport=httpx2.MockTransport(handler), + headers={"Authorization": "Bearer private-api-credential"}, + trust_env=False, + ) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert exchange_authorizations == [None] + assert api_authorizations == ["Bearer safe-token"] + + @pytest.mark.parametrize("response_body", [[], "not-an-object", 42, True]) def test_sync_x509_rejects_non_object_token_responses(response_body: object) -> None: def handler(request: httpx2.Request) -> httpx2.Response: From fe8d0652b487295bf5c76781d1e2a4b28e1b903b Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 18:59:09 +0000 Subject: [PATCH 09/11] fix(client): preserve explicit base URLs on copy --- src/openai/_client.py | 30 ++++++++++++------- ...test_x509_workload_identity_regressions.py | 24 ++++++++++++--- 2 files changed, 40 insertions(+), 14 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index d5c0c33480..c58af6a5d7 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -122,6 +122,7 @@ class OpenAI(SyncAPIClient): _workload_identity_auth: WorkloadIdentityAuth | SyncX509WorkloadIdentityAuth | None _provider: _Provider | None _provider_runtime: _ProviderRuntime | None + _base_url_was_default: bool websocket_base_url: str | httpx2.URL | None """Base URL for WebSocket connections. @@ -267,6 +268,7 @@ def __init__( base_url = provider_runtime.base_url elif base_url is None: base_url = os.environ.get("OPENAI_BASE_URL") + self._base_url_was_default = provider_runtime is None and base_url is None if base_url is None: base_url = MTLS_API_BASE_URL if x509_identity is not None else "https://api.openai.com/v1" @@ -657,6 +659,7 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider + preserve_default_base_url = False auth_options: dict[str, Any] if next_provider is not None: auth_options = { @@ -679,10 +682,9 @@ def copy( next_workload_identity = None current_x509 = is_x509_workload_identity(self.workload_identity) next_x509 = is_x509_workload_identity(next_workload_identity) - current_default_base_url = f"{MTLS_API_BASE_URL}/" if current_x509 else "https://api.openai.com/v1/" - inherited_base_url = ( - None if current_x509 != next_x509 and str(self.base_url) == current_default_base_url else self.base_url - ) + mode_changed = current_x509 != next_x509 + inherited_base_url = None if mode_changed and self._base_url_was_default else self.base_url + preserve_default_base_url = base_url is None and not mode_changed and self._base_url_was_default auth_options = { "api_key": api_key if workload_identity is not None @@ -692,7 +694,7 @@ def copy( "base_url": base_url or inherited_base_url, } - return self.__class__( + copied = self.__class__( organization=organization or inherited_organization, project=project or inherited_project, webhook_secret=webhook_secret or self.webhook_secret, @@ -706,6 +708,9 @@ def copy( **auth_options, **_extra_kwargs, ) + if preserve_default_base_url: + copied._base_url_was_default = True + return copied # Alias for `copy` for nicer inline usage, e.g. # client.with_options(timeout=10).foo.create(...) @@ -757,6 +762,7 @@ class AsyncOpenAI(AsyncAPIClient): _workload_identity_auth: WorkloadIdentityAuth | AsyncX509WorkloadIdentityAuth | None _provider: _Provider | None _provider_runtime: _ProviderRuntime | None + _base_url_was_default: bool websocket_base_url: str | httpx2.URL | None """Base URL for WebSocket connections. @@ -902,6 +908,7 @@ def __init__( base_url = provider_runtime.base_url elif base_url is None: base_url = os.environ.get("OPENAI_BASE_URL") + self._base_url_was_default = provider_runtime is None and base_url is None if base_url is None: base_url = MTLS_API_BASE_URL if x509_identity is not None else "https://api.openai.com/v1" @@ -1304,6 +1311,7 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider + preserve_default_base_url = False auth_options: dict[str, Any] if next_provider is not None: auth_options = { @@ -1326,10 +1334,9 @@ def copy( next_workload_identity = None current_x509 = is_x509_workload_identity(self.workload_identity) next_x509 = is_x509_workload_identity(next_workload_identity) - current_default_base_url = f"{MTLS_API_BASE_URL}/" if current_x509 else "https://api.openai.com/v1/" - inherited_base_url = ( - None if current_x509 != next_x509 and str(self.base_url) == current_default_base_url else self.base_url - ) + mode_changed = current_x509 != next_x509 + inherited_base_url = None if mode_changed and self._base_url_was_default else self.base_url + preserve_default_base_url = base_url is None and not mode_changed and self._base_url_was_default auth_options = { "api_key": api_key if workload_identity is not None @@ -1339,7 +1346,7 @@ def copy( "base_url": base_url or inherited_base_url, } - return self.__class__( + copied = self.__class__( organization=organization or inherited_organization, project=project or inherited_project, webhook_secret=webhook_secret or self.webhook_secret, @@ -1353,6 +1360,9 @@ def copy( **auth_options, **_extra_kwargs, ) + if preserve_default_base_url: + copied._base_url_was_default = True + return copied # Alias for `copy` for nicer inline usage, e.g. # client.with_options(timeout=10).foo.create(...) diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index 73ed0fe54b..bb47fd5d7b 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -114,7 +114,7 @@ async def handler(request: httpx2.Request) -> httpx2.Response: @pytest.mark.parametrize("method", ["copy", "with_options"]) -@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1", "https://api.openai.com/v1"]) def test_sync_copy_can_switch_from_api_key_to_x509_identity(method: str, base_url: str | None) -> None: api_authorizations: list[str] = [] @@ -141,7 +141,7 @@ def handler(request: httpx2.Request) -> httpx2.Response: @pytest.mark.parametrize("method", ["copy", "with_options"]) -@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1", "https://api.openai.com/v1"]) async def test_async_copy_can_switch_from_api_key_to_x509_identity(method: str, base_url: str | None) -> None: api_authorizations: list[str] = [] @@ -170,7 +170,7 @@ async def handler(request: httpx2.Request) -> httpx2.Response: @pytest.mark.parametrize("method", ["copy", "with_options"]) -@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1", "https://mtls.api.openai.com/v1"]) def test_sync_copy_can_switch_from_x509_identity_to_api_key(method: str, base_url: str | None) -> None: requests: list[httpx2.Request] = [] @@ -197,7 +197,7 @@ def handler(request: httpx2.Request) -> httpx2.Response: @pytest.mark.parametrize("method", ["copy", "with_options"]) -@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1"]) +@pytest.mark.parametrize("base_url", [None, "https://custom.example/v1", "https://mtls.api.openai.com/v1"]) async def test_async_copy_can_switch_from_x509_identity_to_api_key(method: str, base_url: str | None) -> None: requests: list[httpx2.Request] = [] @@ -225,6 +225,22 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert requests[0].headers["Authorization"] == "Bearer replacement-api-key" +def test_sync_copy_preserves_implicit_base_url_provenance_across_chained_copies() -> None: + http_client = httpx2.Client(trust_env=False) + with OpenAI(api_key="original-api-key", http_client=http_client) as client: + copied = client.copy(timeout=1).copy(workload_identity=_identity()) + + assert str(copied.base_url) == "https://mtls.api.openai.com/v1/" + + +async def test_async_copy_preserves_implicit_base_url_provenance_across_chained_copies() -> None: + http_client = httpx2.AsyncClient(trust_env=False) + async with AsyncOpenAI(api_key="original-api-key", http_client=http_client) as client: + copied = client.copy(timeout=1).copy(workload_identity=_identity()) + + assert str(copied.base_url) == "https://mtls.api.openai.com/v1/" + + @pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: urls: list[str] = [] From eb39b75108bac2c6e17f24c5a58068ad9d89f41b Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 19:12:51 +0000 Subject: [PATCH 10/11] fix(auth): reject oversized X.509 timing values --- src/openai/auth/_x509.py | 28 +++++++++++++++++----------- tests/test_x509_workload_identity.py | 4 ++-- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 02be50d1a2..a13aeaa297 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -29,6 +29,16 @@ _ALLOWED_IDENTITY_FIELDS = {"type", "identity_provider_id", "service_account_id", "refresh_buffer_seconds"} +def _as_finite_float(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + try: + value = float(value) + except OverflowError: + return None + return value if math.isfinite(value) else None + + def is_x509_workload_identity( identity: WorkloadIdentity | X509WorkloadIdentity | None, ) -> TypeIs[X509WorkloadIdentity]: @@ -46,13 +56,10 @@ def _validate_identity(identity: X509WorkloadIdentity) -> None: raise OpenAIError("X.509 workload identity requires identity-provider and service-account IDs") refresh_buffer = cast(object, identity.get("refresh_buffer_seconds")) - if refresh_buffer is not None and ( - isinstance(refresh_buffer, bool) - or not isinstance(refresh_buffer, (int, float)) - or not math.isfinite(refresh_buffer) - or refresh_buffer < 0 - ): - raise OpenAIError("X.509 workload identity requires a finite, non-negative refresh buffer") + if refresh_buffer is not None: + refresh_buffer_value = _as_finite_float(refresh_buffer) + if refresh_buffer_value is None or refresh_buffer_value < 0: + raise OpenAIError("X.509 workload identity requires a finite, non-negative refresh buffer") def _exchange_payload(identity: X509WorkloadIdentity) -> dict[str, str]: @@ -167,11 +174,10 @@ def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]: @override def _validate_expires_in(self, expires_in: object) -> float: - if isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)): - raise OpenAIError("X.509 token exchange response did not include a positive, finite expires_in") - if not math.isfinite(expires_in) or expires_in <= 0: + expires_in_value = _as_finite_float(expires_in) + if expires_in_value is None or expires_in_value <= 0: raise OpenAIError("X.509 token exchange response did not include a positive, finite expires_in") - return float(expires_in) + return expires_in_value @override def _can_retry_request(self, request: httpx2.Request) -> bool: diff --git a/tests/test_x509_workload_identity.py b/tests/test_x509_workload_identity.py index 97ecde50ad..83a56eaf3a 100644 --- a/tests/test_x509_workload_identity.py +++ b/tests/test_x509_workload_identity.py @@ -76,7 +76,7 @@ def test_x509_rejects_certificate_and_token_material_without_leaking_it(invalid_ assert secret not in str(error.value) -@pytest.mark.parametrize("refresh_buffer", [-1.0, float("inf"), float("nan"), True]) +@pytest.mark.parametrize("refresh_buffer", [-1.0, float("inf"), float("nan"), True, 10**400]) def test_x509_rejects_invalid_refresh_buffer(refresh_buffer: float) -> None: with pytest.raises(OpenAIError, match="finite, non-negative refresh buffer"): OpenAI(workload_identity=_identity(refresh_buffer_seconds=refresh_buffer)) @@ -207,7 +207,7 @@ def test_api_key_clients_keep_ordinary_api_endpoint() -> None: assert str(client.base_url) == "https://api.openai.com/v1/" -@pytest.mark.parametrize("expires_in", [0, -1, True, "3600", None]) +@pytest.mark.parametrize("expires_in", [0, -1, True, "3600", None, 10**400]) def test_x509_rejects_nonpositive_or_nonnumeric_expiration(expires_in: object) -> None: def handler(request: httpx2.Request) -> httpx2.Response: return _token_response(request, expires_in=expires_in) From edeba35fe8ade3bb296576080b35835cbe4063fd Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 19:25:07 +0000 Subject: [PATCH 11/11] fix(auth): isolate X.509 exchange requests --- src/openai/_client.py | 22 +++- src/openai/auth/_x509.py | 26 ++--- ...test_x509_workload_identity_regressions.py | 100 +++++++++++++++--- 3 files changed, 118 insertions(+), 30 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index c58af6a5d7..105193ef4b 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -29,7 +29,7 @@ get_async_library, ) from ._compat import cached_property -from ._httpx2 import is_httpx2_sync_client, is_httpx2_async_client +from ._httpx2 import normalize_httpx_url, is_httpx2_sync_client, is_httpx2_async_client from ._models import SecurityOptions, FinalRequestOptions from ._version import __version__ from ._provider import _Provider, _provider_name, _ProviderRuntime, _configure_provider @@ -132,6 +132,16 @@ class OpenAI(SyncAPIClient): 'wss://example.com' """ + @property + @override + def base_url(self) -> httpx2.URL: + return self._base_url + + @base_url.setter + def base_url(self, url: httpx2.URL | str) -> None: + self._base_url = self._enforce_trailing_slash(normalize_httpx_url(url)) + self._base_url_was_default = False + def __init__( self, *, @@ -772,6 +782,16 @@ class AsyncOpenAI(AsyncAPIClient): 'wss://example.com' """ + @property + @override + def base_url(self) -> httpx2.URL: + return self._base_url + + @base_url.setter + def base_url(self, url: httpx2.URL | str) -> None: + self._base_url = self._enforce_trailing_slash(normalize_httpx_url(url)) + self._base_url_was_default = False + def __init__( self, *, diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index a13aeaa297..9912edc890 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -71,9 +71,13 @@ def _exchange_payload(identity: X509WorkloadIdentity) -> dict[str, str]: } -def _strip_authorization(request: httpx2.Request) -> httpx2.Request: - request.headers.pop("Authorization", None) - return request +def _token_exchange_request(identity: X509WorkloadIdentity) -> httpx2.Request: + return httpx2.Request( + "POST", + _X509_TOKEN_EXCHANGE_URL, + json=_exchange_payload(identity), + extensions={"timeout": httpx2.Timeout(10.0).as_dict()}, + ) def _retry_delay(response: httpx2.Response | None, attempt: int) -> float | None: @@ -213,11 +217,9 @@ def __init__( def _fetch_token_from_exchange(self) -> dict[str, Any]: for attempt in range(self._max_exchange_retries + 1): try: - response = self._http_client.post( - _X509_TOKEN_EXCHANGE_URL, - json=_exchange_payload(self.workload_identity), - auth=_strip_authorization, - timeout=10.0, + response = self._http_client.send( + _token_exchange_request(self.workload_identity), + auth=None, follow_redirects=False, ) except _transport_errors() as error: @@ -258,11 +260,9 @@ async def get_token_async(self) -> str: async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: for attempt in range(self._max_exchange_retries + 1): try: - response = await self._http_client.post( - _X509_TOKEN_EXCHANGE_URL, - json=_exchange_payload(self.workload_identity), - auth=_strip_authorization, - timeout=10.0, + response = await self._http_client.send( + _token_exchange_request(self.workload_identity), + auth=None, follow_redirects=False, ) except _transport_errors() as error: diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py index bb47fd5d7b..5a28c1dfe4 100644 --- a/tests/test_x509_workload_identity_regressions.py +++ b/tests/test_x509_workload_identity_regressions.py @@ -241,6 +241,48 @@ async def test_async_copy_preserves_implicit_base_url_provenance_across_chained_ assert str(copied.base_url) == "https://mtls.api.openai.com/v1/" +@pytest.mark.parametrize("starts_with_x509", [False, True]) +def test_sync_copy_preserves_base_url_assigned_through_setter(starts_with_x509: bool) -> None: + http_client = httpx2.Client(trust_env=False) + client = ( + OpenAI(workload_identity=_identity(), http_client=http_client) + if starts_with_x509 + else OpenAI(api_key="original-api-key", http_client=http_client) + ) + with client: + client.base_url = "https://assigned.example/v1" + same_mode_copy = client.copy(timeout=1) + copied = ( + same_mode_copy.copy(api_key="replacement-api-key") + if starts_with_x509 + else same_mode_copy.copy(workload_identity=_identity()) + ) + + assert str(same_mode_copy.base_url) == "https://assigned.example/v1/" + assert str(copied.base_url) == "https://assigned.example/v1/" + + +@pytest.mark.parametrize("starts_with_x509", [False, True]) +async def test_async_copy_preserves_base_url_assigned_through_setter(starts_with_x509: bool) -> None: + http_client = httpx2.AsyncClient(trust_env=False) + client = ( + AsyncOpenAI(workload_identity=_identity(), http_client=http_client) + if starts_with_x509 + else AsyncOpenAI(api_key="original-api-key", http_client=http_client) + ) + async with client: + client.base_url = "https://assigned.example/v1" + same_mode_copy = client.copy(timeout=1) + copied = ( + same_mode_copy.copy(api_key="replacement-api-key") + if starts_with_x509 + else same_mode_copy.copy(workload_identity=_identity()) + ) + + assert str(same_mode_copy.base_url) == "https://assigned.example/v1/" + assert str(copied.base_url) == "https://assigned.example/v1/" + + @pytest.mark.parametrize("authorization", [None, "Bearer caller-override"]) def test_sync_x509_disables_redirects_without_placeholder_authorization(authorization: str | None) -> None: urls: list[str] = [] @@ -323,50 +365,76 @@ async def handler(request: httpx2.Request) -> httpx2.Response: assert api_authorizations == ["Bearer safe-token"] -def test_sync_x509_exchange_does_not_inherit_caller_authorization_header() -> None: - exchange_authorizations: list[str | None] = [] - api_authorizations: list[str | None] = [] +def test_sync_x509_exchange_does_not_inherit_caller_request_state() -> None: + exchange_headers: list[httpx2.Headers] = [] + api_headers: list[httpx2.Headers] = [] def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: - exchange_authorizations.append(request.headers.get("Authorization")) + exchange_headers.append(request.headers) return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) - api_authorizations.append(request.headers.get("Authorization")) + api_headers.append(request.headers) return httpx2.Response(200, request=request, json={"object": "list", "data": []}) http_client = httpx2.Client( transport=httpx2.MockTransport(handler), - headers={"Authorization": "Bearer private-api-credential"}, + headers={ + "Authorization": "Bearer private-api-credential", + "X-API-Key": "private-api-credential", + "Content-Type": "application/private", + }, + cookies={"session": "private-cookie"}, trust_env=False, ) with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: assert client.models.list().object == "list" - assert exchange_authorizations == [None] - assert api_authorizations == ["Bearer safe-token"] + assert len(exchange_headers) == 1 + assert exchange_headers[0].get("Authorization") is None + assert exchange_headers[0].get("X-API-Key") is None + assert exchange_headers[0].get("Cookie") is None + assert exchange_headers[0]["Content-Type"] == "application/json" + assert len(api_headers) == 1 + assert api_headers[0]["Authorization"] == "Bearer safe-token" + assert api_headers[0]["X-API-Key"] == "private-api-credential" + assert api_headers[0]["Cookie"] == "session=private-cookie" + assert api_headers[0]["Content-Type"] == "application/private" -async def test_async_x509_exchange_does_not_inherit_caller_authorization_header() -> None: - exchange_authorizations: list[str | None] = [] - api_authorizations: list[str | None] = [] +async def test_async_x509_exchange_does_not_inherit_caller_request_state() -> None: + exchange_headers: list[httpx2.Headers] = [] + api_headers: list[httpx2.Headers] = [] async def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: - exchange_authorizations.append(request.headers.get("Authorization")) + exchange_headers.append(request.headers) return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) - api_authorizations.append(request.headers.get("Authorization")) + api_headers.append(request.headers) return httpx2.Response(200, request=request, json={"object": "list", "data": []}) http_client = httpx2.AsyncClient( transport=httpx2.MockTransport(handler), - headers={"Authorization": "Bearer private-api-credential"}, + headers={ + "Authorization": "Bearer private-api-credential", + "X-API-Key": "private-api-credential", + "Content-Type": "application/private", + }, + cookies={"session": "private-cookie"}, trust_env=False, ) async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: assert (await client.models.list()).object == "list" - assert exchange_authorizations == [None] - assert api_authorizations == ["Bearer safe-token"] + assert len(exchange_headers) == 1 + assert exchange_headers[0].get("Authorization") is None + assert exchange_headers[0].get("X-API-Key") is None + assert exchange_headers[0].get("Cookie") is None + assert exchange_headers[0]["Content-Type"] == "application/json" + assert len(api_headers) == 1 + assert api_headers[0]["Authorization"] == "Bearer safe-token" + assert api_headers[0]["X-API-Key"] == "private-api-credential" + assert api_headers[0]["Cookie"] == "session=private-cookie" + assert api_headers[0]["Content-Type"] == "application/private" @pytest.mark.parametrize("response_body", [[], "not-an-object", 42, True])