Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license
This commit is contained in:
@@ -0,0 +1,309 @@
|
||||
"""Regression tests for cross-process races refreshing Anthropic OAuth tokens.
|
||||
|
||||
``CredentialPool._refresh_entry`` explicitly documents (see the comment
|
||||
above the ``if self.provider in ("openai-codex", "xai-oauth", "anthropic"):`` branch in
|
||||
``agent/credential_pool.py``) that single-use OAuth refresh tokens require
|
||||
the whole sync -> POST -> write-back sequence to be serialized across
|
||||
Hermes *processes* via the cross-process ``_auth_store_lock`` flock,
|
||||
otherwise "two processes can both adopt the same on-disk token, both POST
|
||||
it, and the loser gets ``refresh_token_reused``".
|
||||
|
||||
Anthropic's OAuth refresh tokens have the identical single-use property --
|
||||
``agent.anthropic_credentials._refresh_oauth_token`` says so explicitly:
|
||||
"Claude Code's OAuth refresh tokens are single-use: a successful refresh
|
||||
rotates the pair and invalidates the old refresh token." Before the PR, ``"anthropic"`` was absent from the
|
||||
``("openai-codex", "xai-oauth")`` tuple that gets the cross-process flock,
|
||||
and the *only* on-failure recovery path
|
||||
(``CredentialPool._sync_anthropic_entry_from_credentials_file``) was
|
||||
hard-scoped to ``entry.source == "claude_code"``. Entries sourced from
|
||||
Hermes's own PKCE login (``manual:hermes_pkce`` / ``hermes_pkce``) got no
|
||||
recovery at all and were marked exhausted on a lost race, even though a
|
||||
fresh, valid token pair existed on disk (written by the winner). The old
|
||||
dashboard source ``manual:dashboard_pkce`` is now retired with the removed
|
||||
dashboard flow.
|
||||
|
||||
These tests reproduce that race deterministically with a fake OAuth server
|
||||
that enforces single-use refresh tokens, run concurrent pool instances against
|
||||
it, and assert the current protection. The process-level, cross-profile
|
||||
Claude Code witness lives in ``test_anthropic_oauth_stress.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import replace as dc_replace
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.credential_pool import (
|
||||
AUTH_TYPE_OAUTH,
|
||||
STATUS_EXHAUSTED,
|
||||
CredentialPool,
|
||||
PooledCredential,
|
||||
)
|
||||
|
||||
|
||||
def _entry(*, id: str, access_token: str, refresh_token: str, source: str) -> PooledCredential:
|
||||
return PooledCredential(
|
||||
provider="anthropic",
|
||||
id=id,
|
||||
label="anthropic oauth",
|
||||
auth_type=AUTH_TYPE_OAUTH,
|
||||
priority=0,
|
||||
source=source,
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_at_ms=0, # already expired -> force path always refreshes
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fake_pool_store(monkeypatch):
|
||||
"""Back write/read_credential_pool with a shared in-memory dict.
|
||||
|
||||
This stands in for ``~/.hermes/auth.json`` so the two "process-local"
|
||||
CredentialPool instances used by the race tests can actually see each
|
||||
other's persisted writes (exactly what the real cross-process recovery
|
||||
path depends on), without touching the real filesystem.
|
||||
"""
|
||||
store: Dict[str, list] = {}
|
||||
|
||||
def _write(provider, entries, *, removed_ids=None):
|
||||
store[provider] = list(entries)
|
||||
|
||||
def _read(provider=None):
|
||||
if provider is None:
|
||||
return dict(store)
|
||||
return list(store.get(provider, []))
|
||||
|
||||
monkeypatch.setattr("agent.credential_pool.write_credential_pool", _write)
|
||||
monkeypatch.setattr("agent.credential_pool.read_credential_pool", _read)
|
||||
return store
|
||||
|
||||
|
||||
class _SingleUseTokenServer:
|
||||
"""Minimal fake of Anthropic's OAuth token endpoint.
|
||||
|
||||
Enforces the real single-use-refresh-token contract: the first caller
|
||||
to present a given refresh_token gets a fresh pair; every subsequent
|
||||
caller presenting that same (now-spent) refresh_token gets an
|
||||
``invalid_grant`` error, mirroring Anthropic's actual behavior.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._spent: set[str] = set()
|
||||
self._rotation = 0
|
||||
self.calls: list[str] = []
|
||||
# Small artificial delay to widen the race window deterministically.
|
||||
self.delay_seconds = 0.05
|
||||
|
||||
def refresh(self, refresh_token: str, *, use_json: bool = False):
|
||||
self.calls.append(refresh_token)
|
||||
time.sleep(self.delay_seconds)
|
||||
with self._lock:
|
||||
if refresh_token in self._spent:
|
||||
raise ValueError("invalid_grant: refresh token already used")
|
||||
self._spent.add(refresh_token)
|
||||
self._rotation += 1
|
||||
rotation = self._rotation
|
||||
return {
|
||||
"access_token": f"sk-ant-oat-rotated-{rotation}",
|
||||
"refresh_token": f"sk-ant-ort-rotated-{rotation}",
|
||||
"expires_at_ms": int(time.time() * 1000) + 3_600_000,
|
||||
}
|
||||
|
||||
|
||||
def test_anthropic_refresh_is_protected_by_cross_process_lock(monkeypatch):
|
||||
"""Structural check: anthropic refresh acquires ``_auth_store_lock``.
|
||||
|
||||
Regression guard for the gap where "anthropic" was missing from the
|
||||
``("openai-codex", "xai-oauth")`` tuple despite Anthropic OAuth refresh
|
||||
tokens being single-use too -- see the module docstring.
|
||||
"""
|
||||
lock_calls: list[str] = []
|
||||
|
||||
class _RecordingLock:
|
||||
def __init__(self, *a, **kw):
|
||||
lock_calls.append("anthropic")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
return False
|
||||
|
||||
monkeypatch.setattr("agent.credential_pool._auth_store_lock", _RecordingLock)
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials.refresh_anthropic_oauth_pure",
|
||||
lambda refresh_token, use_json=False: {
|
||||
"access_token": "sk-ant-oat-new",
|
||||
"refresh_token": "sk-ant-ort-new",
|
||||
"expires_at_ms": int(time.time() * 1000) + 3_600_000,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials.read_claude_code_credentials", lambda: None
|
||||
)
|
||||
|
||||
pool = CredentialPool(
|
||||
"anthropic",
|
||||
[_entry(id="a1", access_token="stale-at", refresh_token="stale-rt", source="manual:hermes_pkce")],
|
||||
)
|
||||
entry = pool.entries()[0]
|
||||
pool._refresh_entry(entry, force=True)
|
||||
|
||||
assert lock_calls, (
|
||||
"regression: CredentialPool._refresh_entry() for provider='anthropic' "
|
||||
"did not acquire the cross-process _auth_store_lock, even though "
|
||||
"Anthropic OAuth refresh tokens are single-use (same property "
|
||||
"openai-codex/xai-oauth are explicitly locked for)."
|
||||
)
|
||||
|
||||
|
||||
def test_concurrent_hermes_pkce_refresh_loses_credential_despite_valid_token_on_disk(monkeypatch):
|
||||
"""Two 'Hermes processes' race to refresh the same stale hermes_pkce
|
||||
refresh token. The winner gets a valid new pair. The loser -- despite a
|
||||
fresh, valid credential now existing -- has no recovery path and is
|
||||
marked exhausted, because ``_sync_anthropic_entry_from_credentials_file``
|
||||
only helps ``entry.source == "claude_code"``.
|
||||
"""
|
||||
server = _SingleUseTokenServer()
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials.refresh_anthropic_oauth_pure",
|
||||
lambda refresh_token, use_json=False: server.refresh(refresh_token, use_json=use_json),
|
||||
)
|
||||
# No Claude Code credential file in play for this scenario.
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials.read_claude_code_credentials", lambda: None
|
||||
)
|
||||
|
||||
shared_stale_entry = _entry(
|
||||
id="pool-entry",
|
||||
access_token="stale-at",
|
||||
refresh_token="stale-rt",
|
||||
source="manual:hermes_pkce",
|
||||
)
|
||||
|
||||
# Simulate two independent OS processes, each with its own in-memory
|
||||
# pool constructed from the SAME on-disk stale entry.
|
||||
pool_process_a = CredentialPool("anthropic", [dc_replace(shared_stale_entry)])
|
||||
pool_process_b = CredentialPool("anthropic", [dc_replace(shared_stale_entry)])
|
||||
|
||||
results: dict[str, object] = {}
|
||||
|
||||
def _run(pool, key):
|
||||
entry = pool.entries()[0]
|
||||
results[key] = pool._refresh_entry(entry, force=True)
|
||||
|
||||
t_a = threading.Thread(target=_run, args=(pool_process_a, "a"))
|
||||
t_b = threading.Thread(target=_run, args=(pool_process_b, "b"))
|
||||
t_a.start()
|
||||
t_b.start()
|
||||
t_a.join(timeout=5)
|
||||
t_b.join(timeout=5)
|
||||
|
||||
# With the cross-process lock in place, the second thread to acquire it
|
||||
# re-syncs from the pool store first and may adopt the winner's
|
||||
# already-persisted token before ever calling the token endpoint --
|
||||
# so only ONE call is expected now (the race is avoided entirely,
|
||||
# which is strictly better than the old behavior of two racing calls).
|
||||
# Both calls (if two do happen) must only ever spend the one stale
|
||||
# token that was actually on disk -- never anything else.
|
||||
assert 1 <= len(server.calls) <= 2, f"unexpected call count: {server.calls!r}"
|
||||
assert set(server.calls) == {"stale-rt"}, f"unexpected tokens posted: {server.calls!r}"
|
||||
|
||||
winner_result = results["a"] or results["b"]
|
||||
loser_key = "b" if results["a"] else "a"
|
||||
loser_result = results[loser_key]
|
||||
loser_pool = pool_process_b if loser_key == "b" else pool_process_a
|
||||
|
||||
assert winner_result is not None and winner_result.access_token.startswith(
|
||||
"sk-ant-oat-rotated-"
|
||||
), "the winning process should obtain a fresh, valid token pair"
|
||||
|
||||
# This is the bug: the loser should be able to recover (a valid token
|
||||
# now exists -- the winner's), not be left exhausted.
|
||||
assert loser_result is not None, (
|
||||
"regression: the losing process's CredentialPool._refresh_entry() "
|
||||
"returned None (unrecoverable) after losing the single-use-token "
|
||||
"race, even though a fresh valid Anthropic credential exists on "
|
||||
"disk. hermes_pkce-sourced entries must adopt the winner's token "
|
||||
"via _sync_anthropic_entry_from_pool_store(), same as claude_code."
|
||||
)
|
||||
loser_entry_after = loser_pool.entries()[0]
|
||||
assert loser_entry_after.last_status != STATUS_EXHAUSTED, (
|
||||
"regression: the losing process marked its Anthropic hermes_pkce "
|
||||
"credential STATUS_EXHAUSTED after a lost refresh race, even "
|
||||
"though the account is not actually exhausted -- a sibling Hermes "
|
||||
"process holds a perfectly valid rotated token. This causes "
|
||||
"spurious 're-authenticate with Anthropic' failures under "
|
||||
"ordinary concurrent Hermes usage (fleet workers, cron jobs, "
|
||||
"multiple CLI sessions)."
|
||||
)
|
||||
|
||||
|
||||
def test_concurrent_claude_code_refresh_recovers_via_credentials_file(monkeypatch):
|
||||
"""Contrast case: entry.source == 'claude_code' DOES recover from a lost
|
||||
race, because ``_sync_anthropic_entry_from_credentials_file`` re-reads
|
||||
``~/.claude/.credentials.json`` on failure. This asymmetry is itself
|
||||
evidence that hermes_pkce/dashboard-sourced credentials were simply
|
||||
never given the same treatment, not that recovery is impossible.
|
||||
"""
|
||||
server = _SingleUseTokenServer()
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials.refresh_anthropic_oauth_pure",
|
||||
lambda refresh_token, use_json=False: server.refresh(refresh_token, use_json=use_json),
|
||||
)
|
||||
|
||||
winner_creds: dict = {}
|
||||
write_lock = threading.Lock()
|
||||
|
||||
def _fake_read_claude_code_credentials():
|
||||
with write_lock:
|
||||
return dict(winner_creds) if winner_creds else None
|
||||
|
||||
def _fake_write_claude_code_credentials(access_token, refresh_token, expires_at_ms, **_kw):
|
||||
with write_lock:
|
||||
winner_creds.update(
|
||||
accessToken=access_token,
|
||||
refreshToken=refresh_token,
|
||||
expiresAt=expires_at_ms,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials.read_claude_code_credentials",
|
||||
_fake_read_claude_code_credentials,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"agent.anthropic_credentials._write_claude_code_credentials",
|
||||
_fake_write_claude_code_credentials,
|
||||
)
|
||||
|
||||
shared_stale_entry = _entry(
|
||||
id="pool-entry",
|
||||
access_token="stale-at",
|
||||
refresh_token="stale-rt",
|
||||
source="claude_code",
|
||||
)
|
||||
pool_process_a = CredentialPool("anthropic", [dc_replace(shared_stale_entry)])
|
||||
pool_process_b = CredentialPool("anthropic", [dc_replace(shared_stale_entry)])
|
||||
|
||||
results: dict[str, object] = {}
|
||||
|
||||
def _run(pool, key):
|
||||
entry = pool.entries()[0]
|
||||
results[key] = pool._refresh_entry(entry, force=True)
|
||||
|
||||
t_a = threading.Thread(target=_run, args=(pool_process_a, "a"))
|
||||
t_b = threading.Thread(target=_run, args=(pool_process_b, "b"))
|
||||
t_a.start()
|
||||
t_b.start()
|
||||
t_a.join(timeout=5)
|
||||
t_b.join(timeout=5)
|
||||
|
||||
# Both processes should end up with a usable credential: the winner via
|
||||
# its own refresh, the loser via the credentials-file sync-and-retry.
|
||||
assert results["a"] is not None, "claude_code process A should recover"
|
||||
assert results["b"] is not None, "claude_code process B should recover"
|
||||
Reference in New Issue
Block a user