Amazon Bedrock AgentCore Identity (3) : Agentから外部サービスへ3LO認証でアクセスする

from typing import Any
from collections import OrderedDict
from strands import Agent, tool
import asyncio
from strands.agent.conversation_manager.null_conversation_manager import NullConversationManager
from bedrock_agentcore.runtime import BedrockAgentCoreApp
from model.load import load_model
from mcp_client.client import get_streamable_http_mcp_client

# --- 3LO (Google Drive) 用の追加 import ---
from bedrock_agentcore.identity.auth import requires_access_token
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError

app = BedrockAgentCoreApp()
log = app.logger

# Define a Streamable HTTP MCP Client
mcp_clients = [get_streamable_http_mcp_client()]

DEFAULT_SYSTEM_PROMPT = """
You are a helpful assistant. Use tools when appropriate.

"""


# Define a collection of tools used by the model
tools = []

_INLINE_FUNCTION_NAMES = set()

# Define a simple function tool
@tool
def add_numbers(a: int, b: int) -> int:
    """Return the sum of two numbers"""
    return a+b
tools.append(add_numbers)


# --- ここから 3LO (Google Drive) ツール ---
SCOPES = ["https://www.googleapis.com/auth/drive.metadata.readonly"]

# 自前 callback サーバーの URL（ワークロードID・Google の両方に登録済み）
CALLBACK_URL = "http://localhost:8090/callback"


def _list_drive_files(access_token: str) -> str:
    """取得したアクセストークンで Drive のファイル一覧（先頭10件）を文字列で返す"""
    creds = Credentials(token=access_token, scopes=SCOPES)
    try:
        service = build("drive", "v3", credentials=creds)
        results = service.files().list(pageSize=10, fields="files(id, name)").execute()
        items = results.get("files", [])
        if not items:
            return "ファイルが見つかりませんでした。"
        lines = [f"- {it['name']} ({it['id']})" for it in items]
        return "Google Drive のファイル（先頭10件）:\n" + "\n".join(lines)
    except HttpError as error:
        return f"Drive API エラー: {error}"


def _on_auth_url(url: str):
    """3LO の同意 URL を目立つ形で出す"""
    msg = (
        "\n" + "=" * 60
        + "\n▼ 次の URL をブラウザで開いて同意してください:\n"
        + url
        + "\n" + "=" * 60 + "\n"
    )
    print(msg, flush=True)
    log.info(f"AUTH_URL: {url}")


@requires_access_token(
    provider_name="google-provider",   # 登録済み Credential Provider 名
    scopes=SCOPES,
    auth_flow="USER_FEDERATION",        # 3LO（ユーザー委譲）
    callback_url=CALLBACK_URL,          # 自前 callback サーバー（localhost:8090）
    on_auth_url=_on_auth_url,           # 同意 URL の提示先
)
async def _get_drive_files_with_token(*, access_token: str) -> str:
    return _list_drive_files(access_token)


@tool
async def read_google_drive() -> str:
    """ユーザーの Google Drive のファイル一覧（先頭10件）を取得して返す。"""
    try:
        return await _get_drive_files_with_token(access_token="")
    except Exception as e:
        import traceback
        log.error(f"DRIVE_TOOL_ERROR: {type(e).__name__}: {e}\n{traceback.format_exc()}")
        return f"[DEBUG] 3LO でエラー: {type(e).__name__}: {e}"
tools.append(read_google_drive)
# --- 3LO ツールここまで ---


# Add MCP client to tools if available
for mcp_client in mcp_clients:
    if mcp_client:
        tools.append(mcp_client)


def _make_conversation_manager():
    return NullConversationManager()

# Reuses one Agent per session_id so each session keeps its own in-process
# conversation history (best-effort; resets on cold start). The cache is bounded
# to 128 sessions with LRU eviction (least-recently-used is dropped and its
# history reset) so a single process serving many sessions cannot leak history
# between them or grow without limit. For durable history, attach a session manager.
def agent_factory():
    cache = OrderedDict()
    def get_or_create_agent(session_id):
        if session_id in cache:
            cache.move_to_end(session_id)
            return cache[session_id]
        if len(cache) >= 128:
            cache.popitem(last=False)
        cache[session_id] = Agent(
            model=load_model(),
            system_prompt=DEFAULT_SYSTEM_PROMPT,
            tools=tools,
            conversation_manager=_make_conversation_manager(),
            hooks=[
            ],
        )
        return cache[session_id]
    return get_or_create_agent
get_or_create_agent = agent_factory()


def strip_trailing_tool_use(messages: Any) -> list[dict]:
    """Strip toolUse blocks from the tail until the last message has none."""
    if not isinstance(messages, list):
        raise ValueError("messages must be a list")

    messages = list(messages)
    while messages:
        last = messages[-1]
        if not isinstance(last, dict):
            raise ValueError("each message must be an object")
        original_content = last.get("content", [])
        if not isinstance(original_content, list) or not all(isinstance(block, dict) for block in original_content):
            raise ValueError("each message content value must be a list of content blocks")

        content = [block for block in original_content if "toolUse" not in block]
        if len(content) == len(original_content):
            break
        if content:
            messages[-1] = {**last, "content": content}
            break
        messages.pop()

    return messages


def _extract_prompt(payload: dict):
    """Accept validated harness messages, tool results, or a plain prompt string."""
    if not isinstance(payload, dict):
        raise ValueError("payload must be a JSON object")
    if "messages" in payload:
        return strip_trailing_tool_use(payload["messages"])
    if "tool_results" in payload:
        tool_results = payload["tool_results"]
        if not isinstance(tool_results, list) or not all(
            isinstance(tool_result, dict) and isinstance(tool_result.get("toolUseId"), str)
            for tool_result in tool_results
        ):
            raise ValueError("tool_results must contain objects with a toolUseId string")
        return [{"role": "user", "content": [{"toolResult": {
            "toolUseId": tr["toolUseId"],
            "status": tr.get("status", "success"),
            "content": tr.get("content", []),
        }} for tr in tool_results]}]
    prompt = payload.get("prompt", "")
    if not isinstance(prompt, str):
        raise ValueError("prompt must be a string")
    return prompt


def _has_inline_function_call(messages) -> bool:
    """Return True if messages contains an assistant toolUse for an inline function tool."""
    if not _INLINE_FUNCTION_NAMES or not isinstance(messages, list):
        return False
    for msg in messages:
        if msg.get("role") == "assistant":
            for block in msg.get("content", []):
                if isinstance(block, dict) and block.get("toolUse", {}).get("name") in _INLINE_FUNCTION_NAMES:
                    return True
    return False


def _is_inline_function_call(event: dict) -> bool:
    """Check if a contentBlockStart event is for an inline function tool."""
    if not _INLINE_FUNCTION_NAMES:
        return False
    cbs = event.get("contentBlockStart", {})
    start = cbs.get("start", {})
    tool_use = start.get("toolUse") if isinstance(start, dict) else None
    return tool_use is not None and tool_use.get("name") in _INLINE_FUNCTION_NAMES



@app.entrypoint
async def invoke(payload, context):
    log.info("Invoking Agent.....")


    session_id = getattr(context, 'session_id', 'default-session')
    agent = get_or_create_agent(session_id)

    prompt = _extract_prompt(payload)


    async for event in agent.stream_async(
        prompt,
    ):
        if not isinstance(event, dict) or "event" not in event:
            continue
        cbs = event["event"].get("contentBlockStart")
        if cbs is not None and not cbs.get("start"):
            continue
        yield event


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

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "drive3lo"
version = "0.1.0"
description = "AgentCore Runtime Application using Strands SDK"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
    "aws-opentelemetry-distro",
    "bedrock-agentcore >= 1.9.1",
    "botocore[crt] >= 1.35.0",
    "mcp ~= 1.24.0",
    "strands-agents >= 1.15.0",
    "google-api-python-client >= 2.100.0",
    "google-auth >= 2.23.0",
]

[tool.hatch.build.targets.wheel]
packages = ["."]

"""
3LO 用ローカル callback サーバー。
Google からの ?session_id=... を受けて CompleteResourceTokenAuth を呼び、
セッションバインディングを完了する。
user_id は app/drive3lo/.agentcore.json から自動で読み込む。
"""
import json
from pathlib import Path
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.parse import urlparse, parse_qs

import boto3

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

# .agentcore.json から user_id を自動読み込み（絶対パスで固定）
AGENTCORE_JSON = Path.home() / "agentcore-runtime-quickstart" / "drive3lo" / "app" / "drive3lo" / ".agentcore.json"
try:
    _cfg = json.loads(AGENTCORE_JSON.read_text())
    USER_ID = _cfg["user_id"]
    WORKLOAD = _cfg.get("workload_identity_name", "?")
    print(f"Loaded from {AGENTCORE_JSON}\n  user_id={USER_ID}  workload={WORKLOAD}")
except Exception as e:
    raise SystemExit(
        f"[callback] .agentcore.json を読めません: {e}\n"
        f"  先に setup_identity.py を実行して .agentcore.json を生成してください。"
    )

agentcore = boto3.client("bedrock-agentcore", region_name=REGION)


class Handler(BaseHTTPRequestHandler):
    def _send(self, code, body, ctype="application/json"):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.end_headers()
        self.wfile.write(body.encode() if isinstance(body, str) else body)

    def do_GET(self):
        parsed = urlparse(self.path)
        if parsed.path != "/callback":
            self._send(404, '{"error":"not found"}')
            return

        qs = parse_qs(parsed.query)
        session_id = qs.get("session_id", [None])[0]
        if not session_id:
            self._send(400, '{"error":"missing session_id"}')
            return

        try:
            print(f"[callback] complete: session={session_id[:30]}... user_id={USER_ID}", flush=True)
            agentcore.complete_resource_token_auth(
                sessionUri=session_id,
                userIdentifier={"userId": USER_ID},
            )
        except Exception as e:
            print(f"[callback] complete failed: {e}", flush=True)
            self._send(500, json.dumps({"error": f"complete failed: {e}"}))
            return

        print(f"[callback] session bound OK: {session_id[:40]}...", flush=True)
        self._send(
            200,
            "<html><body><h2>認証が完了しました。ターミナルに戻ってください。</h2></body></html>",
            ctype="text/html",
        )

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    print(f"Callback server on http://localhost:{PORT}/callback (user_id={USER_ID})")
    HTTPServer(("127.0.0.1", PORT), Handler).serve_forever()

"""
3LO の下準備スクリプト。
- ワークロード ID と user_id を生成し app/drive3lo/.agentcore.json に保存
- そのワークロード ID に、戻り先URL（callbackサーバー & Credential Provider）を登録
これを最初に一度実行すれば、あとは本番の3ターミナル実行だけで動く。
"""
import json
import os
import uuid
from pathlib import Path

import boto3

REGION = os.environ.get("AWS_REGION", "us-west-2")
LOCAL_CALLBACK = "http://localhost:8090/callback"
GOOGLE_CALLBACK_URL = os.environ["GOOGLE_CALLBACK_URL"]  # env.sh で export 済み

CONFIG_PATH = Path("app/drive3lo/.agentcore.json")

cp = boto3.client("bedrock-agentcore-control", region_name=REGION)

# 既に .agentcore.json があれば再利用、無ければ新規作成
if CONFIG_PATH.exists():
    cfg = json.loads(CONFIG_PATH.read_text())
    workload = cfg["workload_identity_name"]
    user_id = cfg["user_id"]
    print(f"既存の設定を利用: workload={workload} user_id={user_id}")
else:
    workload = cp.create_workload_identity(name=f"drive3lo-{uuid.uuid4().hex[:8]}")["name"]
    user_id = uuid.uuid4().hex[:8]
    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    CONFIG_PATH.write_text(json.dumps(
        {"workload_identity_name": workload, "user_id": user_id}, indent=2
    ))
    print(f"作成: workload={workload} user_id={user_id}")

# 戻り先URLを登録（callbackサーバー & Credential Provider の両方）
cp.update_workload_identity(
    name=workload,
    allowedResourceOauth2ReturnUrls=[LOCAL_CALLBACK, GOOGLE_CALLBACK_URL],
)
print("戻り先URLを登録しました:")
print(f"  - {LOCAL_CALLBACK}")
print(f"  - {GOOGLE_CALLBACK_URL}")
print("\n下準備完了。次は3ターミナルで本番実行してください。")