f12e6b66d9
Generic plugin-OAuth support that the Solstice provider needs. None of it names a provider: - OAuthPKCEConfig.token_request: a confidential client whose secret lives behind a broker hands the grant to the plugin in place of a token_url POST. The old refresh_token is kept when the response omits one, which Google does on refresh. - Plugin OAuth rows rotate 2 minutes ahead of their stored expires_at_ms (shared with the anthropic rule). Opaque bearers such as Google's ya29.* carry no JWT exp, so this stored expiry is the only clock they have. - A row whose refresh grant just went DEAD stays DEAD on the inference 401 that follows; before this it was re-benched EXHAUSTED and the dead token was replayed every cooldown. - The auxiliary client serves an OAuth plugin that ships create_client from its pooled row on the main runtime's endpoint (compression, titles). It used to return "not directly supported". - A 401 on a plugin OAuth row prints a re-sign-in hint, not the API-key advice. - The Gemini native adapter exposes its method names, auth headers and error builder so a Bearer-auth subclass can reuse it without the API-key remedies.
106 lines
5.1 KiB
Python
106 lines
5.1 KiB
Python
"""Plugin-provider refresh support for the credential pool (#116408).
|
|
|
|
A model-provider plugin makes its pooled OAuth rows refreshable by shipping
|
|
``ProviderProfile.refresh_credential(entry) -> Mapping | None``. The pool owns
|
|
the locking, persistence and failure classification around that hook; the
|
|
plugin owns only the token POST. This sibling module keeps that logic out of
|
|
``agent/credential_pool.py`` (near the size cap).
|
|
|
|
Contract (documented in website/docs/developer-guide/model-provider-plugin.md):
|
|
|
|
* the hook returns a mapping of rotated values — dataclass field names
|
|
(``access_token``, ``refresh_token``, ``expires_at_ms`` …) replace the row's
|
|
fields, every other key (``expires_in``, ``token_type``, ``scope`` — the
|
|
natural token-endpoint shape) lands in ``entry.extra``; ``None``/empty = could not
|
|
rotate and the pool benches the row like a failed refresh POST;
|
|
* raising ``AuthError(..., relogin_required=True)`` (or a grant-dead OAuth code)
|
|
is terminal: the row goes DEAD with a WARNING naming ``hermes auth add``;
|
|
any other exception is transient and only benches the row.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
from dataclasses import fields, replace
|
|
from typing import TYPE_CHECKING, Any, Mapping, Optional, Tuple
|
|
|
|
from hermes_cli.auth import _OAUTH_GRANT_DEAD_CODES
|
|
from hermes_cli.auth_constants import AuthError
|
|
|
|
if TYPE_CHECKING: # pragma: no cover
|
|
from agent.credential_pool import CredentialPool, PooledCredential
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_EXPIRY_SKEW_MS = 120_000
|
|
|
|
|
|
def plugin_row_is_expiring(entry: "PooledCredential") -> bool:
|
|
"""An expiry-stamped OAuth row (a plugin's, or Anthropic's) is due within the skew of ``expires_at_ms``.
|
|
|
|
The token endpoint's ``expires_in`` is the only clock: many providers issue opaque bearers (Google's
|
|
``ya29.*``) with no JWT ``exp`` to decode, so without this a long session sends the dead bearer.
|
|
"""
|
|
return entry.expires_at_ms is not None and int(entry.expires_at_ms) <= int(time.time() * 1000) + _EXPIRY_SKEW_MS
|
|
|
|
|
|
def apply_plugin_refresh_result(entry: "PooledCredential", result: Any) -> "PooledCredential":
|
|
"""Merge a ``refresh_credential`` return value into *entry*.
|
|
|
|
Field names go through ``dataclasses.replace``; everything else is merged into ``extra``
|
|
(mirroring ``PooledCredential.from_dict``). Before this split a token-endpoint-shaped mapping
|
|
with ``expires_in`` raised ``TypeError`` inside ``replace`` and the pool benched a row whose
|
|
single-use refresh token the server had already rotated — the login was lost.
|
|
"""
|
|
if not result:
|
|
return entry
|
|
mapping: Mapping[str, Any] = dict(result)
|
|
field_names = {f.name for f in fields(type(entry))} - {"provider", "extra"}
|
|
field_updates = {k: v for k, v in mapping.items() if k in field_names}
|
|
extra_updates = {k: v for k, v in mapping.items() if k not in field_names and k != "provider"}
|
|
if extra_updates:
|
|
field_updates["extra"] = {**entry.extra, **extra_updates}
|
|
return replace(entry, **field_updates) if field_updates else entry
|
|
|
|
|
|
def is_terminal_plugin_refresh_error(exc: BaseException) -> bool:
|
|
"""True when retrying the same plugin refresh token cannot succeed.
|
|
|
|
Plugins have no per-provider code table, so the predicate is the structural one every built-in
|
|
flow shares: a structured ``AuthError`` that asks for a re-login, or one carrying a grant-dead
|
|
OAuth code. Transport errors, 429/5xx-shaped ``RuntimeError`` and plain ``AuthError`` without
|
|
that signal stay transient (EXHAUSTED for one cooldown).
|
|
"""
|
|
if not isinstance(exc, AuthError):
|
|
return False
|
|
return bool(exc.relogin_required) or (exc.code or "") in _OAUTH_GRANT_DEAD_CODES
|
|
|
|
|
|
def recover_failed_plugin_refresh(
|
|
pool: "CredentialPool", entry: "PooledCredential", exc: Exception,
|
|
) -> Tuple[bool, Optional["PooledCredential"]]:
|
|
"""Recovery for a plugin hook that raised: adopt a peer's rotation, or quarantine a dead grant.
|
|
|
|
Returns ``(handled, result)``; ``handled=False`` means the caller should bench the row as a
|
|
transient failure. A peer process may have rotated the single-use token between our in-lock
|
|
sync and the hook's POST — adopt that pair before classifying the failure.
|
|
"""
|
|
from agent.credential_pool import _MARK_OK
|
|
|
|
synced = pool._sync_entry_from_pool_store(entry)
|
|
if synced.refresh_token != entry.refresh_token and (synced.access_token or "").strip():
|
|
logger.debug("%s refresh failed but the pool store has newer tokens — adopting", pool.provider)
|
|
return True, pool._adopt(synced, **_MARK_OK)
|
|
if is_terminal_plugin_refresh_error(exc):
|
|
# WARNING, not debug: this is the moment a login is lost. Benching for a TTL would replay
|
|
# the dead token every cooldown at DEBUG with no trace for the user.
|
|
logger.warning(
|
|
"%s refresh token for %s is terminally invalid (%s); the credential leaves rotation. "
|
|
"Re-run 'hermes auth add %s' to sign in again.",
|
|
pool.provider, entry.label or entry.id[:8], exc, pool.provider,
|
|
)
|
|
pool._mark_dead_refresh_grant(entry, exc)
|
|
return True, None
|
|
return False, None
|