From 8f39bccf8f1bef89de07b56165990886ae785474 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Furga=C5=82a?= <83299832+00200200@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:17:42 +0200 Subject: [PATCH] feat(bedrock): add explicit Mantle auth mode --- src/anthropic/lib/bedrock/_mantle.py | 64 +++++++++++++++++++++----- tests/lib/test_bedrock_mantle.py | 68 ++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 11 deletions(-) diff --git a/src/anthropic/lib/bedrock/_mantle.py b/src/anthropic/lib/bedrock/_mantle.py index 85e261b05..2d7c55fca 100644 --- a/src/anthropic/lib/bedrock/_mantle.py +++ b/src/anthropic/lib/bedrock/_mantle.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from typing import Any, Union, Mapping, TypeVar, Sequence +from typing import Any, Union, Literal, Mapping, TypeVar, Sequence from typing_extensions import Self, override import httpx2 @@ -34,6 +34,7 @@ from ...resources.beta.messages import Messages as BetaMessages, AsyncMessages as AsyncBetaMessages DEFAULT_SERVICE_NAME = "bedrock-mantle" +MantleAuthMode = Literal["auto", "api_key", "sigv4"] _MANTLE_API_KEY_ENV_VARS = ("AWS_BEARER_TOKEN_BEDROCK", "ANTHROPIC_AWS_API_KEY") @@ -102,6 +103,7 @@ def _resolve_mantle_config( aws_region: str | None, aws_profile: str | None, skip_auth: bool, + auth_mode: MantleAuthMode, base_url: str | httpx2.URL | None, default_headers: Mapping[str, str] | None, ) -> tuple[str | None, str | httpx2.URL, bool, dict[str, str]]: @@ -109,26 +111,44 @@ def _resolve_mantle_config( Returns (resolved_api_key, resolved_base_url, use_sigv4, merged_headers). """ + if auth_mode not in ("auto", "api_key", "sigv4"): + raise ValueError("`auth_mode` must be one of `auto`, `api_key`, or `sigv4`") + + if skip_auth and auth_mode != "auto": + raise ValueError("`skip_auth` is mutually exclusive with an explicit `auth_mode`") + if skip_auth: use_sigv4 = False resolved_api_key = None else: validate_credentials(aws_access_key=aws_access_key, aws_secret_key=aws_secret_key) - use_sigv4 = resolve_auth_mode( - api_key=api_key, - aws_access_key=aws_access_key, - aws_secret_key=aws_secret_key, - aws_profile=aws_profile, - api_key_env_vars=_MANTLE_API_KEY_ENV_VARS, - ) + if auth_mode == "auto": + use_sigv4 = resolve_auth_mode( + api_key=api_key, + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_profile=aws_profile, + api_key_env_vars=_MANTLE_API_KEY_ENV_VARS, + ) + else: + use_sigv4 = auth_mode == "sigv4" + + if auth_mode == "sigv4" and api_key is not None: + raise ValueError("`api_key` cannot be used with `auth_mode='sigv4'`") resolved_api_key = resolve_api_key( - api_key=api_key, + api_key=api_key if not use_sigv4 else None, use_sigv4=use_sigv4, api_key_env_vars=_MANTLE_API_KEY_ENV_VARS, ) + if auth_mode == "api_key" and resolved_api_key is None: + raise ValueError( + "`auth_mode='api_key'` requires an API key. Set `api_key` or one of " + "`AWS_BEARER_TOKEN_BEDROCK` / `ANTHROPIC_AWS_API_KEY`." + ) + resolved_region = resolve_region(aws_region) if base_url is None: @@ -159,6 +179,7 @@ class AnthropicBedrockMantle(BaseMantleClient[httpx2.Client, Stream[Any]], SyncA aws_session_token: str | None aws_profile: str | None skip_auth: bool + auth_mode: MantleAuthMode _use_sigv4: bool @@ -172,6 +193,7 @@ def __init__( aws_profile: str | None = None, api_key: str | None = None, skip_auth: bool = False, + auth_mode: MantleAuthMode = "auto", base_url: str | httpx2.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, max_retries: int = DEFAULT_MAX_RETRIES, @@ -188,6 +210,7 @@ def __init__( aws_region=aws_region, aws_profile=aws_profile, skip_auth=skip_auth, + auth_mode=auth_mode, base_url=base_url, default_headers=default_headers, ) @@ -213,6 +236,7 @@ def __init__( self.aws_session_token = aws_session_token self.aws_profile = aws_profile self.skip_auth = skip_auth + self.auth_mode = auth_mode self._use_sigv4 = use_sigv4 self.messages = Messages(self) @@ -278,6 +302,7 @@ def copy( aws_region: str | None = None, aws_profile: str | None = None, skip_auth: bool | None = None, + auth_mode: MantleAuthMode | None = None, base_url: str | httpx2.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx2.Client | None = None, @@ -310,14 +335,20 @@ def copy( elif set_default_query is not None: params = set_default_query + resolved_auth_mode = auth_mode if auth_mode is not None else self.auth_mode + resolved_api_key = api_key or self.api_key + if auth_mode == "sigv4" and api_key is None: + resolved_api_key = None + return self.__class__( - api_key=api_key or self.api_key, + api_key=resolved_api_key, aws_access_key=aws_access_key or self.aws_access_key, aws_secret_key=aws_secret_key or self.aws_secret_key, aws_session_token=aws_session_token or self.aws_session_token, aws_region=aws_region or self.aws_region, aws_profile=aws_profile or self.aws_profile, skip_auth=skip_auth if skip_auth is not None else self.skip_auth, + auth_mode=resolved_auth_mode, base_url=base_url or self.base_url, timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, @@ -352,6 +383,7 @@ class AsyncAnthropicBedrockMantle(BaseMantleClient[httpx2.AsyncClient, AsyncStre aws_session_token: str | None aws_profile: str | None skip_auth: bool + auth_mode: MantleAuthMode _use_sigv4: bool @@ -365,6 +397,7 @@ def __init__( aws_profile: str | None = None, api_key: str | None = None, skip_auth: bool = False, + auth_mode: MantleAuthMode = "auto", base_url: str | httpx2.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, max_retries: int = DEFAULT_MAX_RETRIES, @@ -381,6 +414,7 @@ def __init__( aws_region=aws_region, aws_profile=aws_profile, skip_auth=skip_auth, + auth_mode=auth_mode, base_url=base_url, default_headers=default_headers, ) @@ -406,6 +440,7 @@ def __init__( self.aws_session_token = aws_session_token self.aws_profile = aws_profile self.skip_auth = skip_auth + self.auth_mode = auth_mode self._use_sigv4 = use_sigv4 self.messages = AsyncMessages(self) @@ -471,6 +506,7 @@ def copy( aws_region: str | None = None, aws_profile: str | None = None, skip_auth: bool | None = None, + auth_mode: MantleAuthMode | None = None, base_url: str | httpx2.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx2.AsyncClient | None = None, @@ -503,14 +539,20 @@ def copy( elif set_default_query is not None: params = set_default_query + resolved_auth_mode = auth_mode if auth_mode is not None else self.auth_mode + resolved_api_key = api_key or self.api_key + if auth_mode == "sigv4" and api_key is None: + resolved_api_key = None + return self.__class__( - api_key=api_key or self.api_key, + api_key=resolved_api_key, aws_access_key=aws_access_key or self.aws_access_key, aws_secret_key=aws_secret_key or self.aws_secret_key, aws_session_token=aws_session_token or self.aws_session_token, aws_region=aws_region or self.aws_region, aws_profile=aws_profile or self.aws_profile, skip_auth=skip_auth if skip_auth is not None else self.skip_auth, + auth_mode=resolved_auth_mode, base_url=base_url or self.base_url, timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, diff --git a/tests/lib/test_bedrock_mantle.py b/tests/lib/test_bedrock_mantle.py index e8a6695a4..b457aad1a 100644 --- a/tests/lib/test_bedrock_mantle.py +++ b/tests/lib/test_bedrock_mantle.py @@ -202,6 +202,43 @@ def test_skip_auth_returns_empty_auth_headers(self) -> None: ) assert client.auth_headers == {} + def test_sigv4_mode_ignores_ambient_api_key(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "ambient-key") + monkeypatch.setenv("ANTHROPIC_AWS_API_KEY", "fallback-key") + + client = AnthropicBedrockMantle( + auth_mode="sigv4", + aws_region="us-east-1", + ) + + assert client.auth_mode == "sigv4" + assert client._use_sigv4 is True + assert client.api_key is None + + def test_api_key_mode_uses_key_even_with_aws_credentials(self) -> None: + client = AnthropicBedrockMantle( + auth_mode="api_key", + api_key="my-key", + aws_access_key="AKID", + aws_secret_key="secret", + aws_region="us-east-1", + ) + + assert client.auth_mode == "api_key" + assert client._use_sigv4 is False + assert client.auth_headers == {"Authorization": "Bearer my-key"} + + def test_api_key_mode_requires_a_key(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("ANTHROPIC_AWS_API_KEY", raising=False) + + with pytest.raises(ValueError, match="auth_mode='api_key'.*requires an API key"): + AnthropicBedrockMantle(auth_mode="api_key", base_url="https://example.com") + + def test_skip_auth_cannot_be_combined_with_explicit_mode(self) -> None: + with pytest.raises(ValueError, match="skip_auth.*mutually exclusive"): + AnthropicBedrockMantle(skip_auth=True, auth_mode="sigv4", base_url="https://example.com") + class TestSkipAuth: def test_skip_auth_does_not_sign_request(self, get_auth_headers_recorder: GetAuthHeadersRecorder) -> None: @@ -278,6 +315,14 @@ def fake_get_auth_headers(**_: object) -> dict[str, str]: assert signing_threads[0] != threading.get_ident() assert request.headers["Authorization"] == "AWS4-HMAC-SHA256 stub" + def test_explicit_auth_mode(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "ambient-key") + client = AsyncAnthropicBedrockMantle(auth_mode="sigv4", aws_region="us-east-1") + + assert client.auth_mode == "sigv4" + assert client._use_sigv4 is True + assert client.api_key is None + class TestCopy: def test_copy_preserves_config(self) -> None: @@ -297,6 +342,29 @@ def test_copy_overrides_region(self) -> None: copied = client.copy(aws_region="us-west-2") assert copied.aws_region == "us-west-2" + def test_copy_preserves_auth_mode(self) -> None: + client = AnthropicBedrockMantle(auth_mode="sigv4", aws_region="us-east-1") + copied = client.copy() + + assert copied.auth_mode == "sigv4" + assert copied._use_sigv4 is True + + def test_copy_can_switch_to_sigv4(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "ambient-key") + client = AnthropicBedrockMantle(api_key="my-key", aws_region="us-east-1") + copied = client.copy(auth_mode="sigv4") + + assert copied.auth_mode == "sigv4" + assert copied._use_sigv4 is True + assert copied.api_key is None + + def test_async_copy_preserves_auth_mode(self) -> None: + client = AsyncAnthropicBedrockMantle(auth_mode="sigv4", aws_region="us-east-1") + copied = client.copy() + + assert copied.auth_mode == "sigv4" + assert copied._use_sigv4 is True + def test_copy_x_stainless_helper_header_appends(self) -> None: # `x-stainless-helper` accumulates across copies instead of being clobbered client = AnthropicBedrockMantle(