Amazon Bedrock AgentCore Identityでユーザー委任型認証（3LO）を実装してGoogle Driveにアクセスしてみた

# プロジェクト初期化
uv init

# 必要なライブラリを追加
uv add bedrock-agentcore boto3 pyjwt strands-agents bedrock-agentcore-starter-toolkit google-api-python-client google-auth-httplib2 google-auth-oauthlib fastapi uvicorn pyyaml

# boto3 クライアント経由で AWS API を呼び出し
dp_client.get_workload_access_token_for_jwt(workloadName, userToken)  # JWT を使う場合
dp_client.get_workload_access_token_for_user_id(workloadName, userId) # user_id を使う場合

@requires_access_token(
      provider_name="my-google-provider",
      scopes=["..."],
      auth_flow="USER_FEDERATION",
      callback_url="http://localhost:9090/callback",
 )

identity_client.complete_resource_token_auth(
    session_uri=session_id,
    user_identifier=UserTokenIdentifier(user_token=cognito_access_token)
)

import boto3

REGION = "us-west-2"
PROVIDER_NAME = "google-drive-provider"
GOOGLE_CLIENT_ID = "your-client-id.apps.googleusercontent.com"
GOOGLE_CLIENT_SECRET = "your-client-secret"

def create_google_provider():
    """Google OAuth Credential Provider を作成"""
    client = boto3.client("bedrock-agentcore-control", region_name=REGION)

    try:
        response = client.create_oauth2_credential_provider(
            name=PROVIDER_NAME,
            credentialProviderVendor="GoogleOauth2",
            oauth2ProviderConfigInput={
                "googleOauth2ProviderConfig": {
                    "clientId": GOOGLE_CLIENT_ID,
                    "clientSecret": GOOGLE_CLIENT_SECRET
                }
            },
        )
        print(f"Created provider: {PROVIDER_NAME}")
        print(f"Callback URL: {response.get('callbackUrl', '')}")
        return response

    except client.exceptions.ConflictException:
        # 既に存在する場合は取得
        response = client.get_oauth2_credential_provider(name=PROVIDER_NAME)
        print(f"Provider already exists: {PROVIDER_NAME}")
        print(f"Callback URL: {response.get('callbackUrl', '')}")
        return response

if __name__ == "__main__":
    create_google_provider()

python create_credential_provider.py

Created provider: google-drive-provider
Callback URL: https://identity.agentcore.us-west-2.amazonaws.com/oauth2/callback/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx

import secrets
import string
import boto3

REGION = "us-west-2"
POOL_NAME = "AgentCore3LODemoPool"
CLIENT_NAME = "AgentCore3LOClient"

def generate_password(length: int = 16) -> str:
    """パスワードを生成"""
    alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
    password = "".join(secrets.choice(alphabet) for _ in range(length))
    password += secrets.choice(string.digits)
    return password

def create_cognito_resources():
    """Cognito リソースを作成"""
    client = boto3.client("cognito-idp", region_name=REGION)

    # 1. User Pool 作成
    pool_response = client.create_user_pool(
        PoolName=POOL_NAME,
        AutoVerifiedAttributes=["email"],
        UsernameAttributes=["email"],
    )
    user_pool_id = pool_response["UserPool"]["Id"]

    # 2. User Pool Client 作成
    client_response = client.create_user_pool_client(
        UserPoolId=user_pool_id,
        ClientName=CLIENT_NAME,
        GenerateSecret=True,
        ExplicitAuthFlows=["ALLOW_USER_PASSWORD_AUTH", "ALLOW_REFRESH_TOKEN_AUTH"],
    )
    client_id = client_response["UserPoolClient"]["ClientId"]
    client_secret = client_response["UserPoolClient"]["ClientSecret"]

    # テストユーザー作成
    username = f"testuser{secrets.randbelow(10000):04d}@example.com"
    password = generate_password()

    client.admin_create_user(
        UserPoolId=user_pool_id, Username=username, MessageAction="SUPPRESS"
    )
    client.admin_set_user_password(
        UserPoolId=user_pool_id, Username=username, Password=password, Permanent=True
    )

    # 出力
    discovery_url = f"https://cognito-idp.{REGION}.amazonaws.com/{user_pool_id}/.well-known/openid-configuration"

    print(f"Client ID: {client_id}")
    print(f"Client Secret: {client_secret}")
    print(f"Discovery URL: {discovery_url}")
    print(f"Test Username: {username}")
    print(f"Test Password: {password}")

    print("\n# 環境変数として設定:")
    print(f"export USER_POOL_ID='{user_pool_id}'")
    print(f"export CLIENT_ID='{client_id}'")
    print(f"export CLIENT_SECRET='{client_secret}'")
    print(f"export COGNITO_USERNAME='{username}'")
    print(f"export COGNITO_PASSWORD='{password}'")

if __name__ == "__main__":
    create_cognito_resources()

python create_cognito.py

Client ID: xxxxxxxxxxxxxxxxxxxxxxxxxx
Client Secret: xxxxxxxxxxxxxxxxxxxxxxxxxx
Discovery URL: https://cognito-idp.us-west-2.amazonaws.com/us-west-2_XXXXXX/.well-known/openid-configuration
Test Username: testuser1234@example.com
Test Password: xxxxxxxxxxxxxxxx

# 環境変数として設定:
export USER_POOL_ID='us-west-2_XXXXXX'
...

import os
import hmac
import hashlib
import base64
import boto3

REGION = "us-west-2"

def get_access_token(
    user_pool_id: str,
    client_id: str,
    client_secret: str,
    username: str,
    password: str,
) -> dict:
    """Cognito から Access Token を取得"""
    # SECRET_HASH の計算
    message = username + client_id
    dig = hmac.new(
        client_secret.encode("utf-8"),
        message.encode("utf-8"),
        hashlib.sha256,
    ).digest()
    secret_hash = base64.b64encode(dig).decode()

    client = boto3.client("cognito-idp", region_name=REGION)

    response = client.initiate_auth(
        ClientId=client_id,
        AuthFlow="USER_PASSWORD_AUTH",
        AuthParameters={
            "USERNAME": username,
            "PASSWORD": password,
            "SECRET_HASH": secret_hash,
        },
    )

    return response["AuthenticationResult"]

if __name__ == "__main__":
    result = get_access_token(
        user_pool_id=os.environ["USER_POOL_ID"],
        client_id=os.environ["CLIENT_ID"],
        client_secret=os.environ["CLIENT_SECRET"],
        username=os.environ["COGNITO_USERNAME"],
        password=os.environ["COGNITO_PASSWORD"],
    )

    access_token = result["AccessToken"]
    print(f"Access Token: {access_token[:50]}...")

    # ファイルに保存
    with open(".access_token.txt", "w") as f:
        f.write(access_token)
    print("Saved to .access_token.txt")

python refresh_cognito_token.py

export ACCESS_TOKEN=$(cat .access_token.txt)

"""
Google Drive 3LO Agent

AgentCore Runtime にデプロイするエージェント。
BedrockAgentCoreApp を使用して Runtime コンテキストを受け取ります。
"""

import json
from typing import Any, AsyncGenerator, Dict

from bedrock_agentcore.identity.auth import requires_access_token
from bedrock_agentcore.runtime import BedrockAgentCoreApp
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError

# クライアントアプリケーションの名前を設定
PROVIDER_NAME = "google-drive-provider"
SCOPES = ["https://www.googleapis.com/auth/drive.metadata.readonly"]
CALLBACK_URL = "http://localhost:9090/callback"  # 後で設定

# BedrockAgentCoreApp を使用
app = BedrockAgentCoreApp()

class AuthRequiredException(Exception):
    """認証が必要な場合にスローされる例外"""
    def __init__(self, auth_url: str):
        self.auth_url = auth_url
        super().__init__(f"認証が必要です: {auth_url}")

def raise_auth_required(url: str):
    """認証 URL を受け取ったら例外をスロー"""
    raise AuthRequiredException(url)

@requires_access_token(
    provider_name=PROVIDER_NAME,
    scopes=SCOPES,
    auth_flow="USER_FEDERATION",
    on_auth_url=raise_auth_required,
    force_authentication=False,
    callback_url=CALLBACK_URL,
)
async def access_google_drive(*, access_token: str) -> dict:
    """Google Drive にアクセスしてファイル一覧を取得"""
    creds = Credentials(token=access_token, scopes=SCOPES)

    try:
        service = build("drive", "v3", credentials=creds)
        results = (
            service.files()
            .list(pageSize=10, fields="nextPageToken, files(id, name, mimeType)")
            .execute()
        )
        return {"status": "success", "files": results.get("files", [])}

    except HttpError as error:
        return {"status": "error", "message": str(error)}

@app.entrypoint
async def agent_invocation(payload: Dict[str, Any]) -> AsyncGenerator[str, None]:
    """エージェントのエントリーポイント"""

    try:
        result = await access_google_drive(access_token="")
        yield json.dumps(result, ensure_ascii=False)

    except AuthRequiredException as e:
        response = {
            "status": "auth_required",
            "auth_url": e.auth_url,
            "message": "以下の URL をブラウザで開いて Google 認証を完了してください",
        }
        yield json.dumps(response, ensure_ascii=False)

    except Exception as e:
        yield json.dumps({"status": "error", "message": str(e)}, ensure_ascii=False)

if __name__ == "__main__":
    app.run()

response = {
            "status": "auth_required",
            "auth_url": e.auth_url,
            "message": "以下の URL をブラウザで開いて Google 認証を完了してください",
}

async def access_google_drive(*, access_token: str) -> dict:
    """Google Drive にアクセスしてファイル一覧を取得"""
    creds = Credentials(token=access_token, scopes=SCOPES)

    try:
        service = build("drive", "v3", credentials=creds)
        results = (
            service.files()
            .list(pageSize=10, fields="nextPageToken, files(id, name, mimeType)")
            .execute()
        )
        return {"status": "success", "files": results.get("files", [])}

    except HttpError as error:
        return {"status": "error", "message": str(error)}

uv run agentcore configure -e agent.py

Configuring Bedrock AgentCore...
✓ Using file: agent.py

🏷️  Inferred agent name: agent
Press Enter to use this name, or type a different one (alphanumeric without '-')
Agent name [agent]: agent_sample_test_0123
✓ Using agent name: agent_sample_test_0123

🔍 Detected dependency file: pyproject.toml
Press Enter to use this file, or type a different path (use Tab for autocomplete):
Path or Press Enter to use detected dependency file: pyproject.toml
✓ Using requirements file: pyproject.toml

🚀 Deployment Configuration
Select deployment type:
  1. Direct Code Deploy (recommended) - Python only, no Docker required
  2. Container - For custom runtimes or complex dependencies
Choice [1]: 1

Select Python runtime version:
  1. PYTHON_3_10
  2. PYTHON_3_11
  3. PYTHON_3_12
  4. PYTHON_3_13
Choice [3]: 3
✓ Deployment type: Direct Code Deploy (python.3.12)

🔐 Execution Role
Press Enter to auto-create execution role, or provide execution role ARN/name to use existing
Execution role ARN/name (or press Enter to auto-create):
✓ Will auto-create execution role

🏗️  S3 Bucket
Press Enter to auto-create S3 bucket, or provide S3 URI/path to use existing
S3 URI/path (or press Enter to auto-create):
✓ Will auto-create S3 bucket

🔐 Authorization Configuration
By default, Bedrock AgentCore uses IAM authorization.
Configure OAuth authorizer instead? (yes/no) [no]: yes

📋 OAuth Configuration
Enter OAuth discovery URL: https://cognito-idp.us-west-2.amazonaws.com/us-west-2_XXXXXXXXX/.well-known/openid-configuration
Enter allowed OAuth client IDs (comma-separated): xxxxxxxxxxxxxxxxxxxxxxxxxx
Enter allowed OAuth audience (comma-separated):
Enter allowed OAuth allowed scopes (comma-separated):
Enter allowed OAuth custom claims as JSON string (comma-separated):
✓ OAuth authorizer configuration created

🔒 Request Header Allowlist
Configure which request headers are allowed to pass through to your agent.
Common headers: Authorization, X-Amzn-Bedrock-AgentCore-Runtime-Custom-*
Configure request header allowlist? (yes/no) [no]: no
✓ Using default request header configuration
Configuring BedrockAgentCore agent: agent_sample_test_0123

agentcore deploy

from bedrock_agentcore.services.identity import IdentityClient

# agent_id を知っている場合
agent_id = "agent_sample_test_0123-5PWohQFJ8u"

identity_client = IdentityClient(region="us-west-2")

# Workload Identity を取得
workload_identity = identity_client.get_workload_identity(name=agent_id)
print(workload_identity)

print("=== Workload Identity ===")
print(f"Name: {workload_identity.get('name')}")
print(f"ARN: {workload_identity.get('arn')}")
print(
    f"Allowed Callback URLs: {workload_identity.get('allowedResourceOauth2ReturnUrls')}"
)

identity_client.update_workload_identity(
    name=agent_id,
    allowed_resource_oauth_2_return_urls=["http://localhost:9090/callback"],
) 

=== Workload Identity ===
Name: my-agent-XXXXXXXXXX
ARN: arn:aws:bedrock-agentcore:us-west-2:123456789012:workload-identity-directory/default/workload-identity/my-agent-XXXXXXXXXX
Allowed Callback URLs: []

=== Workload Identity ===
Name: my-agent-XXXXXXXXXX
ARN: arn:aws:bedrock-agentcore:us-west-2:123456789012:workload-identity-directory/default/workload-identity/my-agent-XXXXXXXXXX
Allowed Callback URLs: ['http://localhost:9090/callback']

"""
OAuth2 Callback Server

Google 認証完了後、AgentCore Identity の CompleteResourceTokenAuth を呼び出して
Session Binding を完了します。
"""

import os
import uvicorn
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from bedrock_agentcore.services.identity import IdentityClient, UserTokenIdentifier

PORT = 9090
REGION = "us-west-2"

app = FastAPI()
identity_client = IdentityClient(region=REGION)

@app.get("/callback")
async def oauth2_callback(session_id: str):
    """OAuth2 コールバック処理"""

    # 環境変数から Cognito Access Token を取得
    token = os.environ.get("ACCESS_TOKEN", "")

    if not token:
        return HTMLResponse(
            content="<h1>Error</h1><p>ACCESS_TOKEN not set</p>",
            status_code=500
        )

    try:
        # Session Binding を完了
        identity_client.complete_resource_token_auth(
            session_uri=session_id,
            user_identifier=UserTokenIdentifier(user_token=token)
        )

        html = """
        <!DOCTYPE html>
        <html>
        <head><title>認証完了</title></head>
        <body>
            <h1>認証完了！</h1>
            <p>Token Vault にアクセストークンが保存されました。</p>
            <p>ターミナルに戻って、もう一度 invoke してください。</p>
        </body>
        </html>
        """
        return HTMLResponse(content=html, status_code=200)

    except Exception as e:
        return HTMLResponse(
            content=f"<h1>Error</h1><pre>{str(e)}</pre>",
            status_code=500
        )

if __name__ == "__main__":
    print(f"Callback Server running on http://localhost:{PORT}/callback")
    uvicorn.run(app, host="127.0.0.1", port=PORT)

identity_client.complete_resource_token_auth(
            session_uri=session_id,
            user_identifier=UserTokenIdentifier(user_token=token)
)

export ACCESS_TOKEN=$(cat .access_token.txt)
python callback_server.py

#　ログ
Callback Server running on http://localhost:9090/callback
INFO:     Started server process [77981]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:9090 (Press CTRL+C to quit)

uv run agentcore invoke --agent agent_sample_test_0123 --bearer-token $ACCESS_TOKEN {"prompt": ""}

Using bearer token for OAuth authentication
Using JWT authentication
{"status": "auth_required", "auth_url":
"https://bedrock-agentcore.us-west-2.amazonaws.com/identities/oauth2/authorize?request_uri=urn%3Aietf%3Apa
rams%3Aoauth%3Arequest_uri%sample", "message": "以下の URL
をブラウザで開いて Google 認証を完了してください"}

agentcore invoke --agent agent_sample_test_0123 --bearer-token $ACCESS_TOKEN

{
  "status": "success",
  "files": [
    {"id": "1abc...", "name": "サンプル1", "mimeType": "application/vnd.google-apps.document"},
    {"id": "2def...", "name": "サンプル2", "mimeType": "application/vnd.google-apps.spreadsheet"},
    {"id": "3ghi...", "name": "サンプル3", "mimeType": "application/vnd.google-apps.presentation"},
    ...
  ]
}

# NG: BedrockAgentCoreApp を使わない
async def main():
    result = await access_google_drive(access_token="")

# OK: BedrockAgentCoreApp を使う
app = BedrockAgentCoreApp()

@app.entrypoint
async def agent_invocation(payload):
    ...

# 全てOK
allowed_resource_oauth_2_return_urls=[]

# 指定したURLのみ許可
# http://localhost:9090/callbackのみ許可
allowed_resource_oauth_2_return_urls=["http://localhost:9090/callback"]