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
brancheschar/aliceai-t5-moeall refs →
3.4 KiB118 linesraw
1#!/usr/bin/env python3
2"""
3Minimal MCP server that writes notification + response in a single write() with no flush.
4This reproduces the buffering bug where read_message() can strand the response.
5"""
6import json
7import sys
8import os
9
10TOOLS = [
11    {
12        "name": "echo",
13        "description": "Echo back the input message",
14        "inputSchema": {
15            "type": "object",
16            "properties": {
17                "message": {"type": "string"}
18            },
19            "required": ["message"]
20        }
21    }
22]
23
24def handle_initialize(params, req_id):
25    return {
26        "jsonrpc": "2.0",
27        "id": req_id,
28        "result": {
29            "protocolVersion": "2024-11-05",
30            "capabilities": {"tools": {}},
31            "serverInfo": {"name": "burst-test", "version": "1.0"}
32        }
33    }
34
35def handle_tools_list(params, req_id):
36    return {
37        "jsonrpc": "2.0",
38        "id": req_id,
39        "result": {"tools": TOOLS}
40    }
41
42def handle_tools_call(params, req_id):
43    tool_name = params.get("name")
44    arguments = params.get("arguments", {})
45
46    if tool_name == "echo":
47        message = arguments.get("message", "")
48        notif = {
49            "jsonrpc": "2.0",
50            "method": "notifications/progress",
51            "params": {"progress": 50, "total": 100}
52        }
53        response = {
54            "jsonrpc": "2.0",
55            "id": req_id,
56            "result": {
57                "content": [{"type": "text", "text": f"echo: {message}"}]
58            }
59        }
60        # Single os.write() call: both lines land in one pipe packet atomically.
61        # This is the key difference from mcp_malformed_server.py which flushes between writes.
62        data = (json.dumps(notif) + "\n" + json.dumps(response) + "\n").encode("utf-8")
63        os.write(sys.stdout.fileno(), data)
64        return None  # already written
65    else:
66        response = {
67            "jsonrpc": "2.0",
68            "id": req_id,
69            "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
70        }
71        return response
72
73HANDLERS = {
74    "initialize": handle_initialize,
75    "tools/list": handle_tools_list,
76    "tools/call": handle_tools_call,
77}
78
79def main():
80    # Use line-buffered text mode for regular responses, but the burst write
81    # uses os.write() directly to guarantee a single kernel write().
82    sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
83    sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
84
85    for line in sys.stdin:
86        line = line.strip()
87        if not line:
88            continue
89        try:
90            request = json.loads(line)
91        except json.JSONDecodeError:
92            continue
93
94        method = request.get("method")
95        req_id = request.get("id")
96        params = request.get("params", {})
97
98        # JSON-RPC 2.0: a message without an id is a notification and must not receive a response
99        if req_id is None:
100            continue
101
102        handler = HANDLERS.get(method)
103        if handler:
104            response = handler(params, req_id)
105            if response is not None:
106                sys.stdout.write(json.dumps(response) + "\n")
107                sys.stdout.flush()
108        else:
109            response = {
110                "jsonrpc": "2.0",
111                "id": req_id,
112                "error": {"code": -32601, "message": f"Method not found: {method}"}
113            }
114            sys.stdout.write(json.dumps(response) + "\n")
115            sys.stdout.flush()
116
117if __name__ == "__main__":
118    main()