Aether/tests/services/test_provider_keys_quota_refresh_strategies.py

from __future__ import annotations

import json
import sys
import types
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock

import pytest

from src.config import config as app_config
from src.services.provider_keys.codex_usage_parser import (
    CodexUsageParseError,
    parse_codex_wham_usage_response,
)
from src.services.provider_keys.quota_refresh.antigravity_refresher import (
    refresh_antigravity_key_quota,
)
from src.services.provider_keys.quota_refresh.codex_refresher import refresh_codex_key_quota
from src.services.provider_keys.quota_refresh.kiro_refresher import refresh_kiro_key_quota
from src.services.request.execution_runtime_client import ExecutionRuntimeSyncResult


class _FakeDB:
    def __init__(self) -> None:
        self.commit_count = 0

    def commit(self) -> None:
        self.commit_count += 1


class _FakeResponse:
    def __init__(
        self,
        status_code: int,
        payload: Any = None,
        json_exc: Exception | None = None,
        *,
        headers: dict[str, str] | None = None,
        text: str | None = None,
    ) -> None:
        self.status_code = status_code
        self._payload = payload
        self._json_exc = json_exc
        self.headers = headers or {}
        if text is not None:
            self.text = text
        elif isinstance(payload, (dict, list)):
            self.text = json.dumps(payload, ensure_ascii=False)
        else:
            self.text = ""

    def json(self) -> Any:
        if self._json_exc:
            raise self._json_exc
        return self._payload


class _FakeAsyncClient:
    def __init__(self, response: _FakeResponse, **kwargs: Any) -> None:
        self._response = response
        self.kwargs = kwargs
        self.last_url: str | None = None
        self.last_headers: dict[str, str] | None = None

    async def __aenter__(self) -> "_FakeAsyncClient":
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: Any,
    ) -> bool:
        _ = exc_type, exc, tb
        return False

    async def get(self, url: str, headers: dict[str, str]) -> _FakeResponse:
        self.last_url = url
        self.last_headers = headers
        return self._response


def _install_module(monkeypatch: pytest.MonkeyPatch, name: str, attrs: dict[str, Any]) -> None:
    module = types.ModuleType(name)
    for key, value in attrs.items():
        setattr(module, key, value)
    monkeypatch.setitem(sys.modules, name, module)


def _patch_rust_codex_response(
    monkeypatch: pytest.MonkeyPatch,
    module: Any,
    response: Any | None,
) -> None:
    monkeypatch.setattr(
        module,
        "_try_rust_codex_quota_response",
        AsyncMock(return_value=response),
    )


@pytest.mark.asyncio
async def test_codex_refresher_endpoint_missing_returns_error() -> None:
    key = SimpleNamespace(id="k1", name="K1")
    provider = SimpleNamespace(proxy=None)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=None,
        codex_wham_usage_url="https://example.test",
        metadata_updates={},
        state_updates={},
    )

    assert result["status"] == "error"
    assert "openai:cli" in result["message"]


@pytest.mark.asyncio
async def test_codex_refresher_http_non_200_returns_error(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    response = _FakeResponse(status_code=503, payload={"x": 1})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates={},
        state_updates={},
    )

    assert result["status"] == "error"
    assert result["status_code"] == 503


@pytest.mark.asyncio
async def test_codex_refresher_prefers_rust_executor(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module
    from src.services.request import execution_runtime_client as runtime_module

    monkeypatch.setattr(app_config, "executor_backend", "rust")

    key = SimpleNamespace(
        id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
    )
    provider = SimpleNamespace(id="provider-1", proxy=None)
    endpoint = SimpleNamespace(id="endpoint-1")
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}
    captured: dict[str, Any] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: {"enabled": True, "url": "http://proxy.test:8080"},
            "build_proxy_client_kwargs": lambda proxy, timeout: {"timeout": timeout},
            "get_system_proxy_config_async": AsyncMock(return_value=None),
            "resolve_delegate_config_async": AsyncMock(return_value=None),
            "build_proxy_url_async": AsyncMock(return_value="http://proxy.test:8080"),
            "resolve_proxy_info_async": AsyncMock(return_value={"mode": "http", "label": "proxy"}),
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    monkeypatch.setattr(
        module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 12.5}
    )

    async def _fake_execute_sync_json(self: object, plan: Any) -> ExecutionRuntimeSyncResult:
        captured["plan"] = plan
        return ExecutionRuntimeSyncResult(
            status_code=200,
            response_json={"ok": True},
            headers={"content-type": "application/json"},
        )

    monkeypatch.setattr(
        runtime_module.ExecutionRuntimeClient,
        "execute_sync_json",
        _fake_execute_sync_json,
    )

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test/wham/usage",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "success"
    assert captured["plan"].method == "GET"
    assert captured["plan"].url == "https://example.test/wham/usage"
    assert captured["plan"].proxy is not None
    assert captured["plan"].proxy.url == "http://proxy.test:8080"
    assert metadata_updates["k1"]["codex"]["used_percent"] == 12.5


@pytest.mark.asyncio
async def test_codex_refresher_http_401_marks_auth_invalid_without_disabling_key(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    response = _FakeResponse(status_code=401, payload={"error": {"message": "token expired"}})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "auth_invalid"
    assert result["status_code"] == 401
    assert result["auto_disabled"] is False
    assert metadata_updates == {}
    assert str(state_updates["k1"]["oauth_invalid_reason"]).startswith("[OAUTH_EXPIRED]")


@pytest.mark.asyncio
async def test_codex_refresher_http_402_sets_quota_exhausted_metadata(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc-key",
        auth_type="oauth",
        auth_config="enc-config",
        proxy=None,
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(
        module.crypto_service,
        "decrypt",
        lambda value: (
            "sk-test"
            if value == "enc-key"
            else json.dumps({"plan_type": "team", "account_id": "acc-1"})
        ),
    )
    response = _FakeResponse(status_code=402, payload={"error": {"message": "payment required"}})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "quota_exhausted"
    assert result["status_code"] == 402
    codex_meta = metadata_updates["k1"]["codex"]
    assert codex_meta["plan_type"] == "team"
    assert codex_meta["primary_used_percent"] == 100.0
    assert codex_meta["secondary_used_percent"] == 100.0
    assert state_updates["k1"]["oauth_invalid_at"] is None
    assert state_updates["k1"]["oauth_invalid_reason"] is None


@pytest.mark.asyncio
async def test_codex_refresher_success_preserves_refresh_failed_marker(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc-key",
        auth_type="oauth",
        auth_config="enc-config",
        proxy=None,
        oauth_invalid_at="sentinel-invalid-at",
        oauth_invalid_reason="[REFRESH_FAILED] Token 续期失败 (400): refresh_token_reused",
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(
        module.crypto_service,
        "decrypt",
        lambda value: (
            "sk-test"
            if value == "enc-key"
            else json.dumps({"plan_type": "team", "account_id": "acc-1"})
        ),
    )
    monkeypatch.setattr(
        module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 10.0}
    )
    response = _FakeResponse(status_code=200, payload={"ok": True})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "success"
    assert state_updates["k1"]["oauth_invalid_at"] == "sentinel-invalid-at"
    assert (
        state_updates["k1"]["oauth_invalid_reason"]
        == "[REFRESH_FAILED] Token 续期失败 (400): refresh_token_reused"
    )


@pytest.mark.asyncio
async def test_codex_refresher_quota_exhausted_preserves_refresh_failed_marker(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc-key",
        auth_type="oauth",
        auth_config="enc-config",
        proxy=None,
        oauth_invalid_at="sentinel-invalid-at",
        oauth_invalid_reason="[REFRESH_FAILED] Token 续期失败 (400): refresh_token_reused",
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(
        module.crypto_service,
        "decrypt",
        lambda value: (
            "sk-test"
            if value == "enc-key"
            else json.dumps({"plan_type": "team", "account_id": "acc-1"})
        ),
    )
    response = _FakeResponse(status_code=402, payload={"error": {"message": "payment required"}})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "quota_exhausted"
    assert state_updates["k1"]["oauth_invalid_at"] == "sentinel-invalid-at"
    assert (
        state_updates["k1"]["oauth_invalid_reason"]
        == "[REFRESH_FAILED] Token 续期失败 (400): refresh_token_reused"
    )
    codex_meta = metadata_updates["k1"]["codex"]
    assert codex_meta["secondary_used_percent"] == 100.0


@pytest.mark.asyncio
async def test_codex_refresher_http_403_token_invalidated_marks_oauth_expired(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    response = _FakeResponse(
        status_code=403,
        payload={"error": {"message": "Authentication token has been invalidated."}},
    )
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "forbidden"
    assert result["status_code"] == 403
    assert result["auto_disabled"] is False
    assert str(state_updates["k1"]["oauth_invalid_reason"]).startswith("[OAUTH_EXPIRED]")


@pytest.mark.asyncio
async def test_codex_refresher_http_403_generic_marks_soft_request_failed(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    response = _FakeResponse(
        status_code=403,
        payload={"error": {"message": "Access forbidden for this account."}},
    )
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "forbidden"
    assert result["status_code"] == 403
    assert result["auto_disabled"] is False
    assert str(state_updates["k1"]["oauth_invalid_reason"]).startswith("[REQUEST_FAILED]")


@pytest.mark.asyncio
async def test_codex_refresher_success_updates_metadata(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc",
        auth_type="api_key",
        auth_config=None,
        proxy=None,
        oauth_invalid_at="old",
        oauth_invalid_reason="old",
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    monkeypatch.setattr(
        module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 10.0}
    )
    response = _FakeResponse(status_code=200, payload={"ok": True})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "success"
    assert metadata_updates == {"k1": {"codex": {"used_percent": 10.0}}}
    assert state_updates == {"k1": {"oauth_invalid_at": None, "oauth_invalid_reason": None}}


@pytest.mark.asyncio
async def test_codex_refresher_parse_error_is_diagnostic(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
    monkeypatch.setattr(
        module,
        "parse_codex_wham_usage_response",
        lambda _data: (_ for _ in ()).throw(
            CodexUsageParseError("rate_limit.primary_window 类型错误")
        ),
    )
    response = _FakeResponse(status_code=200, payload={"ok": True})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates={},
        state_updates={},
    )

    assert result["status"] == "error"
    assert "响应结构异常" in result["message"]
    assert "rate_limit.primary_window 类型错误" in result["message"]


@pytest.mark.asyncio
async def test_codex_refresher_oauth_missing_plan_type_adds_account_header(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc",
        auth_type="oauth",
        auth_config="enc-config",
        proxy=None,
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    response = _FakeResponse(status_code=200, payload={"ok": True})
    client_ref: dict[str, _FakeAsyncClient] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return SimpleNamespace(auth_header="Authorization", auth_value="Bearer oauth-token")

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(
        module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 1.0}
    )
    monkeypatch.setattr(
        module.crypto_service, "decrypt", lambda _v: json.dumps({"account_id": "acc-1"})
    )

    async def _fake_rust_response(**kwargs: Any) -> Any:
        client_ref["client"] = SimpleNamespace(last_headers=dict(kwargs["headers"]))
        return response

    monkeypatch.setattr(module, "_try_rust_codex_quota_response", _fake_rust_response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates={},
        state_updates={},
    )

    assert result["status"] == "success"
    assert client_ref["client"].last_headers is not None
    assert client_ref["client"].last_headers.get("chatgpt-account-id") == "acc-1"


@pytest.mark.asyncio
async def test_codex_refresher_oauth_uppercase_free_does_not_add_account_header(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc",
        auth_type="oauth",
        auth_config="enc-config",
        proxy=None,
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    response = _FakeResponse(status_code=200, payload={"ok": True})
    client_ref: dict[str, _FakeAsyncClient] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return SimpleNamespace(auth_header="Authorization", auth_value="Bearer oauth-token")

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(
        module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 1.0}
    )
    monkeypatch.setattr(
        module.crypto_service,
        "decrypt",
        lambda _v: json.dumps({"account_id": "acc-1", "plan_type": "FREE"}),
    )

    async def _fake_rust_response(**kwargs: Any) -> Any:
        client_ref["client"] = SimpleNamespace(last_headers=dict(kwargs["headers"]))
        return response

    monkeypatch.setattr(module, "_try_rust_codex_quota_response", _fake_rust_response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates={},
        state_updates={},
    )

    assert result["status"] == "success"
    assert client_ref["client"].last_headers is not None
    assert "chatgpt-account-id" not in client_ref["client"].last_headers


def test_parse_codex_usage_plan_type_case_insensitive_free_window_semantics() -> None:
    parsed = parse_codex_wham_usage_response(
        {
            "plan_type": "FREE",
            "rate_limit": {
                "primary_window": {
                    "used_percent": "12.5",
                    "reset_after_seconds": "120",
                    "reset_at": "1700000000",
                    "limit_window_seconds": "604800",
                }
            },
        }
    )
    assert parsed is not None
    assert parsed["plan_type"] == "free"
    assert parsed["primary_used_percent"] == 12.5
    assert parsed["primary_window_minutes"] == 10080
    assert "secondary_used_percent" not in parsed


def test_parse_codex_usage_missing_plan_type_infers_paid_windows() -> None:
    parsed = parse_codex_wham_usage_response(
        {
            "rate_limit": {
                "primary_window": {
                    "used_percent": 25,
                    "reset_after_seconds": 600,
                    "reset_at": 1700000000,
                    "limit_window_seconds": 18000,
                },
                "secondary_window": {
                    "used_percent": 80,
                    "reset_after_seconds": 3600,
                    "reset_at": 1700003600,
                    "limit_window_seconds": 604800,
                },
            }
        }
    )
    assert parsed is not None
    assert parsed["primary_used_percent"] == 80.0
    assert parsed["secondary_used_percent"] == 25.0
    assert parsed["primary_window_minutes"] == 10080
    assert parsed["secondary_window_minutes"] == 300


@pytest.mark.parametrize("plan_type", ["plus", "enterprise"])
def test_parse_codex_usage_paid_plan_windows_mapping(plan_type: str) -> None:
    parsed = parse_codex_wham_usage_response(
        {
            "plan_type": plan_type,
            "rate_limit": {
                "primary_window": {
                    "used_percent": 25,
                    "reset_after_seconds": 600,
                    "reset_at": 1700000000,
                    "limit_window_seconds": 18000,
                },
                "secondary_window": {
                    "used_percent": 80,
                    "reset_after_seconds": 3600,
                    "reset_at": 1700003600,
                    "limit_window_seconds": 604800,
                },
            },
        }
    )
    assert parsed is not None
    assert parsed["plan_type"] == plan_type
    assert parsed["primary_used_percent"] == 80.0
    assert parsed["secondary_used_percent"] == 25.0
    assert parsed["primary_window_minutes"] == 10080
    assert parsed["secondary_window_minutes"] == 300


def test_parse_codex_usage_blank_credits_balance_is_ignored() -> None:
    parsed = parse_codex_wham_usage_response(
        {
            "plan_type": "team",
            "rate_limit": {
                "primary_window": {
                    "used_percent": 10,
                    "reset_after_seconds": 100,
                    "reset_at": 1700000000,
                    "limit_window_seconds": 18000,
                }
            },
            "credits": {
                "has_credits": False,
                "balance": "",
                "unlimited": "false",
            },
        }
    )
    assert parsed is not None
    assert parsed["has_credits"] is False
    assert parsed["credits_unlimited"] is False
    assert "credits_balance" not in parsed


def test_parse_codex_usage_invalid_type_raises_diagnostic_error() -> None:
    with pytest.raises(CodexUsageParseError, match="rate_limit.primary_window 类型错误"):
        parse_codex_wham_usage_response({"rate_limit": {"primary_window": []}})


@pytest.mark.asyncio
async def test_antigravity_refresher_forbidden_collects_updates_without_commit(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import antigravity_refresher as module

    class _Forbidden(Exception):
        def __init__(self, reason: str) -> None:
            super().__init__(reason)
            self.reason = reason
            self.message = reason

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return SimpleNamespace(auth_value="Bearer tk", decrypted_auth_config={"pid": "p1"})

    async def _fetch_models_for_key(_ctx: Any, timeout_seconds: float) -> Any:
        _ = timeout_seconds
        raise _Forbidden("forbidden-by-test")

    class _UpstreamModelsFetchContext:  # noqa: D101
        def __init__(self, **kwargs: Any) -> None:
            self.kwargs = kwargs

    _install_module(
        monkeypatch,
        "src.services.model.upstream_fetcher",
        {
            "UpstreamModelsFetchContext": _UpstreamModelsFetchContext,
            "fetch_models_for_key": _fetch_models_for_key,
        },
    )
    _install_module(
        monkeypatch,
        "src.services.provider.adapters.antigravity.client",
        {"AntigravityAccountForbiddenException": _Forbidden},
    )
    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {"resolve_effective_proxy": lambda provider_proxy, key_proxy: None},
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)

    db = _FakeDB()
    provider = SimpleNamespace(proxy=None)
    key = SimpleNamespace(
        id="k1",
        name="K1",
        proxy=None,
        is_active=True,
        oauth_invalid_at=None,
        oauth_invalid_reason=None,
        upstream_metadata={},
    )
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    result = await refresh_antigravity_key_quota(
        db=cast(Any, db),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "forbidden"
    assert result["auto_disabled"] is False
    assert key.is_active is True
    assert key.oauth_invalid_reason is None
    assert "is_active" not in state_updates["k1"]
    assert state_updates["k1"]["oauth_invalid_reason"].startswith("账户访问被禁止")
    assert metadata_updates["k1"]["antigravity"]["is_forbidden"] is True
    assert db.commit_count == 0


@pytest.mark.asyncio
async def test_antigravity_refresher_success_resets_forbidden_flag(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import antigravity_refresher as module

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return SimpleNamespace(auth_value="Bearer tk", decrypted_auth_config={})

    async def _fetch_models_for_key(_ctx: Any, timeout_seconds: float) -> Any:
        _ = timeout_seconds
        return (
            [],
            [],
            True,
            {"antigravity": {"is_forbidden": True, "forbidden_reason": "x", "forbidden_at": 1}},
        )

    class _UpstreamModelsFetchContext:  # noqa: D101
        def __init__(self, **kwargs: Any) -> None:
            self.kwargs = kwargs

    _install_module(
        monkeypatch,
        "src.services.model.upstream_fetcher",
        {
            "UpstreamModelsFetchContext": _UpstreamModelsFetchContext,
            "fetch_models_for_key": _fetch_models_for_key,
        },
    )
    _install_module(
        monkeypatch,
        "src.services.provider.adapters.antigravity.client",
        {"AntigravityAccountForbiddenException": RuntimeError},
    )
    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {"resolve_effective_proxy": lambda provider_proxy, key_proxy: None},
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)

    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}
    result = await refresh_antigravity_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, SimpleNamespace(proxy=None)),
        key=cast(
            Any,
            SimpleNamespace(
                id="k1",
                name="K1",
                proxy=None,
                oauth_invalid_at="old",
                oauth_invalid_reason="old",
            ),
        ),
        endpoint=cast(Any, SimpleNamespace()),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "success"
    assert metadata_updates["k1"]["antigravity"]["is_forbidden"] is False
    assert metadata_updates["k1"]["antigravity"]["forbidden_reason"] is None
    assert metadata_updates["k1"]["antigravity"]["forbidden_at"] is None
    assert state_updates["k1"]["oauth_invalid_at"] is None
    assert state_updates["k1"]["oauth_invalid_reason"] is None


@pytest.mark.asyncio
async def test_kiro_refresher_runtime_401_marks_key_invalid(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class _Banned(Exception):
        pass

    async def _fetch_limits(auth_config: dict[str, Any], proxy_config: object) -> Any:
        _ = auth_config, proxy_config
        raise RuntimeError("401 token expired")

    def _parse_usage(_usage: Any) -> dict[str, Any]:
        return {"quota": 1}

    _install_module(
        monkeypatch,
        "src.services.provider.adapters.kiro.usage",
        {
            "KiroAccountBannedException": _Banned,
            "fetch_kiro_usage_limits": _fetch_limits,
            "parse_kiro_usage_response": _parse_usage,
        },
    )
    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {"resolve_effective_proxy": lambda provider_proxy, key_proxy: None},
    )
    monkeypatch.setattr(
        "src.services.provider_keys.quota_refresh.kiro_refresher.crypto_service.decrypt",
        lambda _v: "{}",
    )

    db = _FakeDB()
    key = SimpleNamespace(
        id="k1",
        name="K1",
        auth_config="enc",
        proxy=None,
        is_active=True,
        oauth_invalid_at=None,
        oauth_invalid_reason=None,
        upstream_metadata={},
    )
    state_updates: dict[str, dict[str, Any]] = {}

    result = await refresh_kiro_key_quota(
        db=cast(Any, db),
        provider=cast(Any, SimpleNamespace(proxy=None)),
        key=cast(Any, key),
        endpoint=None,
        codex_wham_usage_url="https://example.test",
        metadata_updates={},
        state_updates=state_updates,
    )

    assert result["status"] == "error"
    assert "401" in result["message"]
    assert key.is_active is True
    assert key.oauth_invalid_reason is None
    assert "is_active" not in state_updates["k1"]
    assert state_updates["k1"]["oauth_invalid_reason"] == "Kiro Token 无效或已过期"
    assert db.commit_count == 0


@pytest.mark.asyncio
async def test_kiro_refresher_success_updates_metadata_and_auth_config(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class _Banned(Exception):
        pass

    async def _fetch_limits(auth_config: dict[str, Any], proxy_config: object) -> Any:
        _ = auth_config, proxy_config
        return {"usage_data": {"x": 1}, "updated_auth_config": {"token": "new"}}

    def _parse_usage(_usage: Any) -> dict[str, Any]:
        return {"quota": 1}

    _install_module(
        monkeypatch,
        "src.services.provider.adapters.kiro.usage",
        {
            "KiroAccountBannedException": _Banned,
            "fetch_kiro_usage_limits": _fetch_limits,
            "parse_kiro_usage_response": _parse_usage,
        },
    )
    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {"resolve_effective_proxy": lambda provider_proxy, key_proxy: None},
    )
    monkeypatch.setattr(
        "src.services.provider_keys.quota_refresh.kiro_refresher.crypto_service.decrypt",
        lambda _v: json.dumps({"seed": 1}),
    )
    monkeypatch.setattr(
        "src.services.provider_keys.quota_refresh.kiro_refresher.crypto_service.encrypt",
        lambda raw: f"ENC:{raw}",
    )

    key = SimpleNamespace(
        id="k1",
        name="K1",
        auth_config="enc",
        proxy=None,
        oauth_invalid_at="old",
        oauth_invalid_reason="old",
        upstream_metadata={},
    )
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    result = await refresh_kiro_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, SimpleNamespace(proxy=None)),
        key=cast(Any, key),
        endpoint=None,
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "success"
    assert metadata_updates["k1"]["kiro"]["is_banned"] is False
    assert metadata_updates["k1"]["kiro"]["quota"] == 1
    assert key.auth_config == "enc"
    assert state_updates["k1"]["oauth_invalid_at"] is None
    assert state_updates["k1"]["oauth_invalid_reason"] is None
    assert state_updates["k1"]["auth_config"].startswith("ENC:")


@pytest.mark.asyncio
async def test_codex_refresher_http_402_workspace_deactivated_marks_account_block(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.services.provider_keys.quota_refresh import codex_refresher as module

    key = SimpleNamespace(
        id="k1",
        name="K1",
        api_key="enc-key",
        auth_type="oauth",
        auth_config="enc-config",
        proxy=None,
    )
    provider = SimpleNamespace(proxy=None)
    endpoint = SimpleNamespace()
    metadata_updates: dict[str, dict[str, Any]] = {}
    state_updates: dict[str, dict[str, Any]] = {}

    async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
        return None

    _install_module(
        monkeypatch,
        "src.services.proxy_node.resolver",
        {
            "resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
        },
    )
    monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
    monkeypatch.setattr(
        module.crypto_service,
        "decrypt",
        lambda value: (
            "sk-test"
            if value == "enc-key"
            else json.dumps({"plan_type": "team", "account_id": "acc-1"})
        ),
    )
    response = _FakeResponse(status_code=402, payload={"detail": {"code": "deactivated_workspace"}})
    _patch_rust_codex_response(monkeypatch, module, response)

    result = await refresh_codex_key_quota(
        db=cast(Any, _FakeDB()),
        provider=cast(Any, provider),
        key=cast(Any, key),
        endpoint=cast(Any, endpoint),
        codex_wham_usage_url="https://example.test",
        metadata_updates=metadata_updates,
        state_updates=state_updates,
    )

    assert result["status"] == "workspace_deactivated"
    assert result["status_code"] == 402
    assert metadata_updates["k1"]["codex"]["account_disabled"] is True
    assert metadata_updates["k1"]["codex"]["reason"] == "deactivated_workspace"
    assert state_updates["k1"]["oauth_invalid_at"] is not None
    assert str(state_updates["k1"]["oauth_invalid_reason"]).startswith("[ACCOUNT_BLOCK]")