fix(bedrock): surface web identity token aud/iss on InvalidIdentityToken

litellm/llms/bedrock/base_aws_llm.py

@@ -1,3 +1,4 @@
+import base64
 import hashlib
 import json
 import os
@@ -19,7 +20,7 @@
 )
 
 import httpx
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 from litellm._logging import verbose_logger
 from litellm.caching.caching import DualCache
@@ -56,6 +57,11 @@ class Boto3CredentialsInfo(BaseModel):
     aws_bedrock_runtime_endpoint: Optional[str]
 
 
+class _WebIdentityTokenClaims(BaseModel):
+    aud: Optional[Union[str, list[str]]] = None
+    iss: Optional[str] = None
+
+
 class AwsAuthError(Exception):
     def __init__(self, status_code, message):
         self.status_code = status_code
@@ -817,6 +823,25 @@ def _is_already_running_as_role(
 
         return False
 
+    @staticmethod
+    def _unverified_web_identity_audience(oidc_token: str) -> Optional[str]:
+        """Return the public ``aud``/``iss`` claims of a web identity JWT
+        without verifying its signature, so a rejected-token error can name
+        the audience LiteLLM actually sent. The signature is never read, so no
+        secret is exposed."""
+        segments = oidc_token.split(".")
+        if len(segments) != 3:
+            return None
+        payload = segments[1]
+        try:
+            decoded = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))
+            claims = _WebIdentityTokenClaims.model_validate_json(decoded)
+        except (ValueError, ValidationError):
+            return None
+        if claims.aud is None and claims.iss is None:
+            return None
+        return f"aud={claims.aud!r}, iss={claims.iss!r}"
+
     @tracer.wrap()
     def _auth_with_web_identity_token(
         self,
@@ -925,7 +950,21 @@ def _auth_with_web_identity_token(
         if aws_external_id is not None:
             assume_role_params["ExternalId"] = aws_external_id
 
-        sts_response = sts_client.assume_role_with_web_identity(**assume_role_params)
+        try:
+            sts_response = sts_client.assume_role_with_web_identity(
+                **assume_role_params
+            )
+        except sts_client.exceptions.InvalidIdentityTokenException as e:
+            audience = (
+                self._unverified_web_identity_audience(oidc_token)
+                if isinstance(oidc_token, str)
+                else None
+            )
+            detail = f" Token {audience}" if audience else ""
+            raise AwsAuthError(
+                status_code=401,
+                message=f"AWS STS rejected the web identity token: {e}.{detail}",
+            ) from e
 
         iam_creds_dict = {
             "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"],

tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py

@@ -29,6 +29,7 @@
 action.
 """
 
+import base64
 import json
 from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, patch
@@ -157,6 +158,78 @@ def test_no_aws_external_anthropic_statement_collision(self):
         )
 
 
+def _make_jwt(payload: dict) -> str:
+    def _segment(data: dict) -> str:
+        return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode()
+
+    return f"{_segment({'alg': 'RS256', 'typ': 'JWT'})}.{_segment(payload)}.signature"
+
+
+class TestInvalidIdentityTokenSurfacesAudience:
+    """LIT-4026: when STS rejects the web identity token with
+    ``InvalidIdentityToken`` (the "Incorrect token audience" case), the raised
+    error must name the ``aud``/``iss`` the token actually carries so an
+    operator can diagnose the mismatch without enabling LITELLM_LOG=DEBUG on a
+    prod instance."""
+
+    _AUD = "https://guidepoint.litellm-prod.ai"
+    _ISS = "https://accounts.google.com"
+    _STS_MESSAGE = (
+        "An error occurred (InvalidIdentityToken) when calling the "
+        "AssumeRoleWithWebIdentity operation: Incorrect token audience"
+    )
+
+    def _raise_invalid_identity_token(self) -> Exception:
+        from litellm.llms.bedrock.base_aws_llm import AwsAuthError, BaseAWSLLM
+
+        token = _make_jwt({"aud": self._AUD, "iss": self._ISS, "sub": "svc-account"})
+
+        mock_sts = MagicMock()
+
+        class _InvalidIdentityTokenException(Exception):
+            pass
+
+        mock_sts.exceptions.InvalidIdentityTokenException = (
+            _InvalidIdentityTokenException
+        )
+        mock_sts.assume_role_with_web_identity.side_effect = (
+            _InvalidIdentityTokenException(self._STS_MESSAGE)
+        )
+
+        with (
+            patch("boto3.client", return_value=mock_sts),
+            patch(
+                "litellm.llms.bedrock.base_aws_llm.get_secret",
+                return_value=token,
+            ),
+            pytest.raises(AwsAuthError) as exc_info,
+        ):
+            BaseAWSLLM()._auth_with_web_identity_token(
+                aws_web_identity_token="oidc/google/" + self._AUD,
+                aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role",
+                aws_session_name="test-session",
+                aws_region_name="us-east-1",
+                aws_sts_endpoint=None,
+            )
+        return exc_info.value
+
+    def test_error_names_token_audience(self):
+        err = self._raise_invalid_identity_token()
+        assert self._AUD in str(err)
+
+    def test_error_names_token_issuer(self):
+        err = self._raise_invalid_identity_token()
+        assert self._ISS in str(err)
+
+    def test_error_preserves_original_sts_reason(self):
+        err = self._raise_invalid_identity_token()
+        assert "Incorrect token audience" in str(err)
+
+    def test_error_is_401(self):
+        err = self._raise_invalid_identity_token()
+        assert err.status_code == 401
+
+
 class TestPolicyTransportConditions:
     def test_bedrock_statement_keeps_secure_transport_condition(self):
         policy = _captured_policy()