hermes-agent/tests/agent/test_account_usage.py

from types import SimpleNamespace

import httpx
import pytest

from agent import account_usage


class _FakeResponse:
    def __init__(self, payload, status_code=200):
        self._payload = payload
        self.status_code = status_code

    def raise_for_status(self):
        if self.status_code >= 400:
            request = httpx.Request("GET", "https://chatgpt.com/backend-api/wham/usage")
            response = httpx.Response(self.status_code, request=request)
            raise httpx.HTTPStatusError("request failed", request=request, response=response)

    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),
    )
    # 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(
            account_usage.AuthError("no singleton auth", provider="openai-codex", code="codex_auth_missing")
        ),
    )

    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"
    # 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 _explicit_creds_snapshot(monkeypatch, payload):
    calls = []
    monkeypatch.setattr(account_usage.httpx, "Client", lambda timeout: _FakeClient(calls, payload))
    snapshot = account_usage.fetch_account_usage(
        "openai-codex", base_url="https://chatgpt.com/backend-api/codex", api_key="live-agent-token",
    )
    return snapshot, calls


def test_codex_weekly_only_primary_window_is_labeled_weekly(monkeypatch):
    """#65387: a lone 604800s primary_window is the weekly limit, not the session one."""
    payload = {"plan_type": "pro", "rate_limit": {
        "primary_window": {"used_percent": 1, "limit_window_seconds": 604800},
        "secondary_window": None,
    }}
    snapshot, _ = _explicit_creds_snapshot(monkeypatch, payload)
    assert [(w.label, w.used_percent) for w in snapshot.windows] == [("Weekly", 1.0)]


def test_codex_window_labels_follow_duration_with_positional_fallback(monkeypatch):
    # Swapped positions: labels must follow limit_window_seconds.
    payload = {"rate_limit": {
        "primary_window": {"used_percent": 4, "limit_window_seconds": 604800},
        "secondary_window": {"used_percent": 21, "limit_window_seconds": 18000},
    }}
    snapshot, _ = _explicit_creds_snapshot(monkeypatch, payload)
    assert [w.label for w in snapshot.windows] == ["Weekly", "Session"]
    # Missing / unrecognized durations keep the legacy positional labels.
    payload = {"rate_limit": {
        "primary_window": {"used_percent": 4},
        "secondary_window": {"used_percent": 21, "limit_window_seconds": 12345},
    }}
    snapshot, _ = _explicit_creds_snapshot(monkeypatch, payload)
    assert [w.label for w in snapshot.windows] == ["Session", "Weekly"]


def test_codex_snapshot_exposes_exact_raw_payload_with_one_get(monkeypatch, codex_usage_payload):
    """#79695: the decoded body rides along untouched (unknown fields included), from the single GET."""
    codex_usage_payload["future_field"] = {"nested": [1, 2]}
    snapshot, calls = _explicit_creds_snapshot(monkeypatch, codex_usage_payload)
    assert snapshot.raw == codex_usage_payload
    assert snapshot.raw["future_field"] == {"nested": [1, 2]}
    assert len(calls) == 1
    assert [w.label for w in snapshot.windows] == ["Session", "Weekly"]  # normalized limits unchanged
    # Additive: existing constructor calls stay valid and default to no raw body.
    assert account_usage.AccountUsageSnapshot(provider="anthropic", source="x", fetched_at=snapshot.fetched_at).raw is None


def test_codex_invalid_payload_fails_closed(monkeypatch):
    snapshot, _ = _explicit_creds_snapshot(monkeypatch, ["not", "a", "dict"])
    assert snapshot is None


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_retries_401_with_forced_refresh(monkeypatch, codex_usage_payload):
    credential_calls = []
    request_calls = []
    responses = [_FakeResponse({}, status_code=401), _FakeResponse(codex_usage_payload)]

    def resolve(**kwargs):
        credential_calls.append(kwargs)
        token = "fresh-token" if kwargs.get("force_refresh") else "revoked-token"
        return {"api_key": token, "base_url": "https://chatgpt.com/backend-api/codex"}

    class Client:
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        def get(self, url, headers):
            request_calls.append(headers["Authorization"])
            return responses.pop(0)

    monkeypatch.setattr(account_usage, "resolve_codex_runtime_credentials", resolve)
    monkeypatch.setattr(account_usage, "_read_codex_tokens", lambda: {"tokens": {}})
    monkeypatch.setattr(account_usage.httpx, "Client", lambda timeout: Client())

    snapshot = account_usage.fetch_account_usage("openai-codex")

    assert snapshot is not None
    assert snapshot.windows[0].label == "Session"
    assert credential_calls == [
        {"refresh_if_expiring": True},
        {"refresh_if_expiring": True, "force_refresh": True},
    ]
    assert request_calls == ["Bearer revoked-token", "Bearer fresh-token"]


# ── Banked rate-limit reset credits (`/usage reset`) ─────────────────────────


class _FakeResetClient:
    """GET returns the usage payload; POST returns the consume payload."""

    def __init__(self, calls, usage_payload, consume_payload=None):
        self.calls = calls
        self.usage_payload = usage_payload
        self.consume_payload = consume_payload or {}

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        return False

    def get(self, url, headers):
        self.calls.append({"method": "GET", "url": url, "headers": headers})
        return _FakeResponse(self.usage_payload)

    def post(self, url, headers=None, json=None):
        self.calls.append({"method": "POST", "url": url, "headers": headers, "json": json})
        return _FakeResponse(self.consume_payload)


def _usage_payload_with_resets(primary_used, secondary_used, banked):
    return {
        "plan_type": "plus",
        "rate_limit": {
            "primary_window": {"used_percent": primary_used, "reset_at": 1779846359},
            "secondary_window": {"used_percent": secondary_used, "reset_at": 1780230796},
        },
        "rate_limit_reset_credits": {"available_count": banked},
        "credits": {"has_credits": False},
    }
















def test_redeem_retries_401_with_forced_refresh(monkeypatch):
    credential_calls = []
    request_calls = []
    client_count = 0
    payload = _usage_payload_with_resets(100, 40, 1)

    def resolve(base_url, api_key, *, force_refresh=False):
        credential_calls.append(force_refresh)
        token = "fresh-token" if force_refresh else "revoked-token"
        return token, "https://chatgpt.com/backend-api/codex", None

    class Client(_FakeResetClient):
        def get(self, url, headers):
            request_calls.append(("GET", headers["Authorization"]))
            if headers["Authorization"] == "Bearer revoked-token":
                return _FakeResponse({}, status_code=401)
            return _FakeResponse(payload)

        def post(self, url, headers=None, json=None):
            request_calls.append(("POST", headers["Authorization"]))
            return _FakeResponse({"code": "reset", "windows_reset": 2})

    def client_factory(timeout):
        nonlocal client_count
        client_count += 1
        return Client([], payload)

    monkeypatch.setattr(account_usage, "_resolve_codex_usage_credentials", resolve)
    monkeypatch.setattr(account_usage.httpx, "Client", client_factory)

    result = account_usage.redeem_codex_reset_credit()

    assert result.status == "reset"
    assert credential_calls == [False, True]
    assert request_calls == [
        ("GET", "Bearer revoked-token"),
        ("GET", "Bearer fresh-token"),
        ("POST", "Bearer fresh-token"),
    ]
    assert client_count == 2


def test_redeem_missing_credentials_reports_unavailable(monkeypatch):
    monkeypatch.setattr(
        account_usage,
        "_resolve_codex_usage_credentials",
        lambda base_url, api_key, **kwargs: (_ for _ in ()).throw(RuntimeError("no creds")),
    )

    result = account_usage.redeem_codex_reset_credit()

    assert result.status == "unavailable"
    assert "hermes auth" in result.message


def test_codex_usage_401_retry_refreshes_the_explicit_credential_not_another_account(monkeypatch, codex_usage_payload):
    """A live agent on pool entry B hands its own api_key in; after a 401 the retry must refresh B,
    not re-resolve and render the singleton/pool account A's usage."""
    request_calls = []
    refresh_hints = []
    responses = [_FakeResponse({}, status_code=401), _FakeResponse(codex_usage_payload)]

    class Pool:
        def try_refresh_matching(self, api_key_hint=None, credential_id=None):
            refresh_hints.append(api_key_hint)
            return SimpleNamespace(runtime_api_key="pool-B-fresh", runtime_base_url=None)

    class Client:
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        def get(self, url, headers):
            request_calls.append(headers["Authorization"])
            return responses.pop(0)

    monkeypatch.setattr(account_usage, "resolve_codex_runtime_credentials",
                        lambda **kwargs: pytest.fail("must not re-resolve another account's credential"))
    monkeypatch.setattr(account_usage, "_read_codex_tokens", lambda: {"tokens": {"access_token": "singleton-A"}})
    monkeypatch.setattr("agent.credential_pool.load_pool", lambda provider: Pool())
    monkeypatch.setattr(account_usage.httpx, "Client", lambda timeout: Client())

    snapshot = account_usage.fetch_account_usage(
        "openai-codex", base_url="https://chatgpt.com/backend-api/codex", api_key="pool-B-revoked")

    assert snapshot is not None
    assert refresh_hints == ["pool-B-revoked"]
    assert request_calls == ["Bearer pool-B-revoked", "Bearer pool-B-fresh"]