🏢

AgentCore Runtime のマルチテナント設計: CognitoとJWTによるテナントID伝播

に公開

はじめに

マルチテナント SaaS アプリケーションを構築する際、最も重要な課題の一つが「各リクエストをどのテナントのものか識別し、適切にデータへのアクセスを分離する」ことです。AI エージェントを活用したアプリケーションでは、エージェントの実行時にもテナント情報を正しく伝播させる必要があります。

本記事では、Amazon Bedrock AgentCore Runtime を使用したマルチテナントのアプリケーションにおいて、Amazon Cognito と連携してテナントIDを安全に伝播させる方法を紹介します。

要約

お忙しい方向けに先にまとめを記載します。
以下の方法で実現できます。

  • Cognito でカスタム属性として tenant_id を定義
  • Pre Token Generation Lambda トリガーでアクセストークンにカスタム属性を含める
  • Bedrock AgentCore Runtime で requestHeaderAllowlist を設定し、JWT Propagation を行う
  • Bedrock AgentCore Runtime で JWT の claims から tenant_id を取得する

課題と解決策: エージェント実行時のテナント識別

従来の Web アプリケーションでは、HTTP リクエストのヘッダーから JWT トークンを取得し、claims からテナントIDを取得できます。しかし、Bedrock AgentCore Runtime でエージェントを実行する場合、以下の課題があります。

  • エージェントコードにテナント情報をどのように渡すか
  • セキュリティを保ちながら、改ざんされないテナント情報をどう保証するか

本記事では上記の問題に対し、Amazon Bedrock AgentCore Runtime の JWT Propagation 機能を使用して、アクセストークンを AgentCore Runtime に渡すことで解決します。

それでは具体的な手順を見ていきましょう。

1. Cognito のカスタム属性にテナントIDを追加

まず Cognito で作成したユーザープールにカスタム属性としてテナントIDを設定します。

カスタム属性は マネジメントコンソール上では Cognito の「認証 > サインアップ」から追加できます。なお、カスタム属性には以下の制約があることに注意が必要です。

  • 作成後に名前の変更や削除ができない
  • 属性名の前に custom: プレフィックスが自動的に付与される

ユーザー登録時やユーザー情報更新時に、所属するテナントの ID を付与するようにします。

2. Pre Token Generation Lambda トリガーでアクセストークンにテナントIDを埋め込む

Cognito のユーザープールに設定したカスタム属性は、デフォルトでは ID トークンの claims にしか含まれず、アクセストークンの claims には含まれません。

そのため、このままでは AgentCore Runtime の認証・認可で使用されるアクセストークンからテナント情報を取得することができません。この問題を解決するために、Pre Token Generation Lambda トリガーを使用します。

Pre Token Generation Lambda トリガーは、ID トークンやアクセストークンの claims をカスタマイズできる機能です。名前の通りトークン生成の前に呼び出され、Lambda 関数によって claims の内容を編集することができます。

マネジメントコンソール上では、 Cognito のユーザープールの「拡張機能」から設定できます。

Lambda トリガー作成時は「認証 > トークン生成前トリガー」を選択します。トリガーイベントバージョンが複数ありますが、アクセストークンをカスタマイズするにはバージョン 2 以上を選択する必要があります。

以下が Lambda 関数の実装例です。ユーザー属性からテナントIDを取得し、アクセストークンに設定しています。

カスタム属性には custom: プレフィックスが付与されることに注意してください。

def lambda_handler(event, context):
    # ユーザー属性からテナントIDを取得
    tenant = event.get('request', {}).get('userAttributes', {}).get('custom:tenant_id')

    # アクセストークンにテナントIDを追加
    event['response'] = event.get('response', {})
    event['response']['claimsAndScopeOverrideDetails'] = {
        'accessTokenGeneration': {
            'claimsToAddOrOverride': {'tenant_id': tenant}
        }
    }

    return event

以上の対応でアクセストークンにテナントIDを付与することができました。

しかし、これだけでは AgentCore Runtime でテナント情報を取得することはできません。次のステップが必要です。

3. AgentCore Runtime で Authorization ヘッダーを許可

AgentCore Runtime では、Authorization ヘッダーに含まれるアクセストークンをデフォルトでは取得できません。トークンを取得するには、AgentCore Runtime の requestHeaderAllowlistAuthorization を含める必要があります。この設定により、エージェントコード内でトークンにアクセスできるようになります。

以下は bedrock_agentcore_starter_toolkit の Runtime を使用した場合のサンプルコードです。

from bedrock_agentcore_starter_toolkit import Runtime

agentcore_runtime = Runtime()
response = agentcore_runtime.configure(
    entrypoint="agent.py",
    auto_create_execution_role=True,
    auto_create_ecr=True,
    requirements_file="requirements.txt",
    region=region,
    agent_name="agent",

    # JWT 認証の設定
    authorizer_configuration={
        "customJWTAuthorizer": {
            "discoveryUrl": discovery_url,  # Cognito のディスカバリーエンドポイント
            "allowedClients": [client_id]   # 許可する Cognito クライアント ID
        }
    },

    # Authorization ヘッダーの許可(JWT Propagation)
    request_header_configuration={
        "requestHeaderAllowlist": ["Authorization"]
    }
)

4. エージェントコード内で JWT からテナントIDを抽出

これまでの手順により、AgentCore Runtime にアクセストークンを伝播させることができました (この仕組みは JWT Propagation とも呼ばれます)。

最後のステップとして、エージェントコード内で JWT から claims を読み取り、テナントIDを抽出します。以下がサンプル実装です。PyJWT ライブラリを使用してトークンをデコードし、claims から tenant_id を取得しています。

import jwt
# リクエストヘッダーから Authorization ヘッダーを取得
auth_header = context.request_headers.get("Authorization")

# Bearer プレフィックスを除去
token = auth_header.replace('Bearer ', '') if auth_header.startswith('Bearer ') else auth_header

tenant_id = None
try:
    # AgentCore Runtime が既にトークンを検証済みのため、署名検証はスキップ
    claims = jwt.decode(token, options={"verify_signature": False})
    tenant_id = claims.get("tenant_id")

    app.logger.info(f"Extracted tenant_id: {tenant_id}")

except jwt.InvalidTokenError as e:
    app.logger.exception("Invalid JWT token: %s", e)

まとめ

本記事では、Amazon Bedrock AgentCore Runtime と Cognito を連携させ、JWT Propagation を活用してマルチテナント対応を行うために必要なテナントIDを伝播する方法をお伝えしました。

この手法により、以下が実現できます。

  • JWT の claims からテナントIDを取得することで、署名による改ざん防止を実現
  • マルチテナント構成に必要な tenant_id を AgentCore Runtime で安全に使用

マルチテナント SaaS アプリケーションに AI エージェントを組み込む際、本記事の手法が参考になれば幸いです。

参考資料

エクサウィザーズ Tech Blog

Discussion