Follow-up to the salvaged #111787 commit (@KoNit-K), same mechanism as #75578 (@adikpb). What: - Move the per-model cooldown logic out of the 2.7k-line credential_pool.py facade into agent/credential_pool_model_cooldowns.py (mixin + module helpers), keeping only the select/has_available/next_available_at hooks and the mark_exhausted_and_rotate branch in the facade. - The model cooldown uses the same TTL policy as a credential-wide 429 (_exhausted_ttl: provider reset_at wins, a sole credential keeps its 60s bench) instead of a flat 1h, so a single-credential user is never benched LONGER for the failed model than before. - Drop the `quota_scope == "account"` check: nothing in the tree produces that key. - Also keep billing_unverified 429s credential-wide, matching the credential-wide branch. - _rebind_primary_credential_pool read `rt` that was not in its scope (NameError on every post-fallback restore); pass primary_model from the caller instead. - resolve_anthropic_token(model=...) gates only model-aware callers; model-less diagnostics (usage display, model discovery) keep the key as before. - _anthropic_token_or_raise names the cooled model instead of claiming no credentials exist. - Simplify the salvaged call sites (unconditional select(model=)/resolve_anthropic_token(model=)), widen test stubs that lacked the new kwargs, trim the tests to two invariants on a real temp store. - Docs: credential-pools.md documents per-model Anthropic 429 cooldowns. Why: a generic Anthropic 429 is a per-model rate limit; benching the whole credential took every other Claude model offline while the env/borrowed token path handed the same benched key straight back (#111769, #61451). Co-authored-by: KoNit-K <124019182+KoNit-K@users.noreply.github.com> Co-authored-by: adikpb <67222969+adikpb@users.noreply.github.com>
235 lines
7.7 KiB
Python
235 lines
7.7 KiB
Python
"""Credential pools must never cross provider or custom-endpoint boundaries."""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from agent.credential_pool import (
|
|
credential_pool_matches_provider,
|
|
resolve_runtime_pool_key,
|
|
)
|
|
from hermes_cli import runtime_provider as rp
|
|
|
|
|
|
def test_provider_match_requires_exact_non_custom_identity():
|
|
assert credential_pool_matches_provider("deepseek", "deepseek")
|
|
assert not credential_pool_matches_provider("openai-codex", "deepseek")
|
|
assert not credential_pool_matches_provider("", "deepseek")
|
|
|
|
|
|
def test_custom_pool_match_is_scoped_by_endpoint():
|
|
with patch(
|
|
"agent.credential_pool.get_custom_provider_pool_key",
|
|
return_value="custom:lab",
|
|
):
|
|
assert credential_pool_matches_provider(
|
|
"custom:lab", "custom", base_url="https://lab.example/v1"
|
|
)
|
|
assert not credential_pool_matches_provider(
|
|
"custom:other", "custom", base_url="https://lab.example/v1"
|
|
)
|
|
|
|
|
|
def test_named_custom_pool_match_requires_configured_identity_and_endpoint():
|
|
configured = [
|
|
(
|
|
"gemini-display",
|
|
{
|
|
"name": "Gemini Display",
|
|
"provider_key": "gemini-no-filter",
|
|
"base_url": "https://generativelanguage.googleapis.com/v1beta/",
|
|
},
|
|
)
|
|
]
|
|
with patch("agent.credential_pool._iter_custom_providers", return_value=configured):
|
|
assert credential_pool_matches_provider(
|
|
"custom:gemini-display",
|
|
"gemini-no-filter",
|
|
base_url="https://generativelanguage.googleapis.com/v1beta",
|
|
)
|
|
assert credential_pool_matches_provider(
|
|
"custom:gemini-display",
|
|
"custom:gemini-no-filter",
|
|
base_url="https://generativelanguage.googleapis.com/v1beta",
|
|
)
|
|
assert not credential_pool_matches_provider(
|
|
"custom:gemini-display",
|
|
"gemini-no-filter",
|
|
base_url="https://fallback.example/v1",
|
|
)
|
|
assert not credential_pool_matches_provider(
|
|
"custom:gemini-display",
|
|
"custom:gemini-no-filter",
|
|
base_url="https://fallback.example/v1",
|
|
)
|
|
assert not credential_pool_matches_provider(
|
|
"custom:gemini-display",
|
|
"other-provider",
|
|
base_url="https://generativelanguage.googleapis.com/v1beta",
|
|
)
|
|
|
|
|
|
def test_runtime_pool_key_resolves_all_custom_runtime_identities():
|
|
endpoint = "https://generativelanguage.googleapis.com/v1beta"
|
|
configured = [
|
|
(
|
|
"sibling-display",
|
|
{
|
|
"name": "Sibling Display",
|
|
"provider_key": "sibling-provider",
|
|
"base_url": endpoint,
|
|
},
|
|
),
|
|
(
|
|
"gemini-display",
|
|
{
|
|
"name": "Gemini Display",
|
|
"provider_key": "gemini-no-filter",
|
|
"base_url": endpoint,
|
|
},
|
|
)
|
|
]
|
|
with patch("agent.credential_pool._iter_custom_providers", return_value=configured):
|
|
assert resolve_runtime_pool_key("custom", endpoint) == "sibling-provider"
|
|
assert (
|
|
resolve_runtime_pool_key("gemini-no-filter", endpoint)
|
|
== "gemini-no-filter"
|
|
)
|
|
assert (
|
|
resolve_runtime_pool_key("custom:gemini-no-filter", endpoint)
|
|
== "gemini-no-filter"
|
|
)
|
|
assert (
|
|
resolve_runtime_pool_key(
|
|
"gemini-no-filter",
|
|
"https://fallback.example/v1",
|
|
)
|
|
== "gemini-no-filter"
|
|
)
|
|
|
|
|
|
def test_runtime_pool_key_resolves_modern_provider_in_mixed_config():
|
|
endpoint = "https://generativelanguage.googleapis.com/v1beta"
|
|
config = {
|
|
"custom_providers": [
|
|
{
|
|
"name": "Legacy Provider",
|
|
"base_url": "https://legacy.example/v1",
|
|
}
|
|
],
|
|
"providers": {
|
|
"gemini-no-filter": {
|
|
"name": "Gemini Display",
|
|
"api": endpoint,
|
|
}
|
|
},
|
|
}
|
|
|
|
with patch("agent.credential_pool._load_config_safe", return_value=config):
|
|
assert (
|
|
resolve_runtime_pool_key("gemini-no-filter", endpoint)
|
|
== "gemini-no-filter"
|
|
)
|
|
assert (
|
|
resolve_runtime_pool_key("custom:gemini-no-filter", endpoint)
|
|
== "gemini-no-filter"
|
|
)
|
|
assert (
|
|
resolve_runtime_pool_key(
|
|
"custom:gemini-no-filter",
|
|
"https://fallback.example/v1",
|
|
)
|
|
== "custom:gemini-no-filter"
|
|
)
|
|
|
|
|
|
def test_keyed_provider_pool_matches_runtime_aliases():
|
|
configured = [
|
|
(
|
|
"b.ai",
|
|
{
|
|
"name": "B.AI",
|
|
"provider_key": "b-ai",
|
|
"base_url": "https://api.b.ai/v1",
|
|
},
|
|
)
|
|
]
|
|
with patch("agent.credential_pool._iter_custom_providers", return_value=configured):
|
|
assert credential_pool_matches_provider(
|
|
"b-ai", "b-ai", base_url="https://api.b.ai/v1"
|
|
)
|
|
assert credential_pool_matches_provider(
|
|
"b-ai", "custom", base_url="https://api.b.ai/v1"
|
|
)
|
|
assert credential_pool_matches_provider(
|
|
"b-ai", "custom:b.ai", base_url="https://api.b.ai/v1"
|
|
)
|
|
assert not credential_pool_matches_provider(
|
|
"b-ai", "custom", base_url="https://other.example/v1"
|
|
)
|
|
assert not credential_pool_matches_provider(
|
|
"b-ai", "deepseek", base_url="https://api.b.ai/v1"
|
|
)
|
|
|
|
|
|
def test_runtime_pool_key_prefers_durable_provider_slug():
|
|
endpoint = "https://api.b.ai/v1"
|
|
configured = [
|
|
(
|
|
"b.ai",
|
|
{
|
|
"name": "B.AI",
|
|
"provider_key": "b-ai",
|
|
"base_url": endpoint,
|
|
},
|
|
)
|
|
]
|
|
with patch("agent.credential_pool._iter_custom_providers", return_value=configured):
|
|
assert resolve_runtime_pool_key("b-ai", endpoint) == "b-ai"
|
|
assert resolve_runtime_pool_key("custom", endpoint) == "b-ai"
|
|
assert resolve_runtime_pool_key("custom:b.ai", endpoint) == "b-ai"
|
|
|
|
|
|
def test_runtime_pool_key_preserves_non_custom_identity():
|
|
with patch("agent.credential_pool._iter_custom_providers", return_value=[]):
|
|
assert (
|
|
resolve_runtime_pool_key("openai-codex", "https://chatgpt.com/backend-api")
|
|
== "openai-codex"
|
|
)
|
|
|
|
|
|
def test_runtime_ignores_pool_loaded_for_different_provider(monkeypatch):
|
|
entry = SimpleNamespace(
|
|
provider="openai-codex",
|
|
access_token="wrong-token",
|
|
runtime_api_key="wrong-token",
|
|
runtime_base_url="https://chatgpt.com/backend-api/codex",
|
|
base_url="https://chatgpt.com/backend-api/codex",
|
|
)
|
|
pool = SimpleNamespace(
|
|
provider="openai-codex",
|
|
has_credentials=lambda: True,
|
|
select=lambda **_kwargs: entry,
|
|
)
|
|
monkeypatch.setattr(rp, "load_pool", lambda _provider: pool)
|
|
monkeypatch.setattr(rp, "resolve_provider", lambda *_a, **_kw: "deepseek")
|
|
monkeypatch.setattr(
|
|
rp,
|
|
"_get_model_config",
|
|
lambda: {"provider": "deepseek", "default": "deepseek-chat"},
|
|
)
|
|
monkeypatch.setattr(
|
|
rp,
|
|
"resolve_api_key_provider_credentials",
|
|
lambda _provider: {
|
|
"provider": "deepseek",
|
|
"api_key": "deepseek-key",
|
|
"base_url": "https://api.deepseek.com/v1",
|
|
"source": "env",
|
|
},
|
|
)
|
|
|
|
resolved = rp.resolve_runtime_provider(requested="deepseek")
|
|
|
|
assert resolved["provider"] == "deepseek"
|
|
assert resolved["api_key"] == "deepseek-key"
|
|
assert resolved["base_url"] == "https://api.deepseek.com/v1" |