fix(multiplex): build Bedrock and Entra credential clients from the routed profile
boto3 and azure-identity freeze the credential chain into the client at construction, and the process env under a multiplexed turn belongs to the launch profile. A region-only (bedrock) / config-only (lru_cache) slot therefore signed a served profile's calls with the launch profile's keys and served its account's model list to everyone. Under a HERMES_HOME override the clients are built from the profile's secret scope (AWS_* / AZURE_* from its .env) and cached per (home, service, region) / (home, config); the discovery cache key carries the home. The unscoped path keeps the region slot and the maxsize=1 lru byte-for-byte.
This commit is contained in:
@@ -419,12 +419,17 @@ def build_anthropic_bedrock_client(region: str):
|
||||
``context-1m-2025-08-07`` are attached: without the latter Bedrock caps Opus 4.6/4.7 at 200K.
|
||||
A configured ``bedrock.guardrail`` rides as InvokeModel headers so every client built here
|
||||
(primary, auxiliary, per-request rebuild) enforces it."""
|
||||
from agent.bedrock_adapter import bedrock_guardrail_headers
|
||||
from agent.bedrock_adapter import bedrock_guardrail_headers, scoped_aws_session_kwargs
|
||||
sdk = _require_sdk("the Bedrock provider")
|
||||
if not hasattr(sdk, "AnthropicBedrock"):
|
||||
raise ImportError("anthropic.AnthropicBedrock not available. Upgrade with: pip install 'anthropic>=0.39.0'")
|
||||
# Routed multiplex profile: its own AWS_* from the secret scope (the SDK would otherwise read the
|
||||
# launch profile's process env); unscoped passes nothing and keeps the default chain.
|
||||
scoped = scoped_aws_session_kwargs()
|
||||
aws_kwargs = {"aws_access_key": scoped.get("aws_access_key_id"), "aws_secret_key": scoped.get("aws_secret_access_key"),
|
||||
"aws_session_token": scoped.get("aws_session_token"), "aws_profile": scoped.get("profile_name")}
|
||||
return sdk.AnthropicBedrock(
|
||||
aws_region=region, timeout=_client_timeout(None),
|
||||
aws_region=region, timeout=_client_timeout(None), **{k: v for k, v in aws_kwargs.items() if v},
|
||||
max_retries=0, # retry belongs to hermes's outer loop (honors Retry-After)
|
||||
default_headers={**_beta_header([*_COMMON_BETAS, _CONTEXT_1M_BETA]), **bedrock_guardrail_headers()},
|
||||
)
|
||||
|
||||
@@ -60,10 +60,11 @@ def _require_azure_identity():
|
||||
|
||||
|
||||
def reset_credential_cache() -> None:
|
||||
"""Clear the cached ``DefaultAzureCredential`` (tests, profile switches); tolerates a monkeypatched plain function."""
|
||||
cache_clear = getattr(build_credential, "cache_clear", None)
|
||||
"""Clear the cached credentials (tests, profile switches); tolerates a monkeypatched plain function."""
|
||||
cache_clear = getattr(_default_chain_credential, "cache_clear", None)
|
||||
if callable(cache_clear):
|
||||
cache_clear()
|
||||
_credentials_by_home.clear()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -91,15 +92,49 @@ class EntraIdentityConfig:
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def build_credential(config: EntraIdentityConfig) -> Any:
|
||||
"""Cached ``DefaultAzureCredential``. ``maxsize=1`` is intentional: a process uses one ``model.entra.*``
|
||||
block at a time. Only Hermes knobs are passed as kwargs; the rest comes from ``AZURE_*`` env vars."""
|
||||
def _default_chain_credential(config: EntraIdentityConfig) -> Any:
|
||||
"""Cached ``DefaultAzureCredential`` for the unscoped process. ``maxsize=1`` is intentional: a process uses
|
||||
one ``model.entra.*`` block at a time. Only Hermes knobs are passed as kwargs; the rest comes from ``AZURE_*``
|
||||
env vars."""
|
||||
ai = _require_azure_identity()
|
||||
# SDK default already excludes the browser; only pass the kwarg when opting in.
|
||||
kwargs = {} if config.exclude_interactive_browser else {"exclude_interactive_browser_credential": False}
|
||||
return ai.DefaultAzureCredential(**kwargs)
|
||||
|
||||
|
||||
# Routed multiplex profiles: (home key, config) -> credential. DefaultAzureCredential reads AZURE_* from the
|
||||
# process env, which under an override belongs to the LAUNCH profile, so a served profile's service principal
|
||||
# is built explicitly from its own secret scope (client secret first, then workload identity), falling back
|
||||
# to the default chain only when the profile sets no AZURE_* of its own.
|
||||
_credentials_by_home: Dict[tuple, Any] = {}
|
||||
|
||||
|
||||
def _scoped_credential(ai: Any, config: EntraIdentityConfig) -> Any:
|
||||
from agent.secret_scope import current_secret_scope
|
||||
scope = current_secret_scope() or {}
|
||||
read = lambda name: (scope.get(name) or "").strip() # noqa: E731
|
||||
tenant, client = read("AZURE_TENANT_ID"), read("AZURE_CLIENT_ID")
|
||||
if tenant and client and read("AZURE_CLIENT_SECRET"):
|
||||
return ai.ClientSecretCredential(tenant, client, read("AZURE_CLIENT_SECRET"))
|
||||
if tenant and client and read("AZURE_FEDERATED_TOKEN_FILE"):
|
||||
return ai.WorkloadIdentityCredential(tenant_id=tenant, client_id=client, token_file_path=read("AZURE_FEDERATED_TOKEN_FILE"))
|
||||
kwargs = {} if config.exclude_interactive_browser else {"exclude_interactive_browser_credential": False}
|
||||
return ai.DefaultAzureCredential(**kwargs)
|
||||
|
||||
|
||||
def build_credential(config: EntraIdentityConfig) -> Any:
|
||||
"""Cached Entra credential: the process-wide default chain when unscoped, the routed profile's own
|
||||
credential (built from its secret scope) under a HERMES_HOME override."""
|
||||
from hermes_constants import get_hermes_home_override, hermes_home_key
|
||||
if get_hermes_home_override() is None:
|
||||
return _default_chain_credential(config)
|
||||
key = (hermes_home_key(), config)
|
||||
credential = _credentials_by_home.get(key)
|
||||
if credential is None:
|
||||
credential = _credentials_by_home[key] = _scoped_credential(_require_azure_identity(), config)
|
||||
return credential
|
||||
|
||||
|
||||
def _resolve_config(config: Optional[EntraIdentityConfig], scope: Optional[str], **overrides: Any) -> EntraIdentityConfig:
|
||||
if config is not None:
|
||||
return config
|
||||
|
||||
@@ -37,6 +37,29 @@ except Exception:
|
||||
|
||||
_bedrock_runtime_client_cache: Dict[str, Any] = {}
|
||||
_bedrock_control_client_cache: Dict[str, Any] = {}
|
||||
# Routed multiplex profiles: one client per (profile home, region). boto3 freezes the credential
|
||||
# chain into the client at construction, so a region-only slot would sign profile B's calls with A's keys.
|
||||
_bedrock_clients_by_home: Dict[Tuple[str, str, str], Any] = {}
|
||||
|
||||
# botocore session kwarg <- profile .env variable (the explicit sources of the default chain).
|
||||
_AWS_SCOPED_CREDENTIAL_VARS: Tuple[Tuple[str, str], ...] = (
|
||||
("aws_access_key_id", "AWS_ACCESS_KEY_ID"), ("aws_secret_access_key", "AWS_SECRET_ACCESS_KEY"),
|
||||
("aws_session_token", "AWS_SESSION_TOKEN"), ("profile_name", "AWS_PROFILE"),
|
||||
)
|
||||
|
||||
|
||||
def scoped_aws_session_kwargs() -> Dict[str, str]:
|
||||
"""``boto3.session.Session`` kwargs from the routed profile's secret scope, ``{}`` when unscoped.
|
||||
|
||||
Under a HERMES_HOME override the process env holds the LAUNCH profile's ``AWS_*`` (or nothing), so
|
||||
every Bedrock client for a served profile must be built from that profile's own ``.env`` values.
|
||||
"""
|
||||
from hermes_constants import get_hermes_home_override
|
||||
if get_hermes_home_override() is None:
|
||||
return {}
|
||||
from agent.secret_scope import current_secret_scope
|
||||
scope = current_secret_scope() or {}
|
||||
return {kw: scope[var].strip() for kw, var in _AWS_SCOPED_CREDENTIAL_VARS if (scope.get(var) or "").strip()}
|
||||
|
||||
# Bedrock-hosted GPT-5.x models are served from the Bedrock Mantle OpenAI-compatible endpoint, not
|
||||
# Converse. Narrow allowlist so GPT-OSS models stay on the native path.
|
||||
@@ -70,10 +93,21 @@ def _require_boto3():
|
||||
|
||||
|
||||
def _cached_client(cache: Dict[str, Any], service: str, region: str):
|
||||
"""Get or create a per-region boto3 client using the default credential chain."""
|
||||
if region not in cache:
|
||||
cache[region] = _require_boto3().client(service, region_name=region)
|
||||
return cache[region]
|
||||
"""Get or create a per-region boto3 client. Unscoped: the default credential chain, one client per
|
||||
region. Routed profile: one client per (home, service, region), built from that profile's scoped
|
||||
``AWS_*`` (falling back to the default chain only for what the profile does not set)."""
|
||||
from hermes_constants import get_hermes_home_override, hermes_home_key
|
||||
if get_hermes_home_override() is None:
|
||||
if region not in cache:
|
||||
cache[region] = _require_boto3().client(service, region_name=region)
|
||||
return cache[region]
|
||||
key = (hermes_home_key(), service, region)
|
||||
client = _bedrock_clients_by_home.get(key)
|
||||
if client is None:
|
||||
boto3 = _require_boto3()
|
||||
client = boto3.Session(**scoped_aws_session_kwargs()).client(service, region_name=region)
|
||||
_bedrock_clients_by_home[key] = client
|
||||
return client
|
||||
|
||||
|
||||
def _get_bedrock_runtime_client(region: str):
|
||||
@@ -88,13 +122,18 @@ def reset_client_cache():
|
||||
"""Clear cached boto3 clients. Used in tests and profile switches."""
|
||||
_bedrock_runtime_client_cache.clear()
|
||||
_bedrock_control_client_cache.clear()
|
||||
_bedrock_clients_by_home.clear()
|
||||
|
||||
|
||||
def invalidate_runtime_client(region: str) -> bool:
|
||||
"""Evict one region's cached ``bedrock-runtime`` client (stale HTTP pool); True if evicted."""
|
||||
from hermes_constants import get_hermes_home_override, hermes_home_key
|
||||
if get_hermes_home_override() is not None:
|
||||
return _bedrock_clients_by_home.pop((hermes_home_key(), "bedrock-runtime", region), None) is not None
|
||||
return _bedrock_runtime_client_cache.pop(region, None) is not None
|
||||
|
||||
|
||||
|
||||
# --- Bedrock Mantle / OpenAI Responses support ---
|
||||
|
||||
def is_openai_bedrock_model(model_id: str) -> bool:
|
||||
@@ -149,10 +188,9 @@ class BedrockOpenAISigV4Auth(httpx.Auth):
|
||||
self.service = service
|
||||
|
||||
def auth_flow(self, request): # pragma: no cover - exercised by live call
|
||||
import botocore.session
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
credentials = botocore.session.get_session().get_credentials()
|
||||
credentials = _require_boto3().Session(**scoped_aws_session_kwargs()).get_credentials()
|
||||
if credentials is None:
|
||||
raise RuntimeError(
|
||||
"No AWS credentials available for Bedrock OpenAI Responses. "
|
||||
@@ -967,7 +1005,12 @@ def _list_inference_profiles(client, filter_set: set, models: List[Dict[str, Any
|
||||
def discover_bedrock_models(region: str, provider_filter: Optional[List[str]] = None) -> List[Dict[str, Any]]:
|
||||
"""Foundation models + inference profiles (cached 1h per region/filter), ``global.`` profiles first then
|
||||
by name; [] when the client cannot be built."""
|
||||
# The list is account-scoped (whichever credentials the control client signs with), so a routed
|
||||
# profile gets its own entry; unscoped keeps the region:filter key byte-for-byte.
|
||||
from hermes_constants import get_hermes_home_override, hermes_home_key
|
||||
cache_key = f"{region}:{','.join(sorted(provider_filter or []))}"
|
||||
if get_hermes_home_override() is not None:
|
||||
cache_key = f"{hermes_home_key()}|{cache_key}"
|
||||
cached = _discovery_cache.get(cache_key)
|
||||
if cached and (time.time() - cached["timestamp"]) < _DISCOVERY_CACHE_TTL_SECONDS:
|
||||
return cached["models"]
|
||||
|
||||
90
tests/agent/test_multiplex_cloud_credential_clients.py
Normal file
90
tests/agent/test_multiplex_cloud_credential_clients.py
Normal file
@@ -0,0 +1,90 @@
|
||||
"""Cloud-SDK credential clients under ``gateway.multiplex_profiles``: boto3 and azure-identity freeze the
|
||||
credential chain into the client at construction, so a slot keyed by region / config alone would sign a
|
||||
served profile's calls with the launch profile's keys. Each test warms the client under profile A, reads
|
||||
under routed profile B whose ``.env`` differs (real temp homes, real secret scope; no mocks of the cache).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.secret_scope import build_profile_secret_scope, reset_secret_scope, set_secret_scope
|
||||
from hermes_constants import reset_hermes_home_override, set_hermes_home_override
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def two_profiles(tmp_path, monkeypatch):
|
||||
a = tmp_path / ".hermes"
|
||||
b = a / "profiles" / "b"
|
||||
b.mkdir(parents=True)
|
||||
monkeypatch.setenv("HERMES_HOME", str(a))
|
||||
for var in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_PROFILE",
|
||||
"AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
for name, home in (("A", a), ("B", b)):
|
||||
(home / ".env").write_text(
|
||||
f"AWS_ACCESS_KEY_ID=AKIA{name * 16}\nAWS_SECRET_ACCESS_KEY=secret-{name}\n"
|
||||
f"AZURE_TENANT_ID=tenant-{name}\nAZURE_CLIENT_ID=client-{name}\nAZURE_CLIENT_SECRET=s-{name}\n",
|
||||
encoding="utf-8")
|
||||
return a, b
|
||||
|
||||
|
||||
def _under(home: Path, fn):
|
||||
home_token = set_hermes_home_override(str(home))
|
||||
scope_token = set_secret_scope(build_profile_secret_scope(home))
|
||||
try:
|
||||
return fn()
|
||||
finally:
|
||||
reset_secret_scope(scope_token)
|
||||
reset_hermes_home_override(home_token)
|
||||
|
||||
|
||||
def test_bedrock_clients_sign_with_the_routed_profiles_aws_keys(two_profiles):
|
||||
pytest.importorskip("boto3")
|
||||
from agent import bedrock_adapter as ba
|
||||
|
||||
a, b = two_profiles
|
||||
ba.reset_client_cache()
|
||||
try:
|
||||
client_a = _under(a, lambda: ba._get_bedrock_runtime_client("us-east-1"))
|
||||
client_b = _under(b, lambda: ba._get_bedrock_runtime_client("us-east-1"))
|
||||
assert client_b is not client_a
|
||||
keys = {c._request_signer._credentials.get_frozen_credentials().access_key for c in (client_a, client_b)}
|
||||
assert keys == {"AKIA" + "A" * 16, "AKIA" + "B" * 16}
|
||||
# Per-profile slots stay hot; eviction under B leaves A's client alone.
|
||||
assert _under(a, lambda: ba._get_bedrock_runtime_client("us-east-1")) is client_a
|
||||
assert _under(b, lambda: ba.invalidate_runtime_client("us-east-1")) is True
|
||||
assert _under(a, lambda: ba._get_bedrock_runtime_client("us-east-1")) is client_a
|
||||
finally:
|
||||
ba.reset_client_cache()
|
||||
|
||||
|
||||
def test_azure_entra_credential_is_built_from_the_routed_profiles_scope(two_profiles):
|
||||
from agent import azure_identity_adapter as az
|
||||
|
||||
a, b = two_profiles
|
||||
seen: list[tuple] = []
|
||||
|
||||
class _FakeSDK:
|
||||
def ClientSecretCredential(self, tenant, client, secret):
|
||||
seen.append((tenant, client))
|
||||
return object()
|
||||
|
||||
def DefaultAzureCredential(self, **kwargs):
|
||||
seen.append(("default-chain",))
|
||||
return object()
|
||||
|
||||
az.reset_credential_cache()
|
||||
try:
|
||||
cfg = az.EntraIdentityConfig()
|
||||
import unittest.mock as mock
|
||||
with mock.patch.object(az, "_require_azure_identity", lambda: _FakeSDK()):
|
||||
cred_a = _under(a, lambda: az.build_credential(cfg))
|
||||
cred_b = _under(b, lambda: az.build_credential(cfg))
|
||||
assert cred_b is not cred_a
|
||||
assert seen == [("tenant-A", "client-A"), ("tenant-B", "client-B")]
|
||||
assert _under(a, lambda: az.build_credential(cfg)) is cred_a # cached per profile, not rebuilt
|
||||
finally:
|
||||
az.reset_credential_cache()
|
||||
Reference in New Issue
Block a user