diff --git a/agent/account_usage.py b/agent/account_usage.py index da02af3c478..3e688d919ac 100644 --- a/agent/account_usage.py +++ b/agent/account_usage.py @@ -436,20 +436,53 @@ def _resolve_codex_usage_url(base_url: str) -> str: return normalized + "/api/codex/usage" -def _fetch_codex_account_usage() -> Optional[AccountUsageSnapshot]: - creds = resolve_codex_runtime_credentials(refresh_if_expiring=True) - token_data = _read_codex_tokens() - tokens = token_data.get("tokens") or {} - account_id = str(tokens.get("account_id", "") or "").strip() or None +def _resolve_codex_usage_credentials( + base_url: Optional[str], + api_key: Optional[str], +) -> tuple[str, str, Optional[str]]: + """Resolve Codex quota credentials from the native runtime path. + + Prefer explicit live-agent credentials, then the legacy singleton OAuth + state, then the credential pool. Hermes's native OAuth setup now stores + device-code logins in the pool, so quota diagnostics must not depend only + on the older singleton store. + """ + explicit_key = str(api_key or "").strip() + if explicit_key: + return explicit_key, str(base_url or "").strip(), None + + try: + creds = resolve_codex_runtime_credentials(refresh_if_expiring=True) + token_data = _read_codex_tokens() + tokens = token_data.get("tokens") or {} + account_id = str(tokens.get("account_id", "") or "").strip() or None + return creds["api_key"], str(creds.get("base_url", "") or "").strip(), account_id + except Exception: + pass + + from agent.credential_pool import load_pool + + pool = load_pool("openai-codex") + entry = pool.select() + if entry is None: + raise RuntimeError("No available openai-codex credential in credential pool") + return entry.runtime_api_key, str(entry.runtime_base_url or base_url or "").strip(), None + + +def _fetch_codex_account_usage( + base_url: Optional[str] = None, + api_key: Optional[str] = None, +) -> Optional[AccountUsageSnapshot]: + token, resolved_base_url, account_id = _resolve_codex_usage_credentials(base_url, api_key) headers = { - "Authorization": f"Bearer {creds['api_key']}", + "Authorization": f"Bearer {token}", "Accept": "application/json", "User-Agent": "codex-cli", } if account_id: headers["ChatGPT-Account-Id"] = account_id with httpx.Client(timeout=15.0) as client: - response = client.get(_resolve_codex_usage_url(creds.get("base_url", "")), headers=headers) + response = client.get(_resolve_codex_usage_url(resolved_base_url), headers=headers) response.raise_for_status() payload = response.json() or {} rate_limit = payload.get("rate_limit") or {} @@ -628,7 +661,7 @@ def fetch_account_usage( return None try: if normalized == "openai-codex": - return _fetch_codex_account_usage() + return _fetch_codex_account_usage(base_url=base_url, api_key=api_key) if normalized == "anthropic": return _fetch_anthropic_account_usage() if normalized == "openrouter": diff --git a/tests/agent/test_account_usage.py b/tests/agent/test_account_usage.py new file mode 100644 index 00000000000..27068f3b3aa --- /dev/null +++ b/tests/agent/test_account_usage.py @@ -0,0 +1,110 @@ +from types import SimpleNamespace + +import pytest + +from agent import account_usage + + +class _FakeResponse: + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self._payload + + +class _FakeClient: + def __init__(self, calls, payload): + self.calls = calls + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def get(self, url, headers): + self.calls.append({"url": url, "headers": headers}) + return _FakeResponse(self.payload) + + +@pytest.fixture +def codex_usage_payload(): + return { + "plan_type": "plus", + "rate_limit": { + "primary_window": { + "used_percent": 21, + "reset_at": 1779846359, + }, + "secondary_window": { + "used_percent": 4, + "reset_at": 1780230796, + }, + }, + "credits": {"has_credits": False}, + } + + +def test_codex_usage_prefers_explicit_live_agent_credentials(monkeypatch, codex_usage_payload): + calls = [] + monkeypatch.setattr( + account_usage.httpx, + "Client", + lambda timeout: _FakeClient(calls, codex_usage_payload), + ) + monkeypatch.setattr( + account_usage, + "resolve_codex_runtime_credentials", + lambda **kwargs: (_ for _ in ()).throw(AssertionError("legacy auth should not be used")), + ) + + snapshot = account_usage.fetch_account_usage( + "openai-codex", + base_url="https://chatgpt.com/backend-api/codex", + api_key="live-agent-token", + ) + + assert snapshot is not None + assert snapshot.provider == "openai-codex" + assert snapshot.plan == "Plus" + assert [w.label for w in snapshot.windows] == ["Session", "Weekly"] + assert snapshot.windows[0].used_percent == 21 + assert calls[0]["url"] == "https://chatgpt.com/backend-api/wham/usage" + assert calls[0]["headers"]["Authorization"] == "Bearer live-agent-token" + + +def test_codex_usage_falls_back_to_native_credential_pool(monkeypatch, codex_usage_payload): + calls = [] + monkeypatch.setattr( + account_usage.httpx, + "Client", + lambda timeout: _FakeClient(calls, codex_usage_payload), + ) + monkeypatch.setattr( + account_usage, + "resolve_codex_runtime_credentials", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("no singleton auth")), + ) + + pool_entry = SimpleNamespace( + runtime_api_key="pooled-token", + runtime_base_url="https://chatgpt.com/backend-api/codex", + ) + pool = SimpleNamespace(select=lambda: pool_entry) + + import agent.credential_pool as credential_pool + + monkeypatch.setattr(credential_pool, "load_pool", lambda provider: pool) + + snapshot = account_usage.fetch_account_usage("openai-codex") + + assert snapshot is not None + assert snapshot.windows[0].label == "Session" + assert snapshot.windows[1].label == "Weekly" + assert calls[0]["url"] == "https://chatgpt.com/backend-api/wham/usage" + assert calls[0]["headers"]["Authorization"] == "Bearer pooled-token"