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
6.3 KiB176 linesraw
1#pragma once
2
3#include "server-common.h"
4
5#include <atomic>
6#include <chrono>
7#include <functional>
8#include <map>
9#include <memory>
10#include <mutex>
11#include <string>
12#include <thread>
13#include <vector>
14
15//
16// Configuration (Cursor-compatible "mcpServers" JSON)
17//
18
19struct server_mcp_server_config {
20    std::string name; // config key, e.g. "filesystem"
21    std::string command;
22    std::vector<std::string> args;
23    std::map<std::string, std::string> env; // merged over the parent env
24    std::string cwd;
25    int timeout_ms = 30000; // per-tool-call timeout
26
27    // throw on parse errors; missing "mcpServers" yields an empty list; entries without a "command" are skipped
28    static std::vector<server_mcp_server_config> parse_from_json(const std::string & json_str);
29    static std::vector<server_mcp_server_config> parse_cursor_format(const json & j);
30};
31
32// a tool advertised by an MCP server
33struct server_mcp_tool_def {
34    std::string server_name;
35    std::string name; // bare tool name, no "<server>_" prefix
36    std::string description;
37    json input_schema; // JSON Schema for the arguments, or null
38};
39
40//
41// server_mcp_transport: one MCP server session.
42//
43//   caller --send_rpc--> to_server   --[writer]--> framing --> server
44//   caller <--send_rpc-- from_server <--[reader]-- framing <-- server
45//
46// each queue item is one complete serialized JSON message.
47// subclass owns byte I/O and framing; base owns JSON and the JSON-RPC session (handshake, id correlation).
48//
49
50struct server_mcp_transport {
51    std::string name;
52    int timeout_ms = 30000;
53
54    server_pipe<std::string> to_server;   // serialized messages we send to the server
55    server_pipe<std::string> from_server; // serialized messages read from the server
56
57    virtual ~server_mcp_transport() = default;
58
59    virtual bool start() = 0;
60    virtual void close() = 0; // blocking and idempotent
61    virtual bool is_alive() const = 0; // never blocks behind an in-flight send_rpc()
62
63    // human-readable diagnostics for logging when the transport fails/dies
64    // (example: last RPC error, plus any transport-specific detail)
65    // may run on a different thread than send_rpc(), so last_error is read under rpc_mutex
66    virtual std::string diagnostics() {
67        std::lock_guard<std::mutex> lock(rpc_mutex);
68        return last_error;
69    }
70
71    std::vector<server_mcp_tool_def> list_tools(const std::function<bool()> & should_stop);
72
73    json call_tool(const std::string & tool_name,
74                   const json & arguments,
75                   const std::function<bool()> & should_stop);
76
77protected:
78    // per-transport: send_rpc() holds it across the reply wait, so sharing it would stall every server behind one slow call. guards all members below.
79    std::mutex rpc_mutex;
80    uint64_t next_id = 1; // reset to 1 per (re)spawn
81    bool initialized = false;
82    std::string last_error;
83    std::vector<server_mcp_tool_def> tools;
84
85    // both assume rpc_mutex is already held by the public caller
86    bool ensure_init(const std::function<bool()> & should_stop); // initialize handshake, once
87    json send_rpc(const json & request, const std::function<bool()> & should_stop); // returns the reply or an {"error": ...}
88};
89
90//
91// server_mcp_stdio: child process, NDJSON JSON-RPC over stdio (stderr drained to the debug log)
92//
93
94struct server_mcp_stdio : server_mcp_transport {
95    explicit server_mcp_stdio(const server_mcp_server_config & config);
96    ~server_mcp_stdio() override;
97
98    bool start() override;
99    void close() override;
100    bool is_alive() const override;
101    std::string diagnostics() override;
102
103private:
104    server_mcp_server_config config;
105
106    // defined in the .cpp so <windows.h> stays out of this header
107    struct process_handle;
108    std::unique_ptr<process_handle> proc;
109
110    std::thread reader; // child stdout -> NDJSON de-framing -> from_server
111    std::thread writer; // to_server -> NDJSON framing -> child stdin
112    std::thread errlog; // child stderr -> debug log (must be drained or the child blocks)
113
114    // cleared by close() or by the reader on stdout EOF; read without rpc_mutex
115    std::atomic<bool> running{false};
116
117    // bounded tail of the child's stderr, for diagnostics when it dies
118    std::mutex err_mu;
119    std::string err_tail;
120
121    void reader_loop();
122    void writer_loop();
123    void errlog_loop();
124    void join_pumps();
125};
126
127//
128// server_mcp
129// declare before the HTTP context so it outlives every /tools handler.
130//
131
132class server_mcp {
133public:
134    server_mcp() = default;
135    ~server_mcp();
136
137    // parse the MCP config from params (file and/or inline JSON),
138    // then spawn each server once,  list its tools, and shut it down
139    // throws on config parse errors; spawn failures are logged.
140    void start(const common_params & params);
141
142    // true until start() has parsed at least one server from the config
143    bool empty() const { return configs.empty(); }
144
145    std::vector<server_mcp_tool_def> list_tools() const;
146
147    // lazily (re)spawns the transport. returns the MCP result or an {"error": ...}. should_stop is OR-ed with the manager's cancel flag.
148    json call_tool(const std::string & server_name,
149                   const std::string & tool_name,
150                   const json & arguments,
151                   const std::function<bool()> & should_stop = nullptr);
152
153    // flip the cancel flag so in-flight calls return; blocking teardown is in the destructor. call before the HTTP server drains.
154    // note: multiple calls are idempotent
155    void shutdown();
156
157private:
158    std::vector<server_mcp_server_config> configs;
159
160    mutable std::mutex mutex; // guards transports, dead_servers, registry
161
162    // shared_ptr: call_tool() hands a transport to the caller and drops the lock for the blocking RPC, so a concurrent evict/respawn must not destroy it mid-call
163    std::map<std::string, std::shared_ptr<server_mcp_transport>> transports;
164    std::map<std::string, std::chrono::steady_clock::time_point> dead_servers; // spawn-failure cooldown
165    std::vector<server_mcp_tool_def> registry;
166
167    std::atomic<bool> stopping{false};
168
169    const server_mcp_server_config * find_config(const std::string & name) const;
170
171    // the only place that names a concrete transport
172    std::shared_ptr<server_mcp_transport> create_transport(const server_mcp_server_config & cfg);
173
174    // nullptr during cooldown or shutdown
175    std::shared_ptr<server_mcp_transport> get_or_create(const std::string & name);
176};