mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-05-19 04:52:06 +00:00
Add pooled same-provider credential fallback
This commit is contained in:
parent
934fbe3c06
commit
b17e5c101d
18 changed files with 2872 additions and 195 deletions
|
|
@ -201,60 +201,75 @@ def is_claude_code_token_valid(creds: Dict[str, Any]) -> bool:
|
|||
return now_ms < (expires_at - 60_000)
|
||||
|
||||
|
||||
def _refresh_oauth_token(creds: Dict[str, Any]) -> Optional[str]:
|
||||
"""Attempt to refresh an expired Claude Code OAuth token.
|
||||
|
||||
Uses the same token endpoint and client_id as Claude Code / OpenCode.
|
||||
Only works for credentials that have a refresh token (from claude /login
|
||||
or claude setup-token with OAuth flow).
|
||||
|
||||
Returns the new access token, or None if refresh fails.
|
||||
"""
|
||||
def refresh_anthropic_oauth_pure(refresh_token: str, *, use_json: bool = False) -> Dict[str, Any]:
|
||||
"""Refresh an Anthropic OAuth token without mutating local credential files."""
|
||||
import time
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
|
||||
refresh_token = creds.get("refreshToken", "")
|
||||
if not refresh_token:
|
||||
logger.debug("No refresh token available — cannot refresh")
|
||||
return None
|
||||
raise ValueError("refresh_token is required")
|
||||
|
||||
# Client ID used by Claude Code's OAuth flow
|
||||
CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
|
||||
|
||||
data = urllib.parse.urlencode({
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"client_id": CLIENT_ID,
|
||||
}).encode()
|
||||
client_id = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
|
||||
if use_json:
|
||||
data = json.dumps({
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"client_id": client_id,
|
||||
}).encode()
|
||||
content_type = "application/json"
|
||||
else:
|
||||
data = urllib.parse.urlencode({
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"client_id": client_id,
|
||||
}).encode()
|
||||
content_type = "application/x-www-form-urlencoded"
|
||||
|
||||
req = urllib.request.Request(
|
||||
"https://console.anthropic.com/v1/oauth/token",
|
||||
data=data,
|
||||
headers={
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"Content-Type": content_type,
|
||||
"User-Agent": f"claude-cli/{_CLAUDE_CODE_VERSION} (external, cli)",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
result = json.loads(resp.read().decode())
|
||||
new_access = result.get("access_token", "")
|
||||
new_refresh = result.get("refresh_token", refresh_token)
|
||||
expires_in = result.get("expires_in", 3600) # seconds
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
result = json.loads(resp.read().decode())
|
||||
|
||||
if new_access:
|
||||
import time
|
||||
new_expires_ms = int(time.time() * 1000) + (expires_in * 1000)
|
||||
# Write refreshed credentials back to ~/.claude/.credentials.json
|
||||
_write_claude_code_credentials(new_access, new_refresh, new_expires_ms)
|
||||
logger.debug("Successfully refreshed Claude Code OAuth token")
|
||||
return new_access
|
||||
access_token = result.get("access_token", "")
|
||||
if not access_token:
|
||||
raise ValueError("Anthropic refresh response was missing access_token")
|
||||
next_refresh = result.get("refresh_token", refresh_token)
|
||||
expires_in = result.get("expires_in", 3600)
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": next_refresh,
|
||||
"expires_at_ms": int(time.time() * 1000) + (expires_in * 1000),
|
||||
}
|
||||
|
||||
|
||||
def _refresh_oauth_token(creds: Dict[str, Any]) -> Optional[str]:
|
||||
"""Attempt to refresh an expired Claude Code OAuth token."""
|
||||
refresh_token = creds.get("refreshToken", "")
|
||||
if not refresh_token:
|
||||
logger.debug("No refresh token available — cannot refresh")
|
||||
return None
|
||||
|
||||
try:
|
||||
refreshed = refresh_anthropic_oauth_pure(refresh_token, use_json=False)
|
||||
_write_claude_code_credentials(
|
||||
refreshed["access_token"],
|
||||
refreshed["refresh_token"],
|
||||
refreshed["expires_at_ms"],
|
||||
)
|
||||
logger.debug("Successfully refreshed Claude Code OAuth token")
|
||||
return refreshed["access_token"]
|
||||
except Exception as e:
|
||||
logger.debug("Failed to refresh Claude Code token: %s", e)
|
||||
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _write_claude_code_credentials(access_token: str, refresh_token: str, expires_at_ms: int) -> None:
|
||||
|
|
@ -466,14 +481,8 @@ def _generate_pkce() -> tuple:
|
|||
return verifier, challenge
|
||||
|
||||
|
||||
def run_hermes_oauth_login() -> Optional[str]:
|
||||
"""Run Hermes-native OAuth PKCE flow for Claude Pro/Max subscription.
|
||||
|
||||
Opens a browser to claude.ai for authorization, prompts for the code,
|
||||
exchanges it for tokens, and stores them in ~/.hermes/.anthropic_oauth.json.
|
||||
|
||||
Returns the access token on success, None on failure.
|
||||
"""
|
||||
def run_hermes_oauth_login_pure() -> Optional[Dict[str, Any]]:
|
||||
"""Run Hermes-native OAuth PKCE flow and return credential state."""
|
||||
import time
|
||||
import webbrowser
|
||||
|
||||
|
|
@ -564,10 +573,32 @@ def run_hermes_oauth_login() -> Optional[str]:
|
|||
print("No access token in response.")
|
||||
return None
|
||||
|
||||
# Store credentials
|
||||
expires_at_ms = int(time.time() * 1000) + (expires_in * 1000)
|
||||
_save_hermes_oauth_credentials(access_token, refresh_token, expires_at_ms)
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_token,
|
||||
"expires_at_ms": expires_at_ms,
|
||||
}
|
||||
|
||||
|
||||
def run_hermes_oauth_login() -> Optional[str]:
|
||||
"""Run Hermes-native OAuth PKCE flow for Claude Pro/Max subscription.
|
||||
|
||||
Opens a browser to claude.ai for authorization, prompts for the code,
|
||||
exchanges it for tokens, and stores them in ~/.hermes/.anthropic_oauth.json.
|
||||
|
||||
Returns the access token on success, None on failure.
|
||||
"""
|
||||
result = run_hermes_oauth_login_pure()
|
||||
if not result:
|
||||
return None
|
||||
|
||||
access_token = result["access_token"]
|
||||
refresh_token = result["refresh_token"]
|
||||
expires_at_ms = result["expires_at_ms"]
|
||||
|
||||
# Store credentials
|
||||
_save_hermes_oauth_credentials(access_token, refresh_token, expires_at_ms)
|
||||
# Also write to Claude Code's credential file for backward compat
|
||||
_write_claude_code_credentials(access_token, refresh_token, expires_at_ms)
|
||||
|
||||
|
|
@ -607,44 +638,27 @@ def refresh_hermes_oauth_token() -> Optional[str]:
|
|||
|
||||
Returns the new access token, or None if refresh fails.
|
||||
"""
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
creds = read_hermes_oauth_credentials()
|
||||
if not creds or not creds.get("refreshToken"):
|
||||
return None
|
||||
|
||||
try:
|
||||
data = json.dumps({
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": creds["refreshToken"],
|
||||
"client_id": _OAUTH_CLIENT_ID,
|
||||
}).encode()
|
||||
|
||||
req = urllib.request.Request(
|
||||
_OAUTH_TOKEN_URL,
|
||||
data=data,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": f"claude-cli/{_CLAUDE_CODE_VERSION} (external, cli)",
|
||||
},
|
||||
method="POST",
|
||||
refreshed = refresh_anthropic_oauth_pure(
|
||||
creds["refreshToken"],
|
||||
use_json=True,
|
||||
)
|
||||
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
result = json.loads(resp.read().decode())
|
||||
|
||||
new_access = result.get("access_token", "")
|
||||
new_refresh = result.get("refresh_token", creds["refreshToken"])
|
||||
expires_in = result.get("expires_in", 3600)
|
||||
|
||||
if new_access:
|
||||
new_expires_ms = int(time.time() * 1000) + (expires_in * 1000)
|
||||
_save_hermes_oauth_credentials(new_access, new_refresh, new_expires_ms)
|
||||
# Also update Claude Code's credential file
|
||||
_write_claude_code_credentials(new_access, new_refresh, new_expires_ms)
|
||||
logger.debug("Successfully refreshed Hermes OAuth token")
|
||||
return new_access
|
||||
_save_hermes_oauth_credentials(
|
||||
refreshed["access_token"],
|
||||
refreshed["refresh_token"],
|
||||
refreshed["expires_at_ms"],
|
||||
)
|
||||
_write_claude_code_credentials(
|
||||
refreshed["access_token"],
|
||||
refreshed["refresh_token"],
|
||||
refreshed["expires_at_ms"],
|
||||
)
|
||||
logger.debug("Successfully refreshed Hermes OAuth token")
|
||||
return refreshed["access_token"]
|
||||
except Exception as e:
|
||||
logger.debug("Failed to refresh Hermes OAuth token: %s", e)
|
||||
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from typing import Any, Dict, List, Optional, Tuple
|
|||
|
||||
from openai import OpenAI
|
||||
|
||||
from agent.credential_pool import load_pool
|
||||
from hermes_cli.config import get_hermes_home
|
||||
from hermes_constants import OPENROUTER_BASE_URL
|
||||
|
||||
|
|
@ -96,6 +97,45 @@ _CODEX_AUX_MODEL = "gpt-5.2-codex"
|
|||
_CODEX_AUX_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
||||
|
||||
|
||||
def _select_pool_entry(provider: str) -> Tuple[bool, Optional[Any]]:
|
||||
"""Return (pool_exists_for_provider, selected_entry)."""
|
||||
try:
|
||||
pool = load_pool(provider)
|
||||
except Exception as exc:
|
||||
logger.debug("Auxiliary client: could not load pool for %s: %s", provider, exc)
|
||||
return False, None
|
||||
if not pool or not pool.has_credentials():
|
||||
return False, None
|
||||
try:
|
||||
return True, pool.select()
|
||||
except Exception as exc:
|
||||
logger.debug("Auxiliary client: could not select pool entry for %s: %s", provider, exc)
|
||||
return True, None
|
||||
|
||||
|
||||
def _pool_runtime_api_key(entry: Any) -> str:
|
||||
if entry is None:
|
||||
return ""
|
||||
return str(
|
||||
getattr(entry, "runtime_api_key", None)
|
||||
or getattr(entry, "agent_key", None)
|
||||
or getattr(entry, "access_token", "")
|
||||
or ""
|
||||
).strip()
|
||||
|
||||
|
||||
def _pool_runtime_base_url(entry: Any, fallback: str = "") -> str:
|
||||
if entry is None:
|
||||
return str(fallback or "").strip().rstrip("/")
|
||||
return str(
|
||||
getattr(entry, "runtime_base_url", None)
|
||||
or getattr(entry, "inference_base_url", None)
|
||||
or getattr(entry, "base_url", None)
|
||||
or fallback
|
||||
or ""
|
||||
).strip().rstrip("/")
|
||||
|
||||
|
||||
# ── Codex Responses → chat.completions adapter ─────────────────────────────
|
||||
# All auxiliary consumers call client.chat.completions.create(**kwargs) and
|
||||
# read response.choices[0].message.content. This adapter translates those
|
||||
|
|
@ -439,6 +479,21 @@ def _read_nous_auth() -> Optional[dict]:
|
|||
Returns the provider state dict if Nous is active with tokens,
|
||||
otherwise None.
|
||||
"""
|
||||
pool_present, entry = _select_pool_entry("nous")
|
||||
if pool_present:
|
||||
if entry is None:
|
||||
return None
|
||||
return {
|
||||
"access_token": getattr(entry, "access_token", ""),
|
||||
"refresh_token": getattr(entry, "refresh_token", None),
|
||||
"agent_key": getattr(entry, "agent_key", None),
|
||||
"inference_base_url": _pool_runtime_base_url(entry, _NOUS_DEFAULT_BASE_URL),
|
||||
"portal_base_url": getattr(entry, "portal_base_url", None),
|
||||
"client_id": getattr(entry, "client_id", None),
|
||||
"scope": getattr(entry, "scope", None),
|
||||
"token_type": getattr(entry, "token_type", "Bearer"),
|
||||
}
|
||||
|
||||
try:
|
||||
if not _AUTH_JSON_PATH.is_file():
|
||||
return None
|
||||
|
|
@ -467,6 +522,11 @@ def _nous_base_url() -> str:
|
|||
|
||||
def _read_codex_access_token() -> Optional[str]:
|
||||
"""Read a valid, non-expired Codex OAuth access token from Hermes auth store."""
|
||||
pool_present, entry = _select_pool_entry("openai-codex")
|
||||
if pool_present:
|
||||
token = _pool_runtime_api_key(entry)
|
||||
return token or None
|
||||
|
||||
try:
|
||||
from hermes_cli.auth import _read_codex_tokens
|
||||
data = _read_codex_tokens()
|
||||
|
|
@ -513,6 +573,24 @@ def _resolve_api_key_provider() -> Tuple[Optional[OpenAI], Optional[str]]:
|
|||
if provider_id == "anthropic":
|
||||
return _try_anthropic()
|
||||
|
||||
pool_present, entry = _select_pool_entry(provider_id)
|
||||
if pool_present:
|
||||
api_key = _pool_runtime_api_key(entry)
|
||||
if not api_key:
|
||||
continue
|
||||
|
||||
base_url = _pool_runtime_base_url(entry, pconfig.inference_base_url) or pconfig.inference_base_url
|
||||
model = _API_KEY_PROVIDER_AUX_MODELS.get(provider_id, "default")
|
||||
logger.debug("Auxiliary text client: %s (%s) via pool", pconfig.name, model)
|
||||
extra = {}
|
||||
if "api.kimi.com" in base_url.lower():
|
||||
extra["default_headers"] = {"User-Agent": "KimiCLI/1.0"}
|
||||
elif "api.githubcopilot.com" in base_url.lower():
|
||||
from hermes_cli.models import copilot_default_headers
|
||||
|
||||
extra["default_headers"] = copilot_default_headers()
|
||||
return OpenAI(api_key=api_key, base_url=base_url, **extra), model
|
||||
|
||||
creds = resolve_api_key_provider_credentials(provider_id)
|
||||
api_key = str(creds.get("api_key", "")).strip()
|
||||
if not api_key:
|
||||
|
|
@ -562,6 +640,16 @@ def _get_auxiliary_env_override(task: str, suffix: str) -> Optional[str]:
|
|||
|
||||
|
||||
def _try_openrouter() -> Tuple[Optional[OpenAI], Optional[str]]:
|
||||
pool_present, entry = _select_pool_entry("openrouter")
|
||||
if pool_present:
|
||||
or_key = _pool_runtime_api_key(entry)
|
||||
if not or_key:
|
||||
return None, None
|
||||
base_url = _pool_runtime_base_url(entry, OPENROUTER_BASE_URL) or OPENROUTER_BASE_URL
|
||||
logger.debug("Auxiliary client: OpenRouter via pool")
|
||||
return OpenAI(api_key=or_key, base_url=base_url,
|
||||
default_headers=_OR_HEADERS), _OPENROUTER_MODEL
|
||||
|
||||
or_key = os.getenv("OPENROUTER_API_KEY")
|
||||
if not or_key:
|
||||
return None, None
|
||||
|
|
@ -578,7 +666,10 @@ def _try_nous() -> Tuple[Optional[OpenAI], Optional[str]]:
|
|||
auxiliary_is_nous = True
|
||||
logger.debug("Auxiliary client: Nous Portal")
|
||||
return (
|
||||
OpenAI(api_key=_nous_api_key(nous), base_url=_nous_base_url()),
|
||||
OpenAI(
|
||||
api_key=_nous_api_key(nous),
|
||||
base_url=str(nous.get("inference_base_url") or _nous_base_url()).rstrip("/"),
|
||||
),
|
||||
_NOUS_MODEL,
|
||||
)
|
||||
|
||||
|
|
@ -654,11 +745,19 @@ def _try_custom_endpoint() -> Tuple[Optional[OpenAI], Optional[str]]:
|
|||
|
||||
|
||||
def _try_codex() -> Tuple[Optional[Any], Optional[str]]:
|
||||
codex_token = _read_codex_access_token()
|
||||
if not codex_token:
|
||||
return None, None
|
||||
pool_present, entry = _select_pool_entry("openai-codex")
|
||||
if pool_present:
|
||||
codex_token = _pool_runtime_api_key(entry)
|
||||
if not codex_token:
|
||||
return None, None
|
||||
base_url = _pool_runtime_base_url(entry, _CODEX_AUX_BASE_URL) or _CODEX_AUX_BASE_URL
|
||||
else:
|
||||
codex_token = _read_codex_access_token()
|
||||
if not codex_token:
|
||||
return None, None
|
||||
base_url = _CODEX_AUX_BASE_URL
|
||||
logger.debug("Auxiliary client: Codex OAuth (%s via Responses API)", _CODEX_AUX_MODEL)
|
||||
real_client = OpenAI(api_key=codex_token, base_url=_CODEX_AUX_BASE_URL)
|
||||
real_client = OpenAI(api_key=codex_token, base_url=base_url)
|
||||
return CodexAuxiliaryClient(real_client, _CODEX_AUX_MODEL), _CODEX_AUX_MODEL
|
||||
|
||||
|
||||
|
|
@ -668,14 +767,21 @@ def _try_anthropic() -> Tuple[Optional[Any], Optional[str]]:
|
|||
except ImportError:
|
||||
return None, None
|
||||
|
||||
token = resolve_anthropic_token()
|
||||
pool_present, entry = _select_pool_entry("anthropic")
|
||||
if pool_present:
|
||||
if entry is None:
|
||||
return None, None
|
||||
token = _pool_runtime_api_key(entry)
|
||||
else:
|
||||
entry = None
|
||||
token = resolve_anthropic_token()
|
||||
if not token:
|
||||
return None, None
|
||||
|
||||
# Allow base URL override from config.yaml model.base_url, but only
|
||||
# when the configured provider is anthropic — otherwise a non-Anthropic
|
||||
# base_url (e.g. Codex endpoint) would leak into Anthropic requests.
|
||||
base_url = _ANTHROPIC_DEFAULT_BASE_URL
|
||||
base_url = _pool_runtime_base_url(entry, _ANTHROPIC_DEFAULT_BASE_URL) if pool_present else _ANTHROPIC_DEFAULT_BASE_URL
|
||||
try:
|
||||
from hermes_cli.config import load_config
|
||||
cfg = load_config()
|
||||
|
|
|
|||
456
agent/credential_pool.py
Normal file
456
agent/credential_pool.py
Normal file
|
|
@ -0,0 +1,456 @@
|
|||
"""Persistent multi-credential pool for same-provider failover."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
import os
|
||||
from dataclasses import dataclass, fields
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from hermes_constants import OPENROUTER_BASE_URL
|
||||
import hermes_cli.auth as auth_mod
|
||||
from hermes_cli.auth import (
|
||||
ACCESS_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CODEX_ACCESS_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
DEFAULT_AGENT_KEY_MIN_TTL_SECONDS,
|
||||
PROVIDER_REGISTRY,
|
||||
_agent_key_is_usable,
|
||||
_codex_access_token_is_expiring,
|
||||
_decode_jwt_claims,
|
||||
_is_expiring,
|
||||
_load_auth_store,
|
||||
_load_provider_state,
|
||||
read_credential_pool,
|
||||
write_credential_pool,
|
||||
)
|
||||
|
||||
EXHAUSTED_TTL_SECONDS = 24 * 60 * 60
|
||||
|
||||
|
||||
@dataclass
|
||||
class PooledCredential:
|
||||
provider: str
|
||||
id: str
|
||||
label: str
|
||||
auth_type: str
|
||||
priority: int
|
||||
source: str
|
||||
access_token: str
|
||||
refresh_token: Optional[str] = None
|
||||
last_status: Optional[str] = None
|
||||
last_status_at: Optional[float] = None
|
||||
last_error_code: Optional[int] = None
|
||||
base_url: Optional[str] = None
|
||||
expires_at: Optional[str] = None
|
||||
expires_at_ms: Optional[int] = None
|
||||
last_refresh: Optional[str] = None
|
||||
token_type: Optional[str] = None
|
||||
scope: Optional[str] = None
|
||||
client_id: Optional[str] = None
|
||||
portal_base_url: Optional[str] = None
|
||||
inference_base_url: Optional[str] = None
|
||||
obtained_at: Optional[str] = None
|
||||
expires_in: Optional[int] = None
|
||||
agent_key: Optional[str] = None
|
||||
agent_key_id: Optional[str] = None
|
||||
agent_key_expires_at: Optional[str] = None
|
||||
agent_key_expires_in: Optional[int] = None
|
||||
agent_key_reused: Optional[bool] = None
|
||||
agent_key_obtained_at: Optional[str] = None
|
||||
tls: Optional[Dict[str, Any]] = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, provider: str, payload: Dict[str, Any]) -> "PooledCredential":
|
||||
allowed = {f.name for f in fields(cls) if f.name != "provider"}
|
||||
data = {k: payload.get(k) for k in allowed if k in payload}
|
||||
data.setdefault("id", uuid.uuid4().hex[:6])
|
||||
data.setdefault("label", payload.get("source", provider))
|
||||
data.setdefault("auth_type", "api_key")
|
||||
data.setdefault("priority", 0)
|
||||
data.setdefault("source", "manual")
|
||||
data.setdefault("access_token", "")
|
||||
return cls(provider=provider, **data)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {}
|
||||
for field_def in fields(self):
|
||||
if field_def.name == "provider":
|
||||
continue
|
||||
value = getattr(self, field_def.name)
|
||||
if value is not None:
|
||||
result[field_def.name] = value
|
||||
for key in ("last_status", "last_status_at", "last_error_code"):
|
||||
result.setdefault(key, getattr(self, key))
|
||||
return result
|
||||
|
||||
@property
|
||||
def runtime_api_key(self) -> str:
|
||||
if self.provider == "nous":
|
||||
return str(self.agent_key or self.access_token or "")
|
||||
return str(self.access_token or "")
|
||||
|
||||
@property
|
||||
def runtime_base_url(self) -> Optional[str]:
|
||||
if self.provider == "nous":
|
||||
return self.inference_base_url or self.base_url
|
||||
return self.base_url
|
||||
|
||||
|
||||
def _label_from_token(token: str, fallback: str) -> str:
|
||||
claims = _decode_jwt_claims(token)
|
||||
for key in ("email", "preferred_username", "upn"):
|
||||
value = claims.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return fallback
|
||||
|
||||
|
||||
def _next_priority(entries: List[PooledCredential]) -> int:
|
||||
return max((entry.priority for entry in entries), default=-1) + 1
|
||||
|
||||
|
||||
class CredentialPool:
|
||||
def __init__(self, provider: str, entries: List[PooledCredential]):
|
||||
self.provider = provider
|
||||
self._entries = sorted(entries, key=lambda entry: entry.priority)
|
||||
self._current_id: Optional[str] = None
|
||||
|
||||
def has_credentials(self) -> bool:
|
||||
return bool(self._entries)
|
||||
|
||||
def entries(self) -> List[PooledCredential]:
|
||||
return list(sorted(self._entries, key=lambda entry: entry.priority))
|
||||
|
||||
def current(self) -> Optional[PooledCredential]:
|
||||
if not self._current_id:
|
||||
return None
|
||||
return next((entry for entry in self._entries if entry.id == self._current_id), None)
|
||||
|
||||
def _persist(self) -> None:
|
||||
write_credential_pool(
|
||||
self.provider,
|
||||
[entry.to_dict() for entry in sorted(self._entries, key=lambda item: item.priority)],
|
||||
)
|
||||
|
||||
def _mark_exhausted(self, entry: PooledCredential, status_code: Optional[int]) -> None:
|
||||
entry.last_status = "exhausted"
|
||||
entry.last_status_at = time.time()
|
||||
entry.last_error_code = status_code
|
||||
self._persist()
|
||||
|
||||
def _refresh_entry(self, entry: PooledCredential, *, force: bool) -> Optional[PooledCredential]:
|
||||
if entry.auth_type != "oauth" or not entry.refresh_token:
|
||||
if force:
|
||||
self._mark_exhausted(entry, None)
|
||||
return None
|
||||
|
||||
try:
|
||||
if self.provider == "anthropic":
|
||||
from agent.anthropic_adapter import refresh_anthropic_oauth_pure
|
||||
|
||||
refreshed = refresh_anthropic_oauth_pure(
|
||||
entry.refresh_token,
|
||||
use_json=entry.source.endswith("hermes_pkce"),
|
||||
)
|
||||
entry.access_token = refreshed["access_token"]
|
||||
entry.refresh_token = refreshed["refresh_token"]
|
||||
entry.expires_at_ms = refreshed["expires_at_ms"]
|
||||
elif self.provider == "openai-codex":
|
||||
refreshed = auth_mod.refresh_codex_oauth_pure(
|
||||
entry.access_token,
|
||||
entry.refresh_token,
|
||||
)
|
||||
entry.access_token = refreshed["access_token"]
|
||||
entry.refresh_token = refreshed["refresh_token"]
|
||||
entry.last_refresh = refreshed.get("last_refresh")
|
||||
elif self.provider == "nous":
|
||||
refreshed = auth_mod.refresh_nous_oauth_pure(
|
||||
entry.access_token,
|
||||
entry.refresh_token,
|
||||
entry.client_id or "hermes-cli",
|
||||
entry.portal_base_url or "https://portal.nousresearch.com",
|
||||
entry.inference_base_url or "https://inference-api.nousresearch.com/v1",
|
||||
token_type=entry.token_type or "Bearer",
|
||||
scope=entry.scope or "",
|
||||
obtained_at=entry.obtained_at,
|
||||
expires_at=entry.expires_at,
|
||||
agent_key=entry.agent_key,
|
||||
agent_key_expires_at=entry.agent_key_expires_at,
|
||||
min_key_ttl_seconds=DEFAULT_AGENT_KEY_MIN_TTL_SECONDS,
|
||||
force_refresh=force,
|
||||
force_mint=force,
|
||||
)
|
||||
for key, value in refreshed.items():
|
||||
if hasattr(entry, key):
|
||||
setattr(entry, key, value)
|
||||
else:
|
||||
return entry
|
||||
except Exception:
|
||||
self._mark_exhausted(entry, None)
|
||||
return None
|
||||
|
||||
entry.last_status = "ok"
|
||||
entry.last_status_at = None
|
||||
entry.last_error_code = None
|
||||
self._persist()
|
||||
return entry
|
||||
|
||||
def _entry_needs_refresh(self, entry: PooledCredential) -> bool:
|
||||
if entry.auth_type != "oauth":
|
||||
return False
|
||||
if self.provider == "anthropic":
|
||||
if entry.expires_at_ms is None:
|
||||
return False
|
||||
return int(entry.expires_at_ms) <= int(time.time() * 1000) + 120_000
|
||||
if self.provider == "openai-codex":
|
||||
return _codex_access_token_is_expiring(
|
||||
entry.access_token,
|
||||
CODEX_ACCESS_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
)
|
||||
if self.provider == "nous":
|
||||
if _is_expiring(entry.expires_at, ACCESS_TOKEN_REFRESH_SKEW_SECONDS):
|
||||
return True
|
||||
return not _agent_key_is_usable(
|
||||
{
|
||||
"agent_key": entry.agent_key,
|
||||
"agent_key_expires_at": entry.agent_key_expires_at,
|
||||
},
|
||||
DEFAULT_AGENT_KEY_MIN_TTL_SECONDS,
|
||||
)
|
||||
return False
|
||||
|
||||
def select(self) -> Optional[PooledCredential]:
|
||||
now = time.time()
|
||||
for entry in sorted(self._entries, key=lambda item: item.priority):
|
||||
if entry.last_status == "exhausted":
|
||||
if entry.last_status_at and now - entry.last_status_at < EXHAUSTED_TTL_SECONDS:
|
||||
continue
|
||||
entry.last_status = "ok"
|
||||
entry.last_status_at = None
|
||||
entry.last_error_code = None
|
||||
self._persist()
|
||||
if self._entry_needs_refresh(entry):
|
||||
refreshed = self._refresh_entry(entry, force=False)
|
||||
if refreshed is None:
|
||||
continue
|
||||
entry = refreshed
|
||||
self._current_id = entry.id
|
||||
return entry
|
||||
self._current_id = None
|
||||
return None
|
||||
|
||||
def mark_exhausted_and_rotate(self, *, status_code: Optional[int]) -> Optional[PooledCredential]:
|
||||
entry = self.current() or self.select()
|
||||
if entry is None:
|
||||
return None
|
||||
self._mark_exhausted(entry, status_code)
|
||||
self._current_id = None
|
||||
return self.select()
|
||||
|
||||
def try_refresh_current(self) -> Optional[PooledCredential]:
|
||||
entry = self.current()
|
||||
if entry is None:
|
||||
return None
|
||||
refreshed = self._refresh_entry(entry, force=True)
|
||||
if refreshed is not None:
|
||||
self._current_id = refreshed.id
|
||||
return refreshed
|
||||
|
||||
def reset_statuses(self) -> int:
|
||||
count = 0
|
||||
for entry in self._entries:
|
||||
if entry.last_status or entry.last_status_at or entry.last_error_code:
|
||||
entry.last_status = None
|
||||
entry.last_status_at = None
|
||||
entry.last_error_code = None
|
||||
count += 1
|
||||
if count:
|
||||
self._persist()
|
||||
return count
|
||||
|
||||
def remove_index(self, index: int) -> Optional[PooledCredential]:
|
||||
ordered = sorted(self._entries, key=lambda item: item.priority)
|
||||
if index < 1 or index > len(ordered):
|
||||
return None
|
||||
removed = ordered.pop(index - 1)
|
||||
for new_priority, entry in enumerate(ordered):
|
||||
entry.priority = new_priority
|
||||
self._entries = ordered
|
||||
self._persist()
|
||||
if self._current_id == removed.id:
|
||||
self._current_id = None
|
||||
return removed
|
||||
|
||||
def add_entry(self, entry: PooledCredential) -> PooledCredential:
|
||||
entry.priority = _next_priority(self._entries)
|
||||
self._entries.append(entry)
|
||||
self._persist()
|
||||
return entry
|
||||
|
||||
|
||||
def _upsert_entry(entries: List[PooledCredential], provider: str, source: str, payload: Dict[str, Any]) -> bool:
|
||||
existing = next((entry for entry in entries if entry.source == source), None)
|
||||
if existing is None:
|
||||
payload.setdefault("id", uuid.uuid4().hex[:6])
|
||||
payload.setdefault("priority", _next_priority(entries))
|
||||
payload.setdefault("label", payload.get("label") or source)
|
||||
entries.append(PooledCredential.from_dict(provider, payload))
|
||||
return True
|
||||
|
||||
changed = False
|
||||
for key, value in payload.items():
|
||||
if key in {"id", "priority"} or value is None:
|
||||
continue
|
||||
if key == "label" and existing.label:
|
||||
continue
|
||||
if hasattr(existing, key) and getattr(existing, key) != value:
|
||||
setattr(existing, key, value)
|
||||
changed = True
|
||||
return changed
|
||||
|
||||
|
||||
def _seed_from_env(provider: str, entries: List[PooledCredential]) -> bool:
|
||||
changed = False
|
||||
if provider == "openrouter":
|
||||
token = os.getenv("OPENROUTER_API_KEY", "").strip()
|
||||
if token:
|
||||
changed |= _upsert_entry(
|
||||
entries,
|
||||
provider,
|
||||
"env:OPENROUTER_API_KEY",
|
||||
{
|
||||
"source": "env:OPENROUTER_API_KEY",
|
||||
"auth_type": "api_key",
|
||||
"access_token": token,
|
||||
"base_url": OPENROUTER_BASE_URL,
|
||||
"label": "OPENROUTER_API_KEY",
|
||||
},
|
||||
)
|
||||
return changed
|
||||
|
||||
pconfig = PROVIDER_REGISTRY.get(provider)
|
||||
if not pconfig or pconfig.auth_type != "api_key":
|
||||
return changed
|
||||
|
||||
env_url = ""
|
||||
if pconfig.base_url_env_var:
|
||||
env_url = os.getenv(pconfig.base_url_env_var, "").strip().rstrip("/")
|
||||
|
||||
for env_var in pconfig.api_key_env_vars:
|
||||
token = os.getenv(env_var, "").strip()
|
||||
if not token:
|
||||
continue
|
||||
auth_type = "oauth" if provider == "anthropic" and not token.startswith("sk-ant-api") else "api_key"
|
||||
base_url = env_url or pconfig.inference_base_url
|
||||
changed |= _upsert_entry(
|
||||
entries,
|
||||
provider,
|
||||
f"env:{env_var}",
|
||||
{
|
||||
"source": f"env:{env_var}",
|
||||
"auth_type": auth_type,
|
||||
"access_token": token,
|
||||
"base_url": base_url,
|
||||
"label": env_var,
|
||||
},
|
||||
)
|
||||
return changed
|
||||
|
||||
|
||||
def _seed_from_singletons(provider: str, entries: List[PooledCredential]) -> bool:
|
||||
changed = False
|
||||
auth_store = _load_auth_store()
|
||||
|
||||
if provider == "anthropic":
|
||||
from agent.anthropic_adapter import read_claude_code_credentials, read_hermes_oauth_credentials
|
||||
|
||||
hermes_creds = read_hermes_oauth_credentials()
|
||||
if hermes_creds and hermes_creds.get("accessToken"):
|
||||
changed |= _upsert_entry(
|
||||
entries,
|
||||
provider,
|
||||
"hermes_pkce",
|
||||
{
|
||||
"source": "hermes_pkce",
|
||||
"auth_type": "oauth",
|
||||
"access_token": hermes_creds.get("accessToken", ""),
|
||||
"refresh_token": hermes_creds.get("refreshToken"),
|
||||
"expires_at_ms": hermes_creds.get("expiresAt"),
|
||||
"label": _label_from_token(hermes_creds.get("accessToken", ""), "hermes_pkce"),
|
||||
},
|
||||
)
|
||||
claude_creds = read_claude_code_credentials()
|
||||
if claude_creds and claude_creds.get("accessToken"):
|
||||
changed |= _upsert_entry(
|
||||
entries,
|
||||
provider,
|
||||
"claude_code",
|
||||
{
|
||||
"source": "claude_code",
|
||||
"auth_type": "oauth",
|
||||
"access_token": claude_creds.get("accessToken", ""),
|
||||
"refresh_token": claude_creds.get("refreshToken"),
|
||||
"expires_at_ms": claude_creds.get("expiresAt"),
|
||||
"label": _label_from_token(claude_creds.get("accessToken", ""), "claude_code"),
|
||||
},
|
||||
)
|
||||
|
||||
elif provider == "nous":
|
||||
state = _load_provider_state(auth_store, "nous")
|
||||
if state:
|
||||
changed |= _upsert_entry(
|
||||
entries,
|
||||
provider,
|
||||
"device_code",
|
||||
{
|
||||
"source": "device_code",
|
||||
"auth_type": "oauth",
|
||||
"access_token": state.get("access_token", ""),
|
||||
"refresh_token": state.get("refresh_token"),
|
||||
"expires_at": state.get("expires_at"),
|
||||
"token_type": state.get("token_type"),
|
||||
"scope": state.get("scope"),
|
||||
"client_id": state.get("client_id"),
|
||||
"portal_base_url": state.get("portal_base_url"),
|
||||
"inference_base_url": state.get("inference_base_url"),
|
||||
"agent_key": state.get("agent_key"),
|
||||
"agent_key_expires_at": state.get("agent_key_expires_at"),
|
||||
"label": _label_from_token(state.get("access_token", ""), "device_code"),
|
||||
},
|
||||
)
|
||||
|
||||
elif provider == "openai-codex":
|
||||
state = _load_provider_state(auth_store, "openai-codex")
|
||||
tokens = state.get("tokens") if isinstance(state, dict) else None
|
||||
if isinstance(tokens, dict) and tokens.get("access_token"):
|
||||
changed |= _upsert_entry(
|
||||
entries,
|
||||
provider,
|
||||
"device_code",
|
||||
{
|
||||
"source": "device_code",
|
||||
"auth_type": "oauth",
|
||||
"access_token": tokens.get("access_token", ""),
|
||||
"refresh_token": tokens.get("refresh_token"),
|
||||
"base_url": "https://chatgpt.com/backend-api/codex",
|
||||
"last_refresh": state.get("last_refresh"),
|
||||
"label": _label_from_token(tokens.get("access_token", ""), "device_code"),
|
||||
},
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
|
||||
def load_pool(provider: str) -> CredentialPool:
|
||||
provider = (provider or "").strip().lower()
|
||||
raw_entries = read_credential_pool(provider)
|
||||
entries = [PooledCredential.from_dict(provider, payload) for payload in raw_entries]
|
||||
changed = _seed_from_singletons(provider, entries)
|
||||
changed |= _seed_from_env(provider, entries)
|
||||
if changed:
|
||||
write_credential_pool(
|
||||
provider,
|
||||
[entry.to_dict() for entry in sorted(entries, key=lambda item: item.priority)],
|
||||
)
|
||||
return CredentialPool(provider, entries)
|
||||
Loading…
Add table
Add a link
Reference in a new issue