fix(auth): compare pool-identity callers against all candidate keys

Two callers of get_custom_provider_pool_key compared against its single
preferred key and broke when the pool held the other identity:

- _prune_replaced_custom_model_config_credentials skipped only the
  preferred key, so a keyed provider's own legacy-named pool
  (custom:b.ai) was false-pruned of its current model_config credential
  when the preferred key resolved to the bare slug (b-ai).
- _seed_custom_pool seeded only when the pool key equaled the preferred
  key, so a legacy-named pool stopped being seeded from model.api_key.

Both now compare against the full custom_provider_pool_key_candidates
set. Also drops a redundant get_custom_provider_pool_key call from
_try_resolve_from_custom_pool (it returned candidates[0], doubling the
config traversal) and updates the two test files that monkeypatched the
removed module attribute.

Follow-up to #100413.
This commit is contained in:
kshitijk4poor
2026-09-01 21:53:46 +05:30
committed by kshitij
parent 0bee5ff408
commit b81383ec21
6 changed files with 146 additions and 17 deletions
+12 -3
View File
@@ -3596,9 +3596,18 @@ def _seed_custom_pool(pool_key: str, entries: List[PooledCredential]) -> Tuple[b
model_api_key = v.strip()
break
if model_provider == "custom" and model_base_url and model_api_key:
# Check if this model's base_url matches our custom provider
matched_key = get_custom_provider_pool_key(model_base_url)
if matched_key == pool_key:
# Check if this model's base_url matches our custom provider.
# The pool may be keyed under either the durable
# ``providers.<key>`` slug or the legacy ``custom:<name>``
# namespace, so accept the match against any candidate —
# comparing against the single preferred key silently skips
# seeding when the pool holds the other identity (verified
# regression from PR #100413 review).
matched_keys = {
str(key).strip().lower()
for key in custom_provider_pool_key_candidates(model_base_url)
}
if pool_key in matched_keys:
source = "model_config"
if not _is_suppressed(pool_key, source):
active_sources.add(source)
+15 -7
View File
@@ -100,15 +100,23 @@ def _prune_replaced_custom_model_config_credentials(
try:
from agent.credential_pool import (
CUSTOM_POOL_PREFIX,
get_custom_provider_pool_key,
custom_provider_pool_key_candidates,
)
from hermes_cli.auth import read_credential_pool, write_credential_pool
active_pool_key = get_custom_provider_pool_key(
base_url,
provider_name=provider_name or None,
)
if not active_pool_key:
# A keyed ``providers.<key>`` endpoint stores under the durable slug
# while legacy-named pools keep the ``custom:<display-name>`` key, so
# every identity the active endpoint may occupy must be skipped —
# comparing against a single preferred key false-prunes the provider's
# own legacy-named pool (verified regression from PR #100413 review).
active_pool_keys = {
str(key).strip().lower()
for key in custom_provider_pool_key_candidates(
base_url,
provider_name=provider_name or None,
)
}
if not active_pool_keys:
return
pools = read_credential_pool(None)
if not isinstance(pools, dict):
@@ -117,7 +125,7 @@ def _prune_replaced_custom_model_config_credentials(
if (
not isinstance(pool_key, str)
or not pool_key.startswith(CUSTOM_POOL_PREFIX)
or pool_key == active_pool_key
or pool_key in active_pool_keys
or not isinstance(entries, list)
):
continue
-2
View File
@@ -16,7 +16,6 @@ from agent.credential_pool import (
PooledCredential,
credential_pool_matches_provider,
custom_provider_pool_key_candidates,
get_custom_provider_pool_key,
load_pool,
)
from agent.secret_scope import get_secret as _get_secret
@@ -685,7 +684,6 @@ def _try_resolve_from_custom_pool(
_add(key)
except Exception:
pass
_add(get_custom_provider_pool_key(base_url, provider_name=provider_name))
if not candidates:
return None
@@ -146,7 +146,7 @@ class TestCustomProviderUrlFallback:
def select(self):
return _Entry()
monkeypatch.setattr(rp, "get_custom_provider_pool_key", lambda *a, **k: "custom:my-claude")
monkeypatch.setattr(rp, "custom_provider_pool_key_candidates", lambda *a, **k: ["custom:my-claude"])
monkeypatch.setattr(rp, "load_pool", lambda key: _Pool())
resolved = rp._try_resolve_from_custom_pool(
@@ -100,3 +100,117 @@ def test_keyed_provider_runtime_falls_back_to_legacy_custom_namespace(
resolved = rp.resolve_runtime_provider(requested="b-ai")
assert resolved["api_key"] == LEGACY_KEY
assert resolved["api_key"] != "no-key-required"
def _model_config_entry(entry_id, token):
return {
"id": entry_id,
"label": "model_config",
"auth_type": "api_key",
"priority": 0,
"source": "model_config",
"access_token": token,
}
def test_prune_keeps_active_legacy_pool_for_keyed_provider(tmp_path, monkeypatch):
"""A keyed provider's own legacy-named pool must not be false-pruned.
Regression: with keys stored under ``custom:b.ai`` while the provider is
configured as ``providers.b-ai``, the active pool key resolves to the
durable slug ``b-ai``; comparing with ``==`` let the prune strip the
provider's own current credential from ``custom:b.ai``.
"""
config = {
"model": {"default": "b-ai-model", "provider": "b-ai"},
"providers": {
"b-ai": {
"name": "B.AI",
"base_url": "https://api.b.ai/v1",
}
},
}
pools = {
# legacy-named pool for the (now keyed) b.ai provider holding the
# credential seeded from model.api_key — this is the ACTIVE pool
"custom:b.ai": [_model_config_entry("mc1", "sk-current-b-ai-key")],
# an unrelated stale pool that SHOULD be pruned
"custom:old-endpoint": [_model_config_entry("mc2", "sk-stale-key")],
}
hermes_home = tmp_path / ".hermes"
hermes_home.mkdir(exist_ok=True)
monkeypatch.setenv("HERMES_HOME", str(hermes_home))
(hermes_home / "config.yaml").write_text(yaml.safe_dump(config), encoding="utf-8")
(hermes_home / "auth.json").write_text(
json.dumps(
{
"version": 1,
"providers": {},
"credential_pool": pools,
}
),
encoding="utf-8",
)
from hermes_cli.model_setup_flows import (
_prune_replaced_custom_model_config_credentials,
)
_prune_replaced_custom_model_config_credentials(
"https://api.b.ai/v1", provider_name="B.AI"
)
after = json.loads(
(tmp_path / ".hermes" / "auth.json").read_text(encoding="utf-8")
)
kept = (after.get("credential_pool") or {}).get("custom:b.ai")
pruned = (after.get("credential_pool") or {}).get("custom:old-endpoint")
assert kept, "active provider's own legacy-named pool must keep its credential"
assert kept[0]["access_token"] == "sk-current-b-ai-key"
assert pruned == [], "stale pool for a different endpoint must be pruned"
def test_seed_custom_pool_matches_legacy_named_pool(tmp_path, monkeypatch):
"""A legacy-named pool must still seed from model.api_key.
Regression: with model.provider 'custom' pointing at a keyed provider's
base_url, the pool key ``custom:b.ai`` no longer equals the preferred
candidate (the slug ``b-ai``), silently skipping the model_config seed.
"""
config = {
"model": {
"default": "b-ai-model",
"provider": "custom",
"base_url": "https://api.b.ai/v1",
"api_key": "sk-model-config-key",
},
"providers": {
"b-ai": {
"name": "B.AI",
"base_url": "https://api.b.ai/v1",
}
},
}
hermes_home = tmp_path / ".hermes"
hermes_home.mkdir(exist_ok=True)
monkeypatch.setenv("HERMES_HOME", str(hermes_home))
(hermes_home / "config.yaml").write_text(yaml.safe_dump(config), encoding="utf-8")
(hermes_home / "auth.json").write_text(
json.dumps(
{
"version": 1,
"providers": {},
"credential_pool": {"custom:b.ai": []},
}
),
encoding="utf-8",
)
from agent.credential_pool import load_pool
pool = load_pool("custom:b.ai")
entries = pool.entries() if hasattr(pool, "entries") else []
seeded = [e for e in entries if getattr(e, "source", "") == "model_config"]
assert seeded, "legacy-named pool must still seed model_config from model.api_key"
assert getattr(seeded[0], "access_token", "") == "sk-model-config-key"
@@ -115,7 +115,7 @@ class TestCustomProviderPoolLoopbackNoKeyExemption:
('123') for a local Ollama endpoint must resolve to the same
"no-key-required" placeholder every other local no-auth path uses,
not the raw unusable value."""
monkeypatch.setattr(rp, "get_custom_provider_pool_key", lambda base_url, provider_name=None: "custom:local-ollama")
monkeypatch.setattr(rp, "custom_provider_pool_key_candidates", lambda base_url, provider_name=None: ["custom:local-ollama"])
monkeypatch.setattr(rp, "load_pool", lambda pool_key: self._pool_with("123"))
result = rp._try_resolve_from_custom_pool("http://localhost:11434/v1", "custom", None)
@@ -124,7 +124,7 @@ class TestCustomProviderPoolLoopbackNoKeyExemption:
assert result["api_key"] == "no-key-required"
def test_single_char_placeholder_key_also_exempted(self, monkeypatch):
monkeypatch.setattr(rp, "get_custom_provider_pool_key", lambda base_url, provider_name=None: "custom:local")
monkeypatch.setattr(rp, "custom_provider_pool_key_candidates", lambda base_url, provider_name=None: ["custom:local"])
monkeypatch.setattr(rp, "load_pool", lambda pool_key: self._pool_with("m"))
result = rp._try_resolve_from_custom_pool("http://127.0.0.1:11434/v1", "custom", None)
@@ -136,7 +136,7 @@ class TestCustomProviderPoolLoopbackNoKeyExemption:
remote endpoint with a genuinely-too-short key must NOT get a
free pass. The short value passes through unchanged, so the
downstream has_usable_secret() gate still catches it."""
monkeypatch.setattr(rp, "get_custom_provider_pool_key", lambda base_url, provider_name=None: "custom:remote")
monkeypatch.setattr(rp, "custom_provider_pool_key_candidates", lambda base_url, provider_name=None: ["custom:remote"])
monkeypatch.setattr(rp, "load_pool", lambda pool_key: self._pool_with("xy"))
result = rp._try_resolve_from_custom_pool("https://api.remote-vendor.example/v1", "custom", None)
@@ -147,7 +147,7 @@ class TestCustomProviderPoolLoopbackNoKeyExemption:
"""Sanity: a genuinely usable key for a loopback endpoint (a real
API key happens to be configured for a local proxy, say) must not
be silently overwritten."""
monkeypatch.setattr(rp, "get_custom_provider_pool_key", lambda base_url, provider_name=None: "custom:local")
monkeypatch.setattr(rp, "custom_provider_pool_key_candidates", lambda base_url, provider_name=None: ["custom:local"])
monkeypatch.setattr(rp, "load_pool", lambda pool_key: self._pool_with("sk-genuinely-long-real-key-12345"))
result = rp._try_resolve_from_custom_pool("http://localhost:11434/v1", "custom", None)