diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index bf4d31b4f4d..1cde73419a4 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -1208,7 +1208,42 @@ def restore_primary_runtime(agent) -> bool: api_mode=rt.get("compressor_api_mode", ""), ) - # ── Re-select from the credential pool if one is available ── + # ── Rebind and re-select the primary credential pool ── + # A cross-provider fallback attaches the fallback provider's pool. The + # runtime fields above restore the primary, but leaving that pool in + # place makes the next primary 401/429 hit the provider-mismatch guard + # and disables credential rotation. Reload the primary pool first; if + # auth storage is temporarily unreadable, clear the mismatched pool. + primary_provider = str(rt.get("provider") or "").strip().lower() + pool = getattr(agent, "_credential_pool", None) + pool_provider = str(getattr(pool, "provider", "") or "").strip().lower() + pool_matches_primary = pool_provider == primary_provider + if ( + primary_provider == "custom" + and pool_provider.startswith("custom:") + ): + try: + from agent.credential_pool import get_custom_provider_pool_key + + primary_key = ( + get_custom_provider_pool_key(str(rt.get("base_url") or "")) or "" + ).strip().lower() + pool_matches_primary = bool(primary_key) and primary_key == pool_provider + except Exception: + pool_matches_primary = False + if pool is not None and pool_provider and not pool_matches_primary: + agent._credential_pool = None + try: + from agent.credential_pool import load_pool + + agent._credential_pool = load_pool(primary_provider) + except Exception as exc: + logger.warning( + "Restore could not reload primary credential pool for %s: %s", + primary_provider, + exc, + ) + # The snapshot's api_key was captured at construction time. Across # turns the pool may have rotated (token revocation, billing/rate-limit # exhaustion, cooldown), leaving the snapshot key stale. Restoring it @@ -1222,7 +1257,6 @@ def restore_primary_runtime(agent) -> bool: entry = pool.select() if entry is not None: entry_provider = str(getattr(entry, "provider", "") or "").strip().lower() - primary_provider = str(rt.get("provider") or "").strip().lower() entry_matches_primary = entry_provider == primary_provider # Custom endpoints all carry the generic ``custom`` provider on # the agent while the pool entry is keyed ``custom:`` (see @@ -1858,6 +1892,9 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo old_norm = (old_provider or "").strip().lower() new_norm = (new_provider or "").strip().lower() if old_norm != new_norm or getattr(agent, "_credential_pool", None) is None: + # A pool bound to the old provider is worse than no pool: the + # recovery guard rejects it and every later 401/429 skips rotation. + agent._credential_pool = None try: from agent.credential_pool import load_pool agent._credential_pool = load_pool(new_provider) diff --git a/tests/run_agent/test_primary_runtime_restore.py b/tests/run_agent/test_primary_runtime_restore.py index d1ac56dca6b..0d732277e97 100644 --- a/tests/run_agent/test_primary_runtime_restore.py +++ b/tests/run_agent/test_primary_runtime_restore.py @@ -286,15 +286,47 @@ class TestRestorePrimaryRuntime: agent._credential_pool = _DeepseekPool() agent._swap_credential = MagicMock() - with patch("run_agent.OpenAI", return_value=MagicMock()): + primary_pool = MagicMock() + primary_pool.provider = primary_provider + primary_pool.has_available.return_value = False + with ( + patch("run_agent.OpenAI", return_value=MagicMock()), + patch("agent.credential_pool.load_pool", return_value=primary_pool) as load_pool, + ): result = agent._restore_primary_runtime() assert result is True assert agent.provider == primary_provider assert agent.base_url == primary_base_url assert "deepseek" not in str(agent.base_url) + assert agent._credential_pool is primary_pool + load_pool.assert_called_once_with(primary_provider) agent._swap_credential.assert_not_called() + def test_restore_clears_fallback_pool_when_primary_pool_reload_fails(self): + """A fallback pool must never remain attached to the restored primary.""" + agent = _make_agent( + provider="openai-api", + base_url="https://api.openai.com/v1", + ) + agent._fallback_activated = True + fallback_pool = MagicMock() + fallback_pool.provider = "deepseek" + agent._credential_pool = fallback_pool + + with ( + patch("run_agent.OpenAI", return_value=MagicMock()), + patch( + "agent.credential_pool.load_pool", + side_effect=RuntimeError("auth store unavailable"), + ), + ): + result = agent._restore_primary_runtime() + + assert result is True + assert agent.provider == "openai-api" + assert agent._credential_pool is None + def test_restore_swaps_matching_custom_pool_entry(self): """Custom primary + custom: entry whose base_url resolves to the SAME custom key must swap (legitimate same-endpoint rotation).""" diff --git a/tests/run_agent/test_switch_model_pool_reload_52727.py b/tests/run_agent/test_switch_model_pool_reload_52727.py index a1dba807f9f..e1ef7d651c0 100644 --- a/tests/run_agent/test_switch_model_pool_reload_52727.py +++ b/tests/run_agent/test_switch_model_pool_reload_52727.py @@ -213,6 +213,7 @@ class TestSwitchModelReloadsCredentialPool: ) # The switch itself completed (provider/model updated) even though - # the pool reload failed. + # the pool reload failed, without retaining the old provider's pool. assert agent.provider == "groq" - assert agent.model == "llama-3.3-70b" \ No newline at end of file + assert agent.model == "llama-3.3-70b" + assert agent._credential_pool is None