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..105193ef4b 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, @@ -29,11 +29,17 @@ 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 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, @@ -109,13 +115,14 @@ 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 - _workload_identity_auth: WorkloadIdentityAuth | None + _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. @@ -125,12 +132,22 @@ 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, *, 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, @@ -246,12 +263,24 @@ def __init__( self.websocket_base_url = websocket_base_url + if is_x509_workload_identity(workload_identity): + x509_identity = workload_identity + subject_token_identity = None + elif workload_identity is None: + x509_identity = 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: 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 = 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: @@ -273,9 +302,13 @@ 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 + ) + elif subject_token_identity is not None: self._workload_identity_auth = WorkloadIdentityAuth( - workload_identity=workload_identity, + workload_identity=subject_token_identity, _use_httpx2=is_httpx2_sync_client(self._client), ) @@ -458,27 +491,30 @@ 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) 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 not retried - ): - response.close() - 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) + 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 {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}" + return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @override def _send_request( @@ -572,7 +608,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) @@ -588,7 +624,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, @@ -633,6 +669,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 = { @@ -650,14 +687,24 @@ 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) + 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 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, + "workload_identity": next_workload_identity, + "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, @@ -671,6 +718,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(...) @@ -715,13 +765,14 @@ 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 - _workload_identity_auth: WorkloadIdentityAuth | None + _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. @@ -731,12 +782,22 @@ 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, *, 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, @@ -852,12 +913,24 @@ def __init__( self.websocket_base_url = websocket_base_url + if is_x509_workload_identity(workload_identity): + x509_identity = workload_identity + subject_token_identity = None + elif workload_identity is None: + x509_identity = 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: 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 = 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: @@ -879,9 +952,13 @@ 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 + ) + elif subject_token_identity is not None: self._workload_identity_auth = WorkloadIdentityAuth( - workload_identity=workload_identity, + workload_identity=subject_token_identity, _use_httpx2=is_httpx2_async_client(self._client), ) @@ -1064,27 +1141,30 @@ 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) 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 not retried - ): - await response.aclose() - 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) + 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 {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}" + return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) @override async def _send_request( @@ -1189,7 +1269,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 @@ -1205,7 +1287,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, @@ -1249,6 +1331,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 = { @@ -1266,14 +1349,24 @@ 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) + 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 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, + "workload_identity": next_workload_identity, + "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, @@ -1287,6 +1380,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/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..4f9797f092 100644 --- a/src/openai/auth/_workload.py +++ b/src/openai/auth/_workload.py @@ -2,12 +2,13 @@ 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, NotRequired +from typing_extensions import Literal, TypeAlias, NotRequired, override 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 @@ -42,6 +43,38 @@ class WorkloadIdentity(TypedDict): refresh_buffer_seconds: NotRequired[float] +SubjectTokenWorkloadIdentity: TypeAlias = WorkloadIdentity + + +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] + + +_WorkloadIdentityT = TypeVar("_WorkloadIdentityT", WorkloadIdentity, 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: @@ -171,17 +204,18 @@ 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.workload_identity = workload_identity + ) -> None: + self.workload_identity: _WorkloadIdentityT = 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 @@ -222,14 +256,19 @@ 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 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"] @@ -239,32 +278,7 @@ def _perform_refresh(self) -> 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() - - 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) + raise NotImplementedError("Workload identity authentication must implement token exchange") def _handle_token_response(self, response: httpx2.Response) -> dict[str, Any]: try: @@ -278,24 +292,22 @@ 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: 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 _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 + 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 _token_unusable(self) -> bool: return self._cached_token is None or self._token_expired() @@ -314,3 +326,62 @@ 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 + + +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 new file mode 100644 index 0000000000..9912edc890 --- /dev/null +++ b/src/openai/auth/_x509.py @@ -0,0 +1,280 @@ +from __future__ import annotations + +import math +import time +import email.utils +from typing import Any, NoReturn, cast +from typing_extensions import TypeIs, 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, + X509WorkloadIdentity, + _WorkloadIdentityAuth, +) +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" +_REPLAY_FILE_POSITIONS_EXTENSION = "openai_x509_replay_file_positions" +_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]: + 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.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")) + 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]: + 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 _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: + 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): + 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) + 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) + 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[X509WorkloadIdentity]): + 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: + 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 expires_in_value + + @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: + 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 + 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]: + for attempt in range(self._max_exchange_retries + 1): + try: + response = self._http_client.send( + _token_exchange_request(self.workload_identity), + auth=None, + 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]: + for attempt in range(self._max_exchange_retries + 1): + try: + response = await self._http_client.send( + _token_exchange_request(self.workload_identity), + auth=None, + 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/src/openai/lib/azure.py b/src/openai/lib/azure.py index 5e0b6f5203..b0ddbf4e5d 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 @@ -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 @@ -284,7 +285,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, @@ -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, @@ -608,7 +611,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, @@ -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/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/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_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_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..83a56eaf3a --- /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, 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)) + + +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, 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) + + 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() + await asyncio.gather(waiter, return_exceptions=True) + assert waiter.cancelled() + 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() + await asyncio.gather(owner, return_exceptions=True) + assert owner.cancelled() + 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) diff --git a/tests/test_x509_workload_identity_regressions.py b/tests/test_x509_workload_identity_regressions.py new file mode 100644 index 0000000000..5a28c1dfe4 --- /dev/null +++ b/tests/test_x509_workload_identity_regressions.py @@ -0,0 +1,729 @@ +from __future__ import annotations + +import io +import asyncio +import threading +from typing import cast, 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, 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" +_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("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("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("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: + 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("method", ["copy", "with_options"]) +@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] = [] + + 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", "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] = [] + + 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("method", ["copy", "with_options"]) +@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] = [] + + 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", "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] = [] + + 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" + + +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("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] = [] + + 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] + + +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( + 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] + 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( + 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] + assert api_authorizations == ["Bearer safe-token"] + + +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_headers.append(request.headers) + return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + 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", + "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 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_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_headers.append(request.headers) + return httpx2.Response(200, request=request, json={"access_token": "safe-token", "expires_in": 3600}) + 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", + "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 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]) +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_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_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] = [] + + 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 + 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