Building a Router

class JsonRpcError(Exception):
    """Exception that converts to a JSON-RPC error response."""

    def __init__(
        self,
        code: JsonRpcErrorCode,
        message: str | None = None,
        data: Any = None
    ):
        self.code = code
        self.message = message or ERROR_MESSAGES.get(code, "Unknown error")
        self.data = data
        super().__init__(self.message)

from typing import Callable, cast


class JsonRpcRouter:
    """
    Routes JSON-RPC messages to handler functions.

    Usage:
        router = JsonRpcRouter()

        @router.method("tools/list")
        def handle_tools_list(params):
            return {"tools": [...]}

        response = router.handle(request)
    """

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

    def method(self, name: str):
        """Decorator to register a request handler."""
        def decorator(func: Callable):
            self._request_handlers[name] = func
            return func
        return decorator

    def notification(self, name: str):
        """Decorator to register a notification handler."""
        def decorator(func: Callable):
            self._notification_handlers[name] = func
            return func
        return decorator

    def handle(self, message: dict) -> dict | None:
        """
        Process a JSON-RPC message and return a response.

        Args:
            message: Parsed JSON-RPC message

        Returns:
            Response dict for requests, None for notifications
        """
        if is_notification(message):
            self._handle_notification(message)
            return None

        if is_request(message):
            return self._handle_request(message)

        # Invalid message
        return create_error_response(
            request_id=message.get("id"),
            code=JsonRpcErrorCode.INVALID_REQUEST,
            message="Invalid JSON-RPC message"
        )

    def _handle_request(self, request: dict) -> dict:
        """Handle a request and return a response."""
        request_id = cast(int | str, request.get("id"))
        method = request.get("method", "")
        params = request.get("params", {})

        # Find handler
        handler = self._request_handlers.get(method)
        if not handler:
            return create_error_response(
                request_id=request_id,
                code=JsonRpcErrorCode.METHOD_NOT_FOUND,
                message=f"Method not found: {method}",
                data={"method": method}
            )

        # Execute handler
        try:
            result = handler(params)
            return create_success_response(request_id, result)
        except JsonRpcError as e:
            return create_error_response(
                request_id=request_id,
                code=e.code,
                message=e.message,
                data=e.data
            )
        except Exception as e:
            return create_error_response(
                request_id=request_id,
                code=JsonRpcErrorCode.INTERNAL_ERROR,
                message=str(e)
            )

    def _handle_notification(self, notification: dict):
        """Handle a notification silently (no response)."""
        method = notification.get("method", "")
        params = notification.get("params", {})

        handler = self._notification_handlers.get(method)
        if handler:
            try:
                handler(params)
            except Exception as e:
                # Log but never respond
                print(f"Notification error ({method}): {e}")

# =============================================================================
# Try it out
# =============================================================================

if __name__ == "__main__":
    # Create a router
    router = JsonRpcRouter()

    # Register request handlers using decorators
    @router.method("initialize")
    def handle_initialize(params: dict) -> dict:
        return {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "demo-server", "version": "1.0.0"}
        }

    @router.method("tools/list")
    def handle_tools_list(params: dict) -> dict:
        return {
            "tools": [
                {
                    "name": "calculator",
                    "description": "Basic math operations",
                    "inputSchema": {"type": "object"}
                }
            ]
        }

    @router.method("tools/call")
    def handle_tools_call(params: dict) -> dict:
        tool_name = params.get("name", "")
        if tool_name != "calculator":
            raise JsonRpcError(
                JsonRpcErrorCode.INVALID_PARAMS,
                f"Unknown tool: {tool_name}"
            )
        return {"content": [{"type": "text", "text": "4"}]}

    # Register notification handlers
    @router.notification("notifications/initialized")
    def on_initialized(params: dict):
        print("  [Server] Client initialization complete!")

    # Test the router
    print("=" * 60)
    print("JSON-RPC Router Demo")
    print("=" * 60)

    test_messages = [
        # Request: initialize
        {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
        # Notification: initialized (no response expected)
        {"jsonrpc": "2.0", "method": "notifications/initialized"},
        # Request: tools/list
        {"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
        # Request: unknown method
        {"jsonrpc": "2.0", "id": 3, "method": "unknown/method"},
        # Request: tools/call with invalid params
        {"jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": {"name": "invalid"}},
    ]

    for msg in test_messages:
        print(f"\n-> Incoming: {json.dumps(msg)}")
        response = router.handle(msg)
        if response:
            print(f"<- Response: {json.dumps(response, indent=2)}")
        else:
            print("<- (no response - notification)")

    print("\n" + "=" * 60)
    print("Router demo complete!")
    print("=" * 60)