diff --git a/agent/agent_init.py b/agent/agent_init.py index 80fa34d924..65f5e3d7c2 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -34,6 +34,7 @@ from agent.model_metadata import ( MINIMUM_CONTEXT_LENGTH, fetch_model_metadata, is_local_endpoint ) from agent.agent_init_ollama import _clamp_compressor_to_ollama_num_ctx, _configure_ollama_num_ctx +from agent.primary_runtime_snapshot import _snapshot_primary_runtime from agent.process_bootstrap import _install_safe_stdio from agent.subdirectory_hints import SubdirectoryHintTracker from agent.think_scrubber import StreamingThinkScrubber @@ -2206,38 +2207,6 @@ def _emit_compression_summary(agent, cs): agent._compression_feasibility_checked = False -def _snapshot_primary_runtime(agent): - # Per-turn restoration snapshot: after a fallback, the next turn restores these so the - # preferred model gets a fresh attempt. - _cc = agent.context_compressor - agent._primary_runtime = { - "model": agent.model, - "provider": agent.provider, - "requested_provider": agent.requested_provider, - "base_url": agent.base_url, - "api_mode": agent.api_mode, - "api_key": getattr(agent, "api_key", ""), - "request_overrides": dict(getattr(agent, "request_overrides", {}) or {}), - "client_kwargs": dict(agent._client_kwargs), - "use_prompt_caching": agent._use_prompt_caching, - "use_native_cache_layout": agent._use_native_cache_layout, - "reasoning_echo_flag": getattr(agent, "_reasoning_echo_flag", False), - # Engine state _try_activate_fallback() overwrites (getattr: plugin engines may lack them). - "compressor_model": getattr(_cc, "model", agent.model), - "compressor_base_url": getattr(_cc, "base_url", agent.base_url), - "compressor_api_key": getattr(_cc, "api_key", ""), - "compressor_provider": getattr(_cc, "provider", agent.provider), - "compressor_context_length": _cc.context_length, - "compressor_threshold_tokens": _cc.threshold_tokens, - } - if agent.api_mode == "anthropic_messages": - agent._primary_runtime.update({ - "anthropic_api_key": agent._anthropic_api_key, - "anthropic_base_url": agent._anthropic_base_url, - "is_anthropic_oauth": agent._is_anthropic_oauth, - }) - - def _init_usage_state(agent): from agent.runtime_cwd import scope_terminal_cwd # Prefer the session's explicitly adopted workspace (a Desktop session created under the diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index d80b3403eb..eed1b5d6ba 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -34,6 +34,7 @@ from agent.retry_utils import parse_retry_after_seconds, reset_delay_from_messag from agent.message_metadata import MERGED_TURN_PREFIX from agent.turn_context import drop_stale_api_content from agent.agent_runtime_helpers_placeholders import _INTERRUPTED_PLACEHOLDER +from agent.primary_runtime_snapshot import _build_primary_runtime_snapshot from utils import base_url_host_matches, base_url_hostname, env_var_enabled, atomic_json_write logger = logging.getLogger(__name__) @@ -1453,7 +1454,7 @@ def restore_primary_runtime(agent) -> bool: def _load_primary_pool(): """The primary's pool (entry-scoped for a custom entry); None when absent or provider-mismatched.""" from agent.custom_entry_pool import load_runtime_pool - loaded = load_runtime_pool(primary_provider, primary_runtime_base_url, (rt or {}).get("requested_provider")) + loaded = load_runtime_pool(primary_provider, primary_runtime_base_url, (rt or {}).get("requested_provider"), rt or {}) return loaded if loaded is not None and _matches_primary(loaded) else None if _primary_quota_reopened_early(agent, primary_provider, primary_model, _matches_primary, _load_primary_pool): agent._rate_limited_until = 0 @@ -2383,43 +2384,6 @@ def _update_switch_compressor(agent, custom_providers, effective_context_length, revalidate_compression_feasibility(agent) -def _build_primary_runtime_snapshot(agent, api_mode) -> dict[str, Any]: - """The ``_primary_runtime`` record that persists a switch across turns.""" - cc = getattr(agent, "context_compressor", None) or None - rt = { - "model": agent.model, - "provider": agent.provider, - "requested_provider": agent.requested_provider, - "base_url": agent.base_url, - "api_mode": agent.api_mode, - "api_key": getattr(agent, "api_key", ""), - "client_kwargs": dict(agent._client_kwargs), - "use_prompt_caching": agent._use_prompt_caching, - "use_native_cache_layout": agent._use_native_cache_layout, - "reasoning_config": dict(agent.reasoning_config) if getattr(agent, "reasoning_config", None) else None, - "reasoning_echo_flag": getattr(agent, "_reasoning_echo_flag", False), - # Overrides must travel with the switched-to identity or a later recovery/restore resurrects - # PRE-switch overrides from the stale init snapshot. - # See #75091. - "request_overrides": dict(getattr(agent, "request_overrides", {}) or {}), - "runtime_capabilities": dict(getattr(agent, "runtime_capabilities", {}) or {}), - "compressor_model": getattr(cc, "model", agent.model), - "compressor_base_url": getattr(cc, "base_url", agent.base_url), - "compressor_api_key": getattr(cc, "api_key", ""), - "compressor_provider": getattr(cc, "provider", agent.provider), - "compressor_context_length": cc.context_length if cc else 0, - "compressor_api_mode": getattr(cc, "api_mode", agent.api_mode), - "compressor_threshold_tokens": cc.threshold_tokens if cc else 0, - } - if api_mode == "anthropic_messages": - rt.update({ - "anthropic_api_key": agent._anthropic_api_key, - "anthropic_base_url": agent._anthropic_base_url, - "is_anthropic_oauth": agent._is_anthropic_oauth, - }) - return rt - - def _finish_switch(agent, new_provider, old_norm, new_norm) -> None: """Post-switch bookkeeping: fallback reset/prune, request_overrides, billing route.""" agent._fallback_activated = False diff --git a/agent/custom_entry_pool.py b/agent/custom_entry_pool.py index 1a73ba2e9a..4c716648d3 100644 --- a/agent/custom_entry_pool.py +++ b/agent/custom_entry_pool.py @@ -13,10 +13,15 @@ from __future__ import annotations from typing import Any, Callable, Optional _DECLARED_CREDENTIAL_KEYS = ("api_key", "key_env", "api_key_env", "key_cmd") +# The ``_primary_runtime`` field holding the owning entry's provenance (``runtime_entry_owner``). +SNAPSHOT_OWNER_KEY = "custom_entry" def entry_declares_credential(entry: Any) -> bool: - """Whether a custom entry names its own credential (``api_key`` / ``key_env`` / ``key_cmd``).""" + """Whether a custom entry names its own credential (``api_key`` / ``key_env`` / ``key_cmd``). A + snapshot's provenance carries the verdict itself, since it never copies an inline secret.""" + if isinstance(entry, dict) and "declares_credential" in entry: + return bool(entry["declares_credential"]) return isinstance(entry, dict) and any(str(entry.get(k) or "").strip() for k in _DECLARED_CREDENTIAL_KEYS) @@ -85,17 +90,46 @@ def pool_for_aux_provider(provider: Optional[str], pool_key: str, load: Callable return entry_scoped_pool(pool, pool_key, match[1] if match else None) -def load_runtime_pool(provider: Optional[str], base_url: Optional[str], requested_provider: Optional[str] = None) -> Any: +def _custom_identity(provider: Optional[str], requested_provider: Optional[str]) -> tuple[bool, Optional[str]]: + """``(is_custom, custom: or None)`` for a runtime's provider / requested_provider pair.""" + from agent.credential_pool import CUSTOM_POOL_PREFIX + identities = [str(v or "").strip().lower() for v in (provider, requested_provider)] + custom = any(v == "custom" or v.startswith(CUSTOM_POOL_PREFIX) for v in identities) + return custom, next((v for v in identities if v.startswith(CUSTOM_POOL_PREFIX)), None) + + +def runtime_entry_owner(provider: Optional[str], requested_provider: Optional[str], base_url: Optional[str]) -> Optional[dict[str, Any]]: + """Credential provenance of the entry that owns a custom runtime, for its ``_primary_runtime`` + snapshot: matched in the config the route runs on (the frozen scope, else the profile's), the way + resolution matched it. Names only, never a secret. None for a built-in or an unowned endpoint.""" + custom, name = _custom_identity(provider, requested_provider) + if not custom or (match := custom_entry_for_endpoint(base_url, name)) is None: + return None + norm_name, entry = match + provenance = {k: entry.get(k) for k in ("provider_key", "base_url", "key_env", "api_key_env")} + return {"name": norm_name, **provenance, "declares_credential": entry_declares_credential(entry)} + + +def load_runtime_pool( + provider: Optional[str], base_url: Optional[str], requested_provider: Optional[str] = None, + snapshot: Optional[dict[str, Any]] = None, +) -> Any: """The pool for a runtime snapshot's identity. A custom runtime (``custom`` or ``custom:``) is the entry's, so it gets the entry-scoped view; a built-in identity loads its whole pool. Ownership and the pool key come from the runtime identity, never from the pool slug, which both - share: a ``custom:`` identity keys on its own entry, not the first entry on its URL.""" - from agent.credential_pool import CUSTOM_POOL_PREFIX, _pool_keys_for_custom_entry, load_pool, resolve_runtime_pool_key - identities = [str(v or "").strip().lower() for v in (provider, requested_provider)] - if not any(v == "custom" or v.startswith(CUSTOM_POOL_PREFIX) for v in identities): + share: a ``custom:`` identity keys on its own entry, not the first entry on its URL. + + ``snapshot`` is the ``_primary_runtime`` being restored: the owner it recorded wins over live + config, so an edit can't reassign the endpoint or drop the entry's declared credential. A named + identity whose snapshot recorded no owner gets None and keeps the snapshot's own credential.""" + from agent.credential_pool import _pool_keys_for_custom_entry, load_pool, resolve_runtime_pool_key + custom, name = _custom_identity(provider, requested_provider) + if not custom: key = resolve_runtime_pool_key(provider, base_url) return load_pool(key) if key else None - name = next((v for v in identities if v.startswith(CUSTOM_POOL_PREFIX)), None) - match = custom_entry_for_endpoint(base_url, name) - key = (_pool_keys_for_custom_entry(*match) or [""])[0] if name and match else resolve_runtime_pool_key(provider, base_url) + owner = (snapshot or {}).get(SNAPSHOT_OWNER_KEY) + if snapshot is not None and name and not owner: + return None + match = (owner["name"], owner) if owner else custom_entry_for_endpoint(base_url, name) + key = (_pool_keys_for_custom_entry(*match) or [""])[0] if match and (name or owner) else resolve_runtime_pool_key(provider, base_url) return entry_scoped_pool(load_pool(key), key, match[1] if match else None) if key else None diff --git a/agent/primary_runtime_snapshot.py b/agent/primary_runtime_snapshot.py new file mode 100644 index 0000000000..fe8914020a --- /dev/null +++ b/agent/primary_runtime_snapshot.py @@ -0,0 +1,84 @@ +"""The ``_primary_runtime`` snapshot: the route a fallback turn restores on the next turn. + +Both builders (agent construction, and a deliberate model switch) record the custom entry that owns +the route's credential, resolved against the config the route runs on. Restoring scopes the +credential pool by that record, never by a live ``config.yaml`` lookup, so an entry edited or +removed mid-session can't hand the session another account's key. +""" + +from __future__ import annotations + +from typing import Any + +from agent.custom_entry_pool import SNAPSHOT_OWNER_KEY, runtime_entry_owner + + +def _snapshot_primary_runtime(agent): + # Per-turn restoration snapshot: after a fallback, the next turn restores these so the + # preferred model gets a fresh attempt. + _cc = agent.context_compressor + agent._primary_runtime = { + "model": agent.model, + "provider": agent.provider, + "requested_provider": agent.requested_provider, + SNAPSHOT_OWNER_KEY: runtime_entry_owner(agent.provider, agent.requested_provider, agent.base_url), + "base_url": agent.base_url, + "api_mode": agent.api_mode, + "api_key": getattr(agent, "api_key", ""), + "request_overrides": dict(getattr(agent, "request_overrides", {}) or {}), + "client_kwargs": dict(agent._client_kwargs), + "use_prompt_caching": agent._use_prompt_caching, + "use_native_cache_layout": agent._use_native_cache_layout, + "reasoning_echo_flag": getattr(agent, "_reasoning_echo_flag", False), + # Engine state _try_activate_fallback() overwrites (getattr: plugin engines may lack them). + "compressor_model": getattr(_cc, "model", agent.model), + "compressor_base_url": getattr(_cc, "base_url", agent.base_url), + "compressor_api_key": getattr(_cc, "api_key", ""), + "compressor_provider": getattr(_cc, "provider", agent.provider), + "compressor_context_length": _cc.context_length, + "compressor_threshold_tokens": _cc.threshold_tokens, + } + if agent.api_mode == "anthropic_messages": + agent._primary_runtime.update({ + "anthropic_api_key": agent._anthropic_api_key, + "anthropic_base_url": agent._anthropic_base_url, + "is_anthropic_oauth": agent._is_anthropic_oauth, + }) + + +def _build_primary_runtime_snapshot(agent, api_mode) -> dict[str, Any]: + """The ``_primary_runtime`` record that persists a switch across turns.""" + cc = getattr(agent, "context_compressor", None) or None + rt = { + "model": agent.model, + "provider": agent.provider, + "requested_provider": agent.requested_provider, + SNAPSHOT_OWNER_KEY: runtime_entry_owner(agent.provider, agent.requested_provider, agent.base_url), + "base_url": agent.base_url, + "api_mode": agent.api_mode, + "api_key": getattr(agent, "api_key", ""), + "client_kwargs": dict(agent._client_kwargs), + "use_prompt_caching": agent._use_prompt_caching, + "use_native_cache_layout": agent._use_native_cache_layout, + "reasoning_config": dict(agent.reasoning_config) if getattr(agent, "reasoning_config", None) else None, + "reasoning_echo_flag": getattr(agent, "_reasoning_echo_flag", False), + # Overrides must travel with the switched-to identity or a later recovery/restore resurrects + # PRE-switch overrides from the stale init snapshot. + # See #75091. + "request_overrides": dict(getattr(agent, "request_overrides", {}) or {}), + "runtime_capabilities": dict(getattr(agent, "runtime_capabilities", {}) or {}), + "compressor_model": getattr(cc, "model", agent.model), + "compressor_base_url": getattr(cc, "base_url", agent.base_url), + "compressor_api_key": getattr(cc, "api_key", ""), + "compressor_provider": getattr(cc, "provider", agent.provider), + "compressor_context_length": cc.context_length if cc else 0, + "compressor_api_mode": getattr(cc, "api_mode", agent.api_mode), + "compressor_threshold_tokens": cc.threshold_tokens if cc else 0, + } + if api_mode == "anthropic_messages": + rt.update({ + "anthropic_api_key": agent._anthropic_api_key, + "anthropic_base_url": agent._anthropic_base_url, + "is_anthropic_oauth": agent._is_anthropic_oauth, + }) + return rt diff --git a/agent/voice_turn_route.py b/agent/voice_turn_route.py index 1726f20dd8..64d96fa73f 100644 --- a/agent/voice_turn_route.py +++ b/agent/voice_turn_route.py @@ -105,7 +105,7 @@ def _voice_reasoning(agent: Any, effort: Optional[dict[str, Any]]) -> Optional[d def _capture(agent: Any) -> dict[str, Any]: - from agent.agent_runtime_helpers import _build_primary_runtime_snapshot + from agent.primary_runtime_snapshot import _build_primary_runtime_snapshot return { "snapshot": _build_primary_runtime_snapshot(agent, agent.api_mode), **{name: getattr(agent, name, None) for name in _EXTRA_FIELDS}, diff --git a/tests/agent/test_primary_runtime_restore.py b/tests/agent/test_primary_runtime_restore.py index 865b88fdfa..81482f82ae 100644 --- a/tests/agent/test_primary_runtime_restore.py +++ b/tests/agent/test_primary_runtime_restore.py @@ -403,13 +403,6 @@ class TestRestorePrimaryRuntime: fallback_pool = MagicMock() fallback_pool.provider = "openrouter" - agent = _make_agent( - provider="custom:gemini-no-filter", - base_url="https://generativelanguage.googleapis.com/v1beta", - ) - agent._fallback_activated = True - agent._credential_pool = fallback_pool - agent._swap_credential = MagicMock() config = { "custom_providers": [ { @@ -424,6 +417,15 @@ class TestRestorePrimaryRuntime: } }, } + # The session was built on this config: its snapshot records the entry as the owner. + with patch("agent.credential_pool._load_config_safe", return_value=config): + agent = _make_agent( + provider="custom:gemini-no-filter", + base_url="https://generativelanguage.googleapis.com/v1beta", + ) + agent._fallback_activated = True + agent._credential_pool = fallback_pool + agent._swap_credential = MagicMock() with ( patch("agent.credential_pool._load_config_safe", return_value=config), diff --git a/tests/hermes_cli/test_custom_entry_shadow_key_env.py b/tests/hermes_cli/test_custom_entry_shadow_key_env.py index b073043c1d..c11895ee6b 100644 --- a/tests/hermes_cli/test_custom_entry_shadow_key_env.py +++ b/tests/hermes_cli/test_custom_entry_shadow_key_env.py @@ -422,3 +422,122 @@ def test_aux_named_entry_skips_an_unusable_inline_key_for_the_declared_key_env(t from hermes_cli.runtime_provider import _get_named_custom_provider assert _named_custom_api_key(_get_named_custom_provider("custom:mydeep"), "mydeep", ENDPOINT) == ENTRY_KEY + + +# ── review round 3: restore reads the snapshot's owner, never live config ─────────────────────── + +def _init_snapshot(agent): + from agent.agent_init import _snapshot_primary_runtime + _snapshot_primary_runtime(agent) + return agent._primary_runtime + + +def _switch_snapshot(agent): + from agent.agent_runtime_helpers import _build_primary_runtime_snapshot + return _build_primary_runtime_snapshot(agent, agent.api_mode) + + +SNAPSHOT_BUILDERS = [pytest.param(_init_snapshot, id="agent_init"), pytest.param(_switch_snapshot, id="model_switch")] + + +def _snapshotted_agent(builder, initial, requested, api_key=ENTRY_KEY): + """An agent built on the resolved primary ``initial``, snapshotted by a real builder, now on a fallback.""" + from unittest.mock import MagicMock + + from run_agent import AIAgent + + built = AIAgent.__new__(AIAgent) + built.model, built.provider, built.requested_provider = "m", initial["provider"], requested + built.base_url, built.api_mode, built.api_key = initial["base_url"], "chat_completions", api_key + built._client_kwargs = {"api_key": api_key, "base_url": initial["base_url"]} + built._use_prompt_caching = built._use_native_cache_layout = False + built.context_compressor = MagicMock(context_length=128000, threshold_tokens=100000) + snapshot = builder(built) + agent = _restorable_agent({"provider": initial["provider"], "requested_provider": requested, + "base_url": initial["base_url"], "api_key": api_key}) + agent._primary_runtime = snapshot + return agent + + +def _move_beta_endpoint(home): + cfg = yaml.safe_load((home / "config.yaml").read_text()) + cfg["providers"]["beta"]["api"] = "http://127.0.0.1:18081/v1" + (home / "config.yaml").write_text(yaml.safe_dump(cfg), encoding="utf-8") + + +def _drop_declaration(home, key): + cfg = yaml.safe_load((home / "config.yaml").read_text()) + del cfg["providers"]["deepseek"][key] + (home / "config.yaml").write_text(yaml.safe_dump(cfg), encoding="utf-8") + + +@pytest.mark.parametrize("builder", SNAPSHOT_BUILDERS) +def test_restore_keeps_the_frozen_owner_after_its_endpoint_moves(tmp_path, monkeypatch, builder): + """beta's live endpoint moves off the shared URL: the named lookup misses, and the URL-only + fallback used to hand the session alpha's ``hermes auth add`` key.""" + home = _two_entries_on_one_url(tmp_path, monkeypatch) + from hermes_cli import runtime_provider as rp + + initial = rp.resolve_runtime_provider(requested="custom:beta") + assert initial["api_key"] == ENTRY_KEY + agent = _snapshotted_agent(builder, initial, "custom:beta") + _move_beta_endpoint(home) + assert agent._restore_primary_runtime() is True + assert (agent.api_key, agent.base_url) == (ENTRY_KEY, ENDPOINT) + + +@pytest.mark.parametrize("builder", SNAPSHOT_BUILDERS) +@pytest.mark.parametrize("declared", [{"key_env": "ENTRY_KEY_VAR"}, {"api_key": ENTRY_KEY}], ids=["key_env", "api_key"]) +def test_restore_keeps_the_frozen_declaration_after_it_is_removed(tmp_path, monkeypatch, builder, declared): + """The live entry stops declaring a credential: the built-in's env row used to win, moving the + session onto DEEPSEEK_API_KEY and api.deepseek.com.""" + home = _home(tmp_path, monkeypatch, "deepseek", "DEEPSEEK_API_KEY", entry_extra=declared) + from hermes_cli import runtime_provider as rp + + initial = rp.resolve_runtime_provider(requested="custom:deepseek") + assert initial["api_key"] == ENTRY_KEY + agent = _snapshotted_agent(builder, initial, "custom:deepseek") + # The snapshot already holds the session's key; ownership adds no other copy of a secret. + rest = {k: v for k, v in agent._primary_runtime.items() if k not in {"api_key", "client_kwargs", "compressor_api_key"}} + assert ENTRY_KEY not in json.dumps(rest, default=str) + _drop_declaration(home, next(iter(declared))) + assert agent._restore_primary_runtime() is True + assert (agent.api_key, agent.base_url) == (ENTRY_KEY, ENDPOINT) + + +@pytest.mark.parametrize("edit", ["endpoint", "declaration"]) +def test_a_snapshot_without_recorded_ownership_keeps_its_own_credential(tmp_path, monkeypatch, edit): + """A snapshot from before ownership was recorded (the reviewer's probe shape): a named identity + can't prove its owner from live config, so it keeps the snapshot's key rather than another's.""" + if edit == "endpoint": + home, name = _two_entries_on_one_url(tmp_path, monkeypatch), "beta" + else: + home, name = _home(tmp_path, monkeypatch, "deepseek", "DEEPSEEK_API_KEY"), "deepseek" + from hermes_cli import runtime_provider as rp + + initial = rp.resolve_runtime_provider(requested=f"custom:{name}") + agent = _restorable_agent({"provider": initial["provider"], "requested_provider": f"custom:{name}", + "base_url": ENDPOINT, "api_key": ENTRY_KEY}) + if edit == "endpoint": + _move_beta_endpoint(home) + else: + _drop_declaration(home, "key_env") + assert agent._restore_primary_runtime() is True + assert (agent.api_key, agent.base_url) == (ENTRY_KEY, ENDPOINT) + + +@pytest.mark.parametrize("builder", SNAPSHOT_BUILDERS) +def test_restore_still_reselects_the_owners_own_pool_row(tmp_path, monkeypatch, builder): + """Recorded ownership scopes the pool, it doesn't drop it: the entry's ``hermes auth add`` row + replaces a stale snapshot key on restore, and the built-in's row stays out.""" + manual = {"id": "m1", "label": "manual", "auth_type": "api_key", "priority": 0, "source": "manual", + "access_token": MANUAL_KEY} + _home(tmp_path, monkeypatch, "deepseek", "DEEPSEEK_API_KEY", pool={"deepseek": [manual]}) + from hermes_cli import runtime_provider as rp + + initial = rp.resolve_runtime_provider(requested="custom:deepseek") + agent = _snapshotted_agent(builder, initial, "custom:deepseek", api_key="stale-snapshot-key-0123") + agent._swap_credential = lambda entry: setattr(agent, "api_key", entry.runtime_api_key) + assert agent._restore_primary_runtime() is True + assert agent.api_key == MANUAL_KEY + assert {e.access_token for e in agent._credential_pool.entries()} == {MANUAL_KEY}