CalcMCP/mcp_server.py

"""
Scientific Calculator MCP Server - Production Ready
完全符合 MCP 规范的科学计算器服务器
支持通过标准输入输出（STDIO）进行 JSON-RPC 2.0 通信

This implementation strictly follows:
https://modelcontextprotocol.io/docs/develop/build-server
"""

import json
import sys
import logging
from typing import Any, Dict, List
from dataclasses import dataclass, asdict
from calculator import CALCULATOR_TOOLS
import numpy as np

# 配置日志到 stderr（不污染 stdout JSON-RPC 输出）
logging.basicConfig(
    level=logging.WARNING,
    format='[%(levelname)s] %(message)s',
    stream=sys.stderr,
    force=True
)
logger = logging.getLogger(__name__)


@dataclass
class JSONRPCRequest:
    """JSON-RPC 2.0 请求"""
    jsonrpc: str = "2.0"
    method: str = None
    params: Dict[str, Any] = None
    id: Any = None


@dataclass
class JSONRPCResponse:
    """JSON-RPC 2.0 响应"""
    jsonrpc: str = "2.0"
    result: Any = None
    error: Dict[str, Any] = None
    id: Any = None


class MCPCalculatorServer:
    """符合 MCP 规范的计算器服务器"""
    
    def __init__(self):
        self.version = "1.0.0"
        self.name = "Scientific Calculator"
        self.tools = self._build_tool_definitions()
        
    def _build_tool_definitions(self) -> List[Dict[str, Any]]:
        """构建工具定义列表"""
        tools = []
        
        # Consolidated tool definitions (3 tools)
        tool_definitions = {
            "symbolic_tool": {
                "description": (
                    "Exact symbolic math (SymPy). Set operation to one of: simplify, expand, factor, "
                    "derivative, integral, limit, solve, dsolve, taylor, matrix. "
                    "Write expressions in Python/SymPy syntax, e.g. 'x**2 * exp(-x)'; use 'oo' for infinity. "
                    "integral: optional lower_limit/upper_limit for definite integrals. "
                    "solve: pass the equation, '=' is allowed (e.g. 'x**2 = 4' or 'x**3 - 6*x**2 + 11*x - 6'). "
                    "dsolve: solves an ODE symbolically — write the unknown function as y(x), "
                    "e.g. equation='Derivative(y(x), x, 2) + y(x) = 0'. "
                    "taylor: expression + variable + point + order. "
                    "matrix: matrix_data (JSON) with variable=determinant|inverse|rank|trace. "
                    "Results are exact (e.g. '2', 'E', 'C1*sin(x) + C2*cos(x)')."
                ),
                "params": [
                    {"name": "operation", "type": "string", "description": "One of: simplify, expand, factor, derivative, integral, limit, solve, dsolve, taylor, matrix"},
                    {"name": "expression", "type": "string", "description": "Expression in SymPy syntax, e.g. 'sin(x)/x'", "optional": True},
                    {"name": "variable", "type": "string", "description": "Variable name, e.g. 'x' (for matrix: the operation name)", "optional": True},
                    {"name": "order", "type": "integer", "description": "Order for derivative/taylor (default 1 / 5)", "optional": True},
                    {"name": "lower_limit", "type": "string", "description": "Lower bound for definite integral, e.g. '0'", "optional": True},
                    {"name": "upper_limit", "type": "string", "description": "Upper bound for definite integral, e.g. 'oo'", "optional": True},
                    {"name": "equation", "type": "string", "description": "Equation for solve/dsolve; '=' allowed", "optional": True},
                    {"name": "point", "type": "string", "description": "Point for limit/taylor, e.g. '0' or 'oo'", "optional": True},
                    {"name": "matrix_data", "type": "string", "description": "Matrix as JSON, e.g. '[[1,2],[3,4]]'", "optional": True}
                ]
            },
            "numpy_tool": {
                "description": (
                    "Numeric arrays and linear algebra (NumPy). Set operation to one of: "
                    "sum, mean, std, max, min (array_data, optional axis); "
                    "eigenvalues, eigenvectors, determinant, inverse, solve, norm, rank, trace, "
                    "svd, qr, cholesky, matmul, dot, hadamard (matrix_a, and matrix_b where two matrices are needed); "
                    "sin, cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh (values, optional use_degrees); "
                    "poly_eval, poly_derivative, poly_integral (coefficients highest-degree first, optional x_values). "
                    "All array/matrix inputs are JSON strings, e.g. matrix_a='[[2,1],[1,2]]'. Results are numeric."
                ),
                "params": [
                    {"name": "operation", "type": "string", "description": "Operation name, e.g. 'eigenvalues', 'matmul', 'mean', 'poly_eval'"},
                    {"name": "array_data", "type": "string", "description": "Array as JSON, e.g. '[1,2,3]' or '[[1,2],[3,4]]'", "optional": True},
                    {"name": "axis", "type": "string", "description": "Axis for reductions: '0' or '1'", "optional": True},
                    {"name": "matrix_a", "type": "string", "description": "First matrix as JSON", "optional": True},
                    {"name": "matrix_b", "type": "string", "description": "Second matrix as JSON (solve/matmul/dot/hadamard)", "optional": True},
                    {"name": "values", "type": "string", "description": "Input values for trig functions as JSON, e.g. '[0, 90, 180]'", "optional": True},
                    {"name": "use_degrees", "type": "boolean", "description": "Treat trig inputs as degrees (default radians)", "optional": True},
                    {"name": "coefficients", "type": "string", "description": "Polynomial coefficients as JSON, highest degree first", "optional": True},
                    {"name": "x_values", "type": "string", "description": "x values for poly_eval as JSON", "optional": True}
                ]
            },
            "scipy_tool": {
                "description": (
                    "Numerical methods (SciPy). Set operation to one of: "
                    "integrate (numeric definite integral: expression, variable, lower, upper); "
                    "optimize_minimize, optimize_root (expression, initial_value, optional method); "
                    "interpolate_linear, interpolate_cubic, interpolate_spline (x_data, y_data, x_eval); "
                    "special (function name + parameters JSON, e.g. gamma, besselj); "
                    "solve_ode (NUMERIC simulation of dy/dt = f(t,y): expression e.g. '-2*y + t', "
                    "initial_conditions, t_values; for exact general solutions use symbolic_tool with operation='dsolve'); "
                    "statistics ops mean, std, describe, ttest, pearsonr (function + data, optional params); "
                    "fft, rfft (data_series or data); matrix_eigensystem (matrix_a). "
                    "All list inputs are JSON strings. Results are numeric."
                ),
                "params": [
                    {"name": "operation", "type": "string", "description": "Operation name, e.g. 'integrate', 'optimize_minimize', 'solve_ode', 'fft'"},
                    {"name": "expression", "type": "string", "description": "Expression/function of the variable, e.g. 'x**2 * exp(-x)'", "optional": True},
                    {"name": "variable", "type": "string", "description": "Variable name, e.g. 'x'", "optional": True},
                    {"name": "lower", "type": "string", "description": "Lower integration bound (number)", "optional": True},
                    {"name": "upper", "type": "string", "description": "Upper integration bound (number or 'oo')", "optional": True},
                    {"name": "initial_value", "type": "string", "description": "Starting point for optimization/root finding", "optional": True},
                    {"name": "method", "type": "string", "description": "Algorithm, e.g. 'BFGS', 'brentq'", "optional": True},
                    {"name": "x_data", "type": "string", "description": "Interpolation x data as JSON", "optional": True},
                    {"name": "y_data", "type": "string", "description": "Interpolation y data as JSON", "optional": True},
                    {"name": "x_eval", "type": "string", "description": "Points to evaluate the interpolant at, JSON", "optional": True},
                    {"name": "function", "type": "string", "description": "Special-function or statistics operation name", "optional": True},
                    {"name": "parameters", "type": "string", "description": "Special-function parameters as JSON", "optional": True},
                    {"name": "initial_conditions", "type": "string", "description": "ODE initial value(s) as JSON, e.g. '[1.0]'", "optional": True},
                    {"name": "t_values", "type": "string", "description": "ODE time points as JSON, e.g. '[0, 0.1, 0.2]'", "optional": True},
                    {"name": "variables", "type": "string", "description": "ODE dependent variable name, default 'y'", "optional": True},
                    {"name": "data", "type": "string", "description": "Statistics/FFT data as JSON", "optional": True},
                    {"name": "params", "type": "string", "description": "Extra statistics parameter as JSON: for ttest the population mean ('2.3' or '{\"popmean\": 2.3}'), for pearsonr the second data series", "optional": True},
                    {"name": "data_series", "type": "string", "description": "FFT data as JSON", "optional": True},
                    {"name": "matrix_a", "type": "string", "description": "Matrix as JSON for matrix_eigensystem", "optional": True}
                ]
            },
        }
        
        # 构建工具列表
        for tool_name, definition in tool_definitions.items():
            tools.append({
                "name": tool_name,
                "description": definition["description"],
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        param["name"]: {
                            "type": param["type"],
                            "description": param["description"]
                        }
                        for param in definition["params"]
                    },
                    "required": [
                        p["name"] for p in definition["params"] 
                        if not p.get("optional", False)
                    ]
                }
            })
        
        return tools
    
    def call_tool(self, tool_name: str, params: Dict[str, Any]) -> str:
        """调用计算器工具"""
        if tool_name not in CALCULATOR_TOOLS:
            raise ValueError(f"Unknown tool: {tool_name}")
        
        try:
            # 获取工具函数
            tool_func = CALCULATOR_TOOLS[tool_name]
            
            # 调用工具函数
            result = tool_func(**params)
            
            return str(result)
            
        except TypeError as e:
            return f"Parameter Error: {str(e)}"
        except Exception as e:
            return f"Execution Error: {str(e)}"
    
    def handle_rpc_request(self, request_dict: Dict[str, Any]) -> Dict[str, Any]:
        """处理 JSON-RPC 请求"""
        try:
            # 验证 JSON-RPC 2.0 格式
            if request_dict.get("jsonrpc") != "2.0":
                return {
                    "jsonrpc": "2.0",
                    "error": {
                        "code": -32600,
                        "message": "Invalid Request"
                    },
                    "id": request_dict.get("id")
                }
            
            method = request_dict.get("method")
            params = request_dict.get("params", {})
            req_id = request_dict.get("id")

            # JSON-RPC notifications (no id) must not be answered
            if req_id is None and isinstance(method, str) and method.startswith("notifications/"):
                return None

            # 处理特殊方法
            if method == "initialize":
                return {
                    "jsonrpc": "2.0",
                    "result": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {},
                        "serverInfo": {
                            "name": self.name,
                            "version": self.version
                        }
                    },
                    "id": req_id
                }
            
            elif method == "tools/list":
                return {
                    "jsonrpc": "2.0",
                    "result": {
                        "tools": self.tools
                    },
                    "id": req_id
                }
            
            elif method == "tools/call":
                tool_name = params.get("name")
                tool_params = params.get("arguments", {})
                
                result = self.call_tool(tool_name, tool_params)

                is_error = isinstance(result, str) and result.startswith(
                    ("Error", "Parameter Error", "Execution Error")
                )
                return {
                    "jsonrpc": "2.0",
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": result
                            }
                        ],
                        "isError": is_error
                    },
                    "id": req_id
                }
            
            else:
                return {
                    "jsonrpc": "2.0",
                    "error": {
                        "code": -32601,
                        "message": f"Method not found: {method}"
                    },
                    "id": req_id
                }
        
        except Exception as e:
            logger.error(f"RPC Error: {e}")
            return {
                "jsonrpc": "2.0",
                "error": {
                    "code": -32603,
                    "message": f"Internal error: {str(e)}"
                },
                "id": request_dict.get("id")
            }
    
    def run(self):
        """运行 MCP 服务器（STDIO 模式）"""
        logger.info(f"Starting {self.name} v{self.version}")
        logger.info(f"Registered {len(self.tools)} tools")
        
        try:
            while True:
                # 从标准输入读取 JSON-RPC 请求
                line = sys.stdin.readline()
                
                if not line:
                    # EOF
                    break
                
                try:
                    request = json.loads(line)
                    response = self.handle_rpc_request(request)

                    # Notifications return None — nothing must be written back
                    if response is not None:
                        sys.stdout.write(json.dumps(response, ensure_ascii=False) + '\n')
                        sys.stdout.flush()
                    
                except json.JSONDecodeError as e:
                    # 发送 JSON 解析错误
                    error_response = {
                        "jsonrpc": "2.0",
                        "error": {
                            "code": -32700,
                            "message": "Parse error"
                        }
                    }
                    sys.stdout.write(json.dumps(error_response) + '\n')
                    sys.stdout.flush()
        
        except KeyboardInterrupt:
            logger.info("Server interrupted")
        except Exception as e:
            logger.error(f"Server error: {e}")
            sys.exit(1)


def main():
    """启动 MCP 服务器"""
    server = MCPCalculatorServer()
    server.run()


if __name__ == "__main__":
    main()