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.
This commit is contained in:
kshitij 2026-07-24 13:09:56 +05:00 committed by kshitij
parent 73c4b5a045
commit c769081803
3 changed files with 24 additions and 30 deletions

View file

@ -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

View file

@ -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

View file

@ -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 = (