From 820a0525750b16cf22eafa550439d8f2b1e81739 Mon Sep 17 00:00:00 2001
From: openhands <openhands@all-hands.dev>
Date: Wed, 1 Jul 2026 12:43:38 +0000
Subject: [PATCH] fix(agent): keep primary runtime restore on matching
 credential pool (#56374)

---
 agent/agent_runtime_helpers.py                | 22 +++++++++-
 .../run_agent/test_primary_runtime_restore.py | 40 +++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py
index 10228d5f1ef..9c7a046299e 100644
--- a/agent/agent_runtime_helpers.py
+++ b/agent/agent_runtime_helpers.py
@@ -1184,11 +1184,23 @@ def restore_primary_runtime(agent) -> bool:
         if pool is not None and pool.has_available():
             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
+                if primary_provider == "custom" and entry_provider.startswith("custom:"):
+                    primary_base_url = str(rt.get("base_url") or "").strip().rstrip("/").lower()
+                    entry_base_url = str(
+                        getattr(entry, "runtime_base_url", None)
+                        or getattr(entry, "base_url", None)
+                        or ""
+                    ).strip().rstrip("/").lower()
+                    entry_matches_primary = bool(primary_base_url and entry_base_url == primary_base_url)
+
                 entry_key = (
                     getattr(entry, "runtime_api_key", None)
                     or getattr(entry, "access_token", "")
                 )
-                if entry_key:
+                if entry_key and entry_matches_primary:
                     # ``_swap_credential`` rebuilds the OpenAI/Anthropic client,
                     # reapplies base-url-scoped headers, and carries the
                     # accumulated base_url / OAuth-detection fixes (#33163).
@@ -1198,6 +1210,14 @@ def restore_primary_runtime(agent) -> bool:
                         getattr(entry, "id", "?"),
                         getattr(entry, "label", "?"),
                     )
+                elif entry_key:
+                    logger.info(
+                        "Restore skipped pool entry %s (%s): provider %s does not match primary provider %s",
+                        getattr(entry, "id", "?"),
+                        getattr(entry, "label", "?"),
+                        entry_provider or "?",
+                        primary_provider or "?",
+                    )
 
         # ── Reset fallback chain for the new turn ──
         agent._fallback_activated = False
diff --git a/tests/run_agent/test_primary_runtime_restore.py b/tests/run_agent/test_primary_runtime_restore.py
index 7aee1418782..06af15bbc46 100644
--- a/tests/run_agent/test_primary_runtime_restore.py
+++ b/tests/run_agent/test_primary_runtime_restore.py
@@ -201,6 +201,46 @@ class TestRestorePrimaryRuntime:
 
         assert agent._use_prompt_caching == original_caching
 
+    def test_restore_skips_cross_provider_pool_entry(self):
+        """Restore must not swap in a fallback provider credential for the primary runtime."""
+
+        class _Entry:
+            provider = "openrouter"
+            id = "fallback-entry"
+            label = "fallback"
+            runtime_api_key = "fallback-key"
+            runtime_base_url = "https://openrouter.ai/api/v1"
+            access_token = "fallback-key"
+
+        class _Pool:
+            provider = "openrouter"
+
+            def has_available(self):
+                return True
+
+            def select(self):
+                return _Entry()
+
+        agent = _make_agent(
+            provider="custom",
+            base_url="https://primary.example.com/v1",
+            fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
+        )
+        original_base_url = agent.base_url
+        mock_client = _mock_resolve()
+        with patch("agent.auxiliary_client.resolve_provider_client", return_value=(mock_client, None)):
+            agent._try_activate_fallback()
+        agent._credential_pool = _Pool()
+        agent._swap_credential = MagicMock()
+
+        with patch("run_agent.OpenAI", return_value=MagicMock()):
+            result = agent._restore_primary_runtime()
+
+        assert result is True
+        assert agent.provider == "custom"
+        assert agent.base_url == original_base_url
+        agent._swap_credential.assert_not_called()
+
     def test_restore_survives_exception(self):
         """If client rebuild fails, the method returns False gracefully."""
         agent = _make_agent()