Skip to content

fix(xai-oauth): serialize refreshes across profiles #81383

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
alexgunsberg wants to merge 2 commits into
base: main
Choose a base branch
from
+502 −58
Open

fix(xai-oauth): serialize refreshes across profiles #81383

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
64aa1fc
fix(xai-oauth): serialize refreshes across profiles
alexgunsberg Aug 7, 2026
93fb5e6
fix(xai-oauth): persist rotations to resolved source
alexgunsberg Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 64 additions & 31 deletions agent/credential_pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -589,9 +589,9 @@ def _write_through_provider_state_to_global_root(
``invalid_grant`` once its access token expires.

Only updates ``providers.<provider_id>`` in the root store; never touches
the profile store (the caller already saved that). Swallows all errors — a
failed write-through degrades to the pre-existing behavior (root stale), it
must never break the profile's own successful save. Mirrors
the profile store. Persistence errors propagate to the caller: a rotating
provider has already consumed the old refresh token, so reporting success
without durably storing the replacement would lose the token chain. Mirrors
``hermes_cli.auth._write_through_xai_oauth_to_global_root`` (which covers
the non-pool xAI refresh path) for the credential-pool refresh path.
"""
Expand All @@ -615,19 +615,12 @@ def _write_through_provider_state_to_global_root(
return
except Exception:
return
try:
auth_mod._persist_provider_state_to_store(
provider_id,
state,
global_path,
set_active=False,
)
except Exception as exc: # pragma: no cover - best effort
logger.debug(
"%s pool refresh: write-through to global root failed: %s",
provider_id,
exc,
)
auth_mod._persist_provider_state_to_store(
provider_id,
state,
global_path,
set_active=False,
)


class CredentialPool:
Expand Down Expand Up @@ -994,7 +987,9 @@ def _sync_xai_oauth_entry_from_auth_store(self, entry: PooledCredential) -> Pool
try:
with _auth_store_lock():
auth_store = _load_auth_store()
state = _load_provider_state(auth_store, "xai-oauth")
state, _source_path = auth_mod._load_xai_oauth_singleton_state_with_source(
auth_store
)
if not isinstance(state, dict):
return entry
tokens = state.get("tokens")
Expand Down Expand Up @@ -1211,8 +1206,8 @@ def _sync_device_code_entry_to_auth_store(self, entry: PooledCredential) -> None
if not isinstance(state, dict):
return
elif self.provider == "xai-oauth":
state, source_path = _load_provider_state_with_source(
auth_store, "xai-oauth"
state, source_path = auth_mod._load_xai_oauth_singleton_state_with_source(
auth_store
)
if not isinstance(state, dict):
return
Expand Down Expand Up @@ -1286,6 +1281,8 @@ def _sync_device_code_entry_to_auth_store(self, entry: PooledCredential) -> None
_save_auth_store(auth_store)
except Exception as exc:
logger.debug("Failed to sync %s pool entry back to auth store: %s", self.provider, exc)
if self.provider == "xai-oauth":
raise

def _refresh_entry(self, entry: PooledCredential, *, force: bool) -> Optional[PooledCredential]:
if entry.auth_type != AUTH_TYPE_OAUTH or not entry.refresh_token:
Expand All @@ -1303,14 +1300,27 @@ def _refresh_entry(self, entry: PooledCredential, *, force: bool) -> Optional[Po
# the lock, the in-lock re-sync below picks up the rotated token the
# winner persisted and skips the POST.
if self.provider in ("openai-codex", "xai-oauth"):
sync_entry = (
self._sync_codex_entry_from_auth_store
if self.provider == "openai-codex"
else self._sync_xai_oauth_entry_from_pool_store
)
with _auth_store_lock(
timeout_seconds=self._single_use_refresh_lock_timeout()
):
if self.provider == "openai-codex":
sync_entry = self._sync_codex_entry_from_auth_store
refresh_transaction = _auth_store_lock(
timeout_seconds=self._single_use_refresh_lock_timeout()
)
else:
# Singleton-seeded xAI entries can resolve from the global
# root while each profile persists its own pool mirror. The
# profile pool row is therefore not authoritative across
# profiles; re-read the singleton while holding both the
# profile and root locks. Manual entries remain profile-local
# and sync from their exact pool row as before.
sync_entry = (
self._sync_xai_oauth_entry_from_auth_store
if entry.source == "device_code"
else self._sync_xai_oauth_entry_from_pool_store
)
refresh_transaction = auth_mod._xai_oauth_refresh_transaction(
timeout_seconds=self._single_use_refresh_lock_timeout()
)
with refresh_transaction:
synced = sync_entry(entry)
if self.provider == "openai-codex":
if synced is not entry:
Expand Down Expand Up @@ -1500,7 +1510,12 @@ def _refresh_entry_impl(
try:
with _auth_store_lock():
auth_store = _load_auth_store()
state = _load_provider_state(auth_store, "xai-oauth") or {}
state, source_path = (
auth_mod._load_xai_oauth_singleton_state_with_source(
auth_store
)
)
state = dict(state or {})
if isinstance(state, dict):
tokens = state.get("tokens") or {}
if isinstance(tokens, dict):
Expand All @@ -1518,8 +1533,26 @@ def _refresh_entry_impl(
"relogin_required": True,
"at": datetime.now(timezone.utc).isoformat(),
}
_save_provider_state(auth_store, "xai-oauth", state)
_save_auth_store(auth_store)
global_root = _global_auth_file_path()
if (
source_path is not None
and global_root is not None
and _same_path(source_path, global_root)
):
auth_mod._persist_provider_state_to_store(
"xai-oauth",
state,
global_root,
set_active=False,
)
else:
_store_provider_state(
auth_store,
"xai-oauth",
state,
set_active=False,
)
_save_auth_store(auth_store)
except Exception as clear_exc:
logger.debug(
"Failed to clear terminal xAI OAuth state: %s", clear_exc
Expand Down Expand Up @@ -2791,7 +2824,7 @@ def _env_val(key: str) -> str:
# (``providers["xai-oauth"]``). Surface them in the pool too so
# ``hermes auth list`` reflects the logged-in state and so the pool
# is the single source of truth for refresh during runtime resolution.
state = _load_provider_state(auth_store, "xai-oauth")
state, _source_path = auth_mod._load_xai_oauth_singleton_state_with_source(auth_store)
tokens = state.get("tokens") if isinstance(state, dict) else None
if isinstance(tokens, dict) and tokens.get("access_token"):
# Device code is the only supported xAI OAuth flow; the singleton is
Expand Down
166 changes: 139 additions & 27 deletions hermes_cli/auth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1396,6 +1396,37 @@ def _provider_state_transaction(provider_id: str):
yield auth_store, source_state, source_path


@contextmanager
def _xai_oauth_refresh_transaction(
*,
timeout_seconds: float = AUTH_LOCK_TIMEOUT_SECONDS,
):
"""Serialize one xAI refresh across profiles sharing the global root.

xAI refresh tokens rotate on every use. A profile lock alone is not a
shared boundary: two profiles can each hold their own ``auth.lock`` while
spending the same root-fallback refresh token. Hold the active profile
lock and, when distinct, the global-root auth lock for the entire
re-read -> refresh POST -> persist sequence.

The profile-before-root order matches the existing write-through paths,
avoiding a lock inversion during rolling upgrades with older processes.
Profiles with independent xAI grants are conservatively serialized too;
refreshes are rare and correctness is more important than parallelism.
"""
with _auth_store_lock(timeout_seconds=timeout_seconds):
active_path = _auth_file_path()
global_path = _global_auth_file_path()
if global_path is None or _same_path(global_path, active_path):
yield
return
with _auth_store_lock(
timeout_seconds=timeout_seconds,
target_path=global_path,
):
yield


def _load_provider_state(auth_store: Dict[str, Any], provider_id: str) -> Optional[Dict[str, Any]]:
"""Return a provider's persisted state.

Expand Down Expand Up @@ -4418,9 +4449,24 @@ def _entry_usable(entry: Dict[str, Any]) -> bool:
# xAI Grok OAuth — tokens stored in ~/.hermes/auth.json
# =============================================================================

def _xai_oauth_provider_state_from_store(
auth_store: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
"""Return the local ``providers.xai-oauth`` block without fallback."""
providers = auth_store.get("providers")
raw_state = (
providers.get("xai-oauth") if isinstance(providers, dict) else None
)
return dict(raw_state) if isinstance(raw_state, dict) else None


def _xai_oauth_state_from_store(auth_store: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Return usable xAI OAuth state from provider state or credential pool."""
state = _load_provider_state(auth_store, "xai-oauth")
"""Return usable xAI OAuth state from this store's provider or pool."""
# Inspect only the store passed by the caller. Using
# ``_load_provider_state`` here would silently fall back to the global
# root, making it impossible for refresh code to identify where the
# rotating token chain actually came from.
state = _xai_oauth_provider_state_from_store(auth_store)
tokens = state.get("tokens") if isinstance(state, dict) else None
if isinstance(tokens, dict):
access_token = str(tokens.get("access_token", "") or "").strip()
Expand Down Expand Up @@ -4465,17 +4511,63 @@ def _xai_oauth_state_has_usable_tokens(state: Optional[Dict[str, Any]]) -> bool:
)


def _load_xai_oauth_state_with_source(
auth_store: Dict[str, Any],
) -> tuple[Optional[Dict[str, Any]], Optional[Path]]:
"""Resolve usable xAI state and the exact store that owns its token chain.

Unlike generic provider shadowing, an unusable local xAI provider block
must not hide a usable global grant. Refresh tokens rotate after one use,
so selecting the global grant but later saving by local key presence would
strand the rotated pair in the profile and leave root stale.
"""
local_state = _xai_oauth_state_from_store(auth_store)
if _xai_oauth_state_has_usable_tokens(local_state):
return local_state, _auth_file_path()

global_path = _global_auth_file_path()
global_state = _xai_oauth_state_from_store(_load_global_auth_store())
if _xai_oauth_state_has_usable_tokens(global_state):
return global_state, global_path

# Preserve a local invalid-shape state for useful diagnostics and for an
# intentional profile-scoped login. An unusable root block is not a
# refresh source and must not redirect a fresh profile login into root.
if isinstance(local_state, dict):
return local_state, _auth_file_path()
return None, None


def _load_xai_oauth_singleton_state_with_source(
auth_store: Dict[str, Any],
) -> tuple[Optional[Dict[str, Any]], Optional[Path]]:
"""Resolve a singleton-seeded xAI pool entry to its provider store.

Profile pools contain mirrors of root singleton credentials. Those local
rows must never become the apparent source after refresh, or sync-back
writes the rotated chain into the profile and leaves the root stale.
"""
local_state = _xai_oauth_provider_state_from_store(auth_store)
if _xai_oauth_state_has_usable_tokens(local_state):
return local_state, _auth_file_path()

global_path = _global_auth_file_path()
global_state = _xai_oauth_provider_state_from_store(_load_global_auth_store())
if _xai_oauth_state_has_usable_tokens(global_state):
return global_state, global_path

if isinstance(local_state, dict):
return local_state, _auth_file_path()
return None, None


def _read_xai_oauth_tokens(*, _lock: bool = True) -> Dict[str, Any]:
if _lock:
with _auth_store_lock():
auth_store = _load_auth_store()
else:
auth_store = _load_auth_store()
state = _xai_oauth_state_from_store(auth_store)
if not _xai_oauth_state_has_usable_tokens(state):
global_state = _xai_oauth_state_from_store(_load_global_auth_store())
if _xai_oauth_state_has_usable_tokens(global_state):
state = global_state
state, _source_path = _load_xai_oauth_state_with_source(auth_store)
if not state:
raise AuthError(
"No xAI OAuth credentials stored. Select xAI Grok OAuth (SuperGrok / Premium+) in `hermes model`.",
Expand Down Expand Up @@ -4536,10 +4628,9 @@ def _write_through_xai_oauth_to_global_root(state: Dict[str, Any]) -> None:
and every other profile reading the stale root grant dies with
``invalid_grant`` once its access token expires.

Only updates ``providers.xai-oauth`` in the root store; never touches the
profile store (the caller already saved that). Swallows all errors — a
failed write-through degrades to the pre-existing behavior (root stale),
it must never break the profile's own successful save.
Only updates ``providers.xai-oauth`` in the root store. Persistence errors
are fatal: xAI has already consumed the one-use refresh token, so reporting
success without durably storing its replacement would lose the token chain.
"""
global_path = _global_auth_file_path()
if global_path is None:
Expand All @@ -4558,15 +4649,12 @@ def _write_through_xai_oauth_to_global_root(state: Dict[str, Any]) -> None:
return
except Exception:
return
try:
_persist_provider_state_to_store(
"xai-oauth",
state,
global_path,
set_active=False,
)
except Exception as exc: # pragma: no cover - best effort
logger.debug("xAI OAuth: write-through to global root failed: %s", exc)
_persist_provider_state_to_store(
"xai-oauth",
state,
global_path,
set_active=False,
)


def _save_xai_oauth_tokens(
Expand Down Expand Up @@ -4600,9 +4688,7 @@ def _save_xai_oauth_tokens(
# unconditionally creates that key below. Use
# _load_provider_state_with_source to learn where the grant was
# resolved from and write back only to that source.
state, source_path = _load_provider_state_with_source(
auth_store, "xai-oauth"
)
state, source_path = _load_xai_oauth_state_with_source(auth_store)
if state is None:
state = {}
state["tokens"] = tokens
Expand Down Expand Up @@ -4994,7 +5080,12 @@ def resolve_xai_oauth_runtime_credentials(
if (not should_refresh) and refresh_if_expiring:
should_refresh = _xai_access_token_is_expiring(access_token, effective_skew)
if should_refresh:
with _auth_store_lock(timeout_seconds=max(float(AUTH_LOCK_TIMEOUT_SECONDS), refresh_timeout_seconds + 5.0)):
with _xai_oauth_refresh_transaction(
timeout_seconds=max(
float(AUTH_LOCK_TIMEOUT_SECONDS),
refresh_timeout_seconds + 5.0,
)
):
data = _read_xai_oauth_tokens(_lock=False)
tokens = dict(data["tokens"])
access_token = str(tokens.get("access_token", "") or "").strip()
Expand Down Expand Up @@ -5027,7 +5118,10 @@ def resolve_xai_oauth_runtime_credentials(
# without a network retry. Mirrors credential_pool.py quarantine.
try:
_q_store = _load_auth_store()
_q_state = _load_provider_state(_q_store, "xai-oauth") or {}
_q_state, _q_source_path = (
_load_xai_oauth_state_with_source(_q_store)
)
_q_state = dict(_q_state or {})
_q_tokens = dict(_q_state.get("tokens") or {})
_q_tokens.pop("access_token", None)
_q_tokens.pop("refresh_token", None)
Expand All @@ -5040,8 +5134,26 @@ def resolve_xai_oauth_runtime_credentials(
"relogin_required": True,
"at": datetime.now(timezone.utc).isoformat(),
}
_store_provider_state(_q_store, "xai-oauth", _q_state, set_active=False)
_save_auth_store(_q_store)
_q_global_path = _global_auth_file_path()
if (
_q_source_path is not None
and _q_global_path is not None
and _same_path(_q_source_path, _q_global_path)
):
_persist_provider_state_to_store(
"xai-oauth",
_q_state,
_q_global_path,
set_active=False,
)
else:
_store_provider_state(
_q_store,
"xai-oauth",
_q_state,
set_active=False,
)
_save_auth_store(_q_store)
except Exception as _save_exc:
logger.debug(
"xAI OAuth: failed to persist quarantined state: %s", _save_exc,
Expand Down
Loading