04_custom_tools.py

"""04: Custom tools via an in-process MCP server."""

import asyncio
from typing import Annotated, Any

from claude_agent_sdk import (
    AssistantMessage,
    ClaudeAgentOptions,
    ResultMessage,
    SystemMessage,
    ToolAnnotations,
    ToolUseBlock,
    UserMessage,
    ToolResultBlock,
    create_sdk_mcp_server,
    query,
    tool,
)
from demo_util import PROJECT

INVENTORY = {"widget": 42, "gadget": 0, "sprocket": 7}


@tool(
    "get_stock",
    "Look up how many units of a product are in stock",
    {"product": Annotated[str, "Product name, lowercase"]},
    annotations=ToolAnnotations(readOnlyHint=True),  # safe to run in parallel
)
async def get_stock(args: dict[str, Any]) -> dict[str, Any]:
    product = args["product"].lower().rstrip("s")  # accept "widgets" for "widget"
    if product not in INVENTORY:
        return {
            "content": [{"type": "text", "text": f"Unknown product: {product}"}],
            "is_error": True,
        }
    return {"content": [{"type": "text", "text": f"{product}: {INVENTORY[product]} units"}]}


@tool(
    "reorder",
    "Place a reorder for a product",
    {
        "type": "object",
        "properties": {
            "product": {"type": "string"},
            "quantity": {"type": "integer", "minimum": 1, "maximum": 1000},
            "priority": {"type": "string", "enum": ["normal", "rush"]},
        },
        "required": ["product", "quantity"],  # priority is optional
    },
)
async def reorder(args: dict[str, Any]) -> dict[str, Any]:
    priority = args.get("priority", "normal")
    INVENTORY[args["product"]] = INVENTORY.get(args["product"], 0) + args["quantity"]
    return {"content": [{"type": "text", "text":
            f"Reorder placed: {args['quantity']} x {args['product']} ({priority}). "
            f"New stock: {INVENTORY[args['product']]}"}]}


inventory_server = create_sdk_mcp_server(name="inventory", version="1.0.0", tools=[get_stock, reorder])


async def main():
    options = ClaudeAgentOptions(
        cwd=PROJECT,
        mcp_servers={"inventory": inventory_server},
        allowed_tools=["mcp__inventory__*"],   # wildcard: every tool on this server
        tools=[],                               # no built-in tools at all
        max_turns=8,
    )
    async for message in query(
        prompt="Check stock for widgets, gadgets and doohickeys. "
               "If anything is out of stock, place a rush reorder for 25 units.",
        options=options,
    ):
        if isinstance(message, SystemMessage) and message.subtype == "init":
            mcp = [t for t in message.data.get("tools", []) if t.startswith("mcp__")]
            print("MCP tools visible at init:", mcp)
            print("MCP servers:", message.data.get("mcp_servers"))
        elif isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, ToolUseBlock):
                    print(f"[tool call] {block.name}({block.input})")
        elif isinstance(message, UserMessage) and isinstance(message.content, list):
            for block in message.content:
                if isinstance(block, ToolResultBlock):
                    print(f"[tool result]{' ERROR' if block.is_error else ''} {block.content}")
        elif isinstance(message, ResultMessage) and message.subtype == "success":
            print("\nFinal answer:", message.result)

    print("\nInventory after run:", INVENTORY)


asyncio.run(main())