"""OpenAI Codex subscription-auth provider: OAuth flow primitives, credential refresh, wire dialect.

Core owns the protocol primitives (PKCE context, authorization URL, code exchange, refresh) and,
because the pinned redirect URI makes every login a localhost callback, the one-shot callback
catcher (`exchange_code_from_callback()`); browser opening and persistent credential storage
belong to applications and harnesses.

The authorization-code + PKCE redirect flow is the only login flow the public Codex client
supports: its registration pins the redirect URI to `http://localhost:1455/auth/callback`
(exact-match, probed live 2026-08-25), and the auth service serves no device-authorization
endpoint. `exchange_code_from_callback()` serves this exact redirect URI and exchanges the
authorization code.
"""

from __future__ import annotations as _annotations

import base64
import json
import os
from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping
from dataclasses import KW_ONLY, dataclass, field
from datetime import datetime, timedelta, timezone
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any, Protocol
from urllib.parse import urlencode

import anyio
import httpx2
from pydantic import Field, StrictFloat, StrictInt, TypeAdapter, ValidationError
from typing_extensions import Self

from pydantic_ai._http import create_async_httpx2_client
from pydantic_ai.exceptions import ModelAPIError, UserError
from pydantic_ai.profiles import ModelProfile
from pydantic_ai.profiles.openai_codex import openai_codex_model_profile

from ._oauth import OAuthFlow
from ._openai_compatible import (
    AsyncHTTPClient as _OpenAIHTTPClient,
    OpenAICompatibleProvider as _OpenAICompatibleProvider,
)

if TYPE_CHECKING:
    from openai import AsyncOpenAI

try:
    from openai import AsyncOpenAI, OpenAIError
except ImportError as _import_error:  # pragma: no cover
    raise ImportError(
        'Please install the `openai` package to use the OpenAI Codex provider, '
        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
    ) from _import_error

__all__ = (
    'CredentialsPersistenceError',
    'CredentialsRefreshError',
    'OpenAICodexCredentialSource',
    'OpenAICodexCredentials',
    'OpenAICodexOAuthFlow',
    'OpenAICodexProvider',
)

_CODEX_BASE_URL = 'https://chatgpt.com/backend-api/codex'
_CODEX_HOST = httpx2.URL(_CODEX_BASE_URL).host
_AUTHORIZE_URL = 'https://auth.openai.com/oauth/authorize'
_TOKEN_URL = 'https://auth.openai.com/oauth/token'
# The public Codex CLI client: OpenAI's registration pins the redirect URI to exactly this
# localhost URI (probed: alternates rejected pre-login); override `redirect_uri=` only with your own client.
_PUBLIC_CLIENT_ID = 'app_EMoamEEZ73f0CkXaXp7hrann'
_REDIRECT_URI = 'http://localhost:1455/auth/callback'
_DEFAULT_SCOPE = 'openid profile email offline_access'
_ORIGINATOR = 'pydantic-ai'
# 30s pre-expiry refresh hint (unverified JWT `exp` is a hint, not an authority).
_TOKEN_EXPIRY_BUFFER = timedelta(seconds=30)


class _CredentialsError(ModelAPIError):
    """Base for Codex credential failures.

    Subclasses [`ModelAPIError`][pydantic_ai.exceptions.ModelAPIError] so the standard handling of
    provider failures (e.g. [`FallbackModel`][pydantic_ai.models.fallback.FallbackModel]) applies;
    the auth layer runs below any specific model, so `model_name` is the provider name.
    """

    def __init__(self, message: str):
        super().__init__(model_name='openai-codex', message=message)

    def __reduce__(self) -> tuple[type, tuple[Any, ...]]:
        return self.__class__, (self.message,)


class CredentialsRefreshError(_CredentialsError):
    """Refreshing Codex credentials against the token endpoint failed.

    When the underlying error is `invalid_grant`, the stored grant is no longer usable and a fresh
    authorization is required (locally: rerun `codex login`; in an app: rerun your connect flow).
    """


# The OpenAI SDK retries arbitrary transport exceptions, but propagates `OpenAIError` unchanged.
class CredentialsPersistenceError(_CredentialsError, OpenAIError):
    """Rotated credentials were updated in memory but the persistence callback raised.

    The in-memory credentials are current and were handed to the callback before it failed; the
    error surfaces so callers do not mistake a failed save for durability.
    """


@dataclass
class OpenAICodexCredentials:
    """Codex subscription credentials.

    The tokens are excluded from `repr` so a logged or traceback-embedded instance does not leak them.
    """

    _: KW_ONLY
    access_token: str = field(repr=False)
    refresh_token: str = field(repr=False)
    account_id: str

    @classmethod
    def from_codex_cli_auth(cls, data: Mapping[str, Any]) -> Self:
        """Parse the Codex CLI `~/.codex/auth.json` shape (`{'tokens': {...}}`)."""
        try:
            tokens = _codex_cli_auth_ta.validate_python(data).tokens
        except ValidationError as e:
            raise UserError(
                f'Malformed Codex CLI credentials. Run `codex login` to regenerate them.\n\n{e.json(include_input=False)}'
            ) from None
        return cls(
            access_token=tokens.access_token,
            refresh_token=tokens.refresh_token,
            account_id=tokens.account_id,
        )


# Dataclasses for the wire payloads consulted above and below, validated through `TypeAdapter`:
# extra fields are ignored, and a `ValidationError` (a `ValueError`) also rejects non-object payloads.


@dataclass
class _CodexCliTokens:
    """The `tokens` entry of the Codex CLI's `auth.json`."""

    access_token: Annotated[str, Field(min_length=1)]
    refresh_token: Annotated[str, Field(min_length=1)]
    account_id: Annotated[str, Field(min_length=1)]


@dataclass
class _CodexCliAuth:
    """The subset of the Codex CLI's `auth.json` that credentials are built from."""

    tokens: _CodexCliTokens


@dataclass
class _JwtAuthClaim:
    """The nested OpenAI claim carrying the ChatGPT account id."""

    chatgpt_account_id: str | None = None


@dataclass
class _JwtPayload:
    """The unverified JWT claims consulted for expiry and account-id hints."""

    # Strict, so a numeric string is not coerced into an expiry hint.
    exp: StrictInt | StrictFloat | None = None
    # The Codex id_token nests the account id under this claim.
    auth: Annotated[_JwtAuthClaim | None, Field(validation_alias='https://api.openai.com/auth')] = None
    chatgpt_account_id: str | None = None
    account_id: str | None = None


@dataclass
class _TokenResponse:
    """The fields of an OAuth token-endpoint response that credentials are built from.

    `refresh_token` is required because the flow always requests the `offline_access` scope, and
    credentials without it could not survive their first expiry.
    """

    access_token: Annotated[str, Field(min_length=1)]
    refresh_token: Annotated[str, Field(min_length=1)]
    id_token: str | None = None
    account_id: str | None = None


@dataclass
class _TokenErrorResponse:
    """An OAuth token-endpoint error body."""

    error: str | None = None
    error_description: str | None = None


_codex_cli_auth_ta = TypeAdapter(_CodexCliAuth)
_jwt_payload_ta = TypeAdapter(_JwtPayload)
_token_response_ta = TypeAdapter(_TokenResponse)
_token_error_response_ta = TypeAdapter(_TokenErrorResponse)


def _jwt_payload(token: str) -> _JwtPayload | None:
    """Decode a JWT payload without verifying the signature. Returns `None` for anything malformed."""
    try:
        segment = token.split('.')[1]
    except IndexError:
        return None
    padded = segment + '=' * (-len(segment) % 4)
    try:
        return _jwt_payload_ta.validate_python(json.loads(base64.urlsafe_b64decode(padded)))
    except ValueError:
        return None


def _jwt_expires_at(token: str) -> datetime | None:
    """Best-effort unverified `exp` claim, a refresh hint, never an authority."""
    payload = _jwt_payload(token)
    if payload is None or payload.exp is None:
        return None
    try:
        return datetime.fromtimestamp(payload.exp, tz=timezone.utc)
    except (OverflowError, OSError, ValueError):
        return None


def _account_id_from_id_token(token: str) -> str | None:
    payload = _jwt_payload(token)
    if payload is None:
        return None
    if payload.auth and payload.auth.chatgpt_account_id:
        return payload.auth.chatgpt_account_id
    return payload.chatgpt_account_id or payload.account_id or None


def _credentials_from_token_response(
    data: _TokenResponse, fallback_account_id: str | None = None
) -> OpenAICodexCredentials:
    """Build credentials from an OAuth token-endpoint response."""
    account_id = (
        data.account_id or (_account_id_from_id_token(data.id_token) if data.id_token else None) or fallback_account_id
    )
    if not account_id:
        raise CredentialsRefreshError('Could not determine the ChatGPT account id from the token response.')
    return OpenAICodexCredentials(
        access_token=data.access_token, refresh_token=data.refresh_token, account_id=account_id
    )


async def _post_token_request(
    url: str, form: Mapping[str, str], http_client: httpx2.AsyncClient | None = None
) -> _TokenResponse:
    """POST a form-urlencoded OAuth token request and decode the JSON response.

    When `http_client` is given the request goes through it (so custom transports and proxies apply
    to refreshes too); otherwise an ephemeral client is used.
    """
    if http_client is None:
        async with httpx2.AsyncClient(timeout=httpx2.Timeout(timeout=30, connect=5)) as client:
            response = await client.post(url, data=dict(form), headers={'Accept': 'application/json'})
    else:
        response = await http_client.post(url, data=dict(form), headers={'Accept': 'application/json'})
    if response.status_code != 200:
        try:
            body = _token_error_response_ta.validate_python(response.json())
        except ValueError:
            body = _TokenErrorResponse()
        detail = body.error_description or body.error or response.text[:200]
        hint = '; the grant was rejected, rerun the authorization flow' if body.error == 'invalid_grant' else ''
        raise CredentialsRefreshError(
            f'Token request to {url} failed with status {response.status_code}: {detail}{hint}'
        )
    try:
        return _token_response_ta.validate_python(response.json())
    except ValueError as e:
        detail = e.json(include_input=False) if isinstance(e, ValidationError) else str(e)
        raise CredentialsRefreshError(f'Token endpoint {url} returned an unexpected response.\n\n{detail}') from None


async def _refresh_credentials(
    credentials: OpenAICodexCredentials, *, http_client: httpx2.AsyncClient | None = None
) -> OpenAICodexCredentials:
    """Exchange the refresh token for a new credential set against the public Codex client.

    Raises [`CredentialsRefreshError`][pydantic_ai.providers.openai_codex.CredentialsRefreshError]
    when the token endpoint rejects the grant (`invalid_grant` means a fresh authorization is
    required) or returns a malformed response.
    """
    data = await _post_token_request(
        _TOKEN_URL,
        {
            'grant_type': 'refresh_token',
            'refresh_token': credentials.refresh_token,
            'client_id': _PUBLIC_CLIENT_ID,
        },
        http_client=http_client,
    )
    return _credentials_from_token_response(data, fallback_account_id=credentials.account_id)


class OpenAICodexCredentialSource(Protocol):
    """Application-owned storage for the credentials, so refreshed tokens outlive the process.

    The provider owns the credential lifecycle (expiry checks, refresh, single-flight) and calls
    this only to read and write the stored set: `load()` on first use, and `save()` after tokens
    are refreshed.

    Refresh tokens are single-use, so before refreshing, the provider re-reads storage with
    `load()` and adopts newer credentials if another process using the same store already
    refreshed them. That is best-effort, not mutual exclusion: two processes can still race
    between `load()` and `save()`.

    Conformance is structural, but implementations are encouraged to subclass the protocol
    explicitly so type checkers verify the method signatures.
    """

    async def load(self) -> OpenAICodexCredentials:
        """Return the currently stored credentials."""
        ...

    async def save(self, credentials: OpenAICodexCredentials) -> None:
        """Durably replace the stored credentials with a freshly rotated set."""
        ...


def _read_codex_cli_credentials() -> OpenAICodexCredentials:
    """Read-only load of the Codex CLI's `auth.json` (honors `CODEX_HOME`). Never writes it."""
    code_home = Path(os.getenv('CODEX_HOME') or Path.home() / '.codex')
    path = code_home / 'auth.json'
    try:
        text = path.read_text()
    except FileNotFoundError:
        raise UserError(
            f'No Codex CLI credentials found at `{path}`. Run `codex login` first, or pass '
            '`credentials=` / use `OpenAICodexOAuthFlow` explicitly.'
        ) from None
    except OSError as e:
        raise UserError(f'Could not read Codex CLI credentials at `{path}`: {e}') from e
    try:
        data = json.loads(text)
    except ValueError as e:
        raise UserError(f'Malformed Codex CLI credentials at `{path}`: {e}') from e
    return OpenAICodexCredentials.from_codex_cli_auth(data)


class OpenAICodexOAuthFlow(OAuthFlow[OpenAICodexCredentials]):
    """Pure authorization-code + PKCE context for the OpenAI Codex public client.

    This is the only login flow the public client supports (no device flow; redirect URI pinned to
    `localhost:1455`, probed exact-match). Construction does no I/O: build the context anywhere,
    send the user to `authorization_url()`, then let `exchange_code_from_callback()` receive the
    redirect on localhost and exchange its code. The browser and credential storage stay
    caller-owned.
    """

    def __init__(self, *, redirect_uri: str = _REDIRECT_URI, state: str | None = None) -> None:
        """Create a new flow context. Construction does no I/O.

        Args:
            redirect_uri: Where the authorization code is delivered. The public client's
                registration pins this to `http://localhost:1455/auth/callback` (exact-match),
                so leave the default unchanged when using the public client.
            state: The CSRF token bound to the callback; auto-generated when `None`.
        """
        super().__init__(redirect_uri=redirect_uri, state=state)

    def authorization_url(self, *, scope: str | None = None, extra_params: Mapping[str, str] | None = None) -> str:
        """The URL to send the user to. Note the public client pins redirects to localhost.

        Args:
            scope: The OAuth scopes to request; `None` means the standard Codex login scopes.
            extra_params: Additional query parameters, merged over the defaults (so they can also
                override them), except `client_id` and `redirect_uri`: `exchange_code()` always
                posts the public client id and the flow's `redirect_uri`, so overriding either
                here would make the authorization code unusable. The production Codex login's
                `id_token_add_organizations=true` and `codex_cli_simplified_flow=true` are sent by
                default: without the former, the `id_token` can omit the account id for multi-org
                accounts (live-verified 2026-08-25).
        """
        params: dict[str, str] = {
            'response_type': 'code',
            'client_id': _PUBLIC_CLIENT_ID,
            'redirect_uri': self.redirect_uri,
            'scope': scope if scope is not None else _DEFAULT_SCOPE,
            'state': self.state,
            'code_challenge': self.code_challenge,
            'code_challenge_method': 'S256',
            'id_token_add_organizations': 'true',
            'codex_cli_simplified_flow': 'true',
        }
        return f'{_AUTHORIZE_URL}?{urlencode(self._merge_extra_params(params, extra_params))}'

    async def exchange_code(self, code: str) -> OpenAICodexCredentials:
        """Exchange an authorization code for credentials (call this in your callback handler)."""
        data = await _post_token_request(
            _TOKEN_URL,
            {
                'grant_type': 'authorization_code',
                'code': code,
                'code_verifier': self.code_verifier,
                'redirect_uri': self.redirect_uri,
                'client_id': _PUBLIC_CLIENT_ID,
            },
        )
        return _credentials_from_token_response(data)


class _OpenAICodexAuth(httpx2.Auth):
    """httpx auth injecting Codex subscription headers, with single-flight refresh-and-replay.

    Injects `Authorization: Bearer …`, `chatgpt-account-id`, and `originator`, but only on
    HTTPS requests to the Codex host, so a caller-supplied client reused for other destinations
    (or downgraded to plaintext) never leaks credentials. On a 401 it performs at most one refresh-and-replay; non-expiry 401s
    therefore cannot loop. The proactive expiry check treats the unverified JWT `exp` as a hint only.
    """

    def __init__(self, provider: OpenAICodexProvider) -> None:
        self._provider = provider

    def _apply_headers(self, request: httpx2.Request, credentials: OpenAICodexCredentials) -> None:
        request.headers['Authorization'] = f'Bearer {credentials.access_token}'
        request.headers['chatgpt-account-id'] = credentials.account_id
        request.headers['originator'] = _ORIGINATOR

    def sync_auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
        # `httpx.Auth`'s default would send the request unauthenticated; refresh-and-replay is async.
        raise UserError('`OpenAICodexProvider` requires an async HTTP client to inject credentials.')

    async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
        if request.url.scheme != 'https' or request.url.host != _CODEX_HOST:
            # Never send subscription credentials to a foreign destination or over plaintext:
            # a caller-supplied client may be reused for arbitrary requests.
            yield request
            return
        # Buffer the outgoing body so the 401 replay can resend it: a one-shot stream (an async
        # generator upload) would otherwise raise `StreamConsumed` on the second send. Codex
        # payloads are JSON already held in memory, and foreign-host requests bypass this above.
        await request.aread()
        # The two classes are deliberately coupled in one module; the provider owns the state and
        # the auth is its wire-side skin. The replay closure carries whatever context the
        # provider's mode (in-memory single-flight, or application credential source) needs.
        credentials, replay = await self._provider._prepare_request_credentials()  # pyright: ignore[reportPrivateUsage]
        self._apply_headers(request, credentials)
        response = yield request
        if response.status_code != 401:
            return
        # Release the connection before replaying.
        await response.aread()
        self._apply_headers(request, await replay())
        yield request  # replay exactly once; its response goes back to the caller


class OpenAICodexProvider(_OpenAICompatibleProvider):
    """Provider for OpenAI Codex subscription authentication.

    Wraps the standard `OpenAIProvider` machinery pointed at the Codex backend, injecting Codex
    OAuth credentials instead of API keys. One provider instance carries one set of credentials
    (there is no process-global cache). The instance binds its refresh lock to the first event
    loop that awaits a request, so do not reuse it across loops.

    ```python {test="skip" lint="skip"}
    provider = OpenAICodexProvider(credential_source=YourCredentialStore())
    agent = Agent('openai-codex:gpt-5.6-luna', provider=provider)
    ```
    """

    @property
    def name(self) -> str:
        return 'openai-codex'

    @property
    def base_url(self) -> str:
        return _CODEX_BASE_URL

    @property
    def client(self) -> AsyncOpenAI:
        return self._client

    @property
    def credentials(self) -> OpenAICodexCredentials:
        """The credentials currently held in memory, rotated in place by refreshes."""
        if self._credentials is None:
            raise UserError(
                '`credentials` is unavailable: the provider either wraps an existing `openai_client`, '
                'which opts out of credential injection entirely, or has a `credential_source` it has '
                'not loaded from yet.'
            )
        return self._credentials

    @staticmethod
    def model_profile(model_name: str) -> ModelProfile | None:
        return openai_codex_model_profile(model_name)

    def __init__(
        self,
        credentials: OpenAICodexCredentials | None = None,
        *,
        credential_source: OpenAICodexCredentialSource | None = None,
        openai_client: AsyncOpenAI | None = None,
        http_client: _OpenAIHTTPClient | None = None,
    ) -> None:
        """Create a new OpenAI Codex provider.

        Args:
            credentials: The subscription credentials to inject. If both this and
                `credential_source` are omitted, they are loaded **read-only** from the Codex
                CLI's `auth.json` (honors `CODEX_HOME`), which never writes the file: refreshed
                tokens then live in memory only. Pydantic AI never falls back to `OPENAI_API_KEY`.
            credential_source: Application-owned storage for the credentials, so refreshed tokens
                are persisted between runs; see
                [`OpenAICodexCredentialSource`][pydantic_ai.providers.openai_codex.OpenAICodexCredentialSource].
                Mutually exclusive with `credentials`.
            openai_client: An existing `AsyncOpenAI` client to use as-is. Opts out of credential
                injection entirely; `credentials`, `credential_source`, and `http_client` must
                be `None`.
            http_client: An existing `httpx2.AsyncClient` to use. Must be dedicated to this
                provider (no auth of its own): the provider attaches its credential-injecting auth to it, and
                sharing a client between providers would mix their credentials. The auth only
                injects credentials on HTTPS requests to the Codex host, so the client can safely
                be reused for other destinations.
        """
        self._credential_source = credential_source
        self._credentials: OpenAICodexCredentials | None = None
        if openai_client is not None:
            assert credentials is None, 'Cannot provide both `openai_client` and `credentials`'
            assert credential_source is None, 'Cannot provide both `openai_client` and `credential_source`'
            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
            self._client = openai_client
            return

        if credential_source is not None:
            # Loaded lazily on first use: construction stays synchronous and does no I/O.
            assert credentials is None, 'Cannot provide both `credentials` and `credential_source`'
        else:
            self._credentials = credentials if credentials is not None else _read_codex_cli_credentials()
        self._revision = 0
        self._refresh_failures = 0
        self._last_refresh_error: tuple[int, Exception] | None = None
        self._auth = _OpenAICodexAuth(self)
        if http_client is None:
            http_client = create_async_httpx2_client()
            self._own_http_client = http_client
            self._http_client_factory = self._create_http_client
        else:
            if not isinstance(http_client, httpx2.AsyncClient):
                raise UserError(
                    '`OpenAICodexProvider` requires an `httpx2` client for `http_client`: the legacy '
                    '`httpx.AsyncClient` cannot carry its credential-injecting auth.'
                )
            if http_client.auth is not None:
                raise UserError(
                    'The `http_client` already has auth configured (it may belong to another provider); '
                    'pass a dedicated client so credentials cannot mix between providers.'
                )
        http_client.auth = self._auth
        self._http_client = http_client
        self._client = AsyncOpenAI(
            base_url=_CODEX_BASE_URL,
            # The SDK merges its own bearer header into requests; `_OpenAICodexAuth` replaces it.
            api_key='codex-subscription-auth',
            http_client=http_client,
        )

    def _create_http_client(self) -> httpx2.AsyncClient:
        """Factory used when a closed provider-owned client is reopened."""
        client = create_async_httpx2_client()
        client.auth = self._auth
        self._http_client = client
        return client

    def _set_http_client(self, http_client: _OpenAIHTTPClient) -> None:
        http_client.auth = self._auth  # pyright: ignore[reportAttributeAccessIssue]
        self._client._client = http_client  # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue]

    async def _prepare_request_credentials(
        self,
    ) -> tuple[OpenAICodexCredentials, Callable[[], Awaitable[OpenAICodexCredentials]]]:
        """The credentials for an outgoing request, plus a replay resolver for a 401 on it."""
        await self._load_if_needed()
        await self._refresh_if_stale()
        revision_used = self._revision
        # Only share failures that happen after this request starts, not failures from earlier requests.
        refresh_failures = self._refresh_failures

        async def replay() -> OpenAICodexCredentials:
            await self._refresh_for_401(revision_used, refresh_failures=refresh_failures)
            return self.credentials

        return self.credentials, replay

    async def _load_if_needed(self) -> None:
        """First use with a `credential_source`: load the stored credentials."""
        if self._credentials is not None:
            return
        async with self._refresh_lock:
            if self._credentials is None:  # recheck after acquiring: load once, not once per task
                assert self._credential_source is not None
                self._credentials = await self._credential_source.load()

    @cached_property
    def _refresh_lock(self) -> anyio.Lock:
        # Like the base provider's enter lock: bind lazily so we attach to whatever loop first
        # awaits a request rather than construction time.
        return anyio.Lock()

    async def _refresh_if_stale(self) -> None:
        """Proactive refresh from the unverified-JWT `exp` hint.

        Failures are swallowed here: the hint is best-effort, and the 401 path retries with real
        errors surfaced.
        """
        if not self._is_stale():
            return
        try:
            async with self._refresh_lock:
                if self._is_stale():  # recheck after acquiring: single-flight, not just serialized
                    await self._refresh_locked()
        except CredentialsPersistenceError:
            raise  # the refresh itself succeeded; a failed save must never be silent
        except Exception:
            # Transport failures and rejected grants alike fall through to the 401 path, which
            # retries with the still-current token and surfaces real errors.
            pass

    async def _refresh_for_401(self, revision_used: int, *, refresh_failures: int) -> None:
        """Single-flight refresh after a 401 carrying `revision_used`.

        If another task already replaced the credentials since the failed request was sent, no
        network refresh happens: the caller replays with the fresh set directly.
        """
        if self._revision != revision_used:
            return
        async with self._refresh_lock:
            if self._revision != revision_used:  # recheck after acquiring
                return
            if (
                (last := self._last_refresh_error) is not None
                and last[0] == revision_used
                and self._refresh_failures != refresh_failures
            ):
                raise last[1]  # share the single-flight failure instead of re-running it per waiter
            try:
                await self._refresh_locked()
            except Exception as e:
                self._refresh_failures += 1
                self._last_refresh_error = (revision_used, e)
                raise

    async def _refresh_locked(self) -> None:
        # The caller must hold `_refresh_lock`.
        assert self._refresh_lock.locked()
        rejected = self.credentials
        if (source := self._credential_source) is not None:
            # Refresh tokens are single-use, so spending ours when another process already refreshed
            # would invalidate theirs. Re-read storage first and adopt a newer set if one is there.
            if (stored := await source.load()) != rejected:
                self._replace(stored)
                return
        # The provider's own client, so custom transports and proxies apply to refreshes too;
        # the auth flow ignores non-Codex hosts, so this cannot recurse or leak the bearer.
        self._replace(await _refresh_credentials(rejected, http_client=self._http_client))
        if source is not None:
            try:
                await source.save(self.credentials)
            except Exception as e:
                raise CredentialsPersistenceError(
                    'Credentials were refreshed in memory but saving them to the credential source raised.'
                ) from e

    def _replace(self, credentials: OpenAICodexCredentials) -> None:
        # Atomic replace of the complete set, then bump the revision so concurrent 401s observe it.
        self._credentials = credentials
        self._revision += 1

    def _is_stale(self) -> bool:
        # The unverified JWT `exp` claim is a refresh hint, not an authority; refresh proactively
        # once the token is within the pre-expiry buffer.
        expires_at = _jwt_expires_at(self.credentials.access_token)
        return expires_at is not None and datetime.now(timezone.utc) >= expires_at - _TOKEN_EXPIRY_BUFFER