454 lines
19 KiB
Rust
454 lines
19 KiB
Rust
// MCP connectors: the bot can use the tools of Model Context Protocol servers (the original's
|
||
// local-exec-daemon has an MCP client with its own permission layer: shouldBlockMcp /
|
||
// isMcpFullyAllowlisted, and per-server custom instructions). Here: stdio servers (a command
|
||
// AnyBot starts) and HTTP servers (Streamable HTTP, JSON or SSE answers).
|
||
|
||
use serde::{Deserialize, Serialize};
|
||
use serde_json::{json, Value};
|
||
use std::collections::{BTreeMap, HashMap};
|
||
use std::io::{BufRead, BufReader, Write};
|
||
use std::process::{Child, ChildStdin, Command, Stdio};
|
||
use std::sync::mpsc::{channel, Receiver};
|
||
use std::sync::{Arc, Mutex, OnceLock};
|
||
use std::time::Duration;
|
||
|
||
pub const PROTOCOL: &str = "2025-06-18";
|
||
/// Longest wait for one answer (estimated).
|
||
pub const CALL_TIMEOUT_SECS: u64 = 120;
|
||
/// Tool names the model sees: mcp__<server>__<tool>, at most 64 characters.
|
||
pub const PREFIX: &str = "mcp__";
|
||
|
||
#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
|
||
pub struct Server {
|
||
pub name: String,
|
||
/// stdio: the command; empty when `url` is used.
|
||
#[serde(default)]
|
||
pub command: String,
|
||
#[serde(default)]
|
||
pub args: Vec<String>,
|
||
#[serde(default)]
|
||
pub env: BTreeMap<String, String>,
|
||
/// HTTP (Streamable HTTP) address.
|
||
#[serde(default)]
|
||
pub url: String,
|
||
/// e.g. "Authorization: Bearer …" for HTTP servers.
|
||
#[serde(default)]
|
||
pub headers: BTreeMap<String, String>,
|
||
#[serde(default = "yes")]
|
||
pub enabled: bool,
|
||
/// Notes for the bot on how to use this server (the original's MCP custom instructions).
|
||
#[serde(default)]
|
||
pub instructions: String,
|
||
}
|
||
|
||
fn yes() -> bool {
|
||
true
|
||
}
|
||
|
||
#[derive(Clone, Debug)]
|
||
pub struct ToolInfo {
|
||
pub server: String,
|
||
pub name: String,
|
||
pub description: String,
|
||
pub schema: Value,
|
||
}
|
||
|
||
enum Link {
|
||
Stdio { child: Child, stdin: ChildStdin, rx: Receiver<Value> },
|
||
Http { url: String, headers: BTreeMap<String, String>, session: Option<String> },
|
||
}
|
||
|
||
/// The part a tool call mutates; locked per server, so one slow call (up to two minutes)
|
||
/// never blocks other servers, the UI's status line, or the app's own sync.
|
||
struct Live {
|
||
link: Link,
|
||
next_id: u64,
|
||
}
|
||
|
||
struct Conn {
|
||
live: Mutex<Live>,
|
||
tools: Vec<ToolInfo>,
|
||
config: Server,
|
||
}
|
||
|
||
/// One lock for the table, one lock per server's link.
|
||
fn conns() -> &'static Mutex<HashMap<String, Arc<Conn>>> {
|
||
static C: OnceLock<Mutex<HashMap<String, Arc<Conn>>>> = OnceLock::new();
|
||
C.get_or_init(|| Mutex::new(HashMap::new()))
|
||
}
|
||
|
||
pub fn safe(s: &str) -> String {
|
||
s.chars().map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' }).collect()
|
||
}
|
||
|
||
/// A server's part of tool names: its name when it's plain ASCII, else a short stable code.
|
||
pub fn server_key(server: &str) -> String {
|
||
if !server.is_empty() && server.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == ' ') {
|
||
return safe(server);
|
||
}
|
||
let h = server.bytes().fold(0xcbf2_9ce4_8422_2325u64, |h, b| (h ^ b as u64).wrapping_mul(0x100_0000_01b3));
|
||
format!("s{:06x}", h & 0xff_ffff)
|
||
}
|
||
|
||
pub fn tool_name(server: &str, tool: &str) -> String {
|
||
let n = format!("{PREFIX}{}__{}", server_key(server), safe(tool));
|
||
n.chars().take(64).collect()
|
||
}
|
||
|
||
pub fn is_mcp_tool(name: &str) -> bool {
|
||
name.starts_with(PREFIX)
|
||
}
|
||
|
||
fn send_stdio(stdin: &mut ChildStdin, msg: &Value) -> Result<(), String> {
|
||
let mut line = msg.to_string();
|
||
line.push('\n');
|
||
stdin.write_all(line.as_bytes()).and_then(|_| stdin.flush()).map_err(|e| format!("送不出去：{e}"))
|
||
}
|
||
|
||
/// Read the answer with this id (skipping notifications and requests from the server).
|
||
fn wait_id(rx: &Receiver<Value>, id: u64, secs: u64) -> Result<Value, String> {
|
||
let deadline = std::time::Instant::now() + Duration::from_secs(secs);
|
||
loop {
|
||
let left = deadline.saturating_duration_since(std::time::Instant::now());
|
||
if left.is_zero() {
|
||
return Err("MCP 伺服器沒有回應".into());
|
||
}
|
||
match rx.recv_timeout(left) {
|
||
Ok(v) if v["id"].as_u64() == Some(id) && v.get("method").is_none() => {
|
||
if let Some(e) = v.get("error") {
|
||
return Err(format!("MCP 錯誤：{}", e["message"].as_str().unwrap_or(&e.to_string())));
|
||
}
|
||
return Ok(v["result"].clone());
|
||
}
|
||
Ok(_) => continue,
|
||
Err(_) => return Err("MCP 伺服器沒有回應或已經結束".into()),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Answers from Streamable HTTP: plain JSON, or SSE with `data:` lines.
|
||
pub fn parse_http_answer(body: &str, id: u64) -> Result<Value, String> {
|
||
let candidates: Vec<Value> = if body.trim_start().starts_with('{') || body.trim_start().starts_with('[') {
|
||
match serde_json::from_str::<Value>(body) {
|
||
Ok(Value::Array(a)) => a,
|
||
Ok(v) => vec![v],
|
||
Err(e) => return Err(e.to_string()),
|
||
}
|
||
} else {
|
||
body.lines().filter_map(|l| l.strip_prefix("data:")).filter_map(|d| serde_json::from_str(d.trim()).ok()).collect()
|
||
};
|
||
for v in candidates {
|
||
if v["id"].as_u64() == Some(id) {
|
||
if let Some(e) = v.get("error") {
|
||
return Err(format!("MCP 錯誤：{}", e["message"].as_str().unwrap_or(&e.to_string())));
|
||
}
|
||
return Ok(v["result"].clone());
|
||
}
|
||
}
|
||
Err("MCP 伺服器沒有給答案".into())
|
||
}
|
||
|
||
impl Live {
|
||
fn request(&mut self, method: &str, params: Value) -> Result<Value, String> {
|
||
let id = self.next_id;
|
||
self.next_id += 1;
|
||
let msg = json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params});
|
||
match &mut self.link {
|
||
Link::Stdio { stdin, rx, .. } => {
|
||
send_stdio(stdin, &msg)?;
|
||
wait_id(rx, id, CALL_TIMEOUT_SECS)
|
||
}
|
||
Link::Http { url, headers, session } => {
|
||
let agent: ureq::Agent = ureq::Agent::config_builder().timeout_global(Some(Duration::from_secs(CALL_TIMEOUT_SECS))).http_status_as_error(false).build().into();
|
||
let mut r = agent.post(url.as_str()).header("Content-Type", "application/json").header("Accept", "application/json, text/event-stream").header("MCP-Protocol-Version", PROTOCOL);
|
||
for (k, v) in headers.iter() {
|
||
r = r.header(k.as_str(), v.as_str());
|
||
}
|
||
if let Some(s) = session.as_ref() {
|
||
r = r.header("Mcp-Session-Id", s.as_str());
|
||
}
|
||
let mut resp = r.send(msg.to_string().as_bytes()).map_err(|e| format!("連不上 MCP 伺服器：{e}"))?;
|
||
if let Some(s) = resp.headers().get("mcp-session-id").and_then(|v| v.to_str().ok()) {
|
||
*session = Some(s.to_string());
|
||
}
|
||
let st = resp.status().as_u16();
|
||
let body = resp.body_mut().read_to_string().unwrap_or_default();
|
||
if !(200..300).contains(&st) {
|
||
return Err(format!("MCP 伺服器回 HTTP {st}：{}", body.chars().take(200).collect::<String>()));
|
||
}
|
||
parse_http_answer(&body, id)
|
||
}
|
||
}
|
||
}
|
||
|
||
fn notify(&mut self, method: &str) {
|
||
let msg = json!({"jsonrpc": "2.0", "method": method});
|
||
match &mut self.link {
|
||
Link::Stdio { stdin, .. } => {
|
||
let _ = send_stdio(stdin, &msg);
|
||
}
|
||
Link::Http { url, headers, session } => {
|
||
let mut r = ureq::post(url.as_str()).header("Content-Type", "application/json").header("Accept", "application/json, text/event-stream");
|
||
for (k, v) in headers.iter() {
|
||
r = r.header(k.as_str(), v.as_str());
|
||
}
|
||
if let Some(s) = session.as_ref() {
|
||
r = r.header("Mcp-Session-Id", s.as_str());
|
||
}
|
||
let _ = r.send(msg.to_string().as_bytes());
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
fn connect(s: &Server) -> Result<Conn, String> {
|
||
let link = if !s.url.trim().is_empty() {
|
||
Link::Http { url: s.url.trim().to_string(), headers: s.headers.clone(), session: None }
|
||
} else {
|
||
if s.command.trim().is_empty() {
|
||
return Err("請給指令或網址".into());
|
||
}
|
||
let mut child = Command::new(s.command.trim())
|
||
.args(&s.args)
|
||
.envs(&s.env)
|
||
.env("PATH", crate::cloudbox::tool_path())
|
||
.stdin(Stdio::piped())
|
||
.stdout(Stdio::piped())
|
||
.stderr(Stdio::null())
|
||
.spawn()
|
||
.map_err(|e| format!("無法啟動「{}」：{e}", s.command))?;
|
||
let stdin = child.stdin.take().ok_or("沒有標準輸入")?;
|
||
let stdout = child.stdout.take().ok_or("沒有標準輸出")?;
|
||
let (tx, rx) = channel();
|
||
std::thread::spawn(move || {
|
||
for line in BufReader::new(stdout).lines().map_while(Result::ok) {
|
||
if let Ok(v) = serde_json::from_str::<Value>(&line) {
|
||
if tx.send(v).is_err() {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
});
|
||
Link::Stdio { child, stdin, rx }
|
||
};
|
||
let mut live = Live { link, next_id: 1 };
|
||
live.request("initialize", json!({"protocolVersion": PROTOCOL, "capabilities": {}, "clientInfo": {"name": "AnyBot", "version": "1"}}))?;
|
||
live.notify("notifications/initialized");
|
||
let mut tools = vec![];
|
||
let mut cursor: Option<String> = None;
|
||
for _ in 0..20 {
|
||
let r = live.request("tools/list", match &cursor { Some(x) => json!({"cursor": x}), None => json!({}) })?;
|
||
for t in r["tools"].as_array().into_iter().flatten() {
|
||
tools.push(ToolInfo { server: s.name.clone(), name: t["name"].as_str().unwrap_or("").into(), description: t["description"].as_str().unwrap_or("").into(), schema: t["inputSchema"].clone() });
|
||
}
|
||
cursor = r["nextCursor"].as_str().map(String::from);
|
||
if cursor.is_none() {
|
||
break;
|
||
}
|
||
}
|
||
Ok(Conn { live: Mutex::new(live), tools, config: s.clone() })
|
||
}
|
||
|
||
/// Start (or keep) the enabled servers; stop the rest. Returns errors per server.
|
||
/// Connecting happens outside the table's lock (a slow server must not freeze the UI's
|
||
/// status line, which reads that lock every second).
|
||
pub fn sync(servers: &[Server]) -> Vec<(String, String)> {
|
||
let mut errs = vec![];
|
||
let mut to_connect: Vec<Server> = vec![];
|
||
{
|
||
let mut map = conns().lock().unwrap();
|
||
let wanted: Vec<&Server> = servers.iter().filter(|s| s.enabled).collect();
|
||
map.retain(|name, c| {
|
||
let keep = wanted.iter().any(|s| &s.name == name && **s == c.config);
|
||
if !keep {
|
||
if let Link::Stdio { child, .. } = &mut c.live.lock().unwrap().link {
|
||
let _ = child.kill();
|
||
}
|
||
}
|
||
keep
|
||
});
|
||
for s in wanted {
|
||
if !map.contains_key(&s.name) {
|
||
to_connect.push(s.clone());
|
||
}
|
||
}
|
||
}
|
||
for s in to_connect {
|
||
match connect(&s) {
|
||
Ok(c) => {
|
||
conns().lock().unwrap().insert(s.name.clone(), Arc::new(c));
|
||
}
|
||
Err(e) => errs.push((s.name.clone(), e)),
|
||
}
|
||
}
|
||
errs
|
||
}
|
||
|
||
/// The connected servers' tools.
|
||
pub fn tools() -> Vec<ToolInfo> {
|
||
conns().lock().unwrap().values().flat_map(|c| c.tools.clone()).collect()
|
||
}
|
||
|
||
pub fn status() -> HashMap<String, usize> {
|
||
conns().lock().unwrap().iter().map(|(k, c)| (k.clone(), c.tools.len())).collect()
|
||
}
|
||
|
||
/// Tool specs for the model.
|
||
pub fn specs() -> Vec<Value> {
|
||
tools()
|
||
.into_iter()
|
||
.map(|t| {
|
||
let schema = if t.schema.is_object() { t.schema.clone() } else { json!({"type": "object", "properties": {}}) };
|
||
json!({"name": tool_name(&t.server, &t.name), "description": format!("[{}] {}", t.server, t.description.chars().take(900).collect::<String>()), "parameters": schema})
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Instructions the servers' owners wrote for the bot.
|
||
pub fn instructions(servers: &[Server]) -> String {
|
||
servers.iter().filter(|s| s.enabled && !s.instructions.trim().is_empty()).map(|s| format!("- {}：{}", s.name, s.instructions.trim())).collect::<Vec<_>>().join("\n")
|
||
}
|
||
|
||
/// Which server and tool a model-side name means.
|
||
pub fn resolve(name: &str) -> Option<(String, String)> {
|
||
tools().into_iter().find(|t| tool_name(&t.server, &t.name) == name).map(|t| (t.server, t.name))
|
||
}
|
||
|
||
/// Call a tool. Returns (text, ok, first picture as base64).
|
||
pub fn call(name: &str, args: &str) -> (String, bool, Option<String>) {
|
||
let Some((server, tool)) = resolve(name) else { return (format!("沒有「{name}」這個 MCP 工具（伺服器可能已關閉）"), false, None) };
|
||
let arguments: Value = serde_json::from_str(if args.trim().is_empty() { "{}" } else { args }).unwrap_or(json!({}));
|
||
// Copy the server's own handle out of the table, then lock its link alone: a two-minute
|
||
// call to one server used to block every other MCP tool and the UI's status line.
|
||
let conn = conns().lock().unwrap().get(&server).cloned();
|
||
let Some(conn) = conn else { return ("MCP 伺服器沒有連線".into(), false, None) };
|
||
let mut live = conn.live.lock().unwrap();
|
||
match live.request("tools/call", json!({"name": tool, "arguments": arguments})) {
|
||
Ok(r) => {
|
||
let mut text = String::new();
|
||
let mut image = None;
|
||
for part in r["content"].as_array().into_iter().flatten() {
|
||
match part["type"].as_str() {
|
||
Some("text") => {
|
||
text.push_str(part["text"].as_str().unwrap_or(""));
|
||
text.push('\n');
|
||
}
|
||
Some("image") if image.is_none() => image = part["data"].as_str().map(String::from),
|
||
Some("resource") => {
|
||
text.push_str(part["resource"]["text"].as_str().unwrap_or(""));
|
||
text.push('\n');
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
if text.trim().is_empty() && r.get("structuredContent").is_some() {
|
||
text = r["structuredContent"].to_string();
|
||
}
|
||
let ok = r["isError"].as_bool() != Some(true);
|
||
(crate::tools::head_tail(text.trim()), ok, image)
|
||
}
|
||
Err(e) => (e, false, None),
|
||
}
|
||
}
|
||
|
||
pub fn shutdown() {
|
||
for c in conns().lock().unwrap().values() {
|
||
if let Link::Stdio { child, .. } = &mut c.live.lock().unwrap().link {
|
||
let _ = child.kill();
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 連接器市集: (name, what it does, command or address, key to fill in — empty when none).
|
||
pub const MARKET: [(&str, &str, &str, &str); 8] = [
|
||
("files", "讀寫你指定資料夾裡的檔案", "npx -y @modelcontextprotocol/server-filesystem ~/Documents", ""),
|
||
("memory", "長期記憶：把人、事、關係記成知識圖", "npx -y @modelcontextprotocol/server-memory", ""),
|
||
("thinking", "一步步拆解難題", "npx -y @modelcontextprotocol/server-sequential-thinking", ""),
|
||
("fetch", "抓網頁內容轉成文字", "uvx mcp-server-fetch", ""),
|
||
("git", "查看和操作 Git 版本庫", "uvx mcp-server-git", ""),
|
||
("time", "時間和時區換算", "uvx mcp-server-time", ""),
|
||
("playwright", "用瀏覽器自動化操作網頁", "npx -y @playwright/mcp@latest", ""),
|
||
("github", "GitHub 的 Issue、PR、程式碼（官方遠端伺服器）", "https://api.githubcopilot.com/mcp/", "Authorization: Bearer 你的 GitHub Token"),
|
||
];
|
||
|
||
/// A marketplace entry as a server (home folder spelled out).
|
||
pub fn market_server(name: &str, cmd: &str) -> Server {
|
||
let parts: Vec<String> = cmd.split_whitespace().map(|w| if w == "~/Documents" { crate::tools::expand_home(w).display().to_string() } else { w.to_string() }).collect();
|
||
Server { name: name.to_string(), command: parts[0].clone(), args: parts[1..].to_vec(), ..serde_json::from_str("{\"name\":\"x\"}").unwrap() }
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
const FAKE_SERVER: &str = r#"
|
||
import sys, json
|
||
for line in sys.stdin:
|
||
m = json.loads(line)
|
||
if "id" not in m:
|
||
continue
|
||
if m["method"] == "initialize":
|
||
r = {"protocolVersion": "2025-06-18", "capabilities": {"tools": {}}, "serverInfo": {"name": "fake"}}
|
||
elif m["method"] == "tools/list":
|
||
r = {"tools": [{"name": "add", "description": "Add two numbers", "inputSchema": {"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}}]}
|
||
elif m["method"] == "tools/call":
|
||
a = m["params"]["arguments"]
|
||
r = {"content": [{"type": "text", "text": "sum=%d" % (a["a"] + a["b"])}]}
|
||
else:
|
||
r = {}
|
||
sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": m["id"], "result": r}) + "\n")
|
||
sys.stdout.flush()
|
||
"#;
|
||
|
||
#[test]
|
||
fn stdio_server_round_trip() {
|
||
let dir = std::env::temp_dir().join(format!("anybot-mcp-{}", std::process::id()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let script = dir.join("server.py");
|
||
std::fs::write(&script, FAKE_SERVER).unwrap();
|
||
let s = Server { name: "計算機".into(), command: "python3".into(), args: vec!["-I".into(), script.display().to_string()], env: BTreeMap::new(), url: String::new(), headers: BTreeMap::new(), enabled: true, instructions: "用來算數學".into() };
|
||
let errs = sync(std::slice::from_ref(&s));
|
||
assert!(errs.is_empty(), "{errs:?}");
|
||
let name = tool_name("計算機", "add");
|
||
assert!(name.starts_with("mcp__s") && name.ends_with("__add"), "{name}");
|
||
assert_eq!(tool_name("github", "list issues"), "mcp__github__list_issues");
|
||
assert!(specs().iter().any(|t| t["name"] == name.as_str()));
|
||
let (out, ok, _) = call(&name, r#"{"a": 2, "b": 40}"#);
|
||
assert!(ok && out.contains("sum=42"), "{out}");
|
||
assert!(instructions(std::slice::from_ref(&s)).contains("用來算數學"));
|
||
sync(&[]);
|
||
assert!(tools().is_empty());
|
||
let _ = std::fs::remove_dir_all(&dir);
|
||
}
|
||
|
||
#[test]
|
||
fn http_answers_json_or_sse() {
|
||
assert_eq!(parse_http_answer(r#"{"jsonrpc":"2.0","id":3,"result":{"ok":1}}"#, 3).unwrap()["ok"], 1);
|
||
let sse = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":4,\"result\":{\"tools\":[]}}\n\n";
|
||
assert!(parse_http_answer(sse, 4).unwrap()["tools"].is_array());
|
||
assert!(parse_http_answer(r#"{"jsonrpc":"2.0","id":5,"error":{"message":"nope"}}"#, 5).unwrap_err().contains("nope"));
|
||
}
|
||
|
||
/// Downloads and starts every marketplace connector that needs no key (needs the network,
|
||
/// Node.js and uv): cargo test --release -- --ignored marketplace --nocapture
|
||
#[test]
|
||
#[ignore]
|
||
fn marketplace_connectors_start() {
|
||
let servers: Vec<Server> = MARKET.iter().filter(|m| m.3.is_empty()).map(|m| market_server(m.0, m.2)).collect();
|
||
let errors = sync(&servers);
|
||
let st = status();
|
||
for s in &servers {
|
||
println!("{:<12} {:?}", s.name, st.get(&s.name));
|
||
}
|
||
assert!(errors.is_empty(), "{errors:?}");
|
||
for s in &servers {
|
||
assert!(st.get(&s.name).copied().unwrap_or(0) > 0, "{} has no tools", s.name);
|
||
}
|
||
// And one real call through the bot's tool name.
|
||
let (out, ok, _) = call(&tool_name("time", "get_current_time"), r#"{"timezone":"Asia/Taipei"}"#);
|
||
println!("time: {out}");
|
||
assert!(ok && out.contains("Asia/Taipei"), "{out}");
|
||
shutdown();
|
||
}
|
||
}
|