feat: 调度器根据上游配额自动跳过已耗尽的 Key (Kiro/Codex/Antigravity)

在 _check_key_availability 末尾增加 _is_key_quota_exhausted 检查，读取 ProviderAPIKey.upstream_metadata 中各 Provider 的配额信息，配额耗尽时跳过该 Key 并返回可读的跳过原因。同时修正同函数内 logger 调用为 loguru {} 占位符风格。

src/services/cache/aware_scheduler.py

@@ -54,6 +54,7 @@ from src.core.model_permissions import (
     get_allowed_models_preview,
     merge_allowed_models,
 )
+from src.core.provider_types import ProviderType, normalize_provider_type
 from src.models.database import (
     ApiKey,
     Model,
@@ -133,6 +134,106 @@ class ProviderCandidate:
         return self._stable_order_key() < other._stable_order_key()
 
 
+# ---------------------------------------------------------------------------
+# Quota-based skipping
+# ---------------------------------------------------------------------------
+
+
+def _pct_is_exhausted(value: object) -> bool:
+    """Return True when used_percent indicates 0% remaining."""
+    try:
+        pct = float(value)  # type: ignore[arg-type]
+    except (TypeError, ValueError):
+        return False
+    # Some upstreams may return values slightly above 100 due to rounding.
+    return pct >= 100.0 - 1e-6
+
+
+def _float_or_none(value: object) -> float | None:
+    try:
+        if value is None:
+            return None
+        return float(value)  # type: ignore[arg-type]
+    except (TypeError, ValueError):
+        return None
+
+
+def _is_key_quota_exhausted(
+    provider_type: str | None,
+    key: ProviderAPIKey,
+    *,
+    model_name: str,
+) -> tuple[bool, str | None]:
+    """Check ProviderAPIKey.upstream_metadata quota and decide whether to skip.
+
+    Requirements:
+    - Kiro: account-level quota. When remaining == 0, skip this key; allow again when remaining > 0.
+    - Codex: only consider weekly quota + 5H quota (ignore code review quota).
+             If either remaining is 0%, skip this key.
+    - Antigravity: quota is per-model; do not disable the account.
+                   When the requested model's quota is 0%, skip this key.
+    """
+    pt = normalize_provider_type(provider_type)
+
+    upstream = getattr(key, "upstream_metadata", None) or {}
+    if not isinstance(upstream, dict):
+        return False, None
+
+    if pt == ProviderType.KIRO:
+        kiro_meta = upstream.get("kiro")
+        if not isinstance(kiro_meta, dict):
+            return False, None
+
+        remaining = _float_or_none(kiro_meta.get("remaining"))
+
+        if remaining is not None and remaining <= 0.0:
+            return True, "Kiro 账号配额剩余 0"
+
+        return False, None
+
+    if pt == ProviderType.CODEX:
+        codex_meta = upstream.get("codex")
+        if not isinstance(codex_meta, dict):
+            return False, None
+
+        weekly_used = codex_meta.get("primary_used_percent")
+        five_hour_used = codex_meta.get("secondary_used_percent")
+
+        exhausted_parts: list[str] = []
+        if _pct_is_exhausted(weekly_used):
+            exhausted_parts.append("周限额剩余 0%")
+        if _pct_is_exhausted(five_hour_used):
+            exhausted_parts.append("5H 限额剩余 0%")
+
+        if exhausted_parts:
+            return True, "Codex " + "，".join(exhausted_parts)
+
+        return False, None
+
+    if pt == ProviderType.ANTIGRAVITY:
+        ag_meta = upstream.get("antigravity")
+        if not isinstance(ag_meta, dict):
+            return False, None
+
+        quota_by_model = ag_meta.get("quota_by_model")
+        if not isinstance(quota_by_model, dict):
+            return False, None
+
+        model_quota = quota_by_model.get(model_name)
+        if not isinstance(model_quota, dict):
+            return False, None
+
+        remaining_fraction = _float_or_none(model_quota.get("remaining_fraction"))
+        if remaining_fraction is not None and remaining_fraction <= 0.0:
+            return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
+        if _pct_is_exhausted(model_quota.get("used_percent")):
+            return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
+
+        return False, None
+
+    return False, None
+
+
 @dataclass
 class ConcurrencySnapshot:
     key_current: int
@@ -991,6 +1092,8 @@ class CacheAwareScheduler:
         capability_requirements: dict[str, bool] | None = None,
         model_mappings: list[str] | None = None,
         candidate_models: set[str] | None = None,
+        *,
+        provider_type: str | None = None,
     ) -> tuple[bool, str | None, str | None]:
         """
         检查 API Key 的可用性
@@ -1030,22 +1133,24 @@ class CacheAwareScheduler:
             )
             if mapping_matched_model:
                 logger.debug(
-                    f"[Scheduler] Key {key.id[:8]}... 模型名匹配: "
+                    "[Scheduler] Key {}... 模型名匹配: model={} -> {}, allowed_models={}",
-                    f"model={model_name} -> {mapping_matched_model}, "
+                    key.id[:8],
-                    f"allowed_models={key.allowed_models}"
+                    model_name,
+                    mapping_matched_model,
+                    key.allowed_models,
                 )
         except TimeoutError:
             # 正则匹配超时（可能是 ReDoS 攻击或复杂模式）
-            logger.warning(f"映射匹配超时: key_id={key.id}, model={model_name}")
+            logger.warning("映射匹配超时: key_id={}, model={}", key.id, model_name)
             return False, "映射匹配超时，请简化配置", None
         except re.error as e:
             # 正则语法错误（配置问题）
-            logger.warning(f"映射规则无效: key_id={key.id}, model={model_name}, error={e}")
+            logger.warning("映射规则无效: key_id={}, model={}, error={}", key.id, model_name, e)
             return False, f"映射规则无效: {str(e)}", None
         except Exception as e:
             # 其他未知异常
             logger.error(
-                f"映射匹配异常: key_id={key.id}, model={model_name}, error={e}", exc_info=True
+                "映射匹配异常: key_id={}, model={}, error={}", key.id, model_name, e, exc_info=True
             )
             # 异常时保守处理：不允许使用该 Key
             return False, "映射匹配失败", None
@@ -1066,6 +1171,16 @@ class CacheAwareScheduler:
         if not is_match:
             return False, skip_reason, None
 
+        effective_model_name = mapping_matched_model or model_name
+
+        quota_exhausted, quota_reason = _is_key_quota_exhausted(
+            provider_type,
+            key,
+            model_name=effective_model_name,
+        )
+        if quota_exhausted:
+            return False, quota_reason, mapping_matched_model
+
         return True, None, mapping_matched_model
 
     async def _build_candidates(
@@ -1281,6 +1396,7 @@ class CacheAwareScheduler:
                             capability_requirements,
                             model_mappings=model_mappings,
                             candidate_models=provider_model_names,
+                            provider_type=getattr(provider, "provider_type", None),
                         )
                     )
 

tests/services/test_aware_scheduler_quota_skipping.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+from src.services.cache.aware_scheduler import CacheAwareScheduler
+
+
+def _make_key(
+    *,
+    upstream_metadata: dict,
+    allowed_models: list[str] | None = None,
+    capabilities: dict[str, bool] | None = None,
+) -> MagicMock:
+    key = MagicMock()
+    key.id = "k1234567890"
+    key.allowed_models = allowed_models
+    key.capabilities = capabilities or {}
+    key.upstream_metadata = upstream_metadata
+    return key
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_kiro_quota_remaining_zero_skips(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(upstream_metadata={"kiro": {"remaining": 0.0}})
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="openai:chat",
+        model_name="any-model",
+        provider_type="kiro",
+    )
+
+    assert ok is False
+    assert reason == "Kiro 账号配额剩余 0"
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_kiro_quota_remaining_positive_allows(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(upstream_metadata={"kiro": {"remaining": 1.0}})
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="openai:chat",
+        model_name="any-model",
+        provider_type="kiro",
+    )
+
+    assert ok is True
+    assert reason is None
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_codex_weekly_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(
+        upstream_metadata={
+            "codex": {
+                "primary_used_percent": 100.0,
+                "secondary_used_percent": 10.0,
+                "code_review_used_percent": 100.0,
+            }
+        }
+    )
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="openai:cli",
+        model_name="any-model",
+        provider_type="codex",
+    )
+
+    assert ok is False
+    assert reason == "Codex 周限额剩余 0%"
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_codex_5h_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(
+        upstream_metadata={
+            "codex": {
+                "primary_used_percent": 10.0,
+                "secondary_used_percent": 100.0,
+            }
+        }
+    )
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="openai:cli",
+        model_name="any-model",
+        provider_type="codex",
+    )
+
+    assert ok is False
+    assert reason == "Codex 5H 限额剩余 0%"
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_codex_ignores_code_review_quota(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(
+        upstream_metadata={
+            "codex": {
+                "primary_used_percent": 10.0,
+                "secondary_used_percent": 20.0,
+                "code_review_used_percent": 100.0,
+            }
+        }
+    )
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="openai:cli",
+        model_name="any-model",
+        provider_type="codex",
+    )
+
+    assert ok is True
+    assert reason is None
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_antigravity_model_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(
+        upstream_metadata={
+            "antigravity": {
+                "quota_by_model": {
+                    "ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
+                    "other": {"remaining_fraction": 1.0, "used_percent": 0.0},
+                }
+            }
+        },
+    )
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="gemini:chat",
+        model_name="ag-model",
+        provider_type="antigravity",
+    )
+
+    assert ok is False
+    assert reason == "Antigravity 模型 ag-model 配额剩余 0%"
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_antigravity_other_model_not_exhausted_allows(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+    key = _make_key(
+        upstream_metadata={
+            "antigravity": {
+                "quota_by_model": {
+                    "ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
+                    "other": {"remaining_fraction": 1.0, "used_percent": 0.0},
+                }
+            }
+        },
+    )
+
+    ok, reason, _mapped = scheduler._check_key_availability(
+        key,
+        api_format="gemini:chat",
+        model_name="other",
+        provider_type="antigravity",
+    )
+
+    assert ok is True
+    assert reason is None
+
+
+@patch(
+    "src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
+    return_value=(True, None),
+)
+def test_antigravity_quota_uses_mapping_matched_model(_mock_cb: MagicMock) -> None:
+    scheduler = CacheAwareScheduler()
+
+    # Request uses GlobalModel.name, but allowed_models only contains provider-side model id.
+    key = _make_key(
+        upstream_metadata={
+            "antigravity": {
+                "quota_by_model": {
+                    "ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
+                }
+            }
+        },
+        allowed_models=["ag-model"],
+    )
+
+    ok, reason, mapped = scheduler._check_key_availability(
+        key,
+        api_format="gemini:chat",
+        model_name="global-model",
+        model_mappings=["ag-.*"],
+        provider_type="antigravity",
+    )
+
+    assert mapped == "ag-model"
+    assert ok is False
+    assert reason == "Antigravity 模型 ag-model 配额剩余 0%"