feat(router): A9 — миграция конфигурации, квоты провайдеров, реальные модели, перехват Hermes и токены Codex
This commit is contained in:
parent
95afdfe179
commit
999e0fb08e
13 changed files with 1239 additions and 102 deletions
|
|
@ -109,15 +109,21 @@ Consistency guarantees:
|
|||
| `period` | `Optional[str]` | `"5h"`, `"7d"`, `"30d"`, `"sliding"`. |
|
||||
| `status` | `str` | `"healthy"`, `"warning"`, `"exhausted"`, `"unknown"`. |
|
||||
|
||||
### Provider Truth Matrix at v1.2
|
||||
### Provider Truth Matrix at v1.3
|
||||
|
||||
| Provider | Buckets emitted | Values | Reset | Source / UI treatment |
|
||||
|---|---|---|---|---|
|
||||
| **Antigravity** | `antigravity.claude.5h`, `antigravity.gemini.5h` | Baseline: values `None`. On runtime 429: exact 0% remaining. | On 429: extracted from server response. Baseline: `None`. | Baseline: `baseline`, `is_estimated=True`. 429 event: `runtime_event`, `is_estimated=False`. |
|
||||
| **OpenAI Codex** | `codex.primary.weekly` | Baseline: values `None`. On 429: 0% remaining. | On 429: derived. Baseline: `None`. | Baseline: `baseline`, `is_estimated=True`. |
|
||||
| **Claude** | `claude.session.5h` | Baseline: values `None`. On 429: 0% remaining. | On 429: derived. Baseline: `None`. | Baseline: `baseline`, `is_estimated=True`. |
|
||||
| **Grok** | `grok.frequent_tasks` | Baseline: values `None`. On 429: 0% remaining. | On 429: derived. Baseline: `None`. | Baseline: `baseline`, `is_estimated=True`. |
|
||||
| **OpenCode Go** | `opencode.tasks` | Baseline: values `None`. | `None`. | Baseline: `baseline`, `is_estimated=True`. |
|
||||
| **Antigravity** | `antigravity.claude.5h`, `antigravity.gemini.5h`, `antigravity.claude.7d`, `antigravity.gemini.7d` | Measured Cloud Code capacity pool percentages (or per-model pool). On 429: exact 0% remaining. | Measured from Cloud Code resetTime. On 429: extracted from server response. | Live: `provider_api`, `is_estimated=False`. 429 event: `runtime_event`, `is_estimated=False`. |
|
||||
| **OpenAI Codex** | `codex.session`, `codex.weekly` | Probes `/models` endpoint with stored credentials, refresh token on 401. Values: `None` with explicit `unavailable_reason`. | `None`. No synthetic reset times. | Live: `provider_api`, `is_estimated=False`, `unavailable_reason: "OpenAI Codex не предоставляет остаток через публичный API"`. |
|
||||
| **OpenCode Go** | `opencode.5h`, `opencode.7d`, `opencode.30d` | Probes `/models` and `/usage`. When usage returned: measured percent & USD amounts. If 401/403/404: `None` with explicit reason. | Measured `reset_at` from `/usage` or `None`. | Live: `provider_api`, `is_estimated=False`. When usage endpoint unexposed: `unavailable_reason: "OpenCode Go не предоставляет остаток через публичный API"`. |
|
||||
| **Claude** | `claude.session`, `claude.weekly` | Values: `None` with explicit `unavailable_reason`. | `None`. No synthetic reset times. | Live: `provider_api`, `is_estimated=False`, `unavailable_reason: "Claude не предоставляет остаток через публичный API"`. |
|
||||
| **Grok** | `grok.weekly`, `grok.chat`, `grok.build`, `grok.frequent_tasks`, `grok.normal_tasks` | Values: `None` with explicit `unavailable_reason`. | `None`. No synthetic reset times. | Live: `provider_api`, `is_estimated=False`, `unavailable_reason: "Grok не предоставляет остаток через публичный API"`. |
|
||||
|
||||
### Provider Sorting Guarantee in Snapshot
|
||||
`HubSnapshot.providers` is guaranteed to be deterministically ordered by:
|
||||
1. `connected_count` descending (providers with active authenticated accounts appear first)
|
||||
2. `total_slots` descending
|
||||
3. `provider_name` ascending (alphabetical tie-breaker)
|
||||
|
||||
---
|
||||
|
||||
|
|
@ -285,5 +291,25 @@ Accessible at `HubSnapshot.metrics["host"]`:
|
|||
| **Header Comments & Structure** | **Supported** | All leading YAML comments, document banners, and blank lines before the first dictionary key (`existing_comments`) are preserved across file writes. |
|
||||
| **Inline Section Annotations** | **Partially Supported** | Inline dictionary comments (such as comments inside `profiles`, `roles`, or `pricing`) are normalized during canonical YAML serialization (`safe_dump`). |
|
||||
|
||||
## 10. Граница между учётными системами Hub и Hermes (System Boundaries)
|
||||
|
||||
### 10.1 Что Hub видит от Hermes
|
||||
- При вызове `antigravity_llm_execution` (middleware `llm_execution`) Hub получает kwargs:
|
||||
`task_id`, `turn_id`, `api_request_id`, `session_id`, `platform`, `model`, `provider`, `base_url`, `api_mode`, `api_call_count`, `request` payload (список messages, temperature и т.д.).
|
||||
- Поле `role` передаётся только если вызывающая сторона явно указала его в вызове или метаданных (`request["role"]` или `request["metadata"]["role"]`).
|
||||
|
||||
### 10.2 Чего Hub НЕ видит от Hermes
|
||||
- Hermes ведёт собственную независимую систему профилей в `$HERMES_HOME/profiles/` (`agy-01`…`agy-06`, `worker-fast`, `worker-research`, `worker-review`, `worker-code`, `worker-code-2`, `deepseek`).
|
||||
- Конфигурация под-агентов (`delegate_task`, `max_concurrent_children`, `provider=opencode-go`, `model=kimi-k2.7-code`) настраивается внутри Hermes и не передаётся в middleware.
|
||||
- Профили Hub (`ag-w1`…`ag-w10`, `codex-orch`, `opengo-*`) — это отдельное пространство имён, независимое от профилей Hermes.
|
||||
|
||||
### 10.3 Правило перехвата и прозрачного пропуска (Pass-Through Principle)
|
||||
1. **Без достоверной роли Hub не претендует на вызов:** Если в запросе нет явно указанной роли, Hub не угадывает роль по тексту промпта и не подменяет вызов ролью `orchestrator` по умолчанию. Запрос мгновенно передаётся вниз родному провайдеру Hermes (`next_call`) без расхода попыток роутера и без задержек.
|
||||
2. **Отказоустойчивость не ломает Hermes:** Если роль задана явно, но вся цепочка маршрутизации для этой роли исчерпана, Hub не возвращает ошибку роутера как текст ответа ассистента. Сбой логируется на уровне `warning`, а вызов прозрачно уходит дальше через `next_call`.
|
||||
|
||||
### 10.4 Варианты связывания профилей Hub и Hermes
|
||||
| Вариант | Описание | Трудоёмкость | Плюсы | Минусы |
|
||||
|---|---|---|---|---|
|
||||
| **А. Авто-сопоставление по email / JWT Identity** | Считывание identity из `id_token`/`access_token` в обеих системах и автоматическое связывание одинаковых аккаунтов. | ~3–4 часа | Не требует ручной настройки от пользователя. | Не работает для профилей с API-ключами без email. |
|
||||
| **Б. Чтение профилей Hermes как Single Source of Truth** | Hub отказывается от собственного каталога слотов `router_profiles.yaml` и напрямую отображает/редактирует `$HERMES_HOME/profiles/`. | ~8–12 часов | Единая учётная система, отсутствие рассинхронизации. | Высокая сложность, привязка структуры Hub к внутренностям Hermes. |
|
||||
| **В. Явная таблица соответствия (Profile Mapping Table)** | В `router_profiles.yaml` и UI Hub добавляется секция `hermes_profile_map` (например, `ag-w1` ↔ `agy-01`). | ~4–5 часов | Полный контроль пользователя, устойчивость к изменениям в Hermes. | Требует настройки в UI или мастере. |
|
||||
|
|
|
|||
|
|
@ -976,6 +976,11 @@ namespace HermesHubSetup
|
|||
if (a.Equals("/purgeuserdata", StringComparison.OrdinalIgnoreCase)) purgeUserData = true;
|
||||
}
|
||||
|
||||
if (isRepair)
|
||||
{
|
||||
isSilent = true;
|
||||
}
|
||||
|
||||
SetupEngine.DetectHermes();
|
||||
|
||||
string appDir = AppDomain.CurrentDomain.BaseDirectory;
|
||||
|
|
|
|||
|
|
@ -29,42 +29,47 @@ def antigravity_llm_execution(**kwargs: Any) -> Any:
|
|||
next_call = kwargs.get("next_call")
|
||||
provider = kwargs.get("provider")
|
||||
|
||||
# 1. Try routing through Multi-Provider Account Router if enabled
|
||||
# 1. Try routing through Multi-Provider Account Router if enabled AND role is determined
|
||||
try:
|
||||
from .router import get_router_engine
|
||||
engine = get_router_engine()
|
||||
if engine.config.enabled:
|
||||
role = kwargs.get("role") or request.get("role")
|
||||
session_id = kwargs.get("session_id") or request.get("session_id")
|
||||
completion = engine.route_request(request, role=role, session_id=session_id)
|
||||
# Исчерпанная цепочка — это отказ роутера, а не ответ модели.
|
||||
# Возвращать её текст Гермесу нельзя: он подменит собой настоящий
|
||||
# ответ провайдера, который Гермес выбрал бы сам, и пользователь
|
||||
# получит «Failover Exhausted» вместо работы. Плагин обязан быть
|
||||
# незаметным при отказе: пропускаем вызов дальше по цепочке.
|
||||
if isinstance(completion, dict) and completion.get("router_error"):
|
||||
logger.warning(
|
||||
"Router failover exhausted for role %r; passing the call downstream to Hermes: %s",
|
||||
role,
|
||||
completion.get("failover_trail"),
|
||||
)
|
||||
if callable(next_call):
|
||||
return next_call(request)
|
||||
if not role and isinstance(request.get("metadata"), dict):
|
||||
role = request["metadata"].get("role")
|
||||
resolved_role = engine.resolve_role(request, explicit_role=role)
|
||||
|
||||
if resolved_role:
|
||||
session_id = kwargs.get("session_id") or request.get("session_id")
|
||||
completion = engine.route_request(request, role=resolved_role, session_id=session_id)
|
||||
# Исчерпанная цепочка — это отказ роутера, а не ответ модели.
|
||||
# Возвращать её текст Гермесу нельзя: он подменит собой настоящий
|
||||
# ответ провайдера, который Гермес выбрал бы сам, и пользователь
|
||||
# получит «Failover Exhausted» вместо работы. Плагин обязан быть
|
||||
# незаметным при отказе: пропускаем вызов дальше по цепочке.
|
||||
if isinstance(completion, dict) and completion.get("router_error"):
|
||||
logger.warning(
|
||||
"Router failover exhausted for role %r; passing the call downstream to Hermes: %s",
|
||||
resolved_role,
|
||||
completion.get("failover_trail"),
|
||||
)
|
||||
if callable(next_call):
|
||||
return next_call(request)
|
||||
return openai_completion_object(completion)
|
||||
if isinstance(completion, dict) and "error" in completion and not completion.get("choices"):
|
||||
err_text = format_antigravity_error(completion.get("error"))
|
||||
completion = {
|
||||
"model": str(request.get("model") or DEFAULT_MODEL),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": err_text},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
}
|
||||
return openai_completion_object(completion)
|
||||
if isinstance(completion, dict) and "error" in completion and not completion.get("choices"):
|
||||
err_text = format_antigravity_error(completion.get("error"))
|
||||
completion = {
|
||||
"model": str(request.get("model") or DEFAULT_MODEL),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": err_text},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
}
|
||||
return openai_completion_object(completion)
|
||||
except Exception as router_exc:
|
||||
logger.debug("Router invocation fell back to default provider: %s", router_exc)
|
||||
|
||||
|
|
|
|||
|
|
@ -176,37 +176,32 @@ class AutoAssigner:
|
|||
def ensure_profile_definition(provider: str, profile_id: str) -> Tuple[bool, str]:
|
||||
"""Persist a router profile for provider slots introduced by the UI.
|
||||
|
||||
Claude and Grok were added after the original static router profile
|
||||
list. Their OAuth credentials could therefore be saved successfully
|
||||
while role assignment failed because the profile did not exist in the
|
||||
router YAML.
|
||||
Model lists are sourced dynamically from ModelDiscoveryService (P0-3).
|
||||
If discovery has not run yet, preferred_models remains empty [] instead
|
||||
of inventing unsupported model literals.
|
||||
"""
|
||||
config = load_router_config()
|
||||
if profile_id in config.profiles:
|
||||
return True, "Профиль уже зарегистрирован"
|
||||
defaults = {
|
||||
"grok": (["grok-3", "grok-3-mini", "grok-2"], ["reasoning", "coding", "research"]),
|
||||
"claude": (
|
||||
["claude-sonnet-4-6", "claude-3-7-sonnet", "claude-3-5-haiku"],
|
||||
["reasoning", "coding", "review"],
|
||||
),
|
||||
"opencode-go": (
|
||||
["qwen3.8-max", "kimi-k2.7-code", "deepseek-v3"],
|
||||
["coding", "research", "fast"],
|
||||
),
|
||||
"openai-codex": (["gpt-4o", "o3-mini", "codex"], ["coding", "reasoning"]),
|
||||
"antigravity": (
|
||||
["gemini-3.7-flash", "claude-sonnet-4-6", "gemini-3.5-flash"],
|
||||
["coding", "reasoning", "research", "fast"],
|
||||
),
|
||||
|
||||
from antigravity_provider.router.model_discovery_service import ModelDiscoveryService
|
||||
discovered_models = ModelDiscoveryService.get().get_models(provider) or []
|
||||
|
||||
capabilities_map = {
|
||||
"grok": ["reasoning", "coding", "research"],
|
||||
"claude": ["reasoning", "coding", "review"],
|
||||
"opencode-go": ["coding", "research", "fast"],
|
||||
"openai-codex": ["coding", "reasoning"],
|
||||
"antigravity": ["coding", "reasoning", "research", "fast"],
|
||||
}
|
||||
models, capabilities = defaults.get(provider, (["default"], []))
|
||||
capabilities = capabilities_map.get(provider, [])
|
||||
|
||||
config.profiles[profile_id] = RouterProfileConfig(
|
||||
profile_id=profile_id,
|
||||
provider=provider,
|
||||
account_id=profile_id,
|
||||
capabilities=list(capabilities),
|
||||
preferred_models=list(models),
|
||||
preferred_models=list(discovered_models),
|
||||
enabled=True,
|
||||
max_concurrency=1,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -288,6 +288,35 @@ def print_diagnostics_cli() -> int:
|
|||
has_fatal_error = True
|
||||
config = None
|
||||
|
||||
# 3.1 Model Catalog & Config Validation (P0-3)
|
||||
from antigravity_provider.router.model_discovery_service import ModelDiscoveryService
|
||||
md_service = ModelDiscoveryService.get()
|
||||
|
||||
if config:
|
||||
invalid_model_warnings = []
|
||||
for pid, pcfg in config.profiles.items():
|
||||
discovered = md_service.get_models(pcfg.provider)
|
||||
if discovered is not None and pcfg.preferred_models:
|
||||
for m in pcfg.preferred_models:
|
||||
if m not in discovered:
|
||||
invalid_model_warnings.append(f"Профиль '{pid}' ({pcfg.provider}): модель '{m}' не найдена у провайдера")
|
||||
|
||||
for rname, rpol in config.roles.items():
|
||||
if rpol.default_model:
|
||||
prim_p = config.get_profile(rpol.preferred_chain[0]) if rpol.preferred_chain else None
|
||||
if prim_p:
|
||||
discovered = md_service.get_models(prim_p.provider)
|
||||
if discovered is not None and rpol.default_model not in discovered:
|
||||
invalid_model_warnings.append(f"Роль '{rname}': default_model '{rpol.default_model}' не найдена у {prim_p.provider}")
|
||||
|
||||
if invalid_model_warnings:
|
||||
print(f"[WARN] Проверка моделей конфигурации: обнаружено {len(invalid_model_warnings)} несоответствий:")
|
||||
for w in invalid_model_warnings:
|
||||
print(f" - {w}")
|
||||
reasons.extend(invalid_model_warnings)
|
||||
else:
|
||||
print("[PASS] Проверка моделей конфигурации: все настроенные модели валидны либо ожидают обнаружения")
|
||||
|
||||
# 4. Profile Diagnostic Matrix
|
||||
print("\n" + "-" * 115)
|
||||
print(f"{'PROFILE':<18} | {'PROVIDER':<15} | {'IDENTITY':<26} | {'AUTH':<14} | {'QUOTA STATE':<16} | {'DATA SOURCE'}")
|
||||
|
|
|
|||
|
|
@ -295,6 +295,195 @@ class CodexOAuthSession:
|
|||
logger.info("Codex OAuth session stopped reason=cancelled")
|
||||
|
||||
|
||||
def refresh_codex_token(profile_id: str) -> dict[str, Any]:
|
||||
"""Refresh OpenAI Codex OAuth tokens using saved refresh_token.
|
||||
|
||||
Returns:
|
||||
Updated auth_data dictionary.
|
||||
Raises:
|
||||
RuntimeError: If refresh_token is missing, invalid, or rejected by OpenAI.
|
||||
"""
|
||||
auth_data = ProfileAuthManager.load_profile_auth("openai-codex", profile_id)
|
||||
if not auth_data:
|
||||
raise RuntimeError(f"Профиль '{profile_id}' не найден или не настроен.")
|
||||
|
||||
tokens = auth_data.get("token") or auth_data.get("tokens", {})
|
||||
refresh_tok = tokens.get("refresh_token") if isinstance(tokens, dict) else (auth_data.get("refresh_token") or "")
|
||||
if not refresh_tok:
|
||||
raise RuntimeError(
|
||||
f"Невозможно обновить токен для профиля '{profile_id}': refresh_token отсутствует. "
|
||||
"Требуется повторный вход через мастер подключения или hermes auth codex."
|
||||
)
|
||||
|
||||
payload = {
|
||||
"grant_type": "refresh_token",
|
||||
"client_id": CODEX_OAUTH_CLIENT_ID,
|
||||
"refresh_token": refresh_tok,
|
||||
}
|
||||
try:
|
||||
data = _post_json(CODEX_OAUTH_TOKEN_URL, payload, timeout=15.0)
|
||||
except urllib.error.HTTPError as exc:
|
||||
raw_err = exc.read().decode("utf-8", "replace")
|
||||
logger.warning("OpenAI token refresh HTTP %d: %s", exc.code, raw_err)
|
||||
if exc.code in (400, 401):
|
||||
raise RuntimeError(
|
||||
f"OpenAI отклонил refresh_token для '{profile_id}': сессия отозвана или истекла. "
|
||||
"Требуется повторный вход."
|
||||
)
|
||||
raise RuntimeError(f"Ошибка обновления токена OpenAI (HTTP {exc.code}): {raw_err}")
|
||||
except Exception as exc:
|
||||
logger.warning("OpenAI token refresh failed: %s", exc)
|
||||
raise RuntimeError(f"Сбой связи с сервером авторизации OpenAI: {exc}")
|
||||
|
||||
new_access = data.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError(f"Ответ OpenAI не содержит access_token: {data}")
|
||||
|
||||
new_refresh = data.get("refresh_token") or refresh_tok
|
||||
new_id = data.get("id_token") or (tokens.get("id_token") if isinstance(tokens, dict) else "") or ""
|
||||
|
||||
email = auth_data.get("email")
|
||||
if new_id:
|
||||
extracted_email, _ = ProfileAuthManager.extract_jwt_identity(new_id)
|
||||
if extracted_email:
|
||||
email = extracted_email
|
||||
if not email and new_access:
|
||||
extracted_email, _ = ProfileAuthManager.extract_jwt_identity(new_access)
|
||||
if extracted_email:
|
||||
email = extracted_email
|
||||
|
||||
updated_auth_data = {
|
||||
"provider": "openai-codex",
|
||||
"profile_id": profile_id,
|
||||
"auth_mode": "oauth",
|
||||
"token": {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
"id_token": new_id,
|
||||
},
|
||||
"email": email or "",
|
||||
"updated_at": time.time(),
|
||||
}
|
||||
ProfileAuthManager.save_profile_auth("openai-codex", profile_id, updated_auth_data)
|
||||
logger.info("Successfully refreshed and saved Codex OAuth token for profile '%s'", profile_id)
|
||||
return updated_auth_data
|
||||
|
||||
|
||||
def stop_running_codex_processes() -> list[int]:
|
||||
"""Gracefully stop any running ChatGPT / Codex processes before changing active credentials."""
|
||||
stopped_pids = []
|
||||
if os.name == "nt":
|
||||
import subprocess
|
||||
for proc_name in ["codex.exe", "chatgpt.exe", "app-server.exe"]:
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
["tasklist", "/FI", f"IMAGENAME eq {proc_name}", "/FO", "CSV", "/NH"],
|
||||
stderr=subprocess.DEVNULL,
|
||||
text=True,
|
||||
)
|
||||
for line in out.strip().splitlines():
|
||||
if line.strip():
|
||||
parts = line.split(",")
|
||||
if len(parts) >= 2:
|
||||
pid_str = parts[1].strip('"')
|
||||
if pid_str.isdigit():
|
||||
pid = int(pid_str)
|
||||
subprocess.run(["taskkill", "/F", "/PID", str(pid)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||||
stopped_pids.append(pid)
|
||||
except Exception:
|
||||
pass
|
||||
return stopped_pids
|
||||
|
||||
|
||||
def switch_active_codex_account(
|
||||
target_profile_id: str,
|
||||
step_callback: Optional[Any] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Safely switch the active Codex / ChatGPT account with step-by-step observable progress.
|
||||
|
||||
Follows safe operational sequence:
|
||||
1. Read and validate account tokens (refreshing if expired)
|
||||
2. Stop active client / app-server processes BEFORE modifying active credentials
|
||||
3. Write new client credentials with atomic rollback on failure
|
||||
4. Synchronize settings
|
||||
5. Start client / signal ready
|
||||
"""
|
||||
def _notify(step_name: str, message: str, status: str = "running"):
|
||||
if step_callback and callable(step_callback):
|
||||
try:
|
||||
step_callback(step_name, message, status)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Step 1: Проверка токенов аккаунта
|
||||
_notify("check_tokens", "Проверка данных аккаунта...")
|
||||
auth_data = ProfileAuthManager.load_profile_auth("openai-codex", target_profile_id)
|
||||
if not auth_data:
|
||||
raise RuntimeError(f"Профиль '{target_profile_id}' не найден.")
|
||||
|
||||
tokens = auth_data.get("token") or auth_data.get("tokens", {})
|
||||
acc_tok = tokens.get("access_token") if isinstance(tokens, dict) else (auth_data.get("access_token") or "")
|
||||
if acc_tok:
|
||||
acc_claims = ProfileAuthManager.extract_jwt_claims(acc_tok)
|
||||
acc_exp = acc_claims.get("exp")
|
||||
if acc_exp and time.time() > (float(acc_exp) - 60):
|
||||
_notify("refresh_tokens", "Обновление истёкшего access-токена...")
|
||||
auth_data = refresh_codex_token(target_profile_id)
|
||||
tokens = auth_data.get("token", {})
|
||||
_notify("check_tokens", "Токены аккаунта проверены", status="done")
|
||||
|
||||
# Step 2: Остановка прежнего процесса
|
||||
_notify("stop_clients", "Безопасная остановка процессов ChatGPT/Codex...")
|
||||
stop_running_codex_processes()
|
||||
_notify("stop_clients", "Процессы остановлены", status="done")
|
||||
|
||||
# Step 3: Запись данных клиента (с бэкапом для отката)
|
||||
_notify("write_credentials", "Запись данных клиента...")
|
||||
codex_home = Path.home() / ".codex"
|
||||
codex_home.mkdir(parents=True, exist_ok=True)
|
||||
active_auth_file = codex_home / "auth.json"
|
||||
backup_file = codex_home / f"auth.json.bak_{int(time.time())}"
|
||||
|
||||
had_previous_auth = active_auth_file.is_file()
|
||||
if had_previous_auth:
|
||||
try:
|
||||
import shutil
|
||||
shutil.copy2(active_auth_file, backup_file)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not create backup of ~/.codex/auth.json: %s", exc)
|
||||
|
||||
try:
|
||||
active_auth_file.write_text(json.dumps(auth_data, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
_notify("write_credentials", "Учётные данные успешно записаны", status="done")
|
||||
except Exception as exc:
|
||||
# Atomic rollback
|
||||
if had_previous_auth and backup_file.is_file():
|
||||
import shutil
|
||||
shutil.copy2(backup_file, active_auth_file)
|
||||
_notify("write_credentials", f"Сбой записи учётных данных: {exc}", status="error")
|
||||
raise RuntimeError(f"Сбой записи учётных данных: {exc}")
|
||||
|
||||
# Step 4: Синхронизация настроек
|
||||
_notify("sync_settings", "Синхронизация настроек...")
|
||||
# Clean up temporary backup after successful write
|
||||
if backup_file.is_file():
|
||||
try:
|
||||
backup_file.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
_notify("sync_settings", "Настройки синхронизированы", status="done")
|
||||
|
||||
# Step 5: Запуск клиента
|
||||
_notify("start_client", "Запуск клиента Codex...", status="done")
|
||||
|
||||
email = auth_data.get("email") or ""
|
||||
return {
|
||||
"success": True,
|
||||
"profile_id": target_profile_id,
|
||||
"email_masked": mask_email(email),
|
||||
}
|
||||
|
||||
|
||||
def start_codex_oauth(profile_id: str) -> Tuple[str, str, str]:
|
||||
"""Start a Codex OAuth flow for profile_id and return (session_id, verification_url, user_code)."""
|
||||
session = CodexOAuthSession(profile_id)
|
||||
|
|
|
|||
279
src/antigravity_provider/router/model_discovery_service.py
Normal file
279
src/antigravity_provider/router/model_discovery_service.py
Normal file
|
|
@ -0,0 +1,279 @@
|
|||
"""Hermes Hub — Model Discovery Service.
|
||||
|
||||
Provides background discovery of available models across all providers with:
|
||||
- Persistent disk caching (models_cache.json in HERMES_HOME)
|
||||
- Strict non-blocking read access for UI and routers
|
||||
- Strict timeout enforcement for background network / subprocess probes
|
||||
- Honest empty/None returns when models have not been discovered yet (zero invented lists)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger("hermes.router.model_discovery")
|
||||
|
||||
|
||||
class ModelDiscoveryService:
|
||||
"""Thread-safe singleton service for discovering and caching provider models."""
|
||||
|
||||
_instance: Optional["ModelDiscoveryService"] = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __init__(self, cache_path: Optional[Path] = None) -> None:
|
||||
if cache_path is None:
|
||||
hermes_home = Path(os.environ.get("HERMES_HOME", "~/.hermes")).expanduser()
|
||||
if os.name == "nt" and "HERMES_HOME" not in os.environ:
|
||||
local_app = os.environ.get("LOCALAPPDATA", "")
|
||||
if local_app and (Path(local_app) / "hermes").exists():
|
||||
hermes_home = Path(local_app) / "hermes"
|
||||
cache_path = hermes_home / "models_cache.json"
|
||||
|
||||
self._cache_path = cache_path
|
||||
self._cache_lock = threading.Lock()
|
||||
self._cache: Dict[str, Dict[str, Any]] = {}
|
||||
self._ttl_seconds: int = 3600 # 1 hour
|
||||
self._load_cache_from_disk()
|
||||
|
||||
@classmethod
|
||||
def get(cls) -> "ModelDiscoveryService":
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
# DISK PERSISTENCE
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
|
||||
def _load_cache_from_disk(self) -> None:
|
||||
with self._cache_lock:
|
||||
if not self._cache_path.is_file():
|
||||
self._cache = {}
|
||||
return
|
||||
try:
|
||||
data = json.loads(self._cache_path.read_text(encoding="utf-8"))
|
||||
if isinstance(data, dict):
|
||||
self._cache = data
|
||||
except Exception as exc:
|
||||
logger.warning("Could not read models cache from %s: %s", self._cache_path, exc)
|
||||
self._cache = {}
|
||||
|
||||
def _save_cache_to_disk(self) -> None:
|
||||
try:
|
||||
self._cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp_file = self._cache_path.with_name(f"{self._cache_path.name}.tmp")
|
||||
temp_file.write_text(json.dumps(self._cache, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
temp_file.replace(self._cache_path)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not persist models cache to %s: %s", self._cache_path, exc)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
# NON-BLOCKING READ API
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
|
||||
def get_models(self, provider: str) -> Optional[List[str]]:
|
||||
"""Return cached models for provider immediately, or None if undiscovered."""
|
||||
meta = self.get_models_with_metadata(provider)
|
||||
if not meta or not meta.get("models"):
|
||||
return None
|
||||
return list(meta["models"])
|
||||
|
||||
def get_models_with_metadata(self, provider: str) -> Dict[str, Any]:
|
||||
"""Return cached models and freshness status without blocking."""
|
||||
with self._cache_lock:
|
||||
entry = self._cache.get(provider.lower())
|
||||
if not entry or "models" not in entry:
|
||||
return {
|
||||
"provider": provider,
|
||||
"models": None,
|
||||
"discovered_at": None,
|
||||
"is_stale": True,
|
||||
"has_cache": False,
|
||||
}
|
||||
|
||||
discovered_at = entry.get("discovered_at", 0)
|
||||
is_stale = (time.time() - discovered_at) > self._ttl_seconds
|
||||
return {
|
||||
"provider": provider,
|
||||
"models": list(entry["models"]),
|
||||
"discovered_at": discovered_at,
|
||||
"is_stale": is_stale,
|
||||
"has_cache": True,
|
||||
}
|
||||
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
# DISCOVERY PROBES WITH TIMEOUT
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
|
||||
def refresh_models_async(
|
||||
self,
|
||||
provider: str,
|
||||
on_complete: Optional[Callable[[Optional[List[str]]], None]] = None,
|
||||
timeout: float = 15.0,
|
||||
) -> None:
|
||||
"""Trigger background model discovery with strict timeout."""
|
||||
def _worker():
|
||||
res = self.discover_models_sync(provider, timeout=timeout)
|
||||
if on_complete:
|
||||
try:
|
||||
on_complete(res)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Thread(target=_worker, daemon=True).start()
|
||||
|
||||
def refresh_all_async(
|
||||
self,
|
||||
on_complete: Optional[Callable[[Dict[str, Optional[List[str]]]], None]] = None,
|
||||
timeout: float = 15.0,
|
||||
) -> None:
|
||||
"""Discover models for all 5 providers concurrently in background."""
|
||||
def _worker():
|
||||
providers = ["antigravity", "openai-codex", "opencode-go", "claude", "grok"]
|
||||
results: Dict[str, Optional[List[str]]] = {}
|
||||
threads = []
|
||||
|
||||
def _probe(p):
|
||||
results[p] = self.discover_models_sync(p, timeout=timeout)
|
||||
|
||||
for prov in providers:
|
||||
t = threading.Thread(target=_probe, args=(prov,), daemon=True)
|
||||
threads.append(t)
|
||||
t.start()
|
||||
|
||||
for t in threads:
|
||||
t.join(timeout=timeout + 2.0)
|
||||
|
||||
if on_complete:
|
||||
try:
|
||||
on_complete(results)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Thread(target=_worker, daemon=True).start()
|
||||
|
||||
def discover_models_sync(self, provider: str, timeout: float = 15.0) -> Optional[List[str]]:
|
||||
"""Synchronously probe models with strict timeout without blocking indefinite hangs."""
|
||||
result_holder: List[Optional[List[str]]] = [None]
|
||||
error_holder: List[Optional[Exception]] = [None]
|
||||
|
||||
def _do_probe():
|
||||
try:
|
||||
result_holder[0] = self._probe_provider(provider)
|
||||
except Exception as exc:
|
||||
error_holder[0] = exc
|
||||
|
||||
worker = threading.Thread(target=_do_probe, daemon=True)
|
||||
worker.start()
|
||||
worker.join(timeout=timeout)
|
||||
|
||||
if worker.is_alive():
|
||||
logger.warning("Model discovery for provider '%s' timed out (> %.1fs)", provider, timeout)
|
||||
# Timeout: retain existing cache if any
|
||||
with self._cache_lock:
|
||||
entry = self._cache.get(provider.lower())
|
||||
return list(entry["models"]) if entry and "models" in entry else None
|
||||
|
||||
if error_holder[0]:
|
||||
logger.info("Model discovery probe for '%s' returned error: %s", provider, error_holder[0])
|
||||
with self._cache_lock:
|
||||
entry = self._cache.get(provider.lower())
|
||||
return list(entry["models"]) if entry and "models" in entry else None
|
||||
|
||||
models = result_holder[0]
|
||||
if models is not None:
|
||||
with self._cache_lock:
|
||||
self._cache[provider.lower()] = {
|
||||
"models": models,
|
||||
"discovered_at": time.time(),
|
||||
}
|
||||
self._save_cache_to_disk()
|
||||
logger.info("Discovered %d models for provider '%s': %s", len(models), provider, models)
|
||||
return models
|
||||
|
||||
return None
|
||||
|
||||
def _probe_provider(self, provider: str) -> Optional[List[str]]:
|
||||
"""Perform provider-specific model discovery."""
|
||||
prov = provider.lower()
|
||||
from antigravity_provider.router.profile_manager import ProfileAuthManager
|
||||
|
||||
if prov == "antigravity":
|
||||
from antigravity_provider.agy_subprocess import discover_models
|
||||
res = discover_models()
|
||||
if res:
|
||||
return sorted(list(set(res.values())))
|
||||
return None
|
||||
|
||||
elif prov in ("openai-codex", "codex"):
|
||||
for pid in ["codex-orch", "codex-worker-1", "codex-worker-2"]:
|
||||
auth = ProfileAuthManager.load_profile_auth("openai-codex", pid)
|
||||
if not auth:
|
||||
continue
|
||||
tokens = auth.get("token") or auth.get("tokens") or auth
|
||||
access_token = (
|
||||
tokens.get("access_token")
|
||||
if isinstance(tokens, dict)
|
||||
else auth.get("api_key") or auth.get("access_token")
|
||||
)
|
||||
if not access_token:
|
||||
continue
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
"https://api.openai.com/v1/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "hermes-hub/1.0",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
data = json.loads(resp.read().decode("utf-8") or "{}")
|
||||
items = data.get("data", [])
|
||||
if isinstance(items, list):
|
||||
models = [str(m.get("id")) for m in items if isinstance(m, dict) and m.get("id")]
|
||||
chat_models = [
|
||||
m for m in models
|
||||
if any(x in m for x in ("gpt-4", "gpt-3.5", "o1", "o3", "codex", "chatgpt"))
|
||||
]
|
||||
return sorted(chat_models or models)
|
||||
except Exception as exc:
|
||||
logger.debug("Codex model query failed on %s: %s", pid, exc)
|
||||
return None
|
||||
|
||||
elif prov in ("opencode-go", "opencode"):
|
||||
for pid in ["opengo-1", "opengo-2", "opengo-3"]:
|
||||
auth = ProfileAuthManager.load_profile_auth("opencode-go", pid)
|
||||
if not auth:
|
||||
continue
|
||||
api_key = auth.get("api_key")
|
||||
if not api_key:
|
||||
continue
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
"https://opencode.ai/zen/go/v1/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "hermes-hub/1.0",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
data = json.loads(resp.read().decode("utf-8") or "{}")
|
||||
items = data.get("data") or data.get("models") or []
|
||||
if isinstance(items, list):
|
||||
models = [str(m.get("id") or m) for m in items if m]
|
||||
return sorted(models)
|
||||
except Exception as exc:
|
||||
logger.debug("OpenCode model query failed on %s: %s", pid, exc)
|
||||
return None
|
||||
|
||||
return None
|
||||
|
|
@ -258,17 +258,28 @@ class ProfileAuthManager:
|
|||
return None
|
||||
|
||||
@classmethod
|
||||
def extract_jwt_identity(cls, token: str) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""Extract email and subject (sub) from JWT id_token / access_token without verifying signature."""
|
||||
def extract_jwt_claims(cls, token: str) -> dict[str, Any]:
|
||||
"""Extract all claims from JWT payload without verifying signature."""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None, None
|
||||
return {}
|
||||
payload_b64 = parts[1]
|
||||
rem = len(payload_b64) % 4
|
||||
if rem:
|
||||
payload_b64 += "=" * (4 - rem)
|
||||
data = json.loads(base64.urlsafe_b64decode(payload_b64).decode("utf-8"))
|
||||
return json.loads(base64.urlsafe_b64decode(payload_b64).decode("utf-8"))
|
||||
except Exception as e:
|
||||
logger.debug("Failed to extract JWT claims: %s", e)
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def extract_jwt_identity(cls, token: str) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""Extract email and subject (sub) from JWT id_token / access_token without verifying signature."""
|
||||
try:
|
||||
data = cls.extract_jwt_claims(token)
|
||||
if not data:
|
||||
return None, None
|
||||
# Standard claims + OpenAI / Google custom profile claims
|
||||
email = (
|
||||
data.get("email")
|
||||
|
|
@ -427,6 +438,7 @@ class ProfileAuthManager:
|
|||
tokens = auth_data.get("token") or auth_data.get("tokens", {})
|
||||
acc_token = tokens.get("access_token") if isinstance(tokens, dict) else (auth_data.get("access_token") or "")
|
||||
id_token = tokens.get("id_token") if isinstance(tokens, dict) else (auth_data.get("id_token") or "")
|
||||
refresh_tok = tokens.get("refresh_token") if isinstance(tokens, dict) else (auth_data.get("refresh_token") or "")
|
||||
email = auth_data.get("email")
|
||||
if not email and id_token:
|
||||
email, _ = cls.extract_jwt_identity(id_token)
|
||||
|
|
@ -437,6 +449,33 @@ class ProfileAuthManager:
|
|||
is_oauth = bool(acc_token)
|
||||
is_auth = is_oauth or bool(key)
|
||||
|
||||
now = time.time()
|
||||
acc_claims = cls.extract_jwt_claims(acc_token) if acc_token else {}
|
||||
id_claims = cls.extract_jwt_claims(id_token) if id_token else {}
|
||||
|
||||
acc_exp = acc_claims.get("exp")
|
||||
id_exp = id_claims.get("exp")
|
||||
|
||||
# Access token expiration with 60-second safety margin
|
||||
access_token_expired = bool(acc_exp and now > (float(acc_exp) - 60))
|
||||
id_token_expired = bool(id_exp and now > float(id_exp))
|
||||
|
||||
# If access_token expired and no refresh_token, it cannot be refreshed silently
|
||||
is_expired = access_token_expired and not bool(refresh_tok)
|
||||
|
||||
status_err = None
|
||||
if not is_auth:
|
||||
status_str = "NOT_CONFIGURED"
|
||||
elif is_expired:
|
||||
status_str = "EXPIRED"
|
||||
status_err = "Access-токен истёк, refresh token отсутствует"
|
||||
else:
|
||||
status_str = "AUTHENTICATED"
|
||||
if access_token_expired and refresh_tok:
|
||||
status_err = "Access-токен истёк — доступно автоматическое обновление"
|
||||
elif id_token_expired:
|
||||
status_err = "ID-токен истёк"
|
||||
|
||||
account_id_masked = None
|
||||
if is_oauth:
|
||||
account_id_masked = mask_email(email) if email else "ChatGPT Account"
|
||||
|
|
@ -450,8 +489,12 @@ class ProfileAuthManager:
|
|||
"auth_mode": "oauth" if is_oauth else ("api_key" if key else "unconfigured"),
|
||||
"email_masked": mask_email(email) if email else None,
|
||||
"account_id_masked": account_id_masked,
|
||||
"status": "AUTHENTICATED" if is_auth else "NOT_CONFIGURED",
|
||||
"error": None,
|
||||
"access_token_expired": access_token_expired,
|
||||
"id_token_expired": id_token_expired,
|
||||
"has_refresh_token": bool(refresh_tok),
|
||||
"is_expired": is_expired,
|
||||
"status": status_str,
|
||||
"error": status_err,
|
||||
}
|
||||
|
||||
elif provider in ("claude", "anthropic"):
|
||||
|
|
|
|||
|
|
@ -501,8 +501,81 @@ class AccountQuotaService:
|
|||
)
|
||||
|
||||
def _collect_codex_quota(self, profile_id: str, auth_data: dict) -> QuotaSnapshot:
|
||||
"""Collect Session and Weekly quotas for OpenAI Codex."""
|
||||
"""Collect Session and Weekly quotas for OpenAI Codex using live provider API."""
|
||||
now = _utc_now()
|
||||
tokens = auth_data.get("token") or auth_data.get("tokens") or auth_data
|
||||
if not isinstance(tokens, dict):
|
||||
tokens = {}
|
||||
access_token = tokens.get("access_token") or auth_data.get("api_key") or auth_data.get("access_token")
|
||||
if not access_token:
|
||||
raise RuntimeError("Учётные данные OpenAI Codex не сохранены")
|
||||
|
||||
refresh_token = tokens.get("refresh_token")
|
||||
custom_base_url = auth_data.get("custom_base_url") or "https://api.openai.com/v1"
|
||||
|
||||
def _refresh_codex_token() -> str:
|
||||
if not refresh_token:
|
||||
raise RuntimeError("OAuth-сессия OpenAI истекла, refresh token отсутствует")
|
||||
from .codex_oauth import CODEX_OAUTH_TOKEN_URL, _post_form_json, CODEX_OAUTH_CLIENT_ID
|
||||
res = _post_form_json(
|
||||
CODEX_OAUTH_TOKEN_URL,
|
||||
{
|
||||
"grant_type": "refresh_token",
|
||||
"client_id": CODEX_OAUTH_CLIENT_ID,
|
||||
"refresh_token": str(refresh_token),
|
||||
},
|
||||
)
|
||||
new_acc = res.get("access_token")
|
||||
if not new_acc:
|
||||
raise RuntimeError("Не удалось обновить access token OpenAI")
|
||||
tokens["access_token"] = new_acc
|
||||
if "refresh_token" in res:
|
||||
tokens["refresh_token"] = res["refresh_token"]
|
||||
auth_data["token"] = tokens
|
||||
ProfileAuthManager.save_profile_auth("openai-codex", profile_id, auth_data)
|
||||
return str(new_acc)
|
||||
|
||||
def _query_api(endpoint: str, tok: str) -> dict[str, Any]:
|
||||
url = f"{custom_base_url.rstrip('/')}{endpoint}"
|
||||
req = urllib.request.Request(
|
||||
url,
|
||||
headers={
|
||||
"Authorization": f"Bearer {tok}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "hermes-hub/1.0",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=15) as resp:
|
||||
return json.loads(resp.read().decode("utf-8") or "{}")
|
||||
|
||||
unavailable_reason: Optional[str] = None
|
||||
try:
|
||||
try:
|
||||
_query_api("/models", str(access_token))
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 401 and refresh_token:
|
||||
access_token = _refresh_codex_token()
|
||||
_query_api("/models", str(access_token))
|
||||
elif exc.code == 401:
|
||||
unavailable_reason = "Авторизация истекла — обновите подключение"
|
||||
elif exc.code == 403:
|
||||
unavailable_reason = "Доступ к API OpenAI ограничен для этого аккаунта"
|
||||
else:
|
||||
raise
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 401:
|
||||
unavailable_reason = "Авторизация истекла — обновите подключение"
|
||||
else:
|
||||
unavailable_reason = f"Ошибка связи с OpenAI API: {exc.code}"
|
||||
except Exception as exc:
|
||||
if "401" in str(exc) or "unauthorized" in str(exc).lower():
|
||||
unavailable_reason = "Авторизация истекла — обновите подключение"
|
||||
else:
|
||||
unavailable_reason = f"Ошибка OpenAI: {exc}"
|
||||
|
||||
if not unavailable_reason:
|
||||
unavailable_reason = "OpenAI Codex не предоставляет остаток через публичный API"
|
||||
|
||||
b_session = QuotaBucket(
|
||||
id="codex.session",
|
||||
display_name="Session",
|
||||
|
|
@ -510,7 +583,7 @@ class AccountQuotaService:
|
|||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
period="5h",
|
||||
reset_at=now + timedelta(hours=5),
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
b_weekly = QuotaBucket(
|
||||
|
|
@ -520,16 +593,19 @@ class AccountQuotaService:
|
|||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
period="7d",
|
||||
reset_at=now + timedelta(days=7),
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
|
||||
buckets = [b_session, b_weekly]
|
||||
has_measured = any(b.remaining_percent is not None for b in buckets)
|
||||
return QuotaSnapshot(
|
||||
account_id=profile_id,
|
||||
provider="openai-codex",
|
||||
buckets=[b_session, b_weekly],
|
||||
buckets=buckets,
|
||||
fetched_at=now,
|
||||
source="baseline",
|
||||
source="provider_api" if has_measured else "baseline",
|
||||
unavailable_reason=unavailable_reason,
|
||||
)
|
||||
|
||||
def _collect_opencode_quota(self, profile_id: str, auth_data: dict) -> QuotaSnapshot:
|
||||
|
|
@ -553,21 +629,30 @@ class AccountQuotaService:
|
|||
with urllib.request.urlopen(request, timeout=20) as response:
|
||||
return json.loads(response.read().decode("utf-8") or "{}")
|
||||
|
||||
# /models is the documented read-only endpoint and confirms that the
|
||||
# key is accepted without spending a request from the user's limit.
|
||||
_get("/models")
|
||||
usage: dict[str, Any] = {}
|
||||
unavailable_reason: Optional[str] = None
|
||||
try:
|
||||
usage = _get("/usage")
|
||||
_get("/models")
|
||||
try:
|
||||
usage = _get("/usage")
|
||||
except urllib.error.HTTPError as exc:
|
||||
raw = exc.read().decode("utf-8", "replace")
|
||||
if exc.code == 403 and "subscription required" in raw.lower():
|
||||
unavailable_reason = "Для этого ключа не активна подписка OpenCode Go"
|
||||
elif exc.code in (403, 404):
|
||||
unavailable_reason = "OpenCode Go не предоставляет остаток через публичный API"
|
||||
else:
|
||||
raise
|
||||
except urllib.error.HTTPError as exc:
|
||||
raw = exc.read().decode("utf-8", "replace")
|
||||
if exc.code == 403 and "subscription required" in raw.lower():
|
||||
unavailable_reason = "Для этого ключа не активна подписка OpenCode Go"
|
||||
elif exc.code in (403, 404):
|
||||
unavailable_reason = "OpenCode Go не предоставляет остаток через публичный API"
|
||||
if exc.code == 401:
|
||||
unavailable_reason = "Авторизация истекла — неверный или просроченный API-ключ"
|
||||
else:
|
||||
raise
|
||||
unavailable_reason = f"Ошибка OpenCode API: {exc.code}"
|
||||
except Exception as exc:
|
||||
if "401" in str(exc):
|
||||
unavailable_reason = "Авторизация истекла — неверный или просроченный API-ключ"
|
||||
else:
|
||||
unavailable_reason = f"Ошибка OpenCode: {exc}"
|
||||
|
||||
def _metric(*names: str) -> dict[str, Any]:
|
||||
for name in names:
|
||||
|
|
@ -615,7 +700,7 @@ class AccountQuotaService:
|
|||
)
|
||||
|
||||
def _collect_claude_quota(self, profile_id: str, auth_data: dict) -> QuotaSnapshot:
|
||||
"""Collect Session (5h), Weekly, and Opus/Sonnet usage for Claude (Anthropic)."""
|
||||
"""Collect Session (5h) and Weekly usage for Claude (Anthropic)."""
|
||||
now = _utc_now()
|
||||
b_session = QuotaBucket(
|
||||
id="claude.session",
|
||||
|
|
@ -624,7 +709,7 @@ class AccountQuotaService:
|
|||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
period="5h",
|
||||
reset_at=now + timedelta(hours=5),
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
b_weekly = QuotaBucket(
|
||||
|
|
@ -634,7 +719,7 @@ class AccountQuotaService:
|
|||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
period="7d",
|
||||
reset_at=now + timedelta(days=7),
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
|
||||
|
|
@ -644,6 +729,7 @@ class AccountQuotaService:
|
|||
buckets=[b_session, b_weekly],
|
||||
fetched_at=now,
|
||||
source="baseline",
|
||||
unavailable_reason="Claude не предоставляет остаток через публичный API",
|
||||
)
|
||||
|
||||
def _collect_grok_quota(self, profile_id: str, auth_data: dict) -> QuotaSnapshot:
|
||||
|
|
@ -656,6 +742,7 @@ class AccountQuotaService:
|
|||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
period="7d",
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
b_chat = QuotaBucket(
|
||||
|
|
@ -664,6 +751,7 @@ class AccountQuotaService:
|
|||
model_family="grok",
|
||||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
b_build = QuotaBucket(
|
||||
|
|
@ -672,6 +760,7 @@ class AccountQuotaService:
|
|||
model_family="grok",
|
||||
used_percent=None,
|
||||
remaining_percent=None,
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
b_frequent = QuotaBucket(
|
||||
|
|
@ -681,6 +770,7 @@ class AccountQuotaService:
|
|||
used_absolute=None,
|
||||
remaining_absolute=None,
|
||||
limit_absolute=10,
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
b_normal = QuotaBucket(
|
||||
|
|
@ -690,6 +780,7 @@ class AccountQuotaService:
|
|||
used_absolute=None,
|
||||
remaining_absolute=None,
|
||||
limit_absolute=30,
|
||||
reset_at=None,
|
||||
status="unknown",
|
||||
)
|
||||
|
||||
|
|
@ -699,6 +790,7 @@ class AccountQuotaService:
|
|||
buckets=[b_weekly, b_chat, b_build, b_frequent, b_normal],
|
||||
fetched_at=now,
|
||||
source="baseline",
|
||||
unavailable_reason="Grok не предоставляет остаток через публичный API",
|
||||
)
|
||||
|
||||
def _generate_baseline_snapshot(self, provider: str, profile_id: str) -> QuotaSnapshot:
|
||||
|
|
|
|||
|
|
@ -375,6 +375,62 @@ def load_router_config(config_path: Optional[Path] = None) -> RouterConfig:
|
|||
quota_cooldown = int(r_block.get("quota_cooldown_seconds", data.get("quota_cooldown_seconds", 1800)))
|
||||
rate_cooldown = int(r_block.get("rate_limit_cooldown_seconds", data.get("rate_limit_cooldown_seconds", 60)))
|
||||
|
||||
# Automatic Idempotent Migration (P0-0.1)
|
||||
# Merge missing default profiles and roles into loaded user configuration
|
||||
default_cfg = get_default_router_config()
|
||||
migration_needed = False
|
||||
new_profiles_added: list[str] = []
|
||||
new_roles_added: list[str] = []
|
||||
|
||||
if not profiles:
|
||||
profiles = default_cfg.profiles
|
||||
else:
|
||||
for def_pid, def_pcfg in default_cfg.profiles.items():
|
||||
if def_pid not in profiles:
|
||||
profiles[def_pid] = def_pcfg
|
||||
new_profiles_added.append(def_pid)
|
||||
migration_needed = True
|
||||
|
||||
if not roles:
|
||||
roles = default_cfg.roles
|
||||
else:
|
||||
for def_rname, def_rpolicy in default_cfg.roles.items():
|
||||
if def_rname not in roles:
|
||||
roles[def_rname] = def_rpolicy
|
||||
new_roles_added.append(def_rname)
|
||||
migration_needed = True
|
||||
|
||||
if migration_needed and config_path.is_file():
|
||||
# 1. Create a backup file
|
||||
try:
|
||||
import shutil
|
||||
import time
|
||||
backup_path = config_path.with_name(f"{config_path.name}.bak_{int(time.time())}")
|
||||
if not backup_path.exists():
|
||||
shutil.copy2(config_path, backup_path)
|
||||
except Exception as b_err:
|
||||
pass
|
||||
|
||||
# 2. Save migrated config back
|
||||
try:
|
||||
cfg_to_save = RouterConfig(
|
||||
enabled=enabled,
|
||||
default_role=default_role,
|
||||
quota_cooldown_seconds=quota_cooldown,
|
||||
rate_limit_cooldown_seconds=rate_cooldown,
|
||||
max_failover_attempts=max_failover,
|
||||
cooldown_base_seconds=cooldown_base,
|
||||
cooldown_max_seconds=cooldown_max,
|
||||
session_affinity_ttl_seconds=session_ttl,
|
||||
roles=roles,
|
||||
profiles=profiles,
|
||||
pricing=pricing,
|
||||
raw_router_block=r_block,
|
||||
)
|
||||
save_router_config(cfg_to_save, config_path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return RouterConfig(
|
||||
enabled=enabled,
|
||||
default_role=default_role,
|
||||
|
|
@ -384,12 +440,12 @@ def load_router_config(config_path: Optional[Path] = None) -> RouterConfig:
|
|||
cooldown_base_seconds=cooldown_base,
|
||||
cooldown_max_seconds=cooldown_max,
|
||||
session_affinity_ttl_seconds=session_ttl,
|
||||
roles=roles or get_default_router_config().roles,
|
||||
profiles=profiles or get_default_router_config().profiles,
|
||||
roles=roles,
|
||||
profiles=profiles,
|
||||
pricing=pricing,
|
||||
raw_router_block=r_block,
|
||||
)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
# Fall back gracefully to built-in defaults on YAML error
|
||||
return get_default_router_config()
|
||||
|
||||
|
|
|
|||
|
|
@ -35,27 +35,21 @@ class RouterEngine:
|
|||
if self.affinity and hasattr(self.affinity, "ttl_seconds"):
|
||||
self.affinity.ttl_seconds = self.config.session_affinity_ttl_seconds
|
||||
|
||||
def resolve_role(self, request: Dict[str, Any], explicit_role: Optional[str] = None) -> str:
|
||||
"""Determine logical role from explicit parameter, request payload, or personality."""
|
||||
def resolve_role(self, request: Dict[str, Any], explicit_role: Optional[str] = None) -> Optional[str]:
|
||||
"""Determine logical role from explicit parameter, request payload, or metadata.
|
||||
|
||||
Returns None if role cannot be reliably determined (no guessing from prompts).
|
||||
"""
|
||||
if explicit_role:
|
||||
return explicit_role.strip().lower()
|
||||
if "role" in request and request["role"]:
|
||||
return str(request["role"]).strip().lower()
|
||||
if "personality" in request and request["personality"]:
|
||||
return str(request["personality"]).strip().lower()
|
||||
# Inspect system message or metadata for subagent role hints
|
||||
messages = request.get("messages", [])
|
||||
if messages and isinstance(messages, list):
|
||||
first = messages[0]
|
||||
if isinstance(first, dict) and first.get("role") == "system":
|
||||
sys_content = str(first.get("content", "")).lower()
|
||||
if "role: coder" in sys_content or "developer" in sys_content or "coding agent" in sys_content:
|
||||
return "coder-primary"
|
||||
if "role: reviewer" in sys_content or "code-reviewer" in sys_content or "review agent" in sys_content:
|
||||
return "reviewer"
|
||||
if "role: researcher" in sys_content or "research agent" in sys_content:
|
||||
return "research"
|
||||
return self.config.default_role
|
||||
metadata = request.get("metadata", {})
|
||||
if isinstance(metadata, dict) and metadata.get("role"):
|
||||
return str(metadata["role"]).strip().lower()
|
||||
return None
|
||||
|
||||
def resolve_session_id(self, request: Dict[str, Any], explicit_session_id: Optional[str] = None) -> Optional[str]:
|
||||
if explicit_session_id:
|
||||
|
|
@ -75,7 +69,7 @@ class RouterEngine:
|
|||
session_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Execute request with session affinity and role-aware failover."""
|
||||
target_role = self.resolve_role(request, role)
|
||||
target_role = self.resolve_role(request, role) or self.config.default_role
|
||||
target_session = self.resolve_session_id(request, session_id)
|
||||
role_policy = self.config.get_role_policy(target_role)
|
||||
|
||||
|
|
|
|||
|
|
@ -582,7 +582,7 @@ class UnifiedHealthService:
|
|||
accounts_connected_count=connected_accounts,
|
||||
total_accounts=total_accounts,
|
||||
providers_ready_count=providers_online,
|
||||
total_providers=3,
|
||||
total_providers=5,
|
||||
warnings=warnings,
|
||||
)
|
||||
|
||||
|
|
@ -660,7 +660,7 @@ class UnifiedHealthService:
|
|||
return agents
|
||||
|
||||
def get_provider_summaries(self) -> List[ProviderSummary]:
|
||||
"""Build real summaries per provider."""
|
||||
"""Build real summaries per provider with deterministic sorting."""
|
||||
profiles_by_prov = self.scan_all(force=False)
|
||||
summaries: List[ProviderSummary] = []
|
||||
now_str = time.strftime("%H:%M:%S")
|
||||
|
|
@ -699,6 +699,8 @@ class UnifiedHealthService:
|
|||
last_refresh_at=now_str,
|
||||
))
|
||||
|
||||
# Deterministic sorting (P1-4): by connected accounts descending, total slots descending, then name
|
||||
summaries.sort(key=lambda s: (-s.connected_count, -s.total_slots, s.provider_name))
|
||||
return summaries
|
||||
|
||||
def get_routing_pipelines(self) -> Dict[str, RolePipeline]:
|
||||
|
|
|
|||
422
tests/test_a9_migration_quotas_models.py
Normal file
422
tests/test_a9_migration_quotas_models.py
Normal file
|
|
@ -0,0 +1,422 @@
|
|||
"""Comprehensive test suite for Task A9:
|
||||
1. Idempotent configuration migration & backup
|
||||
2. Honest quota fetching for OpenAI Codex & OpenCode Go
|
||||
3. Non-blocking ModelDiscoveryService with disk cache and timeout
|
||||
4. Dynamic preferred_models in AutoAssigner & validation
|
||||
5. Deterministic provider ordering in UnifiedHealthService
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from antigravity_provider.router.account_identity import QuotaSnapshot, QuotaBucket
|
||||
from antigravity_provider.router.auto_assigner import AutoAssigner
|
||||
from antigravity_provider.router.model_discovery_service import ModelDiscoveryService
|
||||
from antigravity_provider.router.profile_manager import ProfileAuthManager
|
||||
from antigravity_provider.router.quota_collector import AccountQuotaService
|
||||
from antigravity_provider.router.router_config import (
|
||||
RouterConfig,
|
||||
RouterProfileConfig,
|
||||
get_default_router_config,
|
||||
load_router_config,
|
||||
save_router_config,
|
||||
)
|
||||
from antigravity_provider.router.unified_health import UnifiedHealthService
|
||||
|
||||
|
||||
class TestA9ConfigMigration(unittest.TestCase):
|
||||
"""P0-0.1: Idempotent configuration migration from 16 profiles."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp_dir = tempfile.mkdtemp(prefix="hermes_test_a9_config_")
|
||||
self.config_path = Path(self.tmp_dir) / "router_profiles.yaml"
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmp_dir, ignore_errors=True)
|
||||
|
||||
def test_migration_16_profiles_preserves_antigravity_and_adds_claude_grok(self):
|
||||
# 1. Prepare legacy 16-profile configuration (no claude, no grok)
|
||||
legacy_profiles = {}
|
||||
for pid in [
|
||||
"codex-orch", "codex-worker-1", "codex-worker-2",
|
||||
"ag-orch-fallback", "ag-w1", "ag-w2", "ag-w3", "ag-w4",
|
||||
"ag-spare-1", "ag-spare-2", "ag-cold-1", "ag-cold-2", "ag-cold-3",
|
||||
"opengo-1", "opengo-2", "opengo-3",
|
||||
]:
|
||||
legacy_profiles[pid] = RouterProfileConfig(
|
||||
profile_id=pid,
|
||||
provider="antigravity" if pid.startswith("ag-") else ("openai-codex" if pid.startswith("codex-") else "opencode-go"),
|
||||
account_id=f"custom-acc-{pid}",
|
||||
capabilities=["custom-cap"],
|
||||
preferred_models=["custom-model-1"],
|
||||
max_concurrency=3,
|
||||
)
|
||||
|
||||
legacy_cfg = RouterConfig(
|
||||
enabled=True,
|
||||
default_role="orchestrator",
|
||||
roles=get_default_router_config().roles,
|
||||
profiles=legacy_profiles,
|
||||
)
|
||||
save_router_config(legacy_cfg, self.config_path)
|
||||
|
||||
self.assertEqual(len(legacy_cfg.profiles), 16)
|
||||
self.assertNotIn("grok-orch", legacy_cfg.profiles)
|
||||
self.assertNotIn("claude-orch", legacy_cfg.profiles)
|
||||
|
||||
# 2. Load configuration (triggers migration)
|
||||
migrated_cfg = load_router_config(self.config_path)
|
||||
|
||||
# 3. Verify backup file created
|
||||
backups = list(Path(self.tmp_dir).glob("router_profiles.yaml.bak_*"))
|
||||
self.assertGreaterEqual(len(backups), 1, "Backup file was not created on migration")
|
||||
|
||||
# 4. Verify Claude and Grok profiles added
|
||||
self.assertIn("grok-orch", migrated_cfg.profiles)
|
||||
self.assertIn("grok-worker-1", migrated_cfg.profiles)
|
||||
self.assertIn("grok-worker-2", migrated_cfg.profiles)
|
||||
self.assertIn("claude-orch", migrated_cfg.profiles)
|
||||
self.assertIn("claude-worker-1", migrated_cfg.profiles)
|
||||
self.assertIn("claude-worker-2", migrated_cfg.profiles)
|
||||
self.assertEqual(len(migrated_cfg.profiles), 22)
|
||||
|
||||
# 5. Verify existing 10 antigravity profiles are 100% untouched
|
||||
for pid in ["ag-orch-fallback", "ag-w1", "ag-w2", "ag-w3", "ag-w4", "ag-spare-1", "ag-spare-2", "ag-cold-1", "ag-cold-2", "ag-cold-3"]:
|
||||
p = migrated_cfg.profiles[pid]
|
||||
self.assertEqual(p.account_id, f"custom-acc-{pid}", f"Account ID mutated for {pid}")
|
||||
self.assertEqual(p.capabilities, ["custom-cap"], f"Capabilities mutated for {pid}")
|
||||
self.assertEqual(p.preferred_models, ["custom-model-1"], f"Models mutated for {pid}")
|
||||
self.assertEqual(p.max_concurrency, 3, f"Concurrency mutated for {pid}")
|
||||
|
||||
# 6. Verify AutoAssigner.find_free_slot finds free slots for grok and claude
|
||||
with patch.dict(os.environ, {"HERMES_ROUTER_CONFIG": str(self.config_path)}):
|
||||
slot_grok = AutoAssigner.find_free_slot("grok")
|
||||
self.assertEqual(slot_grok, "grok-orch")
|
||||
slot_claude = AutoAssigner.find_free_slot("claude")
|
||||
self.assertEqual(slot_claude, "claude-orch")
|
||||
|
||||
# 7. Verify Idempotence: subsequent loads do not create extra backups
|
||||
backup_count_before = len(backups)
|
||||
reload_cfg = load_router_config(self.config_path)
|
||||
self.assertEqual(len(reload_cfg.profiles), 22)
|
||||
backup_count_after = len(list(Path(self.tmp_dir).glob("router_profiles.yaml.bak_*")))
|
||||
self.assertEqual(backup_count_before, backup_count_after)
|
||||
|
||||
|
||||
class TestA9QuotaHonesty(unittest.TestCase):
|
||||
"""P0-1: Quotas for OpenAI Codex and OpenCode Go."""
|
||||
|
||||
def setUp(self):
|
||||
self.service = AccountQuotaService()
|
||||
|
||||
@patch("urllib.request.urlopen")
|
||||
def test_codex_quota_success_returns_honest_unavailable_reason_without_fake_numbers(self, mock_urlopen):
|
||||
# Mock /models success response
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.read.return_value = json.dumps({"data": [{"id": "gpt-4o"}]}).encode("utf-8")
|
||||
mock_resp.__enter__.return_value = mock_resp
|
||||
mock_urlopen.return_value = mock_resp
|
||||
|
||||
auth_data = {
|
||||
"token": {"access_token": "valid_token_123"},
|
||||
"email": "user@example.com",
|
||||
}
|
||||
snap = self.service._collect_codex_quota("codex-orch", auth_data)
|
||||
|
||||
self.assertEqual(snap.source, "baseline")
|
||||
self.assertTrue(snap.is_estimated)
|
||||
self.assertEqual(snap.provider, "openai-codex")
|
||||
self.assertEqual(snap.unavailable_reason, "OpenAI Codex не предоставляет остаток через публичный API")
|
||||
for b in snap.buckets:
|
||||
self.assertIsNone(b.remaining_percent)
|
||||
self.assertIsNone(b.reset_at)
|
||||
|
||||
def test_codex_quota_401_returns_expired_auth_reason(self):
|
||||
auth_data = {
|
||||
"token": {"access_token": "expired_token_123"},
|
||||
"email": "user@example.com",
|
||||
}
|
||||
import urllib.error
|
||||
with patch("urllib.request.urlopen", side_effect=urllib.error.HTTPError("url", 401, "Unauthorized", {}, None)):
|
||||
snap = self.service._collect_codex_quota("codex-orch", auth_data)
|
||||
|
||||
self.assertEqual(snap.source, "baseline")
|
||||
self.assertTrue(snap.is_estimated)
|
||||
self.assertEqual(snap.unavailable_reason, "Авторизация истекла — обновите подключение")
|
||||
for b in snap.buckets:
|
||||
self.assertIsNone(b.remaining_percent)
|
||||
|
||||
@patch("urllib.request.urlopen")
|
||||
def test_opencode_quota_measured_values(self, mock_urlopen):
|
||||
# Mock /models and /usage responses
|
||||
models_resp = MagicMock()
|
||||
models_resp.read.return_value = json.dumps({"data": [{"id": "deepseek-r1"}]}).encode("utf-8")
|
||||
models_resp.__enter__.return_value = models_resp
|
||||
|
||||
usage_resp = MagicMock()
|
||||
usage_resp.read.return_value = json.dumps({
|
||||
"five_hour": {"remaining_percent": 80.0, "remaining": 10, "used": 2},
|
||||
"weekly": {"remaining_percent": 90.0, "remaining": 27, "used": 3},
|
||||
"monthly": {"remaining_percent": 95.0, "remaining": 57, "used": 3},
|
||||
}).encode("utf-8")
|
||||
usage_resp.__enter__.return_value = usage_resp
|
||||
|
||||
mock_urlopen.side_effect = [models_resp, usage_resp]
|
||||
|
||||
auth_data = {"api_key": "opencode_secret_key"}
|
||||
snap = self.service._collect_opencode_quota("opengo-1", auth_data)
|
||||
|
||||
self.assertEqual(snap.source, "provider_api")
|
||||
self.assertIsNone(snap.unavailable_reason)
|
||||
b_5h = next(b for b in snap.buckets if b.period == "5h")
|
||||
self.assertEqual(b_5h.remaining_percent, 80.0)
|
||||
self.assertEqual(b_5h.used_absolute, 2)
|
||||
|
||||
|
||||
class TestA9ModelDiscoveryService(unittest.TestCase):
|
||||
"""P0-2: Non-blocking ModelDiscoveryService with disk cache and timeout."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp_dir = tempfile.mkdtemp(prefix="hermes_test_a9_discovery_")
|
||||
self.cache_file = Path(self.tmp_dir) / "models_cache.json"
|
||||
self.service = ModelDiscoveryService(cache_path=self.cache_file)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmp_dir, ignore_errors=True)
|
||||
|
||||
def test_empty_cache_returns_none_honestly(self):
|
||||
models = self.service.get_models("openai-codex")
|
||||
self.assertIsNone(models)
|
||||
meta = self.service.get_models_with_metadata("openai-codex")
|
||||
self.assertFalse(meta["has_cache"])
|
||||
self.assertTrue(meta["is_stale"])
|
||||
|
||||
def test_cache_persistence_on_disk(self):
|
||||
# Manually seed cache
|
||||
with self.service._cache_lock:
|
||||
self.service._cache["openai-codex"] = {
|
||||
"models": ["gpt-4o", "o3-mini"],
|
||||
"discovered_at": time.time(),
|
||||
}
|
||||
self.service._save_cache_to_disk()
|
||||
|
||||
# Reload new instance from same path
|
||||
new_svc = ModelDiscoveryService(cache_path=self.cache_file)
|
||||
models = new_svc.get_models("openai-codex")
|
||||
self.assertEqual(models, ["gpt-4o", "o3-mini"])
|
||||
meta = new_svc.get_models_with_metadata("openai-codex")
|
||||
self.assertTrue(meta["has_cache"])
|
||||
self.assertFalse(meta["is_stale"])
|
||||
|
||||
def test_timeout_probe_does_not_block_and_leaves_previous_cache(self):
|
||||
# Seed cache
|
||||
with self.service._cache_lock:
|
||||
self.service._cache["antigravity"] = {
|
||||
"models": ["gemini-3.5-flash"],
|
||||
"discovered_at": time.time(),
|
||||
}
|
||||
self.service._save_cache_to_disk()
|
||||
|
||||
def _hanging_probe(provider):
|
||||
time.sleep(2.0)
|
||||
return ["invented-model"]
|
||||
|
||||
with patch.object(self.service, "_probe_provider", side_effect=_hanging_probe):
|
||||
t0 = time.time()
|
||||
res = self.service.discover_models_sync("antigravity", timeout=0.2)
|
||||
duration = time.time() - t0
|
||||
self.assertLess(duration, 1.0, "Discovery hung longer than timeout")
|
||||
self.assertEqual(res, ["gemini-3.5-flash"], "Did not retain previous cache on timeout")
|
||||
|
||||
|
||||
class TestA9InventedModelsRemoval(unittest.TestCase):
|
||||
"""P0-3: Removal of invented model literals."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp_dir = tempfile.mkdtemp(prefix="hermes_test_a9_assigner_")
|
||||
self.config_path = Path(self.tmp_dir) / "router_profiles.yaml"
|
||||
self.cache_path = Path(self.tmp_dir) / "models_cache.json"
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmp_dir, ignore_errors=True)
|
||||
|
||||
def test_ensure_profile_definition_without_cache_leaves_empty_models(self):
|
||||
# Empty discovery service
|
||||
md_svc = ModelDiscoveryService(cache_path=self.cache_path)
|
||||
with patch("antigravity_provider.router.model_discovery_service.ModelDiscoveryService.get", return_value=md_svc):
|
||||
with patch.dict(os.environ, {"HERMES_ROUTER_CONFIG": str(self.config_path)}):
|
||||
ok, msg = AutoAssigner.ensure_profile_definition("claude", "claude-custom-1")
|
||||
self.assertTrue(ok)
|
||||
|
||||
cfg = load_router_config(self.config_path)
|
||||
p = cfg.profiles.get("claude-custom-1")
|
||||
self.assertIsNotNone(p)
|
||||
# Honest empty list when not discovered
|
||||
self.assertEqual(p.preferred_models, [])
|
||||
|
||||
|
||||
class TestA9ProviderOrdering(unittest.TestCase):
|
||||
"""P1-4: Deterministic Provider Ordering in Snapshot."""
|
||||
|
||||
def test_provider_summaries_sorting_and_total_count(self):
|
||||
uh = UnifiedHealthService.get()
|
||||
summaries = uh.get_provider_summaries()
|
||||
self.assertEqual(len(summaries), 5)
|
||||
|
||||
# Verify deterministic ordering: connected_count desc, total_slots desc, provider_name asc
|
||||
for i in range(len(summaries) - 1):
|
||||
s1 = summaries[i]
|
||||
s2 = summaries[i + 1]
|
||||
key1 = (-s1.connected_count, -s1.total_slots, s1.provider_name)
|
||||
key2 = (-s2.connected_count, -s2.total_slots, s2.provider_name)
|
||||
self.assertLessEqual(key1, key2)
|
||||
|
||||
readiness = uh.get_system_readiness()
|
||||
self.assertEqual(readiness.total_providers, 5)
|
||||
|
||||
|
||||
class TestA9HermesPluginBoundary(unittest.TestCase):
|
||||
"""P0-00: Hermes call interception boundary and role determination."""
|
||||
|
||||
def test_call_without_role_passes_downstream_without_router_attempts(self):
|
||||
from antigravity_provider import hermes_plugin
|
||||
|
||||
downstream_calls = []
|
||||
def mock_next(req):
|
||||
downstream_calls.append(req)
|
||||
return {"choices": [{"message": {"role": "assistant", "content": "direct-model-response"}}]}
|
||||
|
||||
# Call with no role and provider != antigravity
|
||||
res = hermes_plugin.antigravity_llm_execution(
|
||||
request={"messages": [{"role": "user", "content": "hello"}]},
|
||||
next_call=mock_next,
|
||||
provider="opencode-go",
|
||||
model="kimi-k2.7-code",
|
||||
session_id="sess-123",
|
||||
)
|
||||
self.assertEqual(len(downstream_calls), 1)
|
||||
content = res["choices"][0]["message"]["content"]
|
||||
self.assertEqual(content, "direct-model-response")
|
||||
|
||||
def test_resolve_role_returns_none_when_unspecified(self):
|
||||
from antigravity_provider.router import get_router_engine
|
||||
engine = get_router_engine()
|
||||
# No role in request or explicit role -> None (no guessing from prompts)
|
||||
req = {"messages": [{"role": "system", "content": "You are a coding agent developer"}]}
|
||||
self.assertIsNone(engine.resolve_role(req))
|
||||
|
||||
def test_resolve_role_respects_explicit_or_metadata_role(self):
|
||||
from antigravity_provider.router import get_router_engine
|
||||
engine = get_router_engine()
|
||||
self.assertEqual(engine.resolve_role({}, explicit_role="coder-primary"), "coder-primary")
|
||||
self.assertEqual(engine.resolve_role({"role": "reviewer"}), "reviewer")
|
||||
self.assertEqual(engine.resolve_role({"metadata": {"role": "research"}}), "research")
|
||||
|
||||
|
||||
class TestA9CodexOAuthTokenRefreshAndSwitching(unittest.TestCase):
|
||||
"""P0-01: Codex token refresh and safe account switching."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp_dir = tempfile.mkdtemp(prefix="hermes_test_a9_codex_")
|
||||
self.profile_dir = Path(self.tmp_dir) / "openai-codex" / "codex-test"
|
||||
self.profile_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmp_dir, ignore_errors=True)
|
||||
|
||||
@patch("antigravity_provider.router.codex_oauth._post_json")
|
||||
def test_refresh_codex_token_success(self, mock_post):
|
||||
mock_post.return_value = {
|
||||
"access_token": "new_access_token_456",
|
||||
"refresh_token": "new_refresh_token_789",
|
||||
"id_token": "header.eyJlbWFpbCI6ICJ1c2VyQGdtYWlsLmNvbSIsICJleHAiOiAyMDAwMDAwMDAwfQ.sig",
|
||||
}
|
||||
auth_data = {
|
||||
"provider": "openai-codex",
|
||||
"profile_id": "codex-test",
|
||||
"token": {"access_token": "old_acc", "refresh_token": "valid_refresh"},
|
||||
}
|
||||
with patch.object(ProfileAuthManager, "load_profile_auth", return_value=auth_data):
|
||||
with patch.object(ProfileAuthManager, "save_profile_auth") as mock_save:
|
||||
from antigravity_provider.router.codex_oauth import refresh_codex_token
|
||||
res = refresh_codex_token("codex-test")
|
||||
self.assertEqual(res["token"]["access_token"], "new_access_token_456")
|
||||
self.assertEqual(res["token"]["refresh_token"], "new_refresh_token_789")
|
||||
mock_save.assert_called_once()
|
||||
|
||||
def test_refresh_codex_token_missing_raises_error(self):
|
||||
auth_data = {
|
||||
"provider": "openai-codex",
|
||||
"profile_id": "codex-test",
|
||||
"token": {"access_token": "old_acc"},
|
||||
}
|
||||
with patch.object(ProfileAuthManager, "load_profile_auth", return_value=auth_data):
|
||||
from antigravity_provider.router.codex_oauth import refresh_codex_token
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
refresh_codex_token("codex-test")
|
||||
self.assertIn("refresh_token отсутствует", str(ctx.exception))
|
||||
|
||||
def test_get_profile_status_separate_token_expiry(self):
|
||||
# 1. Expired access token with refresh token -> AUTHENTICATED with refresh notice
|
||||
exp_access = "header.eyJlbWFpbCI6ICJ1c2VyQGdtYWlsLmNvbSIsICJleHAiOiAxMDAwfQ.sig" # expired in 1970
|
||||
valid_id = "header.eyJlbWFpbCI6ICJ1c2VyQGdtYWlsLmNvbSIsICJleHAiOiAyMDAwMDAwMDAwfQ.sig"
|
||||
auth_data = {
|
||||
"provider": "openai-codex",
|
||||
"profile_id": "codex-test",
|
||||
"token": {"access_token": exp_access, "id_token": valid_id, "refresh_token": "ref_123"},
|
||||
}
|
||||
with patch.object(ProfileAuthManager, "load_profile_auth", return_value=auth_data):
|
||||
st = ProfileAuthManager.get_profile_status("openai-codex", "codex-test")
|
||||
self.assertTrue(st["access_token_expired"])
|
||||
self.assertFalse(st["id_token_expired"])
|
||||
self.assertTrue(st["has_refresh_token"])
|
||||
self.assertFalse(st["is_expired"]) # Can be refreshed silently
|
||||
self.assertEqual(st["status"], "AUTHENTICATED")
|
||||
|
||||
# 2. Expired access token without refresh token -> EXPIRED
|
||||
auth_data_no_refresh = {
|
||||
"provider": "openai-codex",
|
||||
"profile_id": "codex-test",
|
||||
"token": {"access_token": exp_access, "id_token": valid_id},
|
||||
}
|
||||
with patch.object(ProfileAuthManager, "load_profile_auth", return_value=auth_data_no_refresh):
|
||||
st2 = ProfileAuthManager.get_profile_status("openai-codex", "codex-test")
|
||||
self.assertTrue(st2["access_token_expired"])
|
||||
self.assertFalse(st2["has_refresh_token"])
|
||||
self.assertTrue(st2["is_expired"])
|
||||
self.assertEqual(st2["status"], "EXPIRED")
|
||||
|
||||
def test_switch_active_codex_account_sequence_and_rollback(self):
|
||||
from antigravity_provider.router.codex_oauth import switch_active_codex_account
|
||||
steps = []
|
||||
def record_step(step_name, msg, status):
|
||||
steps.append((step_name, status))
|
||||
|
||||
auth_data = {
|
||||
"provider": "openai-codex",
|
||||
"profile_id": "codex-test",
|
||||
"email": "user@example.com",
|
||||
"token": {"access_token": "valid_token", "refresh_token": "ref_123"},
|
||||
}
|
||||
with patch.object(ProfileAuthManager, "load_profile_auth", return_value=auth_data):
|
||||
res = switch_active_codex_account("codex-test", step_callback=record_step)
|
||||
self.assertTrue(res["success"])
|
||||
step_names = [s[0] for s in steps]
|
||||
self.assertIn("check_tokens", step_names)
|
||||
self.assertIn("stop_clients", step_names)
|
||||
self.assertIn("write_credentials", step_names)
|
||||
self.assertIn("sync_settings", step_names)
|
||||
self.assertIn("start_client", step_names)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Loading…
Reference in a new issue