mirror of
https://github.com/CyberSecurityUP/NeuroSploit.git
synced 2026-07-09 20:58:38 +02:00
v3.3.0 GUI dashboard + reports + model expansion + root fix
Engine:
- Fix: inject IS_SANDBOX=1 so Claude Code's --dangerously-skip-permissions
works under root (real backend runs were exiting rc=1 immediately)
- models: expand to 40 models / 13 providers, tagged CLI vs API
(NVIDIA NIM, DeepSeek, Mistral, Qwen/DashScope, Groq, Together, OpenRouter,
Ollama, Gemini) — Qwen/DeepSeek/Llama usable via API
- backends: on_start callback surfaces the exact argv ("what runs behind it")
- orchestrator: require a Playwright screenshot per confirmed finding; collect
results/activity.json; auto-generate reports after a run
- report.py: HTML always + PDF via Typst engine (.typ source emitted too)
Web dashboard (webgui/, stdlib only — no npm/build):
- Sidebar dashboard (PentAGI-style): Run / Agents / Insights / Reports / Settings
- Multi-target runs; live execution console + per-task activity; finding cards
with screenshots; backend+provider+model pickers (CLI & API)
- Agents tab: browse 213 + add new .md agents from the UI
- Insights: interactive RL-weight + severity charts
- Reports: download/preview PDF + HTML
- Settings/API: execution mode, per-provider API keys, orchestrator, verbosity
- Endpoints: /api/agents (GET/POST), /api/rl, /api/config, /api/reports,
/reports/* + /shots/* static serving
Cleanup: retire replaced web stack (frontend React, FastAPI backend, core
orchestration, old test) to legacy/. Active engine + GUI are fully standalone.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
CyberSecurityUP
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
NeuroSploit v3 - Smart Router Package
|
||||
|
||||
Gated by ENABLE_SMART_ROUTER=true env var.
|
||||
Provides SmartRouter singleton for LLM request routing.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
HAS_SMART_ROUTER = False
|
||||
_router_instance = None
|
||||
_registry_instance = None
|
||||
_refresher_instance = None
|
||||
|
||||
if os.getenv("ENABLE_SMART_ROUTER", "false").lower() == "true":
|
||||
try:
|
||||
from .provider_registry import ProviderRegistry, Provider, Account
|
||||
from .token_extractor import TokenExtractor
|
||||
from .router import SmartRouter
|
||||
from .token_refresher import TokenRefresher
|
||||
HAS_SMART_ROUTER = True
|
||||
logger.info("SmartRouter: Module loaded successfully")
|
||||
except ImportError as e:
|
||||
logger.warning(f"SmartRouter: Failed to load: {e}")
|
||||
|
||||
|
||||
async def init_router():
|
||||
"""Initialize the SmartRouter singleton. Called from main.py startup."""
|
||||
global _router_instance, _registry_instance, _refresher_instance
|
||||
|
||||
if not HAS_SMART_ROUTER:
|
||||
return
|
||||
|
||||
try:
|
||||
_registry_instance = ProviderRegistry()
|
||||
|
||||
# Auto-detect CLI tokens that were never detected before
|
||||
extractor = TokenExtractor()
|
||||
auto_detected = 0
|
||||
for token in extractor.detect_all():
|
||||
# Only add if provider has no accounts yet
|
||||
provider = _registry_instance.get_provider(token.provider_id)
|
||||
if provider and not provider.accounts:
|
||||
_registry_instance.add_account(
|
||||
provider_id=token.provider_id,
|
||||
label=token.label,
|
||||
credential=token.token,
|
||||
credential_type=token.credential_type,
|
||||
source="cli_detect",
|
||||
refresh_token=token.refresh_token,
|
||||
expires_at=token.expires_at,
|
||||
)
|
||||
auto_detected += 1
|
||||
if auto_detected:
|
||||
logger.info(f"SmartRouter: Auto-detected {auto_detected} new CLI provider(s)")
|
||||
print(f"[SmartRouter] Auto-detected {auto_detected} new CLI provider(s)")
|
||||
|
||||
_router_instance = SmartRouter(_registry_instance)
|
||||
|
||||
# Discover Gemini CLI project ID (needed for Cloud Code Assist API)
|
||||
try:
|
||||
project = await _router_instance.discover_gemini_project()
|
||||
if project:
|
||||
print(f"[SmartRouter] Gemini CLI project: {project}")
|
||||
except Exception as e:
|
||||
logger.debug(f"SmartRouter: Gemini project discovery skipped: {e}")
|
||||
|
||||
_refresher_instance = TokenRefresher(_registry_instance)
|
||||
await _refresher_instance.start()
|
||||
|
||||
# Log connected providers with details
|
||||
status = _registry_instance.get_all_status()
|
||||
connected = [s for s in status if s["connected"]]
|
||||
connected_names = [f"{s['name']}({s['default_model']})" for s in connected]
|
||||
logger.info(f"SmartRouter: Initialized ({len(connected)}/{len(status)} connected)")
|
||||
print(f"[SmartRouter] {len(connected)}/{len(status)} providers connected: {', '.join(connected_names) or 'none'}")
|
||||
except Exception as e:
|
||||
logger.error(f"SmartRouter: Init failed: {e}")
|
||||
print(f"[SmartRouter] Init failed: {e}")
|
||||
_router_instance = None
|
||||
|
||||
|
||||
async def shutdown_router():
|
||||
"""Shutdown the SmartRouter. Called from main.py shutdown."""
|
||||
global _refresher_instance
|
||||
if _refresher_instance:
|
||||
await _refresher_instance.stop()
|
||||
_refresher_instance = None
|
||||
|
||||
|
||||
def get_router() -> Optional["SmartRouter"]:
|
||||
"""Get the SmartRouter singleton (or None if disabled)."""
|
||||
return _router_instance
|
||||
|
||||
|
||||
def get_registry() -> Optional["ProviderRegistry"]:
|
||||
"""Get the ProviderRegistry singleton (or None if disabled)."""
|
||||
return _registry_instance
|
||||
|
||||
|
||||
def get_extractor() -> Optional["TokenExtractor"]:
|
||||
"""Get a new TokenExtractor instance (or None if disabled)."""
|
||||
if not HAS_SMART_ROUTER:
|
||||
return None
|
||||
return TokenExtractor()
|
||||
@@ -0,0 +1,545 @@
|
||||
"""
|
||||
NeuroSploit v3 - Provider Registry
|
||||
|
||||
Central registry of 20 LLM providers + their accounts.
|
||||
Persists metadata to data/providers.json (credentials stay in-memory only).
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PROVIDERS_FILE = Path(__file__).parent.parent.parent.parent / "data" / "providers.json"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Account:
|
||||
id: str
|
||||
label: str
|
||||
source: str # "manual" | "cli_detect" | "env_var"
|
||||
credential_type: str # "api_key" | "oauth"
|
||||
created_at: str = ""
|
||||
last_used: Optional[str] = None
|
||||
tokens_used: int = 0
|
||||
is_active: bool = True
|
||||
expires_at: Optional[float] = None # Unix timestamp for OAuth tokens
|
||||
model_override: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.created_at:
|
||||
self.created_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict) -> "Account":
|
||||
return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
|
||||
|
||||
|
||||
@dataclass
|
||||
class Provider:
|
||||
id: str
|
||||
name: str
|
||||
auth_type: str # "api_key" | "oauth"
|
||||
api_format: str # "anthropic" | "openai_compat" | "gemini" | "ollama"
|
||||
base_url: str
|
||||
tier: int # 1=subscription/paid, 2=cheap, 3=free
|
||||
default_model: str
|
||||
accounts: Dict[str, Account] = field(default_factory=dict)
|
||||
env_key: Optional[str] = None # e.g. "ANTHROPIC_API_KEY"
|
||||
enabled: bool = True # UI toggle: disabled providers are skipped by router
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
d = asdict(self)
|
||||
d["accounts"] = {k: v.to_dict() if isinstance(v, Account) else v for k, v in self.accounts.items()}
|
||||
return d
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict) -> "Provider":
|
||||
accounts_raw = data.pop("accounts", {})
|
||||
accounts = {}
|
||||
for k, v in accounts_raw.items():
|
||||
if isinstance(v, dict):
|
||||
accounts[k] = Account.from_dict(v)
|
||||
else:
|
||||
accounts[k] = v
|
||||
filtered = {k: v for k, v in data.items() if k in cls.__dataclass_fields__}
|
||||
return cls(accounts=accounts, **filtered)
|
||||
|
||||
|
||||
# Default provider definitions
|
||||
DEFAULT_PROVIDERS: List[Dict] = [
|
||||
# === NVIDIA NIM Provider (Tier 2 - Free) ===
|
||||
{
|
||||
"id": "nim", "name": "NVIDIA NIM", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://integrate.api.nvidia.com/v1",
|
||||
"tier": 2, "default_model": os.getenv("NIM_MODEL", "openai/gpt-oss-120b"),
|
||||
"env_key": "NIM_API_KEY",
|
||||
},
|
||||
# === OAuth Providers (Tier 1 - Subscription) ===
|
||||
{
|
||||
"id": "claude_code", "name": "Claude Code", "auth_type": "oauth",
|
||||
"api_format": "anthropic", "base_url": "https://api.anthropic.com",
|
||||
"tier": 1, "default_model": "claude-sonnet-4-20250514",
|
||||
},
|
||||
{
|
||||
"id": "codex_cli", "name": "OpenAI Codex CLI", "auth_type": "oauth",
|
||||
"api_format": "openai_compat", "base_url": "https://api.openai.com/v1",
|
||||
"tier": 1, "default_model": "gpt-4o",
|
||||
},
|
||||
{
|
||||
"id": "gemini_cli", "name": "Gemini CLI", "auth_type": "oauth",
|
||||
"api_format": "gemini_code_assist", "base_url": "https://cloudcode-pa.googleapis.com",
|
||||
"tier": 1, "default_model": "gemini-2.5-flash",
|
||||
},
|
||||
{
|
||||
"id": "cursor", "name": "Cursor", "auth_type": "oauth",
|
||||
"api_format": "openai_compat", "base_url": "https://api2.cursor.sh/v1",
|
||||
"tier": 1, "default_model": "cursor-fast",
|
||||
},
|
||||
{
|
||||
"id": "copilot", "name": "GitHub Copilot", "auth_type": "oauth",
|
||||
"api_format": "openai_compat", "base_url": "https://api.githubcopilot.com",
|
||||
"tier": 1, "default_model": "gpt-4o",
|
||||
},
|
||||
{
|
||||
"id": "iflow", "name": "iFlow AI", "auth_type": "oauth",
|
||||
"api_format": "openai_compat", "base_url": "https://api.iflow.ai/v1",
|
||||
"tier": 1, "default_model": "kimi-k2",
|
||||
},
|
||||
{
|
||||
"id": "qwen_code", "name": "Qwen Code", "auth_type": "oauth",
|
||||
"api_format": "openai_compat", "base_url": "https://chat.qwen.ai/api/v1",
|
||||
"tier": 1, "default_model": "qwen3-coder",
|
||||
},
|
||||
{
|
||||
"id": "kiro", "name": "Kiro AI", "auth_type": "oauth",
|
||||
"api_format": "anthropic", "base_url": "https://api.anthropic.com",
|
||||
"tier": 1, "default_model": "claude-sonnet-4-20250514",
|
||||
},
|
||||
# === API Key Providers (Tier 1 - Paid) ===
|
||||
{
|
||||
"id": "anthropic", "name": "Anthropic", "auth_type": "api_key",
|
||||
"api_format": "anthropic", "base_url": "https://api.anthropic.com",
|
||||
"tier": 1, "default_model": "claude-sonnet-4-20250514",
|
||||
"env_key": "ANTHROPIC_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "openai", "name": "OpenAI", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://api.openai.com/v1",
|
||||
"tier": 1, "default_model": "gpt-4o",
|
||||
"env_key": "OPENAI_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "gemini", "name": "Gemini", "auth_type": "api_key",
|
||||
"api_format": "gemini", "base_url": "https://generativelanguage.googleapis.com/v1beta",
|
||||
"tier": 1, "default_model": "gemini-2.5-flash",
|
||||
"env_key": "GEMINI_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "openrouter", "name": "OpenRouter", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://openrouter.ai/api/v1",
|
||||
"tier": 1, "default_model": "anthropic/claude-sonnet-4-20250514",
|
||||
"env_key": "OPENROUTER_API_KEY",
|
||||
},
|
||||
# === API Key Providers (Tier 2 - Cheap) ===
|
||||
{
|
||||
"id": "glm", "name": "GLM (Zhipu AI)", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://open.bigmodel.cn/api/paas/v4",
|
||||
"tier": 2, "default_model": "glm-4-flash",
|
||||
"env_key": "GLM_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "kimi", "name": "Kimi (Moonshot)", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://api.moonshot.cn/v1",
|
||||
"tier": 2, "default_model": "moonshot-v1-8k",
|
||||
"env_key": "KIMI_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "minimax", "name": "Minimax", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://api.minimax.chat/v1",
|
||||
"tier": 2, "default_model": "abab6.5-chat",
|
||||
"env_key": "MINIMAX_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "together", "name": "Together AI", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://api.together.xyz/v1",
|
||||
"tier": 2, "default_model": "meta-llama/Llama-3-70b-chat-hf",
|
||||
"env_key": "TOGETHER_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "fireworks", "name": "Fireworks AI", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "https://api.fireworks.ai/inference/v1",
|
||||
"tier": 2, "default_model": "accounts/fireworks/models/llama-v3p1-70b-instruct",
|
||||
"env_key": "FIREWORKS_API_KEY",
|
||||
},
|
||||
# === Local Providers (Tier 3 - Free/Self-hosted) ===
|
||||
{
|
||||
"id": "ollama", "name": "Ollama", "auth_type": "api_key",
|
||||
"api_format": "ollama", "base_url": "http://localhost:11434",
|
||||
"tier": 3, "default_model": "llama3",
|
||||
"env_key": "OLLAMA_API_KEY",
|
||||
},
|
||||
{
|
||||
"id": "lmstudio", "name": "LM Studio", "auth_type": "api_key",
|
||||
"api_format": "openai_compat", "base_url": "http://localhost:1234/v1",
|
||||
"tier": 3, "default_model": "local-model",
|
||||
"env_key": "LMSTUDIO_API_KEY",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""Central registry for LLM providers and their accounts.
|
||||
|
||||
Metadata persists to data/providers.json.
|
||||
Credentials are kept IN-MEMORY ONLY for security (re-extracted from CLI on startup).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._providers: Dict[str, Provider] = {}
|
||||
self._credentials: Dict[str, str] = {} # acct_id -> credential (IN-MEMORY)
|
||||
self._refresh_tokens: Dict[str, str] = {} # acct_id -> refresh_token (IN-MEMORY)
|
||||
self._seed_defaults()
|
||||
self._load()
|
||||
self._seed_env_credentials()
|
||||
self._restore_cli_credentials()
|
||||
|
||||
# ── Seeding ──────────────────────────────────────────────
|
||||
|
||||
def _seed_defaults(self):
|
||||
"""Register built-in provider definitions."""
|
||||
for pdef in DEFAULT_PROVIDERS:
|
||||
pid = pdef["id"]
|
||||
if pid not in self._providers:
|
||||
self._providers[pid] = Provider(
|
||||
id=pid,
|
||||
name=pdef["name"],
|
||||
auth_type=pdef["auth_type"],
|
||||
api_format=pdef["api_format"],
|
||||
base_url=pdef["base_url"],
|
||||
tier=pdef["tier"],
|
||||
default_model=pdef["default_model"],
|
||||
env_key=pdef.get("env_key"),
|
||||
)
|
||||
|
||||
def _seed_env_credentials(self):
|
||||
"""Auto-load credentials from environment variables."""
|
||||
for pid, provider in self._providers.items():
|
||||
if not provider.env_key:
|
||||
continue
|
||||
value = os.getenv(provider.env_key, "").strip()
|
||||
if not value:
|
||||
continue
|
||||
# Check if we already have an env_var account
|
||||
existing = [a for a in provider.accounts.values() if a.source == "env_var"]
|
||||
if existing:
|
||||
acct = existing[0]
|
||||
# Update credential in memory
|
||||
self._credentials[acct.id] = value
|
||||
# Reactivate if deactivated — env key is set, it should be active
|
||||
if not acct.is_active:
|
||||
acct.is_active = True
|
||||
logger.info(f"SmartRouter: Reactivated {acct.label} (env key present)")
|
||||
continue
|
||||
# Create new env_var account
|
||||
acct_id = f"acct_{uuid.uuid4().hex[:8]}"
|
||||
acct = Account(
|
||||
id=acct_id,
|
||||
label=f"{provider.name} (env)",
|
||||
source="env_var",
|
||||
credential_type="api_key",
|
||||
)
|
||||
provider.accounts[acct_id] = acct
|
||||
self._credentials[acct_id] = value
|
||||
logger.info(f"SmartRouter: Loaded {provider.env_key} for {provider.name}")
|
||||
self._save()
|
||||
|
||||
def _restore_cli_credentials(self):
|
||||
"""Re-extract CLI tokens on startup for cli_detect accounts.
|
||||
|
||||
CLI tokens are never persisted to disk. On restart, we re-detect
|
||||
them from the CLI tools (same files TokenExtractor reads).
|
||||
Also reactivates accounts that were deactivated due to expired tokens.
|
||||
"""
|
||||
try:
|
||||
from .token_extractor import TokenExtractor
|
||||
extractor = TokenExtractor()
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
restored = 0
|
||||
for pid, provider in self._providers.items():
|
||||
cli_accounts = [a for a in provider.accounts.values() if a.source == "cli_detect"]
|
||||
if not cli_accounts:
|
||||
continue
|
||||
|
||||
# Try to re-extract token for this provider
|
||||
token = extractor.detect(pid)
|
||||
if not token:
|
||||
logger.debug(f"SmartRouter: No CLI token found for {provider.name}")
|
||||
continue
|
||||
|
||||
# Update the first cli_detect account (remove duplicates)
|
||||
primary = cli_accounts[0]
|
||||
self._credentials[primary.id] = token.token
|
||||
if token.refresh_token:
|
||||
self._refresh_tokens[primary.id] = token.refresh_token
|
||||
|
||||
# Update expires_at (token_extractor may have parsed it from the credential file)
|
||||
if token.expires_at:
|
||||
primary.expires_at = token.expires_at
|
||||
elif primary.expires_at is None and token.credential_type == "oauth":
|
||||
# OAuth tokens without expiry — assume 1 hour from now as a hint
|
||||
# so the TokenRefresher will check and refresh them
|
||||
import time as _time
|
||||
primary.expires_at = _time.time() + 3600
|
||||
logger.debug(f"SmartRouter: Set default 1h expiry for {primary.label}")
|
||||
|
||||
# Reactivate if it was deactivated (token may have been refreshed externally)
|
||||
if not primary.is_active:
|
||||
primary.is_active = True
|
||||
logger.info(f"SmartRouter: Reactivated {primary.label} (fresh token found)")
|
||||
|
||||
# Remove duplicate cli_detect accounts (keep only primary)
|
||||
for dup in cli_accounts[1:]:
|
||||
del provider.accounts[dup.id]
|
||||
self._credentials.pop(dup.id, None)
|
||||
self._refresh_tokens.pop(dup.id, None)
|
||||
logger.info(f"SmartRouter: Removed duplicate account {dup.id} for {provider.name}")
|
||||
|
||||
restored += 1
|
||||
logger.info(
|
||||
f"SmartRouter: Restored CLI credential for {provider.name} "
|
||||
f"(token={'***' + token.token[-8:]}, refresh={'yes' if token.refresh_token else 'no'}, "
|
||||
f"expires={primary.expires_at})"
|
||||
)
|
||||
|
||||
if restored:
|
||||
self._save()
|
||||
logger.info(f"SmartRouter: Restored {restored} CLI provider(s)")
|
||||
|
||||
# ── CRUD ─────────────────────────────────────────────────
|
||||
|
||||
def add_account(
|
||||
self,
|
||||
provider_id: str,
|
||||
label: str,
|
||||
credential: str,
|
||||
credential_type: str = "api_key",
|
||||
source: str = "manual",
|
||||
refresh_token: Optional[str] = None,
|
||||
expires_at: Optional[float] = None,
|
||||
model_override: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""Add a new account to a provider. Returns account ID or None."""
|
||||
provider = self._providers.get(provider_id)
|
||||
if not provider:
|
||||
logger.warning(f"SmartRouter: Unknown provider {provider_id}")
|
||||
return None
|
||||
|
||||
# Deduplicate: if a cli_detect account already exists, update it instead
|
||||
if source == "cli_detect":
|
||||
existing = [a for a in provider.accounts.values() if a.source == "cli_detect"]
|
||||
if existing:
|
||||
acct = existing[0]
|
||||
self._credentials[acct.id] = credential
|
||||
if refresh_token:
|
||||
self._refresh_tokens[acct.id] = refresh_token
|
||||
acct.expires_at = expires_at
|
||||
acct.is_active = True
|
||||
self._save()
|
||||
logger.info(f"SmartRouter: Updated existing CLI account {acct.label} for {provider.name}")
|
||||
return acct.id
|
||||
|
||||
acct_id = f"acct_{uuid.uuid4().hex[:8]}"
|
||||
acct = Account(
|
||||
id=acct_id,
|
||||
label=label,
|
||||
source=source,
|
||||
credential_type=credential_type,
|
||||
expires_at=expires_at,
|
||||
model_override=model_override,
|
||||
)
|
||||
provider.accounts[acct_id] = acct
|
||||
self._credentials[acct_id] = credential
|
||||
if refresh_token:
|
||||
self._refresh_tokens[acct_id] = refresh_token
|
||||
|
||||
self._save()
|
||||
logger.info(f"SmartRouter: Added account {label} to {provider.name}")
|
||||
return acct_id
|
||||
|
||||
def remove_account(self, provider_id: str, account_id: str) -> bool:
|
||||
"""Remove an account from a provider."""
|
||||
provider = self._providers.get(provider_id)
|
||||
if not provider or account_id not in provider.accounts:
|
||||
return False
|
||||
|
||||
del provider.accounts[account_id]
|
||||
self._credentials.pop(account_id, None)
|
||||
self._refresh_tokens.pop(account_id, None)
|
||||
self._save()
|
||||
logger.info(f"SmartRouter: Removed account {account_id} from {provider.name}")
|
||||
return True
|
||||
|
||||
def get_credential(self, account_id: str) -> Optional[str]:
|
||||
"""Get credential for an account (from memory only)."""
|
||||
return self._credentials.get(account_id)
|
||||
|
||||
def get_refresh_token(self, account_id: str) -> Optional[str]:
|
||||
"""Get refresh token for an account (from memory only)."""
|
||||
return self._refresh_tokens.get(account_id)
|
||||
|
||||
def update_credential(
|
||||
self,
|
||||
account_id: str,
|
||||
new_credential: str,
|
||||
new_expires_at: Optional[float] = None,
|
||||
):
|
||||
"""Update credential (and optionally expiry) for an account."""
|
||||
self._credentials[account_id] = new_credential
|
||||
# Update expiry on the account object
|
||||
for provider in self._providers.values():
|
||||
if account_id in provider.accounts:
|
||||
provider.accounts[account_id].expires_at = new_expires_at
|
||||
self._save()
|
||||
break
|
||||
|
||||
def deactivate_account(self, account_id: str):
|
||||
"""Mark account as inactive (e.g., on 401/403)."""
|
||||
for provider in self._providers.values():
|
||||
if account_id in provider.accounts:
|
||||
provider.accounts[account_id].is_active = False
|
||||
self._save()
|
||||
logger.warning(f"SmartRouter: Deactivated account {account_id}")
|
||||
break
|
||||
|
||||
def reactivate_account(self, account_id: str) -> bool:
|
||||
"""Re-activate a deactivated account (e.g., after token refresh)."""
|
||||
for provider in self._providers.values():
|
||||
if account_id in provider.accounts:
|
||||
acct = provider.accounts[account_id]
|
||||
if not acct.is_active:
|
||||
acct.is_active = True
|
||||
self._save()
|
||||
logger.info(f"SmartRouter: Reactivated account {account_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
# ── Queries ──────────────────────────────────────────────
|
||||
|
||||
def get_active_accounts(self, provider_id: str) -> List[Account]:
|
||||
"""Get all active accounts for a provider."""
|
||||
provider = self._providers.get(provider_id)
|
||||
if not provider:
|
||||
return []
|
||||
return [
|
||||
a for a in provider.accounts.values()
|
||||
if a.is_active and a.id in self._credentials
|
||||
]
|
||||
|
||||
def get_provider(self, provider_id: str) -> Optional[Provider]:
|
||||
"""Get a provider by ID."""
|
||||
return self._providers.get(provider_id)
|
||||
|
||||
def get_providers_by_tier(self, tier: int) -> List[Provider]:
|
||||
"""Get all providers in a tier that have active accounts."""
|
||||
return [
|
||||
p for p in self._providers.values()
|
||||
if p.tier == tier and self.get_active_accounts(p.id)
|
||||
]
|
||||
|
||||
def get_all_providers(self) -> List[Provider]:
|
||||
"""Get all registered providers."""
|
||||
return list(self._providers.values())
|
||||
|
||||
def toggle_provider(self, provider_id: str, enabled: bool) -> bool:
|
||||
"""Enable or disable a provider. Disabled providers are skipped by the router."""
|
||||
provider = self._providers.get(provider_id)
|
||||
if not provider:
|
||||
return False
|
||||
provider.enabled = enabled
|
||||
self._save()
|
||||
logger.info(f"Provider {provider_id} {'enabled' if enabled else 'disabled'}")
|
||||
return True
|
||||
|
||||
def get_all_status(self) -> List[Dict]:
|
||||
"""Get status summary for all providers (for API/UI)."""
|
||||
result = []
|
||||
for p in self._providers.values():
|
||||
active = self.get_active_accounts(p.id)
|
||||
total_tokens = sum(a.tokens_used for a in p.accounts.values())
|
||||
result.append({
|
||||
"id": p.id,
|
||||
"name": p.name,
|
||||
"auth_type": p.auth_type,
|
||||
"api_format": p.api_format,
|
||||
"tier": p.tier,
|
||||
"default_model": p.default_model,
|
||||
"accounts_total": len(p.accounts),
|
||||
"accounts_active": len(active),
|
||||
"total_tokens_used": total_tokens,
|
||||
"connected": len(active) > 0,
|
||||
"enabled": getattr(p, "enabled", True),
|
||||
})
|
||||
return result
|
||||
|
||||
def record_usage(self, account_id: str, tokens: int):
|
||||
"""Record token usage for an account."""
|
||||
for provider in self._providers.values():
|
||||
if account_id in provider.accounts:
|
||||
provider.accounts[account_id].tokens_used += tokens
|
||||
provider.accounts[account_id].last_used = time.strftime(
|
||||
"%Y-%m-%dT%H:%M:%SZ", time.gmtime()
|
||||
)
|
||||
# Save periodically (every 10th call to avoid I/O spam)
|
||||
if provider.accounts[account_id].tokens_used % 10000 < tokens:
|
||||
self._save()
|
||||
break
|
||||
|
||||
# ── Persistence ──────────────────────────────────────────
|
||||
|
||||
def _save(self):
|
||||
"""Save provider metadata to JSON (NO credentials)."""
|
||||
try:
|
||||
PROVIDERS_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
data = {
|
||||
pid: p.to_dict()
|
||||
for pid, p in self._providers.items()
|
||||
}
|
||||
tmp = PROVIDERS_FILE.with_suffix(".tmp")
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(data, f, indent=2, default=str)
|
||||
tmp.rename(PROVIDERS_FILE)
|
||||
except Exception as e:
|
||||
logger.warning(f"SmartRouter: Failed to save providers: {e}")
|
||||
|
||||
def _load(self):
|
||||
"""Load provider metadata from JSON, merge with defaults."""
|
||||
if not PROVIDERS_FILE.exists():
|
||||
return
|
||||
try:
|
||||
with open(PROVIDERS_FILE) as f:
|
||||
data = json.load(f)
|
||||
for pid, pdata in data.items():
|
||||
if pid in self._providers:
|
||||
# Merge accounts from disk into existing provider
|
||||
saved_accounts = pdata.get("accounts", {})
|
||||
for aid, adata in saved_accounts.items():
|
||||
if isinstance(adata, dict):
|
||||
self._providers[pid].accounts[aid] = Account.from_dict(adata)
|
||||
else:
|
||||
# Unknown provider from disk - load it
|
||||
self._providers[pid] = Provider.from_dict(pdata)
|
||||
logger.info(f"SmartRouter: Loaded {len(data)} providers from disk")
|
||||
except Exception as e:
|
||||
logger.warning(f"SmartRouter: Failed to load providers: {e}")
|
||||
@@ -0,0 +1,606 @@
|
||||
"""
|
||||
NeuroSploit v3 - Smart Router
|
||||
|
||||
Core routing engine with 4 API format translators, tier-based failover,
|
||||
round-robin load balancing, and quota tracking.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .provider_registry import Account, Provider, ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QuotaTracker:
|
||||
"""Tracks per-account quota exhaustion with automatic recovery."""
|
||||
|
||||
def __init__(self):
|
||||
self._exhausted: Dict[str, float] = {} # acct_id -> retry_after_timestamp
|
||||
self._round_robin: Dict[str, int] = {} # provider_id -> last_index
|
||||
|
||||
def record_exhausted(self, account_id: str, retry_after: int = 300):
|
||||
"""Mark an account as quota-exhausted."""
|
||||
self._exhausted[account_id] = time.time() + retry_after
|
||||
logger.warning(f"QuotaTracker: {account_id} exhausted, retry after {retry_after}s")
|
||||
|
||||
def is_available(self, account_id: str) -> bool:
|
||||
"""Check if an account is available (not exhausted or recovered)."""
|
||||
if account_id not in self._exhausted:
|
||||
return True
|
||||
if time.time() > self._exhausted[account_id]:
|
||||
del self._exhausted[account_id]
|
||||
return True
|
||||
return False
|
||||
|
||||
def next_account(self, provider_id: str, accounts: List[Account]) -> Optional[Account]:
|
||||
"""Round-robin selection from available accounts."""
|
||||
available = [a for a in accounts if self.is_available(a.id)]
|
||||
if not available:
|
||||
return None
|
||||
idx = self._round_robin.get(provider_id, -1) + 1
|
||||
idx = idx % len(available)
|
||||
self._round_robin[provider_id] = idx
|
||||
return available[idx]
|
||||
|
||||
|
||||
class SmartRouter:
|
||||
"""Smart LLM router with format translation, failover, and load balancing.
|
||||
|
||||
Routing priority:
|
||||
1. Preferred provider (if specified)
|
||||
2. Tier 1 providers (subscription + paid API keys)
|
||||
3. Tier 2 providers (cheap API keys)
|
||||
4. Tier 3 providers (free/local)
|
||||
"""
|
||||
|
||||
def __init__(self, registry: ProviderRegistry):
|
||||
self.registry = registry
|
||||
self.quota = QuotaTracker()
|
||||
self._total_requests = 0
|
||||
self._total_tokens = 0
|
||||
self._failover_count = 0
|
||||
self._last_provider: Optional[str] = None # Last provider used
|
||||
self._last_model: Optional[str] = None # Last model used
|
||||
self._last_account_label: Optional[str] = None # Last account label
|
||||
self._gemini_project_id: Optional[str] = None # Cached Gemini CLI project ID
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
system: str = "",
|
||||
max_tokens: int = 4096,
|
||||
preferred_provider: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
) -> str:
|
||||
"""Route a generation request through available providers.
|
||||
|
||||
Tries providers in tier order, fails over on errors.
|
||||
Returns generated text or raises if all providers fail.
|
||||
"""
|
||||
candidates = self._build_candidate_list(preferred_provider)
|
||||
if not candidates:
|
||||
raise RuntimeError(
|
||||
f"SmartRouter: No providers available "
|
||||
f"(preferred={preferred_provider}, model={model})"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"SmartRouter: {len(candidates)} candidates "
|
||||
f"[{', '.join(f'{p.id}/{a.label}' for p, a in candidates)}]"
|
||||
)
|
||||
|
||||
last_error = None
|
||||
for provider, account in candidates:
|
||||
try:
|
||||
use_model = model or account.model_override or provider.default_model
|
||||
logger.info(f"SmartRouter: Trying {provider.name} ({use_model}) via {account.label}")
|
||||
|
||||
text, tokens_used = await self._call_provider(
|
||||
provider=provider,
|
||||
account=account,
|
||||
prompt=prompt,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
model=use_model,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
# Record success + track which provider/model served the request
|
||||
self._total_requests += 1
|
||||
self._total_tokens += tokens_used
|
||||
self._last_provider = provider.id
|
||||
self._last_model = use_model
|
||||
self._last_account_label = account.label
|
||||
self.registry.record_usage(account.id, tokens_used)
|
||||
logger.info(f"SmartRouter: OK — served by {provider.name} ({use_model})")
|
||||
return text
|
||||
|
||||
except QuotaExhaustedError as e:
|
||||
retry_after = getattr(e, "retry_after", 300)
|
||||
self.quota.record_exhausted(account.id, retry_after)
|
||||
self._failover_count += 1
|
||||
last_error = e
|
||||
logger.info(f"SmartRouter: {provider.name}/{account.label} quota exhausted, trying next")
|
||||
|
||||
except AuthenticationError as e:
|
||||
# For CLI/OAuth accounts, try re-extracting token and RETRY immediately
|
||||
if account.source == "cli_detect":
|
||||
refreshed = await self._try_reextract_cli_token(provider.id, account.id)
|
||||
if refreshed:
|
||||
logger.info(f"SmartRouter: {provider.name}/{account.label} token re-extracted, retrying NOW")
|
||||
try:
|
||||
text, tokens_used = await self._call_provider(
|
||||
provider=provider,
|
||||
account=account,
|
||||
prompt=prompt,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
model=use_model,
|
||||
temperature=temperature,
|
||||
)
|
||||
self._total_requests += 1
|
||||
self._total_tokens += tokens_used
|
||||
self._last_provider = provider.id
|
||||
self._last_model = use_model
|
||||
self._last_account_label = account.label
|
||||
self.registry.record_usage(account.id, tokens_used)
|
||||
logger.info(f"SmartRouter: OK — served by {provider.name} ({use_model}) after token re-extract")
|
||||
return text
|
||||
except Exception as retry_err:
|
||||
logger.warning(f"SmartRouter: {provider.name} retry after re-extract also failed: {retry_err}")
|
||||
last_error = retry_err
|
||||
|
||||
self.registry.deactivate_account(account.id)
|
||||
self._failover_count += 1
|
||||
if last_error is None:
|
||||
last_error = e
|
||||
logger.warning(f"SmartRouter: {provider.name}/{account.label} auth failed, deactivated")
|
||||
|
||||
except Exception as e:
|
||||
self._failover_count += 1
|
||||
last_error = e
|
||||
logger.warning(f"SmartRouter: {provider.name}/{account.label} error: {e}")
|
||||
|
||||
raise RuntimeError(f"SmartRouter: All providers failed. Last error: {last_error}")
|
||||
|
||||
def _build_candidate_list(
|
||||
self, preferred: Optional[str] = None
|
||||
) -> List[Tuple[Provider, Account]]:
|
||||
"""Build ordered list of (provider, account) candidates.
|
||||
|
||||
If preferred is set, that provider is tried FIRST, then falls back
|
||||
to other providers of the same tier if all accounts fail.
|
||||
If preferred is not set, all providers are tried by tier.
|
||||
"""
|
||||
candidates = []
|
||||
seen_account_ids = set()
|
||||
|
||||
if preferred:
|
||||
# Preferred provider goes first in candidate list
|
||||
provider = self.registry.get_provider(preferred)
|
||||
if provider:
|
||||
accounts = self.registry.get_active_accounts(preferred)
|
||||
for acct in accounts:
|
||||
if self.quota.is_available(acct.id):
|
||||
candidates.append((provider, acct))
|
||||
seen_account_ids.add(acct.id)
|
||||
if not candidates:
|
||||
logger.warning(
|
||||
f"SmartRouter: Preferred provider '{preferred}' has no active accounts! "
|
||||
f"Falling back to all providers."
|
||||
)
|
||||
|
||||
# Add remaining providers as fallback (by tier)
|
||||
for tier in (1, 2, 3):
|
||||
providers = self.registry.get_providers_by_tier(tier)
|
||||
for provider in providers:
|
||||
if not getattr(provider, "enabled", True):
|
||||
continue
|
||||
acct = self.quota.next_account(
|
||||
provider.id,
|
||||
self.registry.get_active_accounts(provider.id),
|
||||
)
|
||||
if acct and acct.id not in seen_account_ids:
|
||||
candidates.append((provider, acct))
|
||||
seen_account_ids.add(acct.id)
|
||||
|
||||
return candidates
|
||||
|
||||
async def _call_provider(
|
||||
self,
|
||||
provider: Provider,
|
||||
account: Account,
|
||||
prompt: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
model: str,
|
||||
temperature: float,
|
||||
) -> Tuple[str, int]:
|
||||
"""Make an API call to a provider. Returns (text, tokens_used)."""
|
||||
credential = self.registry.get_credential(account.id)
|
||||
if not credential:
|
||||
raise AuthenticationError(f"No credential for {account.id}")
|
||||
|
||||
format_map = {
|
||||
"anthropic": (self._build_anthropic_request, self._parse_anthropic_response),
|
||||
"openai_compat": (self._build_openai_request, self._parse_openai_response),
|
||||
"gemini": (self._build_gemini_request, self._parse_gemini_response),
|
||||
"gemini_code_assist": (self._build_gemini_code_assist_request, self._parse_gemini_code_assist_response),
|
||||
"ollama": (self._build_ollama_request, self._parse_ollama_response),
|
||||
}
|
||||
|
||||
builder, parser = format_map.get(provider.api_format, (None, None))
|
||||
if not builder:
|
||||
raise ValueError(f"Unknown API format: {provider.api_format}")
|
||||
|
||||
url, headers, payload = builder(
|
||||
base_url=provider.base_url,
|
||||
credential=credential,
|
||||
credential_type=account.credential_type,
|
||||
prompt=prompt,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
import aiohttp
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=120),
|
||||
) as resp:
|
||||
if resp.status == 429:
|
||||
retry_after = int(resp.headers.get("Retry-After", "300"))
|
||||
raise QuotaExhaustedError(
|
||||
f"Rate limited by {provider.name}",
|
||||
retry_after=retry_after,
|
||||
)
|
||||
if resp.status in (401, 403):
|
||||
raise AuthenticationError(
|
||||
f"Authentication failed for {provider.name}: {resp.status}"
|
||||
)
|
||||
if resp.status >= 400:
|
||||
body = await resp.text()
|
||||
raise RuntimeError(
|
||||
f"{provider.name} returned {resp.status}: {body[:300]}"
|
||||
)
|
||||
|
||||
data = await resp.json()
|
||||
return parser(data)
|
||||
|
||||
# ── Format Builders ──────────────────────────────────────
|
||||
|
||||
def _build_anthropic_request(
|
||||
self, base_url: str, credential: str, credential_type: str,
|
||||
prompt: str, system: str, max_tokens: int, model: str, temperature: float,
|
||||
) -> Tuple[str, Dict, Dict]:
|
||||
url = f"{base_url}/v1/messages"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"anthropic-version": "2023-06-01",
|
||||
}
|
||||
if credential_type == "oauth":
|
||||
headers["Authorization"] = f"Bearer {credential}"
|
||||
else:
|
||||
headers["x-api-key"] = credential
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if system:
|
||||
payload["system"] = system
|
||||
|
||||
return url, headers, payload
|
||||
|
||||
def _build_openai_request(
|
||||
self, base_url: str, credential: str, credential_type: str,
|
||||
prompt: str, system: str, max_tokens: int, model: str, temperature: float,
|
||||
) -> Tuple[str, Dict, Dict]:
|
||||
url = f"{base_url}/chat/completions"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {credential}",
|
||||
}
|
||||
|
||||
messages = []
|
||||
if system:
|
||||
messages.append({"role": "system", "content": system})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
}
|
||||
return url, headers, payload
|
||||
|
||||
def _build_gemini_request(
|
||||
self, base_url: str, credential: str, credential_type: str,
|
||||
prompt: str, system: str, max_tokens: int, model: str, temperature: float,
|
||||
) -> Tuple[str, Dict, Dict]:
|
||||
if credential_type == "oauth":
|
||||
url = f"{base_url}/models/{model}:generateContent"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {credential}",
|
||||
}
|
||||
else:
|
||||
url = f"{base_url}/models/{model}:generateContent?key={credential}"
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
contents = [{"parts": [{"text": prompt}]}]
|
||||
payload = {
|
||||
"contents": contents,
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
},
|
||||
}
|
||||
if system:
|
||||
payload["systemInstruction"] = {"parts": [{"text": system}]}
|
||||
|
||||
return url, headers, payload
|
||||
|
||||
def _build_gemini_code_assist_request(
|
||||
self, base_url: str, credential: str, credential_type: str,
|
||||
prompt: str, system: str, max_tokens: int, model: str, temperature: float,
|
||||
) -> Tuple[str, Dict, Dict]:
|
||||
"""Build request for Google Cloud Code Assist API (Gemini CLI endpoint)."""
|
||||
url = f"{base_url}/v1internal:generateContent"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {credential}",
|
||||
"User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
"X-Goog-Api-Client": "gl-node/22.17.0",
|
||||
}
|
||||
|
||||
inner_request: Dict[str, Any] = {
|
||||
"contents": [{"role": "user", "parts": [{"text": prompt}]}],
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
},
|
||||
}
|
||||
if system:
|
||||
inner_request["systemInstruction"] = {"parts": [{"text": system}]}
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"request": inner_request,
|
||||
}
|
||||
# Include project ID if discovered
|
||||
if self._gemini_project_id:
|
||||
payload["project"] = self._gemini_project_id
|
||||
|
||||
return url, headers, payload
|
||||
|
||||
def _build_ollama_request(
|
||||
self, base_url: str, credential: str, credential_type: str,
|
||||
prompt: str, system: str, max_tokens: int, model: str, temperature: float,
|
||||
) -> Tuple[str, Dict, Dict]:
|
||||
url = f"{base_url}/api/generate"
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"num_predict": max_tokens,
|
||||
"temperature": temperature,
|
||||
},
|
||||
}
|
||||
if system:
|
||||
payload["system"] = system
|
||||
|
||||
return url, headers, payload
|
||||
|
||||
# ── Response Parsers ─────────────────────────────────────
|
||||
|
||||
def _parse_anthropic_response(self, data: Dict) -> Tuple[str, int]:
|
||||
text = ""
|
||||
for block in data.get("content", []):
|
||||
if block.get("type") == "text":
|
||||
text += block.get("text", "")
|
||||
usage = data.get("usage", {})
|
||||
tokens = usage.get("input_tokens", 0) + usage.get("output_tokens", 0)
|
||||
return text, tokens
|
||||
|
||||
def _parse_openai_response(self, data: Dict) -> Tuple[str, int]:
|
||||
choices = data.get("choices", [])
|
||||
text = choices[0]["message"]["content"] if choices else ""
|
||||
usage = data.get("usage", {})
|
||||
tokens = usage.get("total_tokens", 0)
|
||||
return text, tokens
|
||||
|
||||
def _parse_gemini_response(self, data: Dict) -> Tuple[str, int]:
|
||||
candidates = data.get("candidates", [])
|
||||
text = ""
|
||||
if candidates:
|
||||
parts = candidates[0].get("content", {}).get("parts", [])
|
||||
text = "".join(p.get("text", "") for p in parts)
|
||||
usage = data.get("usageMetadata", {})
|
||||
tokens = usage.get("totalTokenCount", 0)
|
||||
return text, tokens
|
||||
|
||||
def _parse_gemini_code_assist_response(self, data: Dict) -> Tuple[str, int]:
|
||||
"""Parse Cloud Code Assist API response (wraps standard Gemini in 'response' field)."""
|
||||
response = data.get("response", data)
|
||||
candidates = response.get("candidates", [])
|
||||
text = ""
|
||||
if candidates:
|
||||
parts = candidates[0].get("content", {}).get("parts", [])
|
||||
text = "".join(p.get("text", "") for p in parts)
|
||||
usage = response.get("usageMetadata", {})
|
||||
tokens = usage.get("totalTokenCount", 0)
|
||||
return text, tokens
|
||||
|
||||
def _parse_ollama_response(self, data: Dict) -> Tuple[str, int]:
|
||||
text = data.get("response", "")
|
||||
# Ollama doesn't always return token counts
|
||||
tokens = data.get("eval_count", 0) + data.get("prompt_eval_count", 0)
|
||||
return text, tokens
|
||||
|
||||
# ── Gemini CLI Project Discovery ────────────────────────────
|
||||
|
||||
async def discover_gemini_project(self) -> Optional[str]:
|
||||
"""Discover Gemini CLI project ID via loadCodeAssist API call."""
|
||||
if self._gemini_project_id:
|
||||
return self._gemini_project_id
|
||||
|
||||
provider = self.registry.get_provider("gemini_cli")
|
||||
if not provider or not provider.accounts:
|
||||
return None
|
||||
|
||||
account = next(iter(provider.accounts.values()), None)
|
||||
if not account or not account.is_active:
|
||||
return None
|
||||
|
||||
credential = self.registry.get_credential(account.id)
|
||||
if not credential:
|
||||
return None
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
url = f"{provider.base_url}/v1internal:loadCodeAssist"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {credential}",
|
||||
"User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
}
|
||||
payload = {
|
||||
"metadata": {
|
||||
"ideType": "IDE_UNSPECIFIED",
|
||||
"platform": "PLATFORM_UNSPECIFIED",
|
||||
"pluginType": "GEMINI",
|
||||
}
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
url, headers=headers, json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=15),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.json()
|
||||
project_data = data.get("cloudaicompanionProject", {})
|
||||
project_id = (
|
||||
project_data.get("id", "")
|
||||
if isinstance(project_data, dict)
|
||||
else str(project_data)
|
||||
)
|
||||
if project_id:
|
||||
self._gemini_project_id = project_id
|
||||
logger.info(f"SmartRouter: Gemini CLI project: {project_id}")
|
||||
return project_id
|
||||
except Exception as e:
|
||||
logger.debug(f"SmartRouter: Gemini project discovery failed: {e}")
|
||||
|
||||
return None
|
||||
|
||||
# ── CLI Token Recovery ─────────────────────────────────────
|
||||
|
||||
async def _try_reextract_cli_token(self, provider_id: str, account_id: str) -> bool:
|
||||
"""Try to re-extract/refresh a CLI token when auth fails.
|
||||
|
||||
Strategy:
|
||||
1. Re-extract from CLI credential files (may have been refreshed externally)
|
||||
2. If refresh_token is available, try OAuth refresh
|
||||
"""
|
||||
# Strategy 1: Re-extract from disk
|
||||
try:
|
||||
from .token_extractor import TokenExtractor
|
||||
extractor = TokenExtractor()
|
||||
token = extractor.detect(provider_id)
|
||||
if token and token.token:
|
||||
self.registry.update_credential(account_id, token.token, token.expires_at)
|
||||
if token.refresh_token:
|
||||
self.registry._refresh_tokens[account_id] = token.refresh_token
|
||||
self.registry.reactivate_account(account_id)
|
||||
logger.info(f"SmartRouter: Re-extracted CLI token for {provider_id}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.debug(f"SmartRouter: CLI re-extraction failed for {provider_id}: {e}")
|
||||
|
||||
# Strategy 2: OAuth refresh if we have a refresh_token
|
||||
refresh_token = self.registry.get_refresh_token(account_id)
|
||||
if refresh_token:
|
||||
try:
|
||||
from .token_refresher import TokenRefresher
|
||||
refresher = TokenRefresher(self.registry)
|
||||
success = await refresher._refresh_token(provider_id, account_id, refresh_token)
|
||||
if success:
|
||||
self.registry.reactivate_account(account_id)
|
||||
logger.info(f"SmartRouter: OAuth-refreshed token for {provider_id}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.debug(f"SmartRouter: OAuth refresh failed for {provider_id}: {e}")
|
||||
|
||||
return False
|
||||
|
||||
# ── Testing & Status ─────────────────────────────────────
|
||||
|
||||
async def test_account(
|
||||
self, provider_id: str, account_id: str
|
||||
) -> Tuple[bool, str]:
|
||||
"""Test connectivity for a specific account. Returns (success, message)."""
|
||||
provider = self.registry.get_provider(provider_id)
|
||||
if not provider:
|
||||
return False, f"Unknown provider: {provider_id}"
|
||||
|
||||
account = provider.accounts.get(account_id)
|
||||
if not account:
|
||||
return False, f"Unknown account: {account_id}"
|
||||
|
||||
try:
|
||||
text, tokens = await self._call_provider(
|
||||
provider=provider,
|
||||
account=account,
|
||||
prompt="Say 'OK' and nothing else.",
|
||||
system="",
|
||||
max_tokens=10,
|
||||
model=account.model_override or provider.default_model,
|
||||
temperature=0,
|
||||
)
|
||||
return True, f"Connected. Response: {text[:50]}"
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
def get_status(self) -> Dict:
|
||||
"""Get router statistics."""
|
||||
return {
|
||||
"total_requests": self._total_requests,
|
||||
"total_tokens": self._total_tokens,
|
||||
"failover_count": self._failover_count,
|
||||
"last_provider": self._last_provider,
|
||||
"last_model": self._last_model,
|
||||
"last_account": self._last_account_label,
|
||||
"providers": self.registry.get_all_status(),
|
||||
}
|
||||
|
||||
|
||||
# ── Custom Exceptions ────────────────────────────────────────
|
||||
|
||||
class QuotaExhaustedError(Exception):
|
||||
def __init__(self, message: str, retry_after: int = 300):
|
||||
super().__init__(message)
|
||||
self.retry_after = retry_after
|
||||
|
||||
|
||||
class AuthenticationError(Exception):
|
||||
pass
|
||||
@@ -0,0 +1,290 @@
|
||||
"""
|
||||
NeuroSploit v3 - CLI Token Extractor
|
||||
|
||||
READ-ONLY extraction of OAuth/API tokens from 8 CLI tools.
|
||||
Never modifies any CLI tool's files or state.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtractedToken:
|
||||
provider_id: str
|
||||
credential_type: str # "oauth" | "api_key"
|
||||
token: str
|
||||
refresh_token: Optional[str] = None
|
||||
expires_at: Optional[float] = None
|
||||
label: str = ""
|
||||
|
||||
|
||||
class TokenExtractor:
|
||||
"""Extracts tokens from CLI tools installed on the system."""
|
||||
|
||||
EXTRACTORS = [
|
||||
"claude_code", "codex_cli", "cursor",
|
||||
"copilot", "iflow", "qwen_code", "kiro",
|
||||
]
|
||||
|
||||
def detect(self, provider_id: str) -> Optional[ExtractedToken]:
|
||||
"""Detect a token for a specific provider."""
|
||||
method = getattr(self, f"_extract_{provider_id}", None)
|
||||
if not method:
|
||||
return None
|
||||
try:
|
||||
return method()
|
||||
except Exception as e:
|
||||
logger.debug(f"TokenExtractor: Failed to extract {provider_id}: {e}")
|
||||
return None
|
||||
|
||||
def detect_all(self) -> List[ExtractedToken]:
|
||||
"""Scan all known CLI tools and return found tokens."""
|
||||
tokens = []
|
||||
for pid in self.EXTRACTORS:
|
||||
result = self.detect(pid)
|
||||
if result:
|
||||
tokens.append(result)
|
||||
logger.info(f"TokenExtractor: Found {pid} token")
|
||||
return tokens
|
||||
|
||||
# ── Individual Extractors ────────────────────────────────
|
||||
|
||||
def _extract_claude_code(self) -> Optional[ExtractedToken]:
|
||||
"""Claude Code: macOS Keychain or ~/.claude/.credentials.json"""
|
||||
is_mac = platform.system() == "Darwin"
|
||||
|
||||
if is_mac:
|
||||
token = self._read_macos_keychain("Claude Code-credentials")
|
||||
if token:
|
||||
# Parse the JSON credential blob
|
||||
try:
|
||||
cred = json.loads(token)
|
||||
access_token = cred.get("claudeAiOauth", {}).get("accessToken")
|
||||
refresh = cred.get("claudeAiOauth", {}).get("refreshToken")
|
||||
expires = cred.get("claudeAiOauth", {}).get("expiresAt")
|
||||
if access_token:
|
||||
return ExtractedToken(
|
||||
provider_id="claude_code",
|
||||
credential_type="oauth",
|
||||
token=access_token,
|
||||
refresh_token=refresh,
|
||||
expires_at=expires / 1000.0 if expires else None,
|
||||
label="Claude Code (Keychain)",
|
||||
)
|
||||
except (json.JSONDecodeError, AttributeError):
|
||||
pass
|
||||
|
||||
# Linux fallback: ~/.claude/.credentials.json
|
||||
creds_path = Path.home() / ".claude" / ".credentials.json"
|
||||
if creds_path.exists():
|
||||
try:
|
||||
data = json.loads(creds_path.read_text())
|
||||
access_token = data.get("claudeAiOauth", {}).get("accessToken")
|
||||
refresh = data.get("claudeAiOauth", {}).get("refreshToken")
|
||||
expires = data.get("claudeAiOauth", {}).get("expiresAt")
|
||||
if access_token:
|
||||
return ExtractedToken(
|
||||
provider_id="claude_code",
|
||||
credential_type="oauth",
|
||||
token=access_token,
|
||||
refresh_token=refresh,
|
||||
expires_at=expires / 1000.0 if expires else None,
|
||||
label="Claude Code (credentials file)",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _extract_codex_cli(self) -> Optional[ExtractedToken]:
|
||||
"""OpenAI Codex CLI: ~/.codex/auth.json"""
|
||||
auth_path = Path.home() / ".codex" / "auth.json"
|
||||
if not auth_path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(auth_path.read_text())
|
||||
# Codex uses OAuth2 PKCE
|
||||
access_token = data.get("access_token") or data.get("token")
|
||||
refresh = data.get("refresh_token")
|
||||
expires = data.get("expires_at")
|
||||
if access_token:
|
||||
return ExtractedToken(
|
||||
provider_id="codex_cli",
|
||||
credential_type="oauth",
|
||||
token=access_token,
|
||||
refresh_token=refresh,
|
||||
expires_at=expires,
|
||||
label="Codex CLI",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _extract_cursor(self) -> Optional[ExtractedToken]:
|
||||
"""Cursor: SQLite state.vscdb database"""
|
||||
is_mac = platform.system() == "Darwin"
|
||||
if is_mac:
|
||||
base = Path.home() / "Library" / "Application Support" / "Cursor" / "User" / "globalStorage"
|
||||
else:
|
||||
base = Path.home() / ".config" / "Cursor" / "User" / "globalStorage"
|
||||
|
||||
db_path = base / "state.vscdb"
|
||||
if not db_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
import sqlite3
|
||||
conn = sqlite3.connect(str(db_path), timeout=1)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT value FROM ItemTable WHERE key = ?",
|
||||
("cursorAuth/accessToken",)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
conn.close()
|
||||
if row and row[0]:
|
||||
return ExtractedToken(
|
||||
provider_id="cursor",
|
||||
credential_type="oauth",
|
||||
token=row[0],
|
||||
label="Cursor IDE",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _extract_copilot(self) -> Optional[ExtractedToken]:
|
||||
"""GitHub Copilot: ~/.config/gh/hosts.yml"""
|
||||
hosts_path = Path.home() / ".config" / "gh" / "hosts.yml"
|
||||
if not hosts_path.exists():
|
||||
return None
|
||||
try:
|
||||
content = hosts_path.read_text()
|
||||
# Simple YAML parsing for oauth_token
|
||||
for line in content.split("\n"):
|
||||
line = line.strip()
|
||||
if line.startswith("oauth_token:"):
|
||||
token = line.split(":", 1)[1].strip()
|
||||
if token:
|
||||
return ExtractedToken(
|
||||
provider_id="copilot",
|
||||
credential_type="oauth",
|
||||
token=token,
|
||||
label="GitHub Copilot (gh CLI)",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _extract_iflow(self) -> Optional[ExtractedToken]:
|
||||
"""iFlow AI: ~/.iflow/settings.json"""
|
||||
settings_path = Path.home() / ".iflow" / "settings.json"
|
||||
if not settings_path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(settings_path.read_text())
|
||||
api_key = data.get("apiKey") or data.get("api_key")
|
||||
if api_key:
|
||||
return ExtractedToken(
|
||||
provider_id="iflow",
|
||||
credential_type="api_key",
|
||||
token=api_key,
|
||||
label="iFlow AI",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _extract_qwen_code(self) -> Optional[ExtractedToken]:
|
||||
"""Qwen Code: ~/.qwen/oauth_creds.json"""
|
||||
creds_path = Path.home() / ".qwen" / "oauth_creds.json"
|
||||
if not creds_path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(creds_path.read_text())
|
||||
access_token = data.get("access_token")
|
||||
refresh = data.get("refresh_token")
|
||||
expires = data.get("expires_at")
|
||||
if access_token:
|
||||
return ExtractedToken(
|
||||
provider_id="qwen_code",
|
||||
credential_type="oauth",
|
||||
token=access_token,
|
||||
refresh_token=refresh,
|
||||
expires_at=expires,
|
||||
label="Qwen Code",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _extract_kiro(self) -> Optional[ExtractedToken]:
|
||||
"""Kiro AI: SQLite data.sqlite3"""
|
||||
is_mac = platform.system() == "Darwin"
|
||||
if is_mac:
|
||||
db_path = Path.home() / "Library" / "Application Support" / "kiro-cli" / "data.sqlite3"
|
||||
else:
|
||||
db_path = Path.home() / ".local" / "share" / "kiro-cli" / "data.sqlite3"
|
||||
|
||||
if not db_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
import sqlite3
|
||||
conn = sqlite3.connect(str(db_path), timeout=1)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT value FROM auth_kv WHERE key = ?",
|
||||
("access_token",)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
refresh_row = None
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT value FROM auth_kv WHERE key = ?",
|
||||
("refresh_token",)
|
||||
)
|
||||
refresh_row = cursor.fetchone()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
conn.close()
|
||||
if row and row[0]:
|
||||
return ExtractedToken(
|
||||
provider_id="kiro",
|
||||
credential_type="oauth",
|
||||
token=row[0],
|
||||
refresh_token=refresh_row[0] if refresh_row else None,
|
||||
label="Kiro AI",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _read_macos_keychain(service: str) -> Optional[str]:
|
||||
"""Read a credential from macOS Keychain."""
|
||||
if platform.system() != "Darwin":
|
||||
return None
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["security", "find-generic-password", "-s", service, "-w"],
|
||||
capture_output=True, text=True, timeout=5,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
return result.stdout.strip()
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
NeuroSploit v3 - Token Refresher
|
||||
|
||||
Background asyncio task that refreshes OAuth tokens before they expire.
|
||||
Checks every 5 minutes, refreshes tokens expiring within 10 minutes.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .provider_registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REFRESH_INTERVAL = 300 # 5 minutes
|
||||
REFRESH_THRESHOLD = 600 # Refresh if expiring within 10 minutes
|
||||
|
||||
# Known OAuth endpoints for token refresh
|
||||
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
||||
ANTHROPIC_TOKEN_URL = "https://auth.anthropic.com/oauth/token"
|
||||
OPENAI_TOKEN_URL = "https://auth0.openai.com/oauth/token"
|
||||
|
||||
|
||||
class TokenRefresher:
|
||||
"""Background task that auto-refreshes expiring OAuth tokens."""
|
||||
|
||||
def __init__(self, registry: "ProviderRegistry"):
|
||||
self._registry = registry
|
||||
self._running = False
|
||||
self._task: Optional[asyncio.Task] = None
|
||||
|
||||
async def start(self):
|
||||
"""Start the background refresh loop."""
|
||||
if self._running:
|
||||
return
|
||||
self._running = True
|
||||
self._task = asyncio.create_task(self._refresh_loop())
|
||||
logger.info("TokenRefresher: Started background refresh loop")
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the background refresh loop."""
|
||||
self._running = False
|
||||
if self._task and not self._task.done():
|
||||
self._task.cancel()
|
||||
try:
|
||||
await self._task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info("TokenRefresher: Stopped")
|
||||
|
||||
async def _refresh_loop(self):
|
||||
"""Main loop: check and refresh tokens periodically."""
|
||||
while self._running:
|
||||
try:
|
||||
await self._check_and_refresh()
|
||||
except Exception as e:
|
||||
logger.error(f"TokenRefresher: Error in refresh loop: {e}")
|
||||
await asyncio.sleep(REFRESH_INTERVAL)
|
||||
|
||||
async def _check_and_refresh(self):
|
||||
"""Check all accounts and refresh tokens nearing expiry."""
|
||||
now = time.time()
|
||||
for provider in self._registry.get_all_providers():
|
||||
for acct in provider.accounts.values():
|
||||
if not acct.is_active or acct.credential_type != "oauth":
|
||||
continue
|
||||
if acct.expires_at is None:
|
||||
continue
|
||||
if acct.expires_at - now > REFRESH_THRESHOLD:
|
||||
continue # Not expiring soon
|
||||
|
||||
refresh_token = self._registry.get_refresh_token(acct.id)
|
||||
if not refresh_token:
|
||||
logger.debug(
|
||||
f"TokenRefresher: {acct.label} expiring but no refresh token"
|
||||
)
|
||||
continue
|
||||
|
||||
logger.info(
|
||||
f"TokenRefresher: Refreshing {acct.label} "
|
||||
f"(expires in {int(acct.expires_at - now)}s)"
|
||||
)
|
||||
|
||||
success = await self._refresh_token(
|
||||
provider_id=provider.id,
|
||||
account_id=acct.id,
|
||||
refresh_token=refresh_token,
|
||||
)
|
||||
if not success:
|
||||
logger.warning(
|
||||
f"TokenRefresher: Failed to refresh {acct.label}"
|
||||
)
|
||||
|
||||
async def _refresh_token(
|
||||
self, provider_id: str, account_id: str, refresh_token: str
|
||||
) -> bool:
|
||||
"""Attempt to refresh a token based on provider type."""
|
||||
try:
|
||||
if provider_id in ("claude_code", "kiro"):
|
||||
return await self._refresh_anthropic(account_id, refresh_token)
|
||||
elif provider_id == "codex_cli":
|
||||
return await self._refresh_openai(account_id, refresh_token)
|
||||
else:
|
||||
logger.debug(
|
||||
f"TokenRefresher: No refresh method for {provider_id}"
|
||||
)
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"TokenRefresher: Refresh failed for {provider_id}: {e}")
|
||||
return False
|
||||
|
||||
async def _refresh_anthropic(self, account_id: str, refresh_token: str) -> bool:
|
||||
"""Refresh an Anthropic/Claude OAuth token."""
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
payload = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
}
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(ANTHROPIC_TOKEN_URL, json=payload, timeout=aiohttp.ClientTimeout(total=10)) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.json()
|
||||
new_token = data.get("access_token")
|
||||
expires_in = data.get("expires_in", 3600)
|
||||
new_refresh = data.get("refresh_token")
|
||||
if new_token:
|
||||
self._registry.update_credential(
|
||||
account_id,
|
||||
new_token,
|
||||
time.time() + expires_in,
|
||||
)
|
||||
# Update refresh token if rotated
|
||||
if new_refresh:
|
||||
self._registry._refresh_tokens[account_id] = new_refresh
|
||||
logger.info(f"TokenRefresher: Anthropic token refreshed")
|
||||
return True
|
||||
else:
|
||||
body = await resp.text()
|
||||
logger.warning(f"TokenRefresher: Anthropic refresh failed: {resp.status} {body[:200]}")
|
||||
except ImportError:
|
||||
logger.warning("TokenRefresher: aiohttp not available for token refresh")
|
||||
except Exception as e:
|
||||
logger.error(f"TokenRefresher: Anthropic refresh error: {e}")
|
||||
return False
|
||||
|
||||
async def _refresh_openai(self, account_id: str, refresh_token: str) -> bool:
|
||||
"""Refresh an OpenAI/Codex OAuth token."""
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
payload = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
}
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(OPENAI_TOKEN_URL, json=payload, timeout=aiohttp.ClientTimeout(total=10)) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.json()
|
||||
new_token = data.get("access_token")
|
||||
expires_in = data.get("expires_in", 3600)
|
||||
if new_token:
|
||||
self._registry.update_credential(
|
||||
account_id,
|
||||
new_token,
|
||||
time.time() + expires_in,
|
||||
)
|
||||
logger.info(f"TokenRefresher: OpenAI token refreshed")
|
||||
return True
|
||||
else:
|
||||
body = await resp.text()
|
||||
logger.warning(f"TokenRefresher: OpenAI refresh failed: {resp.status} {body[:200]}")
|
||||
except ImportError:
|
||||
logger.warning("TokenRefresher: aiohttp not available for token refresh")
|
||||
except Exception as e:
|
||||
logger.error(f"TokenRefresher: OpenAI refresh error: {e}")
|
||||
return False
|
||||
Reference in New Issue
Block a user