c7d0045a8d
The delegate fix admits a MIXED same-provider pool (any one entry serves the child's endpoint) but only guarded the initial lease. A 429/402/401 in the Azure child later ran recover_with_credential_pool -> mark_exhausted_and_rotate -> _swap_credential(next_entry) with no endpoint check, and credential_pool_matches_provider passes a non-custom pool on provider match alone, so the child was rebound to the api.openai.com entry (#68237 via a sibling path). _entry_serves_endpoint moves into agent/credential_pool.py as credential_pool_entry_serves_endpoint (next to credential_pool_matches_provider; tools/delegate_tool_config imports it back), and _rotate_and_swap vetoes a next_entry whose base_url does not serve agent.base_url, returning False exactly like a rotation that yields nothing: the failed entry stays benched, the session keeps its credential and the fallback chain takes over.
57 lines
2.3 KiB
Python
57 lines
2.3 KiB
Python
# Copyright 2025 Nous Research (Licensed under the Apache License, Version 2.0)
|
|
"""Mid-run credential rotation never rebinds a session to a same-provider entry for another endpoint.
|
|
|
|
A mixed ``openai`` pool (public api.openai.com key + Azure resource key) is legitimately shared with
|
|
an Azure-bound child (#68237). ``_swap_credential`` adopts the entry's base_url, so rotating the
|
|
Azure session onto the public entry on a 429 would send its traffic — and the public key — to the
|
|
wrong host. Rotation must be vetoed exactly like a rotation that yields nothing.
|
|
"""
|
|
|
|
from agent.agent_runtime_helpers import recover_with_credential_pool
|
|
from agent.credential_pool import CredentialPool, PooledCredential
|
|
|
|
_AZURE = "https://res.cognitiveservices.azure.com/openai/v1"
|
|
_PUBLIC = "https://api.openai.com/v1"
|
|
|
|
|
|
def _entry(eid, url):
|
|
return PooledCredential(provider="openai", id=eid, label=eid, auth_type="api_key", priority=0,
|
|
source=f"env:{eid}", access_token=f"key-{eid}", base_url=url)
|
|
|
|
|
|
class _AzureChild:
|
|
"""Session bound to the Azure entry of a mixed pool; real pool + real recovery helper."""
|
|
|
|
provider = "openai"
|
|
model = "gpt-5.4"
|
|
base_url = _AZURE
|
|
_fallback_activated = False
|
|
|
|
def __init__(self, pool):
|
|
self._credential_pool = pool
|
|
self._credential_pool_entry_id = "az"
|
|
self.api_key = "key-az"
|
|
|
|
def _swap_credential(self, entry):
|
|
self.api_key = entry.runtime_api_key
|
|
self.base_url = entry.base_url
|
|
self._credential_pool_entry_id = entry.id
|
|
return True
|
|
|
|
def _is_entitlement_failure(self, error_context, status_code):
|
|
return False
|
|
|
|
|
|
def test_rotation_on_mixed_pool_never_rebinds_to_an_entry_for_another_endpoint():
|
|
pool = CredentialPool("openai", [_entry("az", _AZURE), _entry("pub", _PUBLIC)])
|
|
agent = _AzureChild(pool)
|
|
|
|
recovered, _ = recover_with_credential_pool(
|
|
agent, status_code=429, has_retried_429=True, error_context={"message": "Rate limit"},
|
|
)
|
|
|
|
assert recovered is False
|
|
assert (agent.api_key, agent.base_url, agent._credential_pool_entry_id) == ("key-az", _AZURE, "az")
|
|
# The failed Azure entry is still benched; only the swap onto the wrong host was refused.
|
|
assert next(e for e in pool.entries() if e.id == "az").last_status == "exhausted"
|