cleveragents-core/features/mocks/mcp_stub_server_stdio.py

#!/usr/bin/env python3
"""MCP stub server for integration testing of StdioMCPTransport.

Implements the MCP JSON-RPC 2.0 stdio protocol for deterministic testing.
This server speaks MCP 2024-11-05 protocol version.
"""

from __future__ import annotations

import json
import sys
import threading
import time
from typing import Any

MCP_PROTOCOL_VERSION = "2024-11-05"

TOOLS = [
    {
        "name": "test/echo",
        "description": "Echoes back the input arguments",
        "inputSchema": {
            "type": "object",
            "properties": {
                "message": {"type": "string", "description": "Message to echo"},
            },
            "required": ["message"],
        },
    },
    {
        "name": "test/multiply",
        "description": "Multiplies two numbers",
        "inputSchema": {
            "type": "object",
            "properties": {
                "a": {"type": "number", "description": "First number"},
                "b": {"type": "number", "description": "Second number"},
            },
            "required": ["a", "b"],
        },
    },
    {
        "name": "test/slow",
        "description": "A slow tool that takes time to respond",
        "inputSchema": {
            "type": "object",
            "properties": {
                "delay": {"type": "number", "description": "Delay in seconds"},
            },
        },
    },
]

request_id = 0
capabilities: dict[str, Any] = {}
initialized = False
slow_tool_delay = 0.0
lock = threading.Lock()


def send_response(resp_id: int | str | None, result: Any) -> None:
    response = {
        "jsonrpc": "2.0",
        "id": resp_id,
        "result": result,
    }
    sys.stdout.write(json.dumps(response) + "\n")
    sys.stdout.flush()


def send_error(resp_id: int | str | None, code: int, message: str) -> None:
    response = {
        "jsonrpc": "2.0",
        "id": resp_id,
        "error": {"code": code, "message": message},
    }
    sys.stdout.write(json.dumps(response) + "\n")
    sys.stdout.flush()


def handle_initialize(params: dict[str, Any]) -> dict[str, Any]:
    global capabilities, initialized
    capabilities = params.get("capabilities", {})
    initialized = True
    return {
        "protocolVersion": MCP_PROTOCOL_VERSION,
        "capabilities": {"tools": True, "resources": {}},
        "serverInfo": {
            "name": "test-mcp-stdio-server",
            "version": "1.0.0",
        },
    }


def handle_tools_list(params: dict[str, Any]) -> dict[str, Any]:
    return {"tools": TOOLS}


def handle_tools_call(req_id: int | str | None, params: dict[str, Any]) -> None:
    global slow_tool_delay

    tool_name = params.get("name", "")
    arguments = params.get("arguments", {})

    if tool_name == "test/echo":
        message = arguments.get("message", "")
        send_response(
            req_id,
            {"content": [{"type": "text", "text": json.dumps({"echoed": message})}]},
        )
        return

    if tool_name == "test/multiply":
        a = arguments.get("a", 0)
        b = arguments.get("b", 0)
        send_response(
            req_id,
            {"content": [{"type": "text", "text": json.dumps({"result": a * b})}]},
        )
        return

    if tool_name == "test/slow":
        delay = arguments.get("delay", 0.1)
        if delay > slow_tool_delay:
            time.sleep(delay)
        send_response(
            req_id,
            {"content": [{"type": "text", "text": json.dumps({"delayed": delay})}]},
        )
        return

    send_error(req_id, -32602, f"Unknown tool: {tool_name}")


def handle_request(request: dict[str, Any]) -> None:
    global request_id, slow_tool_delay

    method = request.get("method", "")
    params = request.get("params", {})
    req_id = request.get("id")

    if method == "initialize":
        result = handle_initialize(params)
        send_response(req_id, result)
    elif method == "notifications/initialized":
        pass
    elif method == "tools/list":
        if not initialized:
            send_error(req_id, -32602, "Not initialized")
            return
        result = handle_tools_list(params)
        send_response(req_id, result)
    elif method == "tools/call":
        if not initialized:
            send_error(req_id, -32602, "Not initialized")
            return
        handle_tools_call(req_id, params)
    elif method == "test/set_slow_delay":
        slow_tool_delay = params.get("delay", 0.0)
        send_response(req_id, {"ok": True})
    else:
        send_error(req_id, -32601, f"Unknown method: {method}")


def main() -> None:
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            request = json.loads(line)
        except json.JSONDecodeError:
            continue

        if isinstance(request, dict) and request.get("jsonrpc") == "2.0":
            handle_request(request)


if __name__ == "__main__":
    main()