feat(bedrock): add explicit Mantle auth mode

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,33 +103,52 @@ 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]]:
     """Resolve and validate all Mantle client configuration.
 
     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,

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(