fix: address SDE feedback. Add @requires_iam_access_token decorator for AWS STS JWT tokens

- Add new @requires_iam_access_token decorator for AWS IAM JWT federation
- Decorator calls STS:GetWebIdentityToken directly (not via AgentCore Identity)
- Separate from @requires_access_token to avoid confusion between OAuth and STS flows
- Bump boto3 minimum version to support new STS API

Author: Sundar Raghavan

src/bedrock_agentcore/identity/auth.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, Dict, List, Literal, Optional
 
 import boto3
+from botocore.exceptions import ClientError
 
 from bedrock_agentcore.runtime import BedrockAgentCoreContext
 from bedrock_agentcore.services.identity import IdentityClient, TokenPoller
@@ -20,94 +21,41 @@
 
 def requires_access_token(
     *,
-    # OAuth parameters (required for M2M and USER_FEDERATION)
-    provider_name: Optional[str] = None,
-    scopes: Optional[List[str]] = None,
+    provider_name: str,
+    into: str = "access_token",
+    scopes: List[str],
     on_auth_url: Optional[Callable[[str], Any]] = None,
+    auth_flow: Literal["M2M", "USER_FEDERATION"],
     callback_url: Optional[str] = None,
     force_authentication: bool = False,
     token_poller: Optional[TokenPoller] = None,
     custom_state: Optional[str] = None,
     custom_parameters: Optional[Dict[str, str]] = None,
-    # AWS JWT parameters (required for AWS_JWT)
-    audience: Optional[List[str]] = None,
-    signing_algorithm: str = "ES384",
-    duration_seconds: int = 300,
-    tags: Optional[List[Dict[str, str]]] = None,
-    # Common parameters
-    into: str = "access_token",
-    auth_flow: Literal["M2M", "USER_FEDERATION", "AWS_JWT"] = "USER_FEDERATION",
 ) -> Callable:
-    """Decorator that fetches an access token before calling the decorated function.
-
-    Supports three authentication flows:
+    """Decorator that fetches an OAuth2 access token before calling the decorated function.
 
-    1. USER_FEDERATION (OAuth 3LO): User consent required, uses credential provider
-    2. M2M (OAuth client credentials): Machine-to-machine, uses credential provider
-    3. AWS_JWT: Direct AWS STS JWT, no secrets required
-
-    OAuth Parameters (for M2M and USER_FEDERATION):
-        provider_name: The credential provider name (required for OAuth flows)
-        scopes: OAuth2 scopes to request (required for OAuth flows)
-        on_auth_url: Callback for handling authorization URLs (USER_FEDERATION only)
+    Args:
+        provider_name: The credential provider name
+        into: Parameter name to inject the token into
+        scopes: OAuth2 scopes to request
+        on_auth_url: Callback for handling authorization URLs
+        auth_flow: Authentication flow type ("M2M" or "USER_FEDERATION")
         callback_url: OAuth2 callback URL
         force_authentication: Force re-authentication
         token_poller: Custom token poller implementation
-        custom_state: State for callback verification
-        custom_parameters: Additional OAuth parameters
-
-    AWS JWT Parameters (for AWS_JWT):
-        audience: List of intended token recipients (required for AWS_JWT)
-        signing_algorithm: 'ES384' (default) or 'RS256'
-        duration_seconds: Token lifetime 60-3600 (default 300)
-        tags: Custom claims as [{'Key': str, 'Value': str}, ...]
-
-    Common Parameters:
-        into: Parameter name to inject the token into (default: 'access_token')
-        auth_flow: Authentication flow type
+        custom_state: A state that allows applications to verify the validity of callbacks to callback_url
+        custom_parameters: A map of custom parameters to include in authorization request to the credential provider
+                           Note: these parameters are in addition to standard OAuth 2.0 flow parameters
 
     Returns:
         Decorator function
-
-    Examples:
-        # OAuth USER_FEDERATION flow
-        @requires_access_token(
-            provider_name="CognitoProvider",
-            scopes=["openid"],
-            auth_flow="USER_FEDERATION",
-            on_auth_url=lambda url: print(f"Please authorize: {url}")
-        )
-        async def call_oauth_api(*, access_token: str):
-            ...
-
-        # AWS JWT flow (no secrets!)
-        @requires_access_token(
-            auth_flow="AWS_JWT",
-            audience=["https://api.example.com"],
-            signing_algorithm="ES384",
-        )
-        async def call_external_api(*, access_token: str):
-            ...
     """
-    # Validate parameters based on flow
-    if auth_flow in ["M2M", "USER_FEDERATION"]:
-        if not provider_name:
-            raise ValueError(f"provider_name is required for auth_flow='{auth_flow}'")
-        if not scopes:
-            raise ValueError(f"scopes is required for auth_flow='{auth_flow}'")
-    elif auth_flow == "AWS_JWT":
-        if not audience:
-            raise ValueError("audience is required for auth_flow='AWS_JWT'")
-        if signing_algorithm not in ["ES384", "RS256"]:
-            raise ValueError("signing_algorithm must be 'ES384' or 'RS256'")
-        if not (60 <= duration_seconds <= 3600):
-            raise ValueError("duration_seconds must be between 60 and 3600")
 
     def decorator(func: Callable) -> Callable:
         client = IdentityClient(_get_region())
 
-        async def _get_oauth_token() -> str:
-            """Get token via OAuth flow (existing logic)."""
+        async def _get_token() -> str:
+            """Common token fetching logic."""
             return await client.get_token(
                 provider_name=provider_name,
                 agent_identity_token=await _get_workload_access_token(client),
@@ -121,23 +69,6 @@ async def _get_oauth_token() -> str:
                 custom_parameters=custom_parameters,
             )
 
-        async def _get_aws_jwt_token() -> str:
-            """Get token via AWS STS (new logic)."""
-            result = client.get_aws_jwt_token_sync(
-                audience=audience,
-                signing_algorithm=signing_algorithm,
-                duration_seconds=duration_seconds,
-                tags=tags,
-            )
-            return result["token"]
-
-        async def _get_token() -> str:
-            """Route to appropriate token retrieval method."""
-            if auth_flow == "AWS_JWT":
-                return await _get_aws_jwt_token()
-            else:
-                return await _get_oauth_token()
-
         @wraps(func)
         async def async_wrapper(*args: Any, **kwargs_func: Any) -> Any:
             token = await _get_token()
@@ -170,6 +101,122 @@ def sync_wrapper(*args: Any, **kwargs_func: Any) -> Any:
     return decorator
 
 
+def requires_iam_access_token(
+    *,
+    audience: List[str],
+    signing_algorithm: str = "ES384",
+    duration_seconds: int = 300,
+    tags: Optional[List[Dict[str, str]]] = None,
+    into: str = "access_token",
+) -> Callable:
+    """Decorator that fetches an AWS IAM JWT token before calling the decorated function.
+
+    This decorator obtains a signed JWT from AWS STS using the GetWebIdentityToken API.
+    The JWT can be used to authenticate with external services that support OIDC token
+    validation. No client secrets are required - the token is signed by AWS.
+
+    This is separate from @requires_access_token which uses AgentCore Identity for
+    OAuth 2.0 flows. Use this decorator for M2M authentication with services that
+    accept AWS-signed JWTs.
+
+    Args:
+        audience: List of intended token recipients (populates 'aud' claim in JWT).
+                  Must match what the external service expects.
+        signing_algorithm: Algorithm for signing the JWT.
+                        'ES384' (default) or 'RS256'.
+        duration_seconds: Token lifetime in seconds (60-3600, default 300).
+        tags: Optional custom claims as [{'Key': str, 'Value': str}, ...].
+              These are added to the JWT as additional claims.
+        into: Parameter name to inject the token into (default: 'access_token').
+
+    Returns:
+        Decorator function that wraps the target function.
+
+    Raises:
+        ValueError: If parameters are invalid.
+        RuntimeError: If AWS JWT federation is not enabled for the account.
+        ClientError: If the STS API call fails.
+
+    Example:
+        @tool
+        @requires_iam_access_token(
+            audience=["https://api.example.com"],
+            signing_algorithm="ES384",
+            duration_seconds=300,
+        )
+        def call_external_api(query: str, *, access_token: str) -> str:
+            '''Call external API with AWS JWT authentication.'''
+            import requests
+            response = requests.get(
+                "https://api.example.com/data",
+                headers={"Authorization": f"Bearer {access_token}"},
+                params={"q": query},
+            )
+            return response.text
+
+    Note:
+        Before using this decorator, you must:
+        1. Enable AWS IAM Outbound Web Identity Federation for your account
+           (via `agentcore identity setup-aws-jwt` or IAM API)
+        2. Ensure the execution role has `sts:GetWebIdentityToken` permission
+        3. Configure the external service to trust your AWS account's issuer URL
+    """
+    # Validate parameters
+    if not audience:
+        raise ValueError("audience is required")
+    if signing_algorithm not in ["ES384", "RS256"]:
+        raise ValueError("signing_algorithm must be 'ES384' or 'RS256'")
+    if not (60 <= duration_seconds <= 3600):
+        raise ValueError("duration_seconds must be between 60 and 3600")
+
+    logger = logging.getLogger(__name__)
+
+    def _get_iam_jwt_token(region: str) -> str:
+        """Get JWT from AWS STS - NO IdentityClient involved."""
+        logger.info("Getting AWS IAM JWT token from STS...")
+        sts_client = boto3.client("sts", region_name=region)
+
+        params = {
+            "Audience": audience,
+            "SigningAlgorithm": signing_algorithm,
+            "DurationSeconds": duration_seconds,
+        }
+        if tags:
+            params["Tags"] = tags
+
+        try:
+            response = sts_client.get_web_identity_token(**params)
+            logger.info("Successfully obtained AWS IAM JWT token")
+            return response["WebIdentityToken"]
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code", "")
+            if error_code in ["FeatureDisabledException", "FeatureDisabled"]:
+                raise RuntimeError("AWS IAM Outbound Web Identity Federation is not enabled.") from e
+            logger.error("Failed to get AWS IAM JWT token: %s", str(e))
+            raise
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs_func: Any) -> Any:
+            region = _get_region()
+            token = _get_iam_jwt_token(region)
+            kwargs_func[into] = token
+            return await func(*args, **kwargs_func)
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs_func: Any) -> Any:
+            region = _get_region()
+            token = _get_iam_jwt_token(region)
+            kwargs_func[into] = token
+            return func(*args, **kwargs_func)
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        return sync_wrapper
+
+    return decorator
+
+
 def requires_api_key(*, provider_name: str, into: str = "api_key") -> Callable:
     """Decorator that fetches an API key before calling the decorated function.
 

src/bedrock_agentcore/services/identity.py

@@ -249,97 +249,3 @@ async def get_api_key(self, *, provider_name: str, agent_identity_token: str) ->
         req = {"resourceCredentialProviderName": provider_name, "workloadIdentityToken": agent_identity_token}
 
         return self.dp_client.get_resource_api_key(**req)["apiKey"]
-
-    async def get_aws_jwt_token(
-        self,
-        *,
-        audience: List[str],
-        signing_algorithm: str = "ES384",
-        duration_seconds: int = 300,
-        tags: Optional[List[Dict[str, str]]] = None,
-    ) -> Dict[str, Any]:
-        """Get a signed JWT from AWS STS for external service authentication.
-
-        This method calls STS:GetWebIdentityToken directly - no AgentCore Identity
-        service involvement, no secrets required.
-
-        Args:
-            audience: List of intended recipients (populates 'aud' claim)
-            signing_algorithm: Algorithm for signing ('ES384' or 'RS256')
-            duration_seconds: Token lifetime (60-3600, default 300)
-            tags: Optional list of {'Key': str, 'Value': str} for custom claims
-
-        Returns:
-            Dict with 'token' (the JWT string) and 'expiration' (datetime)
-
-        Raises:
-            Exception: If token generation fails
-        """
-        self.logger.info("Getting AWS JWT token...")
-
-        # Create regional STS client (GetWebIdentityToken not available on global endpoint)
-        sts_client = boto3.client("sts", region_name=self.region)
-
-        # Build request parameters
-        params = {
-            "Audience": audience,
-            "SigningAlgorithm": signing_algorithm,
-            "DurationSeconds": duration_seconds,
-        }
-
-        # Add tags if provided
-        if tags:
-            params["Tags"] = tags
-
-        try:
-            response = sts_client.get_web_identity_token(**params)
-
-            self.logger.info("Successfully obtained AWS JWT token")
-
-            return {
-                "token": response["WebIdentityToken"],
-                "expiration": response["Expiration"],
-            }
-
-        except Exception as e:
-            self.logger.error("Failed to get AWS JWT token: %s", str(e))
-            raise
-
-    def get_aws_jwt_token_sync(
-        self,
-        *,
-        audience: List[str],
-        signing_algorithm: str = "ES384",
-        duration_seconds: int = 300,
-        tags: Optional[List[Dict[str, str]]] = None,
-    ) -> Dict[str, Any]:
-        """Synchronous version of get_aws_jwt_token.
-
-        See get_aws_jwt_token for full documentation.
-        """
-        self.logger.info("Getting AWS JWT token (sync)...")
-
-        sts_client = boto3.client("sts", region_name=self.region)
-
-        params = {
-            "Audience": audience,
-            "SigningAlgorithm": signing_algorithm,
-            "DurationSeconds": duration_seconds,
-        }
-
-        if tags:
-            params["Tags"] = tags
-
-        try:
-            response = sts_client.get_web_identity_token(**params)
-
-            self.logger.info("Successfully obtained AWS JWT token")
-
-            return {
-                "token": response["WebIdentityToken"],
-                "expiration": response["Expiration"],
-            }
-
-        except Exception as e:
-            self.logger.error("Failed to get AWS JWT token: %s", str(e))
-            raise