From c7690818033646a8b0fe86eb5b5ceb77aa019cb5 Mon Sep 17 00:00:00 2001
From: kshitij <kshitij@kshitij.dev>
Date: Fri, 24 Jul 2026 13:09:56 +0500
Subject: [PATCH] refactor: extract sync_credential_pool_entry_id helper

Replace 3 duplicated entry_id resolution blocks (try/except +
entry_id_for_api_key + fallback to None) in agent_init.py,
chat_completion_helpers.py, and switch_model with a single
sync_credential_pool_entry_id(agent) function in agent_runtime_helpers.

Follow-up to #70323.
---
 agent/agent_init.py              | 12 ++----------
 agent/agent_runtime_helpers.py   | 30 ++++++++++++++++++++----------
 agent/chat_completion_helpers.py | 12 ++----------
 3 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/agent/agent_init.py b/agent/agent_init.py
index f2cf3e841f02..7373f80f0bfe 100644
--- a/agent/agent_init.py
+++ b/agent/agent_init.py
@@ -1346,16 +1346,8 @@ def init_agent(
     # OAuth refreshes can replace the runtime token before a failed request is
     # recovered, so the mutable API-key value alone cannot reliably attribute
     # the failure to its source entry.
-    agent._credential_pool_entry_id = None
-    if agent._credential_pool is not None:
-        try:
-            agent._credential_pool_entry_id = (
-                agent._credential_pool.entry_id_for_api_key(
-                    getattr(agent, "api_key", None)
-                )
-            )
-        except Exception:
-            agent._credential_pool_entry_id = None
+    from agent.agent_runtime_helpers import sync_credential_pool_entry_id
+    sync_credential_pool_entry_id(agent)
     
     # Provider fallback chain — ordered list of backup providers tried
     # when the primary is exhausted (rate-limit, overload, connection
diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py
index e295a199be92..5da612b92585 100644
--- a/agent/agent_runtime_helpers.py
+++ b/agent/agent_runtime_helpers.py
@@ -850,6 +850,25 @@ def strip_think_blocks(agent, content: str) -> str:
 
 
 
+def sync_credential_pool_entry_id(agent) -> None:
+    """Rebind ``agent._credential_pool_entry_id`` from the current pool + key.
+
+    OAuth refreshes can replace the runtime token before a failed request is
+    recovered, so the mutable API-key value alone cannot reliably attribute
+    the failure to its source entry.  This resolves the stable pool-entry ID
+    for the agent's current ``api_key`` and clears it when no pool is bound.
+    """
+    pool = getattr(agent, "_credential_pool", None)
+    try:
+        agent._credential_pool_entry_id = (
+            pool.entry_id_for_api_key(getattr(agent, "api_key", None))
+            if pool is not None
+            else None
+        )
+    except Exception:
+        agent._credential_pool_entry_id = None
+
+
 def recover_with_credential_pool(
     agent,
     *,
@@ -2254,16 +2273,7 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo
                 shared=True,
             )
 
-        pool = getattr(agent, "_credential_pool", None)
-        if pool is not None:
-            try:
-                agent._credential_pool_entry_id = pool.entry_id_for_api_key(
-                    getattr(agent, "api_key", None)
-                )
-            except Exception:
-                agent._credential_pool_entry_id = None
-        else:
-            agent._credential_pool_entry_id = None
+        sync_credential_pool_entry_id(agent)
     except Exception:
         # Rollback every mutated field to the pre-swap snapshot so the agent
         # is left consistent (old model + old provider + old client) and the
diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py
index f22036ee93fd..d46bdeee7cd1 100644
--- a/agent/chat_completion_helpers.py
+++ b/agent/chat_completion_helpers.py
@@ -1802,16 +1802,8 @@ def try_activate_fallback(agent, reason: "FailoverReason | None" = None) -> bool
                 # not only after a later credential-rotation rebuild.
                 agent._replace_primary_openai_client(reason="fallback_timeout_apply")
 
-        fallback_pool = getattr(agent, "_credential_pool", None)
-        if fallback_pool is not None:
-            try:
-                agent._credential_pool_entry_id = (
-                    fallback_pool.entry_id_for_api_key(agent.api_key)
-                )
-            except Exception:
-                agent._credential_pool_entry_id = None
-        else:
-            agent._credential_pool_entry_id = None
+        from agent.agent_runtime_helpers import sync_credential_pool_entry_id
+        sync_credential_pool_entry_id(agent)
 
         # Re-evaluate prompt caching for the new provider/model
         agent._use_prompt_caching, agent._use_native_cache_layout = (