Bedrock AgentCore RuntimeのStateful MCPサーバーでElicitation・Sampling・Progress Notificationsを動かす

import asyncio
import json
 
from fastmcp import FastMCP, Context
 
mcp = FastMCP(
    "Travel Planner",
    instructions="A stateful travel planning MCP server that demonstrates "
    "elicitation, sampling, and progress notifications.",
)
 
DESTINATIONS = {
    "paris": {
        "city": "Paris",
        "country": "France",
        "highlights": ["Eiffel Tower", "Louvre Museum", "Seine River Cruise"],
        "best_season": "Spring",
    },
    "tokyo": {
        "city": "Tokyo",
        "country": "Japan",
        "highlights": ["Shibuya Crossing", "Senso-ji Temple", "Tsukiji Market"],
        "best_season": "Autumn",
    },
    "new york": {
        "city": "New York",
        "country": "USA",
        "highlights": ["Central Park", "Statue of Liberty", "Broadway"],
        "best_season": "Fall",
    },
    "bali": {
        "city": "Bali",
        "country": "Indonesia",
        "highlights": ["Ubud Rice Terraces", "Tanah Lot Temple", "Seminyak Beach"],
        "best_season": "Dry Season (Apr-Oct)",
    },
}
 
 
@mcp.resource("travel://destinations")
def list_destinations() -> str:
    return json.dumps(DESTINATIONS, indent=2)
 
 
@mcp.resource("travel://destination/{city}")
def get_destination(city: str) -> str:
    dest = DESTINATIONS.get(city.lower())
    if dest:
        return json.dumps(dest, indent=2)
    return json.dumps({"error": f"Destination '{city}' not found"})
 
 
@mcp.prompt()
def packing_list(destination: str, days: int, trip_type: str) -> str:
    return (
        f"Create a detailed {days}-day packing list for a {trip_type} trip "
        f"to {destination}. Include weather-appropriate clothing, essentials, "
        f"and destination-specific items."
    )
 
 
@mcp.prompt()
def local_phrases(destination: str) -> str:
    return (
        f"Teach me 10 essential local phrases for visiting {destination}. "
        f"Include greetings, asking for directions, ordering food, "
        f"and emergency phrases with pronunciation guides."
    )
 
 
@mcp.tool()
async def plan_trip(ctx: Context) -> str:
    """Plan a complete trip using elicitation, sampling, and progress notifications."""
    total_steps = 5
    await ctx.report_progress(progress=0, total=total_steps)
    dest_result = await ctx.elicit(
        message="Where would you like to go?\nOptions: Paris, Tokyo, New York, Bali",
        response_type=str,
    )
    if dest_result.action != "accept":
        return "Trip planning cancelled."
    destination = dest_result.data
    await ctx.report_progress(progress=1, total=total_steps)
 
    days_result = await ctx.elicit(
        message=f"How many days will you spend in {destination}?",
        response_type=int,
    )
    if days_result.action != "accept":
        return "Trip planning cancelled."
    days = days_result.data
    await ctx.report_progress(progress=2, total=total_steps)
 
    type_result = await ctx.elicit(
        message="What type of trip?\nOptions: leisure, business, adventure",
        response_type=str,
    )
    if type_result.action != "accept":
        return "Trip planning cancelled."
    trip_type = type_result.data
    await ctx.report_progress(progress=3, total=total_steps)
 
    response = await ctx.sample(
        messages=f"Give 3 brief tips for a {trip_type} trip to {destination} "
        f"lasting {days} days. Be concise.",
        max_tokens=200,
    )
    ai_tips = response.text
    await ctx.report_progress(progress=5, total=total_steps)
 
    dest_info = DESTINATIONS.get(destination.lower(), {})
    highlights = dest_info.get("highlights", ["No specific highlights available"])
    return json.dumps({
        "destination": destination,
        "days": days,
        "trip_type": trip_type,
        "highlights": highlights,
        "ai_tips": ai_tips,
        "status": "planned",
    }, indent=2)
 
 
@mcp.tool()
async def quick_recommend(ctx: Context) -> str:
    """Get a quick destination recommendation using sampling only."""
    response = await ctx.sample(
        messages="Recommend one travel destination from: Paris, Tokyo, New York, Bali. "
        "Give a one-sentence reason why.",
        max_tokens=100,
    )
    return f"Recommendation: {response.text}"
 
 
@mcp.tool()
async def search_flights(ctx: Context, origin: str, destination: str) -> str:
    """Simulate a flight search with progress notifications."""
    total = 4
    stages = [
        "Searching airlines...",
        "Comparing prices...",
        "Checking availability...",
        "Finalizing results...",
    ]
    for i, stage in enumerate(stages):
        await ctx.report_progress(progress=i + 1, total=total)
        await ctx.info(stage)
        await asyncio.sleep(0.3)
    return json.dumps({
        "origin": origin,
        "destination": destination,
        "flights": [
            {"airline": "AirExample", "price": "$450", "duration": "8h 30m"},
            {"airline": "SkyDemo", "price": "$520", "duration": "7h 15m"},
        ],
    }, indent=2)
 
 
if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="0.0.0.0", port=8000, stateless_http=False)

# 仮想環境の作成と依存インストール
uv venv && source .venv/bin/activate
uv pip install "fastmcp>=2.10.0" mcp
 
# サーバー起動（別ターミナルで実行）
python travel_server.py
# → http://0.0.0.0:8000/mcp で起動

import asyncio
import json
import urllib.parse
 
import boto3
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
import httpx
 
from mcp.client.streamable_http import streamablehttp_client
from mcp.client.session import ClientSession
from mcp.types import CreateMessageResult, ElicitResult, TextContent
 
AGENT_RUNTIME_ARN = "arn:aws:bedrock-agentcore:us-west-2:ACCOUNT_ID:runtime/YOUR_RUNTIME_ID"
REGION = "us-west-2"
 
 
async def elicit_callback(context, params):
    """Elicitationコールバック: サーバーからの質問に応答する。"""
    msg = params.message if hasattr(params, "message") else str(params)
    print(f"  [ELICITATION] Server asks: {msg}")
    # 実運用ではここでユーザーに入力を求める
    response = input("  Your answer: ").strip()
    return ElicitResult(action="accept", content={"value": response})
 
 
async def sampling_callback(context, params):
    """Samplingコールバック: サーバーからのLLM生成要求に応答する。"""
    msg_text = ""
    if hasattr(params, "messages"):
        for m in params.messages:
            if hasattr(m.content, "text"):
                msg_text = m.content.text
                break
    print(f"  [SAMPLING] Server requests: {msg_text[:100]}...")
    # 実運用ではここでLLM APIを呼び出す
    ai_response = "1. Research local customs. 2. Pack light. 3. Book in advance."
    return CreateMessageResult(
        role="assistant",
        content=TextContent(type="text", text=ai_response),
        model="your-model-id",
        stopReason="endTurn",
    )
 
 
class SigV4HttpxAuth(httpx.Auth):
    """AgentCore Runtime接続用のSigV4認証。"""
    def __init__(self, region, service="bedrock-agentcore"):
        self.region = region
        self.service = service
        session = boto3.Session()
        creds = session.get_credentials().get_frozen_credentials()
        self.credentials = Credentials(creds.access_key, creds.secret_key, creds.token)
 
    def auth_flow(self, request):
        aws_request = AWSRequest(
            method=request.method, url=str(request.url),
            headers=dict(request.headers), data=request.content,
        )
        SigV4Auth(self.credentials, self.service, self.region).add_auth(aws_request)
        for key, value in aws_request.headers.items():
            request.headers[key] = value
        yield request
 
 
def create_sigv4_httpx_client(**kwargs):
    kwargs.pop("auth", None)  # SDK由来のauthを除去
    return httpx.AsyncClient(auth=SigV4HttpxAuth(REGION), **kwargs)
 
 
async def main():
    encoded_arn = urllib.parse.quote(AGENT_RUNTIME_ARN, safe="")
    url = f"https://bedrock-agentcore.{REGION}.amazonaws.com/runtimes/{encoded_arn}/invocations?qualifier=DEFAULT"
 
    async with streamablehttp_client(
        url, httpx_client_factory=create_sigv4_httpx_client,
        timeout=120, sse_read_timeout=120, terminate_on_close=False,
    ) as (read_stream, write_stream, get_session_id):
        async with ClientSession(
            read_stream, write_stream,
            elicitation_callback=elicit_callback,
            sampling_callback=sampling_callback,
        ) as session:
            await session.initialize()
            print(f"Session ID: {get_session_id()}")
 
            # ツール一覧
            tools = await session.list_tools()
            for t in tools.tools:
                print(f"  {t.name}: {t.description}")
 
            # ツール実行（Elicitation + Sampling + Progress）
            result = await session.call_tool("plan_trip", {})
            for c in result.content:
                if hasattr(c, "text"):
                    print(c.text)
 
if __name__ == "__main__":
    asyncio.run(main())