Felipe97/llama-cpp-compiled

1#!/usr/bin/env python32"""3MCP server that sleeps before responding, for timeout testing.4"""5import json6import sys7import os8import time9import argparse10 11TOOLS = [12    {13        "name": "sleep",14        "description": "Sleep for a given number of seconds",15        "inputSchema": {16            "type": "object",17            "properties": {18                "seconds": {"type": "number", "description": "Seconds to sleep"}19            },20            "required": ["seconds"]21        }22    }23]24 25def handle_initialize(params, req_id):26    return {27        "jsonrpc": "2.0",28        "id": req_id,29        "result": {30            "protocolVersion": "2024-11-05",31            "capabilities": {"tools": {}},32            "serverInfo": {"name": "slow-test", "version": "1.0"}33        }34    }35 36def handle_tools_list(params, req_id):37    return {38        "jsonrpc": "2.0",39        "id": req_id,40        "result": {"tools": TOOLS}41    }42 43def handle_tools_call(params, req_id):44    tool_name = params.get("name")45    arguments = params.get("arguments", {})46 47    if tool_name == "sleep":48        seconds = arguments.get("seconds", 1)49        time.sleep(seconds)50        return {51            "jsonrpc": "2.0",52            "id": req_id,53            "result": {54                "content": [{"type": "text", "text": f"slept {seconds}s"}]55            }56        }57    else:58        return {59            "jsonrpc": "2.0",60            "id": req_id,61            "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}62        }63 64HANDLERS = {65    "initialize": handle_initialize,66    "tools/list": handle_tools_list,67    "tools/call": handle_tools_call,68}69 70def main():71    parser = argparse.ArgumentParser()72    parser.add_argument("--delay", type=float, default=5.0, help="Delay in seconds for sleep tool")73    args = parser.parse_args()74 75    # Override the sleep duration76    global handle_tools_call77    def handle_tools_call(params, req_id):78        tool_name = params.get("name")79        arguments = params.get("arguments", {})80 81        if tool_name == "sleep":82            seconds = arguments.get("seconds", args.delay)83            time.sleep(seconds)84            return {85                "jsonrpc": "2.0",86                "id": req_id,87                "result": {88                    "content": [{"type": "text", "text": f"slept {seconds}s"}]89                }90            }91        else:92            return {93                "jsonrpc": "2.0",94                "id": req_id,95                "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}96            }97 98    sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)99    sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)100 101    for line in sys.stdin:102        line = line.strip()103        if not line:104            continue105        try:106            request = json.loads(line)107        except json.JSONDecodeError:108            continue109 110        method = request.get("method")111        req_id = request.get("id")112        params = request.get("params", {})113 114        # JSON-RPC 2.0: a message without an id is a notification and must not receive a response115        if req_id is None:116            continue117 118        handler = HANDLERS.get(method)119        if handler:120            response = handler(params, req_id)121        else:122            response = {123                "jsonrpc": "2.0",124                "id": req_id,125                "error": {"code": -32601, "message": f"Method not found: {method}"}126            }127 128        sys.stdout.write(json.dumps(response) + "\n")129        sys.stdout.flush()130 131if __name__ == "__main__":132    main()133