Post Syndicated from Abrom Douglas original https://aws.amazon.com/blogs/security/empower-ai-agents-with-user-context-using-amazon-cognito/

Amazon Cognito is a managed customer identity and access management (CIAM) service that enables seamless user sign-up and sign-in for web and mobile applications. Through user pools, Amazon Cognito provides a user directory with strong authentication features, including passkeys, federation to external identity providers (IdPs), and OAuth 2.0 flows for secure machine-to-machine (M2M) authorization.

Amazon Cognito issues standard JSON Web Tokens (JWTs) and supports the customization of identity and access tokens for user authentication by using the pre token generation Lambda trigger. Learn more about this in How to customize access tokens in Amazon Cognito user pools. Amazon Cognito has extended token customization capabilities to support access token customization for M2M and the ability to pass metadata from the client during M2M authorization. Application builders can use these two features to support multiple use cases, including customizing access tokens based on unique runtime policies, entitlements, environment, or passed metadata. This can simplify and enrich M2M authentication and authorization scenarios and opens up new possibilities for emerging use cases, such as identity and access management for AI agents.

This post demonstrates how Amazon Cognito enables AI agents to perform authorized actions on behalf of users through user-contextualized access tokens for OAuth 2.0-enabled resource servers. AI agents represent a class of autonomous services that require robust identity management and precise access control, especially when acting on behalf of users. By using the Amazon Cognito client credentials flow with access token customization, you can establish distinct identities for AI agents that carry critical information about their capabilities, scope of access, and intended use cases. This approach provides a foundation for more secure, auditable AI agent operations while maintaining clear boundaries around their authorized activities.

The identity of an AI agent can be represented within Amazon Cognito as an app client. The AI agent obtains an access token (JSON Web Token (JWT)) through an OAuth 2.0 client credentials grant. This JWT can be customized to contain claims that represent the authenticated human user whom the AI agent is acting on behalf of. This token can then be used to authorize access to other services that has established trust with the Amazon Cognito user pool by trusting the issuer and audience of the token. For example, this third-party service could be a claims processor, a travel agency service, or a scheduling service acting on behalf of a user. The focus of this post is on foundational building blocks using Amazon Cognito for AI agents and how to obtain a customized access token with user context.

Solution overview and reference architecture

Looking at an example architecture (Figure 1), a user signs in to a web or mobile application using an Amazon Cognito user pool, and tokens for the user are returned to the client. Here, the application could be a serverless digital assistant using an Amazon Bedrock agent that needs to gather and process data residing in a third-party cross-domain service. The AI agent obtains its own access token by performing an OAuth 2.0 client credentials grant while passing the user’s access token as context using the aws_client_metadata request parameter. The AI agent receives the user contextualized access token and calls an external, third-party, or cross-domain service that trusts the issuer and audience of the AI agent’s access token issued from an Amazon Cognito user pool. The cross-domain service can obtain the JSON Web Key Set (JWKS) to verify the token and extract claims presenting both the AI agent and most importantly, the underlying user. Authorization takes place within the cross-domain service using the claims of the customized access token and for fine-grain authorization, Amazon Verified Permissions is used. See Figure 1 for a detailed flow of this example.

Figure 1: AI agent identity reference architecture

1. The user navigates to the application through the client.
2. There is no existing session or token for the user, so the user authentication flow with the Amazon Cognito user pool begins.
3. After a successful sign-in, Amazon Cognito returns access, ID, and refresh tokens to the client for the user.
4. As the user interacts with AI agent through the application, the client sends the user’s access token to an Amazon API Gateway endpoint.
5. The API gateway integrates with the AI agent, which is using an Amazon Bedrock agent. As an example, this can use several AWS Lambda functions interacting with an Amazon Bedrock Knowledge Base or a Retrieval-Augmented Generation (RAG) process.
6. The AI agent obtains its own access token from an Amazon Cognito user pool using an OAuth 2.0 client credentials grant. The user’s access token, obtained in step 1, is sent with the token request in the aws_client_metadata request parameter.

Note: You can use different Amazon Cognito user pools for user authentication and for agent (machine) authentication. This promotes separation and provides the ability to apply different settings and controls on each user pool if needed to meet security requirements.

7. Amazon Cognito validates the client ID and secret from the AI agent and invokes the pre token generation Lambda trigger to customize the access token for the AI agent.

Note: Within the pre token generation Lambda trigger, the user’s access token is verified before returning a customized access token to the AI agent using the aws-jwt-verify library.

8. The customized access token is returned to the AI agent, including custom claims representing the user.
9. The AI agent, using its own access token, calls the cross-domain service to perform the requested action on behalf of the user. For example, this can be a third-party reservation system or a photo sharing service.
10. The resource server in the cross-domain service verifies that the access token from the AI agent is valid. The resource server must be pre-configured to trust the user pool that issued the agent access token.
11. Coarse- and fine-grained authorization can happen either locally in the service code or using Verified Permissions.
12. A response from the cross-domain service flows back to the AI agent, if necessary.
13. A response from the AI agent to the user application or client is returned, if necessary.
14. Actions that take place throughout the flow are logged in AWS CloudTrail, providing end-to-end logging and auditing.

Implementation details

Let’s take a deeper look into the three core components of this scenario:

1. The AI agent obtaining its own OAuth 2.0 access token
2. The Amazon Cognito pre token generation Lambda trigger used to enrich the AI agent’s access token with user context
3. The cross-domain resource server performing fine-grained authorization

AI agent

Figure 2: AI agent obtaining a user access token from the frontend application through API Gateway

Amazon Bedrock Agents is used in this solution with a
custom orchestration configured to use Lambda. When the application interacts with the Amazon Bedrock agent, the custom orchestrator initiation begins with the agent passing the user’s access token to a Lambda function as part of the custom orchestration (shown in Figure 2). The Lambda function validates the user’s token to verify that it’s not expired and hasn’t been tampered with. This custom orchestrator begins the process for the agent to obtain its own OAuth access token and to access downstream and cross-domain resources on behalf of the user. The human user’s access token is included in the call from the application through the client. To learn more about Amazon Bedrock Agents custom orchestrator, see
Getting started with Amazon Bedrock Agents custom orchestrator. The following is an example of what a human user’s decoded access token provided through an API Gateway REST API might look like.

{
  sub: "user-identity-4e4c-example-7cede8e609a2",
  cognito:groups: 
    [
    "exampleChatApplicationAccess"
    ]
  ,
  iss: https://cognito-idp.<region>.amazonaws.com/<region>_example,
  version: 2,
  client_id: "1example23456789",
  origin_jti: "",
  token_use: "access",
  scope: "openid profile email",
  auth_time: 499192140,
  exp: 1445444940,
  iat: 499192140,
  jti: "",
  username: "my-example-username"
}

The following is a Node.js code sample that an AI agent can use to obtain its own access token from Amazon Cognito. This can be the Lambda function part of the custom orchestration for the Amazon Bedrock agent. Notice the clientMetadata variable being set, which will be passed to the Cognito /token endpoint using the aws_client_metadata request parameter. This request parameter is where the user’s access token is provided. In the following code example, you will find an attribute called callerApp, which is set to ExampleChatApplication, which serves as a unique identifier for the application. The callerApp value is preconfigured in the backend of the solution. This unique application identifier is included in the customized access token for the agent and used for additional authorization checks later. It’s a security best practice to use AWS Secrets Manager to store the client ID and client secret and obtain these credentials at runtime. As a security best practice, the user’s access token should be verified prior to passing it to the AI agent backend.

async function getAccessToken() {
    const clientId = 'exampleAiAgentClientId'; // use Secrets Manager
    const clientSecret = 'exampleAiAgentClientSecret'; // use Secrets Manager
    const tokenEndpoint = 'https://mydomain.auth.<region>.amazoncognito.com/oauth2/token';
    const scope = 'crossDomainService/read userData/read';
    const clientMetadata = '{"onBehalfOfToken":"<HUMAN-USER-ACCESS-TOKEN>", "callerApp":"ExampleChatApplication"}';
  
    const basicAuth = Buffer.from(`${clientId}:${clientSecret}`).toString('base64');
  
    const body = new URLSearchParams({
      grant_type: 'client_credentials',
      scope,
      aws_client_metadata: clientMetadata
    });
  
    const res = await fetch(tokenEndpoint, {
      method: 'POST',
      headers: {
        'Authorization': `Basic ${basicAuth}`,
        'Content-Type': 'application/x-www-form-urlencoded'
      },
      body
    });
  
    if (!res.ok) {
      const error = await res.text();
      throw new Error(`Token request failed: ${res.status} ${error}`);
    }
  
    const { access_token } = await res.json();
    console.log('Access Token:', access_token);
  
    return access_token;
  }
  
  getAccessToken().catch(err => console.error('Error:', err.message));

The access token for the AI agent is returned only if the client ID and secret are correct and the provided user access token is valid. However, before it’s returned, the AI agent’s access token is customized by the Amazon Cognito pre token generation Lambda trigger.

Amazon Cognito pre token generation Lambda trigger

Figure 3: AI agent access token customization with Cognito pre token generation Lambda trigger

After the AI agent’s action calls the Amazon Cognito /token endpoint with a valid client ID and secret, Cognito invokes the pre token generation Lambda trigger. The following is an example Lambda function that takes the aws_client_metadata request parameter, which contains the access token of the user and the callerApp attribute that was defined while the user was authenticating. In the following Lambda function, the access token provided from the user is verified (shown in Figure 3). The aws-jwt-verify library is used to verify the token is not expired, the token has not been tampered with by verifying the signature, and it’s making sure that an access token was provided. The Lambda function is also pre-configured to accept user tokens from a specific issuer and audience, this protects against malicious context injection risks. This is also an opportunity to perform additional authorization. For example, check if the user is a member of certain groups.

After the token is verified, the Lambda function customizes the access token to be returned to the AI agent.

import { CognitoJwtVerifier } from "aws-jwt-verify";

// Initialize the JWT verifier to verify the user’s access token
// Provide the user pool ID, token use, and client ID 
const jwtVerifier = CognitoJwtVerifier.create({
  userPoolId: process.env.USER_POOL_ID,  // user pool for user authentication
  clientId: process.env.CLIENT_ID,
  // groups: "exampleChatApplicationAccess", // optional group membership authorization
  tokenUse: 'access'
});

export const handler = async function(event, context) {
  try {
    const onBehalfOfToken = event.request.clientMetadata?.onBehalfOfToken || '';
    // It’s recommended that the provided “callerApp” value from the application is authorized for use with the app client for the AI agent
    const callerApp = event.request.clientMetadata?.callerApp || '';

    // The below console log will display the authenticated user’s JWT
    // Keep this logging with caution in a production environment
    console.log('Original event:', event);

    // Verify the access token from the human user
    // You could optionally also perform some authorization checks here as well
    // Example: check for the membership of a group
    let decodedJWT;
    if (onBehalfOfToken) {
      try {
        decodedJWT = await jwtVerifier.verify(onBehalfOfToken);
        console.log('Decoded JWT:', decodedJWT);
      } catch (err) {
        console.error('Token verification failed:', err);
        throw new Error('Token verification failed');
      }
    }

    // Create the onBehalfOf claim structure
    const behalfOfClaim = decodedJWT ? {
      sub: decodedJWT.sub,
      username: decodedJWT.username,
      groups: decodedJWT['cognito:groups'] || []
    } : {};

    // Customized token returned to client
    event.response = {
      "claimsAndScopeOverrideDetails": {
        "accessTokenGeneration": {
          "claimsToAddOrOverride": {
            "onBehalfOf": behalfOfClaim,
            "callerApp": callerApp
          },
        }
      }
    };

    return event;

  } catch (error) {
    console.error('Error in Lambda execution:', error);
    throw error;
  }
};

Notice in the preceding Lambda function that two custom claims are being dynamically created within the event.response: onBehalfOf and callerApp. The onBehalfOf claim contains nested claims that were extracted from the human user’s access token. The callerApp is carried forward from the frontend application and provided alongside the user access token. It’s recommended for the callerApp value to also be verified against some custom logic to add additional layer of protection. The return AI agent’s access token would look like the following JWT.

{    
	"sub": "agent-identity-4e4c-example-7cede8e609a2",
	"onBehalfOf": {
		"sub": "user-identity-4e4c-example-7cede8e609a2",
		"username": "my-example-username",
		"groups": [
			"readaccess"        
				]    
		},    
		"iss": "https://cognito-idp..amazonaws.com/_example",
		"version": 2,
		"client_id": "1example23456789",
		"callerApp": "ExampleChatApplication",
		"token_use": "access",
		"scope": "crossDomainService123/read userData/read",
		"auth_time": 499192140,
		"exp": 1445444940,
		"iat": 499192140,
		"jti": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
}

Cross-domain resource server authorization check

At this point, shown in Figure 4, the human user has successfully authenticated to the web application, the human user’s access token was sent as context to the backend, an AI agent obtained its own customized access token containing the human user context, and now the agent is ready to call an external cross-domain service.

Figure 4: Cross-domain resource server performing fine-grained authorization with Amazon Verified Permissions

As shown in Figure 4, the cross-domain service is the resource server and therefore needs to perform an authorization check. For this example, we’ll keep things straightforward and make sure that three core things are verified:

1. The AI agent’s OAuth access token is valid
2. The AI agent is authorized to access this service
3. The AI agent is authorized to interact with the user data

Depending on your use case and requirements, you might also need to verify that the user’s consent has been obtained prior to the AI agent acting on their behalf. Ultimately, you want to verify that the AI agent can access a user’s data on their behalf and only for the purpose for which consent has been provided by the user.

For the token verification, use the aws-jwt-verify library again. The following is a Node.js example to verify the AI agent’s access token.

import { CognitoJwtVerifier } from "aws-jwt-verify";

// add custom logic to verify that AI agent is authorized to perform this action on behalf of the user

// Verifier that expects valid access tokens:
const verifier = CognitoJwtVerifier.create({
  userPoolId: "<user_pool_id>", // user pool for AI agent authentication
  tokenUse: "access",
  clientId: "<client_id>",
});

try {
  const payload = await verifier.verify(
    "eyJraWQeyJhdF9oYXNoIjoidk..." //this will be the AI agent's access token
  );
  console.log("Token is valid. Payload:", payload);
} catch {
  console.log("Token not valid!");
}

Fine-grained authorization with Verified Permissions

As a security best practice, the zero trust principle of enforcing fine-grained identity-based authorization should take place using Verified Permissions. The preceding Node.js code sample is a basic validation of the AI agents access token that can happen within the application logic. Instead of keeping authorization logic within the resource server, you can use Verified Permissions to offload the authorization policies to a managed service. The following is an example Cedar policy for this use case.

permit(
    principal == Agent::"agent-identity-4e4c-example-7cede8e609a2",
    action == Action::"readOnly",
    resource == Resource::"crossDomainService123::userData"
)
when {
    resource.scope == Scope::"crossDomainService123/read" &&
    resource.owner == User::" user-identity-4e4c-example-7cede8e609a2" &&
    context.onBehalfOf.sub == " user-identity-4e4c-example-7cede8e609a2" &&
    context.callerApp == "ExampleChatApplication"
};

With the preceding Cedar policy example, you are permitting the AI agent to read userData from the crossDomainService123 resource. This is only permitted when the AI agent’s access token contains the crossDomainService/read scope and when the resource owner and the onBehalfOf user (from the access token) are the same—the human user in this case. There’s also an additional when clause in the policy to make sure that this interaction initiated from ExampleChatApplication.

The cross-domain resource server would use the AI agent’s access token and call the Verified Permissions IsAuthorizedWithToken API. To learn more, see Simplify fine-grained authorization with Amazon Verified Permissions and Amazon Cognito.

The following is a Node.js example using the IsAuthorizedWithToken API from Verified Permissions using the AWS SDK for JavaScript v3.

import { VerifiedPermissionsClient, IsAuthorizedWithTokenCommand } from "@aws-sdk/client-verifiedpermissions";

const client = new VerifiedPermissionsClient({ region: "<region>" });

// Dynamically provided token 
const jwtToken = "eyJraWQiOiJrMWtleSIsInR..."; //AI agent's access token

async function checkAccess() {
  const input = {
    policyStoreId: "ps-abc123example", // your AVP policy store ID
    accessToken: jwtToken,
    action: {
      actionType: "Action",
      actionId: "readOnly"
    },
    resource: {
      entityType: "crossDomainService123",
      entityId: "userData"
    }
  };

  const command = new IsAuthorizedWithTokenCommand(input);

  try {
    const response = await client.send(command);
    console.log("Authorization Decision:", response.decision);
  } catch (err) {
    console.error("Authorization error:", err);
  }
}

Based on the preceding examples of the AI agent’s access token (with user context), the Cedar policy, and the IsAuthorizedWithToken API call, the resource server would get an Allow decision for this action to take place. The following is an example of the authorization decision response.

{
    "decision": "Allow",
    "determiningPolicies": [{
        "determiningPolicyId": "ps-abc123example"
    }],
    "errors": []
}

Before this policy can be evaluated, you must define a schema that includes the relevant entity types (Agent, User, Resource, Scope, and so on), and create corresponding entities in your policy store that match the IDs used in the policy and request.

Bringing it all together, the requested data from the AI agent, on behalf of the user, is returned from the cross-domain service to the AI agent. This additional data can now be used within the context of the AI agent workload. The entire solution can be used for a chat application, such as the one described in Protect sensitive data in RAG applications with Amazon Bedrock.

Conclusion

Amazon Cognito M2M access token customization and support for passing client metadata provides you the extensibility to solve complex use cases and enables emerging ones like AI agent identity and access management. For example, passing contextual client metadata and customizing access tokens at runtime can help software as a service (SaaS) and multi-tenant service providers scale to an unlimited number of resource servers, because these can be dynamically determined at runtime. As organizations increasingly explore the use of AI agents, having a secure, scalable identity management solution becomes crucial for maintaining control and accountability. By using these new features, you can build more secure and scalable solutions with Amazon Cognito to prepare for the future of autonomous AI agent use cases.

Use the comments section to leave feedback about this post. If you have questions about this post, start a new thread on Amazon Cognito re:Post or contact AWS Support.