char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenserver: support MCP stdio (#26062)20455a4ad

char/aliceai-t5-moe
3.5 KiB132 linesraw
1#!/usr/bin/env python3
2"""
3MCP server that sleeps before responding, for timeout testing.
4"""
5import json
6import sys
7import os
8import time
9import argparse
10
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 duration
76    global handle_tools_call
77    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            continue
105        try:
106            request = json.loads(line)
107        except json.JSONDecodeError:
108            continue
109
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 response
115        if req_id is None:
116            continue
117
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()