Edge Cases to Verify Before Using Stateful MCP in Production

# Transactions table
aws dynamodb create-table \
  --table-name mcp-finance-transactions \
  --attribute-definitions \
    AttributeName=user_alias,AttributeType=S \
    AttributeName=transaction_id,AttributeType=S \
  --key-schema \
    AttributeName=user_alias,KeyType=HASH \
    AttributeName=transaction_id,KeyType=RANGE \
  --billing-mode PAY_PER_REQUEST \
  --region us-west-2
 
# Budgets table
aws dynamodb create-table \
  --table-name mcp-finance-budgets \
  --attribute-definitions \
    AttributeName=user_alias,AttributeType=S \
    AttributeName=category,AttributeType=S \
  --key-schema \
    AttributeName=user_alias,KeyType=HASH \
    AttributeName=category,KeyType=RANGE \
  --billing-mode PAY_PER_REQUEST \
  --region us-west-2

import uuid
from datetime import datetime, timezone
from decimal import Decimal
 
import boto3
from boto3.dynamodb.conditions import Key
 
 
class FinanceDB:
    def __init__(self, region_name: str = "us-west-2"):
        dynamodb = boto3.resource("dynamodb", region_name=region_name)
        self.transactions = dynamodb.Table("mcp-finance-transactions")
        self.budgets = dynamodb.Table("mcp-finance-budgets")
 
    def add_transaction(self, user_alias, tx_type, amount, description, category):
        tx_id = str(uuid.uuid4())[:8]
        self.transactions.put_item(Item={
            "user_alias": user_alias,
            "transaction_id": tx_id,
            "type": tx_type,
            "amount": Decimal(str(amount)),
            "description": description,
            "category": category,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })
        return f"Expense of ${abs(amount):.2f} added for {user_alias}"
 
    def get_transactions(self, user_alias):
        return self.transactions.query(
            KeyConditionExpression=Key("user_alias").eq(user_alias)
        ).get("Items", [])
 
    def get_budgets(self, user_alias):
        return self.budgets.query(
            KeyConditionExpression=Key("user_alias").eq(user_alias)
        ).get("Items", [])

import os
from pydantic import BaseModel
from fastmcp import FastMCP, Context
from fastmcp.server.elicitation import AcceptedElicitation
from dynamo_utils import FinanceDB
 
mcp = FastMCP(name="FinanceMCP")
db = FinanceDB(region_name=os.environ.get("AWS_REGION", "us-west-2"))
 
class AmountInput(BaseModel):
    amount: float
 
class DescriptionInput(BaseModel):
    description: str
 
class CategoryInput(BaseModel):
    category: str
 
class ConfirmInput(BaseModel):
    confirm: str
 
 
@mcp.tool()
async def add_expense_interactive(user_alias: str, ctx: Context) -> str:
    """Interactively add a new expense using elicitation."""
    result = await ctx.elicit("How much did you spend?", AmountInput)
    if not isinstance(result, AcceptedElicitation):
        return "Expense entry cancelled."
    amount = result.data.amount
 
    result = await ctx.elicit("What was it for?", DescriptionInput)
    if not isinstance(result, AcceptedElicitation):
        return "Expense entry cancelled."
    description = result.data.description
 
    result = await ctx.elicit(
        "Select a category (food, transport, bills, entertainment, other):",
        CategoryInput,
    )
    if not isinstance(result, AcceptedElicitation):
        return "Expense entry cancelled."
    category = result.data.category
 
    confirm_msg = (
        f"Confirm: add expense of ${amount:.2f} for {description}"
        f" (category: {category})? Reply Yes or No"
    )
    result = await ctx.elicit(confirm_msg, ConfirmInput)
    if not isinstance(result, AcceptedElicitation) or result.data.confirm != "Yes":
        return "Expense entry cancelled."
 
    return db.add_transaction(user_alias, "expense", -abs(amount), description, category)
 
 
@mcp.tool()
async def analyze_spending(user_alias: str, ctx: Context) -> str:
    """Fetch expenses and ask the client's LLM to analyse them."""
    transactions = db.get_transactions(user_alias)
    if not transactions:
        return f"No transactions found for {user_alias}."
 
    lines = "\n".join(
        f"- {t['description']} (${abs(float(t['amount'])):.2f}, {t['category']})"
        for t in transactions
    )
    prompt = (
        f"Here are the recent expenses:\n{lines}\n\n"
        f"Give 3 concise, actionable recommendations. Under 120 words."
    )
 
    ai_analysis = "Analysis unavailable."
    try:
        response = await ctx.sample(messages=prompt, max_tokens=300)
        if hasattr(response, "text") and response.text:
            ai_analysis = response.text
    except Exception:
        pass
 
    return f"Spending Analysis for {user_alias}:\n\n{ai_analysis}"
 
 
@mcp.tool()
async def generate_report(user_alias: str, ctx: Context) -> str:
    """Generate a monthly financial report with progress notifications."""
    total = 5
 
    await ctx.report_progress(progress=1, total=total)
    transactions = db.get_transactions(user_alias)
 
    await ctx.report_progress(progress=2, total=total)
    by_category = {}
    for t in transactions:
        cat = t["category"]
        by_category[cat] = by_category.get(cat, 0) + abs(float(t["amount"]))
 
    await ctx.report_progress(progress=3, total=total)
    budgets = {b["category"]: float(b["monthly_limit"]) for b in db.get_budgets(user_alias)}
 
    await ctx.report_progress(progress=4, total=total)
    lines = []
    for cat, spent in sorted(by_category.items(), key=lambda x: -x[1]):
        limit = budgets.get(cat)
        if limit:
            pct = (spent / limit) * 100
            status = "OVER" if spent > limit else "OK"
            lines.append(f"  {cat:<15} ${spent:>8.2f} / ${limit:.2f}  [{pct:.0f}%] {status}")
        else:
            lines.append(f"  {cat:<15} ${spent:>8.2f}  (no budget set)")
 
    await ctx.report_progress(progress=5, total=total)
    total_spent = sum(by_category.values())
    return (
        f"Monthly Report for {user_alias}\n"
        f"{'=' * 50}\n"
        f"  {'Category':<15} {'Spent':>10}   {'Budget':>8}  Status\n"
        f"{'-' * 50}\n"
        + "\n".join(lines)
        + f"\n{'-' * 50}\n"
        f"  {'TOTAL':<15} ${total_spent:>8.2f}\n"
    )
 
 
if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="0.0.0.0", port=8000, stateless_http=False)

"""Test client — normal flow + edge cases."""
import asyncio
import os
 
os.environ.setdefault("AWS_DEFAULT_REGION", "us-west-2")
 
from fastmcp import Client
from fastmcp.client.elicitation import ElicitResult
from mcp.types import CreateMessageResult, TextContent
from finance_server import mcp
 
 
async def test_normal_flow():
    """Normal: add expenses -> analyze -> report"""
    expenses = [
        [{"amount": 45.50}, {"description": "Lunch at the office"}, {"category": "food"}, {"confirm": "Yes"}],
        [{"amount": 120.00}, {"description": "Electric bill"}, {"category": "bills"}, {"confirm": "Yes"}],
        [{"amount": 15.99}, {"description": "Movie tickets"}, {"category": "entertainment"}, {"confirm": "Yes"}],
        [{"amount": 85.30}, {"description": "Weekly groceries"}, {"category": "food"}, {"confirm": "Yes"}],
    ]
    idx = [0, 0]  # [expense_index, step_index]
 
    async def elicit_handler(message, response_type, params, context):
        resp = expenses[idx[0]][idx[1]]
        print(f"  Server asks: {message}")
        print(f"  Responding:  {resp}")
        idx[1] += 1
        if idx[1] >= len(expenses[idx[0]]):
            idx[1] = 0
            idx[0] += 1
        return resp
 
    async def sampling_handler(messages, params, ctx):
        return CreateMessageResult(
            role="assistant",
            content=TextContent(type="text", text="1. Meal prep. 2. Cancel unused subscriptions. 3. Save energy."),
            model="test-model", stopReason="endTurn",
        )
 
    async def progress_handler(progress, total, message):
        pct = int((progress / total) * 100) if total else 0
        bar = "#" * (pct // 5) + "-" * (20 - pct // 5)
        print(f"  Progress: [{bar}] {pct}% ({int(progress)}/{int(total or 0)})")
 
    async with Client(mcp, elicitation_handler=elicit_handler,
                      sampling_handler=sampling_handler, progress_handler=progress_handler) as client:
        for i in range(4):
            print(f"\n--- Adding expense {i + 1} ---")
            result = await client.call_tool("add_expense_interactive", {"user_alias": "testuser"})
            print(f"  Result: {result.content[0].text}")
 
        print("\n--- Analyzing spending ---")
        result = await client.call_tool("analyze_spending", {"user_alias": "testuser"})
        print(f"  Result:\n{result.content[0].text}")
 
        print("\n--- Generating report ---")
        result = await client.call_tool("generate_report", {"user_alias": "testuser"})
        print(f"  Result:\n{result.content[0].text}")
 
 
async def test_decline():
    """Edge case: decline at 2nd question"""
    call_count = [0]
 
    async def handler(message, response_type, params, context):
        call_count[0] += 1
        print(f"  Server asks: {message}")
        if call_count[0] == 1:
            print("  Responding: {'amount': 99.99}")
            return {"amount": 99.99}
        print("  Responding: DECLINE")
        return ElicitResult(action="decline", content=None)
 
    async with Client(mcp, elicitation_handler=handler) as client:
        result = await client.call_tool("add_expense_interactive", {"user_alias": "testuser"})
        print(f"  Result: {result.content[0].text}")
 
 
async def test_cancel():
    """Edge case: cancel at 1st question"""
    async def handler(message, response_type, params, context):
        print(f"  Server asks: {message}")
        print("  Responding: CANCEL")
        return ElicitResult(action="cancel", content=None)
 
    async with Client(mcp, elicitation_handler=handler) as client:
        result = await client.call_tool("add_expense_interactive", {"user_alias": "testuser"})
        print(f"  Result: {result.content[0].text}")
 
 
async def test_confirm_no():
    """Edge case: answer all, then No at confirmation"""
    responses = iter([{"amount": 50.0}, {"description": "Test"}, {"category": "other"}, {"confirm": "No"}])
 
    async def handler(message, response_type, params, context):
        resp = next(responses)
        print(f"  Server asks: {message}")
        print(f"  Responding: {resp}")
        return resp
 
    async with Client(mcp, elicitation_handler=handler) as client:
        result = await client.call_tool("add_expense_interactive", {"user_alias": "testuser"})
        print(f"  Result: {result.content[0].text}")
 
 
async def test_sampling_failure():
    """Edge case: sampling handler throws exception"""
    async def handler(messages, params, ctx):
        raise RuntimeError("Simulated LLM failure")
 
    async with Client(mcp, sampling_handler=handler) as client:
        result = await client.call_tool("analyze_spending", {"user_alias": "testuser"})
        print(f"  Result: {result.content[0].text}")
 
 
async def test_no_sampling_handler():
    """Edge case: no sampling handler registered"""
    async with Client(mcp) as client:
        result = await client.call_tool("analyze_spending", {"user_alias": "testuser"})
        print(f"  Result: {result.content[0].text}")
 
 
async def main():
    for name, fn in [
        ("Normal Flow", test_normal_flow),
        ("Elicitation Decline", test_decline),
        ("Elicitation Cancel", test_cancel),
        ("Confirm No", test_confirm_no),
        ("Sampling Failure", test_sampling_failure),
        ("No Sampling Handler", test_no_sampling_handler),
    ]:
        print(f"\n{'=' * 50}\n  {name}\n{'=' * 50}")
        await fn()
 
 
if __name__ == "__main__":
    asyncio.run(main())