From 791583704b5000446d40f884117fb4a8a7c71298 Mon Sep 17 00:00:00 2001 From: yoma Date: Sun, 5 Jul 2026 01:44:13 +0800 Subject: [PATCH] fix(auth): prune stale custom model credentials --- hermes_cli/model_setup_flows.py | 58 +++++++++++++ .../test_custom_provider_model_switch.py | 84 +++++++++++++++++++ 2 files changed, 142 insertions(+) diff --git a/hermes_cli/model_setup_flows.py b/hermes_cli/model_setup_flows.py index b6769b69d86..a234fbbee19 100644 --- a/hermes_cli/model_setup_flows.py +++ b/hermes_cli/model_setup_flows.py @@ -27,6 +27,59 @@ import subprocess from hermes_cli.config import clear_model_endpoint_credentials +def _prune_replaced_custom_model_config_credentials( + base_url: str, + *, + provider_name: str = "", +) -> None: + """Drop stale ``model_config`` credentials from inactive custom pools. + + ``model_config`` means "the credential currently stored under + ``model.api_key``". After an explicit custom-endpoint switch, any old + custom pool still carrying that source points at the previous endpoint and + can be selected before the freshly saved config is tried. + """ + try: + from agent.credential_pool import ( + CUSTOM_POOL_PREFIX, + get_custom_provider_pool_key, + ) + 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: + return + pools = read_credential_pool(None) + if not isinstance(pools, dict): + return + for pool_key, entries in pools.items(): + if ( + not isinstance(pool_key, str) + or not pool_key.startswith(CUSTOM_POOL_PREFIX) + or pool_key == active_pool_key + or not isinstance(entries, list) + ): + continue + retained = [] + removed_ids = [] + changed = False + for entry in entries: + if isinstance(entry, dict) and entry.get("source") == "model_config": + changed = True + entry_id = entry.get("id") + if entry_id: + removed_ids.append(str(entry_id)) + continue + retained.append(entry) + if changed: + write_credential_pool(pool_key, retained, removed_ids=removed_ids) + except Exception: + return + + def _prompt_auth_credentials_choice(title: str) -> str: """Prompt for reuse / reauthenticate / cancel with the standard radio UI. @@ -942,6 +995,11 @@ def _model_flow_custom(config): name=display_name, api_mode=api_mode, ) + _prune_replaced_custom_model_config_credentials( + effective_url, + provider_name=display_name, + ) + def _model_flow_azure_foundry(config, current_model=""): """Azure Foundry provider: configure endpoint, auth mode, API mode, and model. diff --git a/tests/hermes_cli/test_custom_provider_model_switch.py b/tests/hermes_cli/test_custom_provider_model_switch.py index 4dfd019ed68..e778a30b2ca 100644 --- a/tests/hermes_cli/test_custom_provider_model_switch.py +++ b/tests/hermes_cli/test_custom_provider_model_switch.py @@ -32,6 +32,90 @@ def config_home(tmp_path, monkeypatch): class TestCustomProviderModelSwitch: """Ensure _model_flow_named_custom always probes and shows menu.""" + def test_custom_endpoint_switch_prunes_stale_model_config_pool_entry( + self, + config_home, + ): + """Switching custom endpoints must not leave the old model.api_key + credential selectable from the previous endpoint's pool.""" + import yaml + from agent.credential_pool import load_pool + from hermes_cli.auth import read_credential_pool, write_credential_pool + from hermes_cli.main import _model_flow_custom + + config_path = config_home / "config.yaml" + config_path.write_text( + "model:\n" + " default: old-model\n" + " provider: custom\n" + " base_url: https://old.example.test/v1\n" + " api_key: sk-old-model-config\n" + "custom_providers:\n" + "- name: Old Endpoint\n" + " base_url: https://old.example.test/v1\n" + " api_key: sk-old-config\n" + " model: old-model\n" + ) + write_credential_pool( + "custom:old-endpoint", + [ + { + "id": "old-model-config", + "source": "model_config", + "auth_type": "api_key", + "access_token": "sk-old-model-config", + "base_url": "https://old.example.test/v1", + "label": "model_config", + }, + { + "id": "old-manual", + "source": "manual", + "auth_type": "api_key", + "access_token": "sk-old-manual", + "base_url": "https://old.example.test/v1", + "label": "manual", + }, + ], + ) + + with patch( + "hermes_cli.models.probe_api_models", + return_value={ + "models": ["new-model"], + "used_fallback": False, + "probed_url": "https://new.example.test/v1/models", + }, + ), \ + patch("hermes_cli.secret_prompt.masked_secret_prompt", return_value="sk-new"), \ + patch("hermes_cli.main._prompt_custom_api_mode_selection", return_value=""), \ + patch( + "builtins.input", + side_effect=[ + "https://new.example.test/v1", + "", + "", + "New Endpoint", + ], + ), \ + patch("builtins.print"): + _model_flow_custom({}) + + auth = read_credential_pool(None) + old_sources = [ + entry.get("source") + for entry in auth.get("custom:old-endpoint", []) + if isinstance(entry, dict) + ] + assert old_sources == ["manual"] + + new_pool = load_pool("custom:new-endpoint") + selected = new_pool.select() + assert selected is not None + assert selected.access_token == "sk-new" + + config = yaml.safe_load(config_path.read_text()) or {} + assert config["model"]["base_url"] == "https://new.example.test/v1" + def test_saved_model_still_probes_endpoint(self, config_home): """When a model is already saved, the function must still call fetch_api_models to probe the endpoint — not skip with early return."""