diff --git a/agent/account_usage.py b/agent/account_usage.py index 3e688d919ac..712e57feda4 100644 --- a/agent/account_usage.py +++ b/agent/account_usage.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional import httpx from agent.anthropic_adapter import _is_oauth_token, resolve_anthropic_token -from hermes_cli.auth import _read_codex_tokens, resolve_codex_runtime_credentials +from hermes_cli.auth import AuthError, _read_codex_tokens, resolve_codex_runtime_credentials from hermes_cli.runtime_provider import resolve_runtime_provider if TYPE_CHECKING: @@ -451,15 +451,40 @@ def _resolve_codex_usage_credentials( if explicit_key: return explicit_key, str(base_url or "").strip(), None + # Tier 2: the native runtime resolver. It ALREADY falls back to the + # credential pool when the singleton is empty (see + # ``resolve_codex_runtime_credentials`` — issue #32992), so in a pool-only + # setup this returns a usable ``source="credential_pool"`` token. + # + # Only ``AuthError`` ("no creds" / rate-limited) is caught so tier 3 can + # run: a broad ``except Exception`` would (a) mask a transient refresh / + # network failure and silently hand back a DIFFERENT pool account's usage, + # and (b) hide genuine programming errors. A refresh/network error must + # propagate — the outer ``fetch_account_usage`` guard fails open (shows + # nothing this turn) rather than reporting the wrong account. + # + # The ``account_id`` (for the ``ChatGPT-Account-Id`` header) is read + # best-effort: a partial/missing singleton token store must not sink an + # otherwise-usable resolver credential and force a header-less pool fallback. 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 + account_id: Optional[str] = None + try: + token_data = _read_codex_tokens() + tokens = token_data.get("tokens") or {} + account_id = str(tokens.get("account_id", "") or "").strip() or None + except AuthError: + # Pool-only creds carry no singleton account_id; header is optional. + logger.debug("codex ▸ /usage account_id read failed (best-effort)", exc_info=True) return creds["api_key"], str(creds.get("base_url", "") or "").strip(), account_id - except Exception: - pass + except AuthError: + logger.debug("codex ▸ /usage runtime resolver returned no creds; trying pool", exc_info=True) + # Tier 3: direct pool select. Reached only when the resolver itself raises + # AuthError (e.g. singleton missing AND its own pool read found nothing at + # resolve time, but a pool entry is usable now). Pool credentials have no + # account_id concept, so the ChatGPT-Account-Id header is intentionally + # omitted here. from agent.credential_pool import load_pool pool = load_pool("openai-codex") diff --git a/tests/agent/test_account_usage.py b/tests/agent/test_account_usage.py index acf1745cbfc..41950c70a12 100644 --- a/tests/agent/test_account_usage.py +++ b/tests/agent/test_account_usage.py @@ -85,10 +85,15 @@ def test_codex_usage_falls_back_to_native_credential_pool(monkeypatch, codex_usa "Client", lambda timeout: _FakeClient(calls, codex_usage_payload), ) + # Pool fallback fires only on AuthError (the documented "no creds" mode of + # the resolver), NOT on arbitrary exceptions — see the transient-error guard + # test below. monkeypatch.setattr( account_usage, "resolve_codex_runtime_credentials", - lambda **kwargs: (_ for _ in ()).throw(RuntimeError("no singleton auth")), + lambda **kwargs: (_ for _ in ()).throw( + account_usage.AuthError("no singleton auth", provider="openai-codex", code="codex_auth_missing") + ), ) pool_entry = SimpleNamespace( @@ -108,6 +113,85 @@ def test_codex_usage_falls_back_to_native_credential_pool(monkeypatch, codex_usa 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" + # Pool creds have no account_id concept — the ChatGPT-Account-Id header must + # be omitted rather than sent stale/wrong. + assert "ChatGPT-Account-Id" not in calls[0]["headers"] + + +def test_codex_usage_does_not_swap_to_pool_on_transient_resolver_error(monkeypatch, codex_usage_payload): + """A transient refresh/network failure (non-AuthError) must NOT silently + downgrade to a possibly-different pool account. It fails open (no snapshot) + instead of reporting the wrong account's usage.""" + 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("refresh endpoint 503")), + ) + + pool_entry = SimpleNamespace( + runtime_api_key="pooled-token-WRONG-ACCOUNT", + runtime_base_url="https://chatgpt.com/backend-api/codex", + ) + pool = SimpleNamespace(select=lambda: pool_entry) + + import agent.credential_pool as credential_pool + + # If the guard regressed, this pool would be consulted and return a snapshot + # for the wrong account. It must NOT be. + monkeypatch.setattr(credential_pool, "load_pool", lambda provider: pool) + + snapshot = account_usage.fetch_account_usage("openai-codex") + + assert snapshot is None + assert calls == [] # HTTP usage endpoint never hit with a wrong-account token + + +def test_codex_usage_account_id_read_failure_keeps_singleton_token(monkeypatch, codex_usage_payload): + """When the resolver succeeds but the separate account_id read raises, the + working singleton token must still be used (best-effort account_id), NOT + abandoned in favor of a header-less pool credential.""" + calls = [] + monkeypatch.setattr( + account_usage.httpx, + "Client", + lambda timeout: _FakeClient(calls, codex_usage_payload), + ) + monkeypatch.setattr( + account_usage, + "resolve_codex_runtime_credentials", + lambda **kwargs: { + "api_key": "singleton-token", + "base_url": "https://chatgpt.com/backend-api/codex", + }, + ) + monkeypatch.setattr( + account_usage, + "_read_codex_tokens", + lambda *a, **k: (_ for _ in ()).throw( + account_usage.AuthError("partial store", provider="openai-codex", code="codex_auth_invalid_shape") + ), + ) + + import agent.credential_pool as credential_pool + + monkeypatch.setattr( + credential_pool, + "load_pool", + lambda provider: (_ for _ in ()).throw(AssertionError("pool must not be consulted")), + ) + + snapshot = account_usage.fetch_account_usage("openai-codex") + + assert snapshot is not None + assert calls[0]["headers"]["Authorization"] == "Bearer singleton-token" + # account_id read failed → header omitted, but the singleton token is kept. + assert "ChatGPT-Account-Id" not in calls[0]["headers"] def test_codex_usage_treats_wham_used_percent_as_used_not_remaining(monkeypatch):