From 2cc6d88c1d947ce6b7becb8b037f3fc61bc60f97 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Sat, 12 Sep 2026 01:03:45 -0700 Subject: [PATCH] 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. --- agent/anthropic_adapter.py | 9 +- agent/azure_identity_adapter.py | 45 ++++++++-- agent/bedrock_adapter.py | 55 ++++++++++-- ...test_multiplex_cloud_credential_clients.py | 90 +++++++++++++++++++ 4 files changed, 186 insertions(+), 13 deletions(-) create mode 100644 tests/agent/test_multiplex_cloud_credential_clients.py diff --git a/agent/anthropic_adapter.py b/agent/anthropic_adapter.py index ecffcabcb6..db0902f1ef 100644 --- a/agent/anthropic_adapter.py +++ b/agent/anthropic_adapter.py @@ -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()}, ) diff --git a/agent/azure_identity_adapter.py b/agent/azure_identity_adapter.py index 8c31bfa907..76d4ded11d 100644 --- a/agent/azure_identity_adapter.py +++ b/agent/azure_identity_adapter.py @@ -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 diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py index 0424a6321b..d2a42ba43e 100644 --- a/agent/bedrock_adapter.py +++ b/agent/bedrock_adapter.py @@ -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"] diff --git a/tests/agent/test_multiplex_cloud_credential_clients.py b/tests/agent/test_multiplex_cloud_credential_clients.py new file mode 100644 index 0000000000..eb8a99b3cc --- /dev/null +++ b/tests/agent/test_multiplex_cloud_credential_clients.py @@ -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()