Add support for Identity OAuth 2.0 federation enhancements

Author: brianlaoaws

src/bedrock_agentcore/identity/auth.py

@@ -28,6 +28,7 @@ def requires_access_token(
     callback_url: Optional[str] = None,
     force_authentication: bool = False,
     token_poller: Optional[TokenPoller] = None,
+    session_state: Optional[str] = None,
 ) -> Callable:
     """Decorator that fetches an OAuth2 access token before calling the decorated function.
 
@@ -59,6 +60,7 @@ async def _get_token() -> str:
                 callback_url=callback_url,
                 force_authentication=force_authentication,
                 token_poller=token_poller,
+                session_state=session_state,
             )
 
         @wraps(func)

src/bedrock_agentcore/services/identity.py

@@ -120,6 +120,7 @@ async def get_token(
         callback_url: Optional[str] = None,
         force_authentication: bool = False,
         token_poller: Optional[TokenPoller] = None,
+        session_state: Optional[str] = None,
     ) -> str:
         """Get an OAuth2 access token for the specified provider.
 
@@ -132,6 +133,7 @@ async def get_token(
             callback_url: OAuth2 callback URL (must be pre-registered)
             force_authentication: Force re-authentication even if token exists in the token vault
             token_poller: Custom token poller implementation
+            session_state: A state that allows applications to verify the validity of call backs to callback_url
 
         Returns:
             The access token string
@@ -155,6 +157,8 @@ async def get_token(
             req["resourceOauth2ReturnUrl"] = callback_url
         if force_authentication:
             req["forceAuthentication"] = force_authentication
+        if session_state:
+            req["includedState"] = session_state
 
         response = self.dp_client.get_resource_oauth2_token(**req)
 
@@ -176,6 +180,9 @@ async def get_token(
             if force_authentication:
                 req["forceAuthentication"] = False
 
+            if "sessionUri" in response:
+                req["sessionUri"] = response["sessionUri"]
+
             # Poll for the token
             active_poller = token_poller or _DefaultApiTokenPoller(
                 auth_url, lambda: self.dp_client.get_resource_oauth2_token(**req).get("accessToken", None)

tests/bedrock_agentcore/services/test_identity.py

@@ -239,6 +239,7 @@ async def test_get_token_with_optional_parameters(self):
             agent_identity_token = "test-agent-token"
             callback_url = "https://example.com/callback"
             force_authentication = True
+            session_state = "myAppIncludedState"
             expected_token = "test-access-token"
 
             mock_client.get_resource_oauth2_token.return_value = {"accessToken": expected_token}
@@ -250,6 +251,7 @@ async def test_get_token_with_optional_parameters(self):
                 auth_flow="USER_FEDERATION",
                 callback_url=callback_url,
                 force_authentication=force_authentication,
+                session_state=session_state,
             )
 
             assert result == expected_token
@@ -260,6 +262,7 @@ async def test_get_token_with_optional_parameters(self):
                 workloadIdentityToken=agent_identity_token,
                 resourceOauth2ReturnUrl=callback_url,
                 forceAuthentication=force_authentication,
+                includedState=session_state,
             )
 
     @pytest.mark.asyncio