mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge pull request #169 from AAEE86/master
feat: 调度器根据上游配额自动跳过已耗尽的 Key (Kiro/Codex/Antigravity)
This commit is contained in:
fawney19
128
src/services/cache/aware_scheduler.py
vendored
128
src/services/cache/aware_scheduler.py
vendored
@@ -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]}... 模型名匹配: "
|
||||
f"model={model_name} -> {mapping_matched_model}, "
|
||||
f"allowed_models={key.allowed_models}"
|
||||
"[Scheduler] Key {}... 模型名匹配: model={} -> {}, allowed_models={}",
|
||||
key.id[:8],
|
||||
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),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
225
tests/services/test_aware_scheduler_quota_skipping.py
Normal file
225
tests/services/test_aware_scheduler_quota_skipping.py
Normal file
@@ -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%"
|
||||
Reference in New Issue
Block a user