LazyBoy2/crates/grokboy-core/src/mcp.rs

632 lines
21 KiB
Rust

//! MCP client shaped after Grok Bot: GetMcpTools then CallMcpTool.
//!
//! Config: `~/.grokboy/mcp.json` (override with `GROKBOY_MCP_CONFIG`).
//! Workspace `.grokboy/mcp.json` is merged on top.
use anyhow::{anyhow, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::Mutex;
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct McpFile {
#[serde(default, rename = "mcpServers")]
pub mcp_servers: HashMap<String, McpServerConfig>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct McpServerConfig {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub command: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub args: Vec<String>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub env: HashMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub url: Option<String>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub headers: HashMap<String, String>,
}
impl McpServerConfig {
fn transport(&self) -> &'static str {
if self.url.as_ref().is_some_and(|u| !u.is_empty()) {
"http"
} else {
"stdio"
}
}
}
pub fn mcp_config_path() -> Result<PathBuf> {
if let Some(path) = std::env::var_os("GROKBOY_MCP_CONFIG").filter(|s| !s.is_empty()) {
return Ok(PathBuf::from(path));
}
let home = std::env::var_os("HOME")
.or_else(|| std::env::var_os("USERPROFILE"))
.map(PathBuf::from)
.ok_or_else(|| anyhow!("cannot resolve home directory"))?;
Ok(home.join(".grokboy").join("mcp.json"))
}
pub fn load_mcp_file(path: &Path) -> Result<McpFile> {
if !path.exists() {
return Ok(McpFile::default());
}
let text = std::fs::read_to_string(path)
.with_context(|| format!("read {}", path.display()))?;
serde_json::from_str(&text).with_context(|| format!("parse {}", path.display()))
}
pub fn save_mcp_file(path: &Path, file: &McpFile) -> Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, serde_json::to_vec_pretty(file)?)
.with_context(|| format!("write {}", path.display()))
}
fn merge_workspace(cwd: &Path, mut file: McpFile) -> McpFile {
let extra = cwd.join(".grokboy").join("mcp.json");
if let Ok(local) = load_mcp_file(&extra) {
file.mcp_servers.extend(local.mcp_servers);
}
file
}
struct StdioSession {
#[allow(dead_code)]
child: Child,
stdin: ChildStdin,
stdout: BufReader<ChildStdout>,
next_id: u64,
}
enum Live {
Stdio(StdioSession),
Http {
url: String,
headers: HashMap<String, String>,
next_id: u64,
},
}
pub struct McpHub {
path: PathBuf,
file: Mutex<McpFile>,
live: Mutex<HashMap<String, Live>>,
}
impl McpHub {
pub fn empty() -> Arc<Self> {
Arc::new(Self {
path: std::env::temp_dir().join("grokboy-mcp-empty.json"),
file: Mutex::new(McpFile::default()),
live: Mutex::new(HashMap::new()),
})
}
pub fn load(cwd: &Path) -> Result<Arc<Self>> {
let path = mcp_config_path()?;
let file = merge_workspace(cwd, load_mcp_file(&path)?);
Ok(Arc::new(Self {
path,
file: Mutex::new(file),
live: Mutex::new(HashMap::new()),
}))
}
pub async fn status(&self, server: Option<&str>) -> Result<Value> {
let file = self.file.lock().await;
let mut rows = Vec::new();
for (name, cfg) in file.mcp_servers.iter() {
if server.is_some_and(|s| s != name) {
continue;
}
rows.push(json!({
"server": name,
"status": "configured",
"transport": cfg.transport(),
"command": cfg.command,
"url": cfg.url,
}));
}
rows.sort_by(|a, b| a["server"].as_str().cmp(&b["server"].as_str()));
Ok(json!({"servers": rows, "config": self.path}))
}
pub async fn add(
&self,
name: &str,
url: Option<&str>,
headers: Option<HashMap<String, String>>,
command: Option<&str>,
args: Option<Vec<String>>,
env: Option<HashMap<String, String>>,
) -> Result<Value> {
if name.trim().is_empty() {
return Err(anyhow!("name is required"));
}
let mut cfg = McpServerConfig::default();
if let Some(url) = url.filter(|s| !s.is_empty()) {
if let Some(err) = validate_remote_url(url) {
return Err(anyhow!(err));
}
cfg.url = Some(url.to_string());
cfg.headers = headers.unwrap_or_default();
} else if let Some(command) = command.filter(|s| !s.is_empty()) {
cfg.command = Some(command.to_string());
cfg.args = args.unwrap_or_default();
cfg.env = env.unwrap_or_default();
} else {
return Err(anyhow!(
"provide url (https MCP endpoint) or command (stdio server)"
));
}
let mut file = self.file.lock().await;
file.mcp_servers.insert(name.to_string(), cfg);
save_mcp_file(&self.path, &file)?;
drop(file);
self.live.lock().await.remove(name);
self.status(Some(name)).await
}
pub async fn remove(&self, name: &str) -> Result<Value> {
let mut file = self.file.lock().await;
let removed = file.mcp_servers.remove(name).is_some();
save_mcp_file(&self.path, &file)?;
drop(file);
self.live.lock().await.remove(name);
Ok(json!({"removed": removed, "server": name}))
}
pub async fn get_tools(
&self,
server: Option<&str>,
tool_name: Option<&str>,
pattern: Option<&str>,
) -> Result<Value> {
let names: Vec<String> = {
let file = self.file.lock().await;
let mut names: Vec<String> = file.mcp_servers.keys().cloned().collect();
names.sort();
if let Some(server) = server {
if !file.mcp_servers.contains_key(server) {
return Err(anyhow!(
"MCP server `{server}` is not installed. GetMcpServerStatus lists identifiers; AddMcpServer installs one."
));
}
names.retain(|n| n == server);
}
names
};
if names.is_empty() {
return Ok(json!({
"mode": "catalog",
"servers": [],
"instruction": "No MCP servers are installed. Use AddMcpServer with a stdio command or an https URL."
}));
}
let re = pattern
.filter(|s| !s.is_empty())
.map(|p| regex::Regex::new(&format!("(?i){p}")))
.transpose()
.map_err(|e| anyhow!("invalid pattern: {e}"))?;
let mut servers = Vec::new();
for name in names {
match self.list_server_tools(&name).await {
Ok(mut tools) => {
if let Some(want) = tool_name.filter(|s| !s.is_empty()) {
tools.retain(|t| t["name"].as_str() == Some(want));
}
if let Some(re) = &re {
tools.retain(|t| {
let blob = format!(
"{} {}",
t["name"].as_str().unwrap_or(""),
t["description"].as_str().unwrap_or("")
);
re.is_match(&blob)
});
}
servers.push(json!({
"server": name,
"serverStatus": "connected",
"tools": tools,
}));
}
Err(err) => servers.push(json!({
"server": name,
"serverStatus": "error",
"serverError": err.to_string(),
"tools": [],
})),
}
}
Ok(json!({
"mode": if tool_name.is_some() { "single_tool" } else { "catalog" },
"servers": servers,
"instruction": "Call a listed tool with CallMcpTool using the server identifier and tool name. Refetch this descriptor if a later call fails."
}))
}
pub async fn call_tool(&self, server: &str, tool: &str, arguments: Value) -> Result<Value> {
if server.trim().is_empty() || tool.trim().is_empty() {
return Err(anyhow!("CallMcpTool requires server and tool_name"));
}
let exists = self.file.lock().await.mcp_servers.contains_key(server);
if !exists {
return Err(anyhow!(
"MCP server `{server}` does not exist. GetMcpTools / GetMcpServerStatus list identifiers."
));
}
let raw = self.rpc(server, "tools/call", json!({"name": tool, "arguments": arguments})).await?;
Ok(json!({
"server": server,
"tool": tool,
"result": raw,
}))
}
async fn list_server_tools(&self, server: &str) -> Result<Vec<Value>> {
let raw = self.rpc(server, "tools/list", json!({})).await?;
let tools = raw["tools"].as_array().cloned().unwrap_or_default();
Ok(tools
.into_iter()
.map(|t| {
json!({
"name": t["name"],
"description": t["description"].as_str().unwrap_or("").chars().take(200).collect::<String>(),
"inputSchema": t.get("inputSchema").cloned().unwrap_or(json!({"type":"object"})),
})
})
.collect())
}
async fn rpc(&self, server: &str, method: &str, params: Value) -> Result<Value> {
self.ensure(server).await?;
let mut live = self.live.lock().await;
let session = live
.get_mut(server)
.ok_or_else(|| anyhow!("MCP server `{server}` failed to start"))?;
match session {
Live::Stdio(stdio) => stdio.request(method, params).await,
Live::Http {
url,
headers,
next_id,
} => http_request(url, headers, next_id, method, params).await,
}
}
async fn ensure(&self, server: &str) -> Result<()> {
if self.live.lock().await.contains_key(server) {
return Ok(());
}
let cfg = self
.file
.lock()
.await
.mcp_servers
.get(server)
.cloned()
.ok_or_else(|| anyhow!("unknown MCP server `{server}`"))?;
let live = connect(&cfg).await?;
self.live.lock().await.insert(server.to_string(), live);
Ok(())
}
}
fn validate_remote_url(raw: &str) -> Option<String> {
let parsed = match reqwest::Url::parse(raw) {
Ok(url) => url,
Err(_) => return Some(format!("\"{raw}\" is not a valid URL.")),
};
if parsed.scheme() != "http" && parsed.scheme() != "https" {
return Some(format!(
"The server URL must be http(s); \"{}\" is not supported.",
parsed.scheme()
));
}
if !parsed.username().is_empty() || parsed.password().is_some() {
return Some(
"Don't put credentials in the server URL — pass them as headers instead.".into(),
);
}
None
}
async fn connect(cfg: &McpServerConfig) -> Result<Live> {
if let Some(url) = cfg.url.as_ref().filter(|s| !s.is_empty()) {
let mut live = Live::Http {
url: url.clone(),
headers: cfg.headers.clone(),
next_id: 1,
};
initialize(&mut live).await?;
return Ok(live);
}
let command = cfg
.command
.as_deref()
.ok_or_else(|| anyhow!("stdio MCP server needs command"))?;
let mut cmd = Command::new(command);
cmd.args(&cfg.args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
for (k, v) in &cfg.env {
cmd.env(k, v);
}
let mut child = cmd.spawn().with_context(|| format!("spawn MCP `{command}`"))?;
let stdin = child.stdin.take().ok_or_else(|| anyhow!("mcp stdin"))?;
let stdout = BufReader::new(child.stdout.take().ok_or_else(|| anyhow!("mcp stdout"))?);
let mut live = Live::Stdio(StdioSession {
child,
stdin,
stdout,
next_id: 1,
});
initialize(&mut live).await?;
Ok(live)
}
async fn initialize(live: &mut Live) -> Result<()> {
let result = match live {
Live::Stdio(stdio) => {
stdio
.request(
"initialize",
json!({
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "grokboy", "version": "0.1.0"}
}),
)
.await?
}
Live::Http {
url,
headers,
next_id,
} => {
http_request(
url,
headers,
next_id,
"initialize",
json!({
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "grokboy", "version": "0.1.0"}
}),
)
.await?
}
};
let _ = result;
match live {
Live::Stdio(stdio) => {
write_frame(
&mut stdio.stdin,
&json!({"jsonrpc":"2.0","method":"notifications/initialized"}),
)
.await
}
Live::Http { .. } => Ok(()),
}
}
impl StdioSession {
async fn request(&mut self, method: &str, params: Value) -> Result<Value> {
let id = self.next_id;
self.next_id += 1;
write_frame(
&mut self.stdin,
&json!({"jsonrpc":"2.0","id":id,"method":method,"params":params}),
)
.await?;
loop {
let msg = read_frame(&mut self.stdout).await?;
if msg.get("id") == Some(&json!(id)) {
if let Some(err) = msg.get("error") {
return Err(anyhow!("MCP {method} error: {err}"));
}
return Ok(msg.get("result").cloned().unwrap_or(json!({})));
}
}
}
}
async fn write_frame(stdin: &mut ChildStdin, value: &Value) -> Result<()> {
let body = serde_json::to_vec(value)?;
let header = format!("Content-Length: {}\r\n\r\n", body.len());
stdin.write_all(header.as_bytes()).await?;
stdin.write_all(&body).await?;
stdin.flush().await?;
Ok(())
}
async fn read_frame(stdout: &mut BufReader<ChildStdout>) -> Result<Value> {
let mut content_length = None;
loop {
let mut line = String::new();
let n = stdout.read_line(&mut line).await?;
if n == 0 {
return Err(anyhow!("MCP server closed stdout"));
}
if line == "\r\n" || line == "\n" {
break;
}
if let Some(rest) = line
.to_ascii_lowercase()
.strip_prefix("content-length:")
.or_else(|| line.strip_prefix("Content-Length:"))
{
content_length = rest.trim().parse().ok();
}
}
let len = content_length.ok_or_else(|| anyhow!("MCP response missing Content-Length"))?;
let mut buf = vec![0u8; len];
stdout.read_exact(&mut buf).await?;
serde_json::from_slice(&buf).context("parse MCP JSON-RPC")
}
async fn http_request(
url: &str,
headers: &HashMap<String, String>,
next_id: &mut u64,
method: &str,
params: Value,
) -> Result<Value> {
let id = *next_id;
*next_id += 1;
let mut req = reqwest::Client::new()
.post(url)
.header("content-type", "application/json")
.json(&json!({"jsonrpc":"2.0","id":id,"method":method,"params":params}));
for (k, v) in headers {
req = req.header(k, v);
}
let parsed: Value = req.send().await?.json().await?;
if let Some(err) = parsed.get("error") {
return Err(anyhow!("MCP {method} error: {err}"));
}
Ok(parsed.get("result").cloned().unwrap_or(json!({})))
}
pub fn parse_headers(value: &Value) -> Option<HashMap<String, String>> {
let obj = value.as_object()?;
let mut out = HashMap::new();
for (k, v) in obj {
if let Some(s) = v.as_str() {
out.insert(k.clone(), s.to_string());
}
}
Some(out)
}
pub fn parse_env(value: &Value) -> Option<HashMap<String, String>> {
parse_headers(value)
}
pub fn parse_args_list(value: &Value) -> Option<Vec<String>> {
value
.as_array()
.map(|a| {
a.iter()
.filter_map(|v| v.as_str().map(str::to_string))
.collect()
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_credentials_in_url() {
assert!(validate_remote_url("https://user:pw@example.com/mcp").is_some());
assert!(validate_remote_url("https://example.com/mcp").is_none());
}
#[test]
fn round_trips_config() {
let dir = std::env::temp_dir().join(format!("gb-mcp-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("mcp.json");
let mut file = McpFile::default();
file.mcp_servers.insert(
"echo".into(),
McpServerConfig {
command: Some("python3".into()),
args: vec!["-c".into(), "pass".into()],
..Default::default()
},
);
save_mcp_file(&path, &file).unwrap();
let loaded = load_mcp_file(&path).unwrap();
assert_eq!(loaded.mcp_servers["echo"].command.as_deref(), Some("python3"));
let _ = std::fs::remove_dir_all(dir);
}
const ECHO_PY: &str = r#"
import json, sys
def read_msg():
headers = {}
while True:
line = sys.stdin.buffer.readline()
if not line or line in (b"\r\n", b"\n"):
break
k, _, v = line.decode().partition(":")
headers[k.strip().lower()] = v.strip()
n = int(headers.get("content-length", "0"))
return json.loads(sys.stdin.buffer.read(n))
def write_msg(obj):
data = json.dumps(obj).encode()
sys.stdout.buffer.write(f"Content-Length: {len(data)}\r\n\r\n".encode() + data)
sys.stdout.buffer.flush()
while True:
msg = read_msg()
method = msg.get("method")
mid = msg.get("id")
if method == "initialize":
write_msg({"jsonrpc":"2.0","id":mid,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"echo","version":"0"}}})
elif method == "notifications/initialized":
continue
elif method == "tools/list":
write_msg({"jsonrpc":"2.0","id":mid,"result":{"tools":[{"name":"echo","description":"echo text","inputSchema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}}]}})
elif method == "tools/call":
text = ((msg.get("params") or {}).get("arguments") or {}).get("text", "")
write_msg({"jsonrpc":"2.0","id":mid,"result":{"content":[{"type":"text","text": text}]}})
"#;
#[tokio::test]
async fn stdio_get_mcp_tools_and_call() {
let _lock = crate::test_env::lock_async().await;
if std::process::Command::new("python3")
.arg("-c")
.arg("pass")
.status()
.ok()
.is_none_or(|s| !s.success())
{
return;
}
let dir = std::env::temp_dir().join(format!("gb-mcp-live-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).unwrap();
let script = dir.join("echo.py");
std::fs::write(&script, ECHO_PY).unwrap();
let cfg = dir.join("mcp.json");
unsafe { std::env::set_var("GROKBOY_MCP_CONFIG", &cfg) };
let mut file = McpFile::default();
file.mcp_servers.insert(
"echo".into(),
McpServerConfig {
command: Some("python3".into()),
args: vec![script.to_string_lossy().into_owned()],
..Default::default()
},
);
save_mcp_file(&cfg, &file).unwrap();
let hub = McpHub::load(&dir).unwrap();
let listed = hub.get_tools(Some("echo"), None, None).await.unwrap();
assert_eq!(listed["servers"][0]["tools"][0]["name"], "echo");
let called = hub
.call_tool("echo", "echo", json!({"text": "hello-mcp"}))
.await
.unwrap();
let blob = called.to_string();
assert!(blob.contains("hello-mcp"), "{blob}");
unsafe { std::env::remove_var("GROKBOY_MCP_CONFIG") };
let _ = std::fs::remove_dir_all(dir);
}
}