Building Calculator Tools

"""
tools.py - Reusable tool registry for MCP servers.

This module builds on jsonrpc.py from Lesson 3.
"""

from typing import Callable

# Import error handling from Lesson 3
from jsonrpc import JsonRpcError, JsonRpcErrorCode


# =============================================================================
# Section 1: Result Helpers
# =============================================================================

def text_result(text: str) -> dict:
    """Create a successful text result."""
    return {
        "content": [{"type": "text", "text": str(text)}]
    }


def error_result(message: str) -> dict:
    """Create an error result (tool executed but failed)."""
    return {
        "content": [{"type": "text", "text": f"Error: {message}"}],
        "isError": True
    }

# =============================================================================
# Section 2: Tool Registry
# =============================================================================

class ToolRegistry:
    """
    Registry for MCP tools with decorator-based registration.

    Usage:
        tools = ToolRegistry()

        @tools.register(
            name="add",
            description="Add two numbers",
            input_schema={"type": "object", "properties": {...}}
        )
        def add_handler(arguments: dict) -> dict:
            return text_result(f"Sum: {arguments['a'] + arguments['b']}")
    """

    def __init__(self):
        self._tools: dict[str, dict] = {}
        self._handlers: dict[str, Callable] = {}

    def register(
        self,
        name: str,
        description: str,
        input_schema: dict
    ) -> Callable:
        """Decorator to register a tool handler."""
        def decorator(func: Callable[[dict], dict]) -> Callable:
            # Store tool definition
            self._tools[name] = {
                "name": name,
                "description": description,
                "inputSchema": input_schema
            }
            # Store handler function
            self._handlers[name] = func
            return func
        return decorator

    def list_tools(self) -> list[dict]:
        """Return all tool definitions for tools/list."""
        return list(self._tools.values())

    def call_tool(self, name: str, arguments: dict) -> dict:
        """Execute a tool by name."""
        if name not in self._handlers:
            raise JsonRpcError(
                JsonRpcErrorCode.INVALID_PARAMS,
                f"Unknown tool: {name}"
            )
        return self._handlers[name](arguments)

# =============================================================================
# Section 3: MCP Router Integration
# =============================================================================

def register_tool_handlers(router, tools: ToolRegistry) -> None:
    """
    Register tools/list and tools/call handlers with an MCP router.

    Args:
        router: A JsonRpcRouter from mcp_core.py
        tools: A ToolRegistry with registered tools
    """

    @router.method("tools/list")
    def handle_tools_list(params: dict) -> dict:
        return {"tools": tools.list_tools()}

    @router.method("tools/call")
    def handle_tools_call(params: dict) -> dict:
        name = params.get("name", "")
        arguments = params.get("arguments", {})
        return tools.call_tool(name, arguments)

# =============================================================================
# Test the module
# =============================================================================

if __name__ == "__main__":
    # Create a registry and register a test tool
    tools = ToolRegistry()

    @tools.register(
        name="greet",
        description="Say hello",
        input_schema={
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"]
        }
    )
    def greet(args: dict) -> dict:
        return text_result(f"Hello, {args['name']}!")

    # Test listing and calling
    print("Registered tools:")
    for tool in tools.list_tools():
        print(f"  - {tool['name']}: {tool['description']}")

    print("\nCalling greet tool:")
    result = tools.call_tool("greet", {"name": "World"})
    print(f"  {result}")

"""
calculator_example.py - Calculator server using tools.py and mcp_core.py

This demonstrates how to build an MCP server by combining modules
from previous lessons.

Run: python calculator_example.py
Test: curl commands below
"""

import json
from http.server import HTTPServer, BaseHTTPRequestHandler

# Import from our modules (Lessons 3, 4, and 5)
from jsonrpc import JsonRpcErrorCode, create_error_response
from mcp_core import (
    router, state, ConnectionState,
    SERVER_NAME, SERVER_VERSION,
    handle_request
)
from tools import ToolRegistry, text_result, register_tool_handlers


# =============================================================================
# Create tool registry and register calculator tools
# =============================================================================

calculator_tools = ToolRegistry()


@calculator_tools.register(
    name="add",
    description="Add two numbers together",
    input_schema={
        "type": "object",
        "properties": {
            "a": {"type": "number", "description": "First number"},
            "b": {"type": "number", "description": "Second number"}
        },
        "required": ["a", "b"]
    }
)
def add(args: dict) -> dict:
    a, b = args["a"], args["b"]
    return text_result(f"{a} + {b} = {a + b}")


@calculator_tools.register(
    name="subtract",
    description="Subtract the second number from the first",
    input_schema={
        "type": "object",
        "properties": {
            "a": {"type": "number", "description": "Number to subtract from"},
            "b": {"type": "number", "description": "Number to subtract"}
        },
        "required": ["a", "b"]
    }
)
def subtract(args: dict) -> dict:
    a, b = args["a"], args["b"]
    return text_result(f"{a} - {b} = {a - b}")


@calculator_tools.register(
    name="multiply",
    description="Multiply two numbers",
    input_schema={
        "type": "object",
        "properties": {
            "a": {"type": "number", "description": "First number"},
            "b": {"type": "number", "description": "Second number"}
        },
        "required": ["a", "b"]
    }
)
def multiply(args: dict) -> dict:
    a, b = args["a"], args["b"]
    return text_result(f"{a} × {b} = {a * b}")


# Register tool handlers with the MCP router
register_tool_handlers(router, calculator_tools)


# =============================================================================
# HTTP Server (reuses mcp_core.handle_request)
# =============================================================================

class MCPHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        print(f"[Calculator] {args[0]}")

    def send_json(self, data: dict, status: int = 200):
        body = json.dumps(data, indent=2).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length)

        try:
            request = json.loads(body)
        except json.JSONDecodeError as e:
            self.send_json(create_error_response(
                None, JsonRpcErrorCode.PARSE_ERROR, str(e)
            ), 400)
            return

        if request.get("jsonrpc") != "2.0":
            self.send_json(create_error_response(
                request.get("id"),
                JsonRpcErrorCode.INVALID_REQUEST,
                "jsonrpc must be '2.0'"
            ), 400)
            return

        response = handle_request(request)
        if response is None:
            self.send_response(204)
            self.end_headers()
        else:
            self.send_json(response)


if __name__ == "__main__":
    print("Calculator MCP Server")
    print("Tools: add, subtract, multiply")
    print("Running on http://localhost:8080")
    print("Press Ctrl+C to stop\n")

    server = HTTPServer(("localhost", 8080), MCPHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")