[Update] Amazon Bedrock AgentCore Runtime now supports Stateful MCP servers

"""
Domestic Travel Planner - Stateful MCP Server
Demonstration of Elicitation / Sampling / Progress Notifications
"""
import asyncio
import logging
from fastmcp import FastMCP, Context
from enum import Enum

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("travel-planner")

mcp = FastMCP("Japan-Travel-Planner")

class TripType(str, Enum):
    SOLO = "ひとり旅"
    COUPLE = "カップル"
    FAMILY = "家族旅行"
    FRIENDS = "友人旅行"

DESTINATIONS = {
    "京都": {
        "name": "京都",
        "transport": 14000,
        "hotel": 12000,
        "highlights": ["伏見稲荷大社", "嵐山竹林", "清水寺"],
        "gourmet": ["湯豆腐", "抹茶スイーツ", "にしんそば"],
    },
    "沖縄": {
        "name": "沖縄",
        "transport": 40000,
        "hotel": 10000,
        "highlights": ["美ら海水族館", "古宇利島", "首里城"],
        "gourmet": ["ソーキそば", "タコライス", "海ぶどう"],
    },
    "北海道": {
        "name": "北海道",
        "transport": 35000,
        "hotel": 9000,
        "highlights": ["富良野ラベンダー畑", "小樽運河", "旭山動物園"],
        "gourmet": ["海鮮丼", "ジンギスカン", "スープカレー"],
    },
    "福岡": {
        "name": "福岡",
        "transport": 22000,
        "hotel": 8000,
        "highlights": ["太宰府天満宮", "中洲屋台", "糸島"],
        "gourmet": ["博多ラーメン", "もつ鍋", "明太子"],
    },
}

@mcp.tool()
async def plan_trip(ctx: Context) -> str:
    """
    Create domestic travel plan (using all MCP features):
    1. Elicitation - gather travel preferences
    2. Progress - search progress for transportation & accommodation
    3. Sampling - generate AI recommendations
    """
    # ---- Phase 1: Elicitation ----
    logger.info("[Phase 1] Elicitation start - asking about destination")
    dest_result = await ctx.elicit(
        message="Where would you like to go?\nOptions: Kyoto, Okinawa, Hokkaido, Fukuoka",
        response_type=str,
    )
    if dest_result.action != "accept":
        logger.info("[Phase 1] User canceled (destination)")
        return "Plan creation canceled."
    dest_key = dest_result.data.strip()
    dest = DESTINATIONS.get(dest_key, DESTINATIONS["京都"])
    logger.info(f"[Phase 1] Destination: {dest['name']}")

    type_result = await ctx.elicit(
        message="What type of trip?\n1. Solo\n2. Couple\n3. Family\n4. Friends",
        response_type=TripType,
    )
    if type_result.action != "accept":
        logger.info("[Phase 1] User canceled (trip type)")
        return "Plan creation canceled."
    trip_type = type_result.data.value if hasattr(type_result.data, "value") else type_result.data
    logger.info(f"[Phase 1] Trip type: {trip_type}")

    days_result = await ctx.elicit(
        message="How many nights? (1-7 nights)",
        response_type=int,
    )
    if days_result.action != "accept":
        return "Plan creation canceled."
    days = max(1, min(7, days_result.data))

    travelers_result = await ctx.elicit(
        message="How many people?",
        response_type=int,
    )
    if travelers_result.action != "accept":
        return "Plan creation canceled."
    travelers = travelers_result.data
    logger.info(f"[Phase 1] Elicitation complete - {dest['name']}/{trip_type}/{days} nights/{travelers} people")

    # ---- Phase 2: Progress Notifications ----
    logger.info("[Phase 2] Progress Notifications start - search process")
    total_steps = 5
    for step in range(1, total_steps + 1):
        await ctx.report_progress(progress=step, total=total_steps)
        await asyncio.sleep(0.4)

    transport_cost = dest["transport"] * travelers
    hotel_cost = dest["hotel"] * days * ((travelers + 1) // 2)
    total_cost = transport_cost + hotel_cost
    logger.info(f"[Phase 2] Progress complete - cost calculation: ¥{total_cost:,}")

    # ---- Phase 3: Sampling ----
    logger.info("[Phase 3] Sampling start - requesting AI inference from client")
    ai_tips = f"Enjoy {dest['name']}!"
    try:
        response = await ctx.sample(
            messages=(
                f"Please give me 3 recommendations for a {trip_type} trip to {dest['name']} "
                f"({travelers} people, {days} nights) in 60 characters or less each. Please answer in Japanese."
            ),
            max_tokens=200,
        )
        if hasattr(response, "text") and response.text:
            ai_tips = response.text
        logger.info(f"[Phase 3] Sampling complete - response: {ai_tips[:50]}...")
    except Exception as e:
        logger.warning(f"[Phase 3] Sampling failed: {e}")
        ai_tips = f"{dest['highlights'][0]} is a must-see. Don't miss {dest['gourmet'][0]} too!"

    # ---- Final Confirmation ----
    logger.info("[Phase 4] Final confirmation Elicitation")
    confirm = await ctx.elicit(
        message=f"""
========== Travel Plan Summary ==========
Destination: {dest['name']}
Type: {trip_type}
Schedule: {days} nights {days + 1} days
People: {travelers}

Estimated cost:
  Transportation: ¥{transport_cost:,}
  Accommodation: ¥{hotel_cost:,} ({(travelers + 1) // 2} rooms × {days} nights)
  Total: ¥{total_cost:,}

Would you like to book this plan?""",
        response_type=["Book", "Cancel"],
    )
    if confirm.action != "accept" or confirm.data == "Cancel":
        logger.info("[Phase 4] User canceled")
        return "Booking canceled. Search results will be stored for 24 hours."

    logger.info(f"[Phase 4] Booking confirmed - {dest['name']}/{trip_type}/{days} nights/{travelers} people/¥{total_cost:,}")
    highlights_str = "\n".join(f"  - {h}" for h in dest["highlights"])
    gourmet_str = "\n".join(f"  - {g}" for g in dest["gourmet"])
    return f"""
{'=' * 50}
Booking Confirmed!
{'=' * 50}
Reservation number: TRV-{ctx.session_id[:8].upper()}

Destination: {dest['name']}
Schedule: {days} nights {days + 1} days / {travelers} people
Type: {trip_type}

Transportation: ¥{transport_cost:,}
Accommodation: ¥{hotel_cost:,} ({(travelers + 1) // 2} rooms × {days} nights)
Total: ¥{total_cost:,}

Recommended spots:
{highlights_str}

Local cuisine:
{gourmet_str}

AI recommendations:
{ai_tips}
{'=' * 50}
"""

if __name__ == "__main__":
    mcp.run(
        transport="streamable-http",
        host="0.0.0.0",
        port=8000,
        stateless_http=False,
    )

"""
Domestic Travel Planner - Test Client
Implementation of Elicitation / Sampling / Progress handlers

Environment variables:
  LOCAL_TEST=true (default) : Connect to local server
  LOCAL_TEST=false          : Connect to AgentCore Runtime (AGENT_ARN, BEARER_TOKEN required)
  USE_BEDROCK=true          : Use Amazon Bedrock for Sampling
  BEDROCK_MODEL_ID          : Bedrock model ID (default: us.anthropic.claude-haiku-4-5-20251001-v1:0)
"""
import asyncio
import os
import sys
import typing
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
from fastmcp.client.elicitation import ElicitResult
from mcp.types import CreateMessageResult, TextContent

# --- Elicitation handler ---

def _extract_options(response_type) -> list[str] | None:
    """Extract option list from response_type.
    fastmcp converts ["Book", "Cancel"] to Literal["Book", "Cancel"] when passing.
    """
    if isinstance(response_type, list):
        return response_type
    if typing.get_origin(response_type) is typing.Literal:
        args = typing.get_args(response_type)
        if args:
            return list(args)
    return None

def _prompt_choice(options: list[str]) -> str:
    """Select by number. Retry on invalid input."""
    for i, opt in enumerate(options, 1):
        print(f"    {i}. {opt}")
    while True:
        raw = input("    Select number: ").strip()
        try:
            idx = int(raw)
            if 1 <= idx <= len(options):
                return options[idx - 1]
        except ValueError:
            pass
        print(f"    ※ Enter a number between 1-{len(options)}")

def _prompt_int(label: str = "Answer (number)") -> int:
    """Input an integer. Retry on invalid input."""
    while True:
        raw = input(f"    {label}: ").strip()
        try:
            return int(raw)
        except ValueError:
            print("    ※ Please enter a number")

async def elicit_handler(message, response_type, params, ctx):
    print(f"\n>>> Question from server: {message}")
    try:
        options = _extract_options(response_type)
        if options:
            response = _prompt_choice(options)
        elif response_type is int:
            response = _prompt_int()
        else:
            response = input("    Answer: ").strip()
        return ElicitResult(action="accept", content={"value": response})
    except (KeyboardInterrupt, EOFError):
        print("\n    (Canceled)")
        return ElicitResult(action="decline", content=None)

# --- Sampling handler ---

def _extract_prompt_text(messages) -> str:
    """Extract prompt text from a list of SamplingMessage objects"""
    if isinstance(messages, str):
        return messages
    if isinstance(messages, list):
        parts = []
        for msg in messages:
            if hasattr(msg, "content") and hasattr(msg.content, "text"):
                parts.append(msg.content.text)
            else:
                parts.append(str(msg))
        return "\n".join(parts)
    return str(messages)

def _invoke_bedrock(prompt: str) -> str:
    """Run inference with Amazon Bedrock"""
    import boto3
    model_id = os.getenv("BEDROCK_MODEL_ID", "us.anthropic.claude-haiku-4-5-20251001-v1:0")
    client = boto3.client("bedrock-runtime")
    response = client.converse(
        modelId=model_id,
        messages=[{"role": "user", "content": [{"text": prompt}]}],
        inferenceConfig={"maxTokens": 200},
    )
    return response["output"]["message"]["content"][0]["text"]

async def sampling_handler(messages, params, ctx):
    prompt_text = _extract_prompt_text(messages)
    print(f"\n>>> AI sampling request")
    print(f"    Prompt: {prompt_text[:120]}...")

    use_bedrock = os.getenv("USE_BEDROCK", "false").lower() == "true"
    try:
        if use_bedrock:
            print("    (Inferring with Bedrock...)")
            ai_response = _invoke_bedrock(prompt_text)
            print(f"    Bedrock response: {ai_response}")
        else:
            ai_response = input("    Enter AI response (Enter for default): ").strip()
            if not ai_response:
                ai_response = "1. Visit popular spots early 2. Enjoy local cuisine 3. Check out local hidden gems"
    except Exception as e:
        print(f"    ※ Sampling error: {e}")
        ai_response = "Could not retrieve recommendations."

    return CreateMessageResult(
        role="assistant",
        content=TextContent(type="text", text=ai_response),
        model="bedrock" if use_bedrock else "manual",
        stopReason="endTurn",
    )

# --- Progress handler ---

async def progress_handler(progress, total, message):
    pct = int((progress / total) * 100) if total else 0
    bar = "#" * (pct // 5) + "-" * (20 - pct // 5)
    print(f"\r    Progress: [{bar}] {pct}%", end="", flush=True)
    if progress == total:
        print(" Complete!")

# --- Main ---

async def main():
    local_test = os.getenv("LOCAL_TEST", "true").lower() == "true"

    if local_test:
        url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000/mcp"
        headers = {}
    else:
        agent_arn = os.getenv("AGENT_ARN")
        token = os.getenv("BEARER_TOKEN")
        if not agent_arn or not token:
            print("ERROR: AGENT_ARN / BEARER_TOKEN not set")
            sys.exit(1)
        encoded_arn = agent_arn.replace(":", "%3A").replace("/", "%2F")
        endpoint = os.getenv(
            "MCP_ENDPOINT",
            f"https://bedrock-agentcore.{os.getenv('AWS_REGION', 'us-west-2')}.amazonaws.com",
        )
        url = f"{endpoint}/runtimes/{encoded_arn}/invocations?qualifier=DEFAULT"
        headers = {"Authorization": f"Bearer {token}"}

    transport = StreamableHttpTransport(url=url, headers=headers)
    client = Client(
        transport,
        elicitation_handler=elicit_handler,
        sampling_handler=sampling_handler,
        progress_handler=progress_handler,
    )

    try:
        await client.__aenter__()
    except Exception as e:
        print(f"\nERROR: Failed to connect to server: {e}")
        sys.exit(1)

    try:
        print("\nTesting plan_trip tool...")
        print("(Running full flow: Elicitation → Progress → Sampling)\n")
        result = await client.call_tool("plan_trip", {})
        print("\n" + "=" * 60)
        print("Result:")
        print("=" * 60)
        print(result.content[0].text)
    except KeyboardInterrupt:
        print("\n\nInterrupted")
    except Exception as e:
        print(f"\nERROR: {e}")
    finally:
        try:
            await client.__aexit__(None, None, None)
        except Exception:
            pass

if __name__ == "__main__":
    asyncio.run(main())