Images of messages (pasted images, screenshots, generated images) were base64 inside the message JSON, so opening a chat downloaded all of them at once (66 MB for a chat with 120 screenshots). They now live in a `blobs` table, addressed by SHA-256; messages keep a reference (blob, size). The UI loads them from /api/blobs/<hash> lazily with an immutable cache header, the runner fills the bytes in before model requests, JSON exports stay self-contained. - Store::add_message moves inline images out; large opaque PNGs become JPEG (q85) when smaller. Older chats are converted once at start-up in the background, unused images are removed (also after deleting a chat) and the database is vacuumed when that frees 32 MB or more. - Responses are compressed with gzip/brotli (not images, SSE streams or binary downloads). Example (copy of production): chat load 66 MB -> 1.5 MB (370 KB brotli), database 90 MB -> 27 MB. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
706 lines
30 KiB
Rust
706 lines
30 KiB
Rust
//! Minimal MCP client: JSON-RPC over stdio or Streamable HTTP.
|
||
//!
|
||
//! Supports tools (list/call, list_changed), roots and elicitation – server
|
||
//! questions are forwarded to the user through the same UI as agent questions.
|
||
|
||
use crate::app::App;
|
||
use crate::config::{McpConfig, McpKind, expand_env};
|
||
use crate::tools::{Tool, ToolCtx, ToolOutput};
|
||
use crate::types::{ImageData, Question, QuestionOption};
|
||
use anyhow::{Context, Result, anyhow, bail};
|
||
use async_trait::async_trait;
|
||
use eventsource_stream::Eventsource;
|
||
use futures::StreamExt;
|
||
use serde::Serialize;
|
||
use serde_json::{Map, Value, json};
|
||
use std::collections::{BTreeMap, HashMap};
|
||
use std::process::Stdio;
|
||
use std::sync::atomic::{AtomicU64, Ordering};
|
||
use std::sync::{Arc, Mutex, Weak};
|
||
use std::time::Duration;
|
||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||
use tokio::sync::{Mutex as AsyncMutex, oneshot};
|
||
use tokio_util::sync::CancellationToken;
|
||
|
||
const PROTOCOL_VERSION: &str = "2025-06-18";
|
||
const CALL_TIMEOUT: Duration = Duration::from_secs(600);
|
||
|
||
#[derive(Debug, Clone, Serialize)]
|
||
pub struct McpToolInfo {
|
||
pub name: String,
|
||
pub description: String,
|
||
pub input_schema: Value,
|
||
/// The server marks the tool as read-only (annotations.readOnlyHint).
|
||
#[serde(default)]
|
||
pub read_only: bool,
|
||
}
|
||
|
||
#[derive(Debug, Clone, Serialize)]
|
||
#[serde(rename_all = "snake_case", tag = "state")]
|
||
pub enum McpState {
|
||
Disabled,
|
||
Connecting,
|
||
Connected { tools: usize },
|
||
Failed { error: String },
|
||
}
|
||
|
||
type Pending = Arc<Mutex<HashMap<u64, oneshot::Sender<Result<Value>>>>>;
|
||
|
||
enum Transport {
|
||
Stdio {
|
||
stdin: AsyncMutex<tokio::process::ChildStdin>,
|
||
pending: Pending,
|
||
_child: AsyncMutex<tokio::process::Child>,
|
||
},
|
||
Http {
|
||
client: reqwest::Client,
|
||
url: String,
|
||
headers: BTreeMap<String, String>,
|
||
session: Mutex<Option<String>>,
|
||
protocol: Mutex<Option<String>>,
|
||
},
|
||
}
|
||
|
||
pub struct McpServer {
|
||
pub name: String,
|
||
transport: Transport,
|
||
next_id: AtomicU64,
|
||
app: Weak<App>,
|
||
/// Session that triggered the current tool call (for elicitation routing).
|
||
current_session: Mutex<Option<String>>,
|
||
pub tools: Mutex<Vec<McpToolInfo>>,
|
||
}
|
||
|
||
fn rpc(id: u64, method: &str, params: Value) -> Value {
|
||
json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params})
|
||
}
|
||
|
||
impl McpServer {
|
||
pub async fn connect(name: &str, cfg: &McpConfig, app: Weak<App>) -> Result<Arc<Self>> {
|
||
let transport = match cfg.kind {
|
||
McpKind::Stdio => {
|
||
if cfg.command.is_empty() {
|
||
bail!("command is required for stdio servers");
|
||
}
|
||
let mut cmd = tokio::process::Command::new(expand_env(&cfg.command));
|
||
cmd.args(cfg.args.iter().map(|a| expand_env(a)))
|
||
.envs(cfg.env.iter().map(|(k, v)| (k.clone(), expand_env(v))))
|
||
.stdin(Stdio::piped())
|
||
.stdout(Stdio::piped())
|
||
.stderr(Stdio::piped())
|
||
.kill_on_drop(true);
|
||
if let Some(app) = app.upgrade() {
|
||
cmd.current_dir(app.workspace());
|
||
}
|
||
let mut child = cmd.spawn().with_context(|| format!("failed to start {}", cfg.command))?;
|
||
let stdin = child.stdin.take().unwrap();
|
||
let stdout = child.stdout.take().unwrap();
|
||
let stderr = child.stderr.take().unwrap();
|
||
let server_name = name.to_string();
|
||
tokio::spawn(async move {
|
||
let mut lines = BufReader::new(stderr).lines();
|
||
while let Ok(Some(l)) = lines.next_line().await {
|
||
tracing::debug!("mcp {server_name} stderr: {l}");
|
||
}
|
||
});
|
||
let pending: Pending = Arc::default();
|
||
let server = Arc::new(Self {
|
||
name: name.to_string(),
|
||
transport: Transport::Stdio {
|
||
stdin: AsyncMutex::new(stdin),
|
||
pending: pending.clone(),
|
||
_child: AsyncMutex::new(child),
|
||
},
|
||
next_id: AtomicU64::new(1),
|
||
app,
|
||
current_session: Mutex::new(None),
|
||
tools: Mutex::new(Vec::new()),
|
||
});
|
||
let weak = Arc::downgrade(&server);
|
||
tokio::spawn(async move {
|
||
let mut lines = BufReader::new(stdout).lines();
|
||
while let Ok(Some(line)) = lines.next_line().await {
|
||
let Ok(msg) = serde_json::from_str::<Value>(&line) else { continue };
|
||
let Some(server) = weak.upgrade() else { break };
|
||
// Server requests (elicitation) may wait for the user: never block the reader.
|
||
tokio::spawn(async move { server.handle_incoming(msg).await });
|
||
}
|
||
// Process ended: fail all pending requests.
|
||
for (_, tx) in pending.lock().unwrap().drain() {
|
||
let _ = tx.send(Err(anyhow!("MCP server exited")));
|
||
}
|
||
});
|
||
server
|
||
}
|
||
McpKind::Http => Arc::new(Self {
|
||
name: name.to_string(),
|
||
transport: Transport::Http {
|
||
client: crate::provider::http_client(),
|
||
url: expand_env(&cfg.url),
|
||
headers: cfg.headers.iter().map(|(k, v)| (k.clone(), expand_env(v))).collect(),
|
||
session: Mutex::new(None),
|
||
protocol: Mutex::new(None),
|
||
},
|
||
next_id: AtomicU64::new(1),
|
||
app,
|
||
current_session: Mutex::new(None),
|
||
tools: Mutex::new(Vec::new()),
|
||
}),
|
||
};
|
||
let server = transport;
|
||
server.initialize().await?;
|
||
server.refresh_tools().await?;
|
||
Ok(server)
|
||
}
|
||
|
||
async fn initialize(&self) -> Result<()> {
|
||
let res = self
|
||
.request(
|
||
"initialize",
|
||
json!({
|
||
"protocolVersion": PROTOCOL_VERSION,
|
||
"capabilities": {"elicitation": {}, "roots": {"listChanged": false}},
|
||
"clientInfo": {"name": "codecraft", "version": codecraft_core::version::VERSION}
|
||
}),
|
||
)
|
||
.await?;
|
||
if let Transport::Http { protocol, .. } = &self.transport {
|
||
*protocol.lock().unwrap() =
|
||
Some(res["protocolVersion"].as_str().unwrap_or(PROTOCOL_VERSION).to_string());
|
||
}
|
||
self.notify("notifications/initialized", json!({})).await?;
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn refresh_tools(&self) -> Result<()> {
|
||
let mut tools = Vec::new();
|
||
let mut cursor: Option<String> = None;
|
||
loop {
|
||
let params = match &cursor {
|
||
Some(c) => json!({"cursor": c}),
|
||
None => json!({}),
|
||
};
|
||
let res = self.request("tools/list", params).await?;
|
||
for t in res["tools"].as_array().cloned().unwrap_or_default() {
|
||
tools.push(McpToolInfo {
|
||
name: t["name"].as_str().unwrap_or_default().to_string(),
|
||
description: t["description"].as_str().unwrap_or_default().to_string(),
|
||
input_schema: t.get("inputSchema").cloned().unwrap_or(json!({"type": "object"})),
|
||
read_only: t["annotations"]["readOnlyHint"].as_bool().unwrap_or(false),
|
||
});
|
||
}
|
||
cursor = res["nextCursor"].as_str().map(String::from);
|
||
if cursor.is_none() {
|
||
break;
|
||
}
|
||
}
|
||
*self.tools.lock().unwrap() = tools;
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn call_tool(&self, session_id: &str, name: &str, args: Value, cancel: CancellationToken) -> Result<ToolOutput> {
|
||
*self.current_session.lock().unwrap() = Some(session_id.to_string());
|
||
let res = tokio::select! {
|
||
r = tokio::time::timeout(CALL_TIMEOUT, self.request("tools/call", json!({"name": name, "arguments": args}))) => {
|
||
r.map_err(|_| anyhow!("MCP tool call timed out"))?
|
||
}
|
||
_ = cancel.cancelled() => bail!("aborted"),
|
||
}?;
|
||
let mut out = ToolOutput { is_error: res["isError"].as_bool().unwrap_or(false), ..Default::default() };
|
||
let mut texts = Vec::new();
|
||
for c in res["content"].as_array().cloned().unwrap_or_default() {
|
||
match c["type"].as_str().unwrap_or_default() {
|
||
"text" => texts.push(c["text"].as_str().unwrap_or_default().to_string()),
|
||
"image" => out.images.push(ImageData {
|
||
blob: None,
|
||
size: None,
|
||
media_type: c["mimeType"].as_str().unwrap_or("image/png").into(),
|
||
data: c["data"].as_str().unwrap_or_default().into(),
|
||
}),
|
||
"resource" => {
|
||
let r = &c["resource"];
|
||
texts.push(match r["text"].as_str() {
|
||
Some(t) => format!("[resource {}]\n{t}", r["uri"].as_str().unwrap_or("")),
|
||
None => format!("[binary resource {}]", r["uri"].as_str().unwrap_or("")),
|
||
});
|
||
}
|
||
"resource_link" => texts.push(format!("[resource link {}]", c["uri"].as_str().unwrap_or(""))),
|
||
other => texts.push(format!("[{other} content]")),
|
||
}
|
||
}
|
||
if texts.is_empty() {
|
||
if let Some(s) = res.get("structuredContent") {
|
||
texts.push(serde_json::to_string_pretty(s)?);
|
||
}
|
||
}
|
||
out.content = crate::tools::truncate(&texts.join("\n"), 80_000);
|
||
Ok(out)
|
||
}
|
||
|
||
// --- transport -------------------------------------------------------
|
||
|
||
async fn request(&self, method: &str, params: Value) -> Result<Value> {
|
||
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
|
||
let msg = rpc(id, method, params);
|
||
match &self.transport {
|
||
Transport::Stdio { stdin, pending, .. } => {
|
||
let (tx, rx) = oneshot::channel();
|
||
pending.lock().unwrap().insert(id, tx);
|
||
self.write_stdio(stdin, &msg).await?;
|
||
let res = tokio::time::timeout(CALL_TIMEOUT, rx)
|
||
.await
|
||
.map_err(|_| anyhow!("MCP request {method} timed out"))?
|
||
.map_err(|_| anyhow!("MCP server closed the connection"))?;
|
||
res
|
||
}
|
||
Transport::Http { .. } => self.http_send(&msg, Some(id)).await?.ok_or_else(|| anyhow!("no response to {method}")),
|
||
}
|
||
}
|
||
|
||
async fn notify(&self, method: &str, params: Value) -> Result<()> {
|
||
let msg = json!({"jsonrpc": "2.0", "method": method, "params": params});
|
||
self.send_raw(&msg).await
|
||
}
|
||
|
||
async fn send_raw(&self, msg: &Value) -> Result<()> {
|
||
match &self.transport {
|
||
Transport::Stdio { stdin, .. } => self.write_stdio(stdin, msg).await,
|
||
Transport::Http { .. } => self.http_send(msg, None).await.map(|_| ()),
|
||
}
|
||
}
|
||
|
||
async fn write_stdio(&self, stdin: &AsyncMutex<tokio::process::ChildStdin>, msg: &Value) -> Result<()> {
|
||
let mut line = serde_json::to_vec(msg)?;
|
||
line.push(b'\n');
|
||
let mut w = stdin.lock().await;
|
||
w.write_all(&line).await?;
|
||
w.flush().await?;
|
||
Ok(())
|
||
}
|
||
|
||
/// POST a message. For requests, waits for the matching response, handling
|
||
/// server requests (elicitation) that arrive on the same SSE stream.
|
||
async fn http_send(&self, msg: &Value, want_id: Option<u64>) -> Result<Option<Value>> {
|
||
let Transport::Http { client, url, headers, session, protocol } = &self.transport else { unreachable!() };
|
||
let mut rb = client
|
||
.post(url)
|
||
.header("Accept", "application/json, text/event-stream")
|
||
.header("Content-Type", "application/json")
|
||
.json(msg);
|
||
for (k, v) in headers {
|
||
rb = rb.header(k, v);
|
||
}
|
||
if let Some(s) = session.lock().unwrap().clone() {
|
||
rb = rb.header("Mcp-Session-Id", s);
|
||
}
|
||
if let Some(p) = protocol.lock().unwrap().clone() {
|
||
rb = rb.header("MCP-Protocol-Version", p);
|
||
}
|
||
let resp = rb.send().await?;
|
||
if let Some(s) = resp.headers().get("mcp-session-id").and_then(|v| v.to_str().ok()) {
|
||
*session.lock().unwrap() = Some(s.to_string());
|
||
}
|
||
let status = resp.status();
|
||
if !status.is_success() {
|
||
let body = resp.text().await.unwrap_or_default();
|
||
bail!("MCP HTTP {status}: {}", body.chars().take(400).collect::<String>());
|
||
}
|
||
let Some(want) = want_id else { return Ok(None) };
|
||
let ctype = resp.headers().get("content-type").and_then(|v| v.to_str().ok()).unwrap_or("").to_string();
|
||
if ctype.contains("text/event-stream") {
|
||
let mut events = resp.bytes_stream().eventsource();
|
||
while let Some(ev) = events.next().await {
|
||
let ev = ev.map_err(|e| anyhow!("MCP stream error: {e}"))?;
|
||
let Ok(v) = serde_json::from_str::<Value>(&ev.data) else { continue };
|
||
if v.get("id").and_then(|i| i.as_u64()) == Some(want) && v.get("method").is_none() {
|
||
return Ok(Some(Self::result_of(v)?));
|
||
}
|
||
// Server→client requests/notifications on the stream.
|
||
Box::pin(self.handle_incoming(v)).await;
|
||
}
|
||
bail!("MCP stream ended without a response");
|
||
}
|
||
let v: Value = resp.json().await?;
|
||
Ok(Some(Self::result_of(v)?))
|
||
}
|
||
|
||
fn result_of(v: Value) -> Result<Value> {
|
||
if let Some(e) = v.get("error") {
|
||
bail!("MCP error {}: {}", e["code"], e["message"].as_str().unwrap_or(""));
|
||
}
|
||
Ok(v.get("result").cloned().unwrap_or(Value::Null))
|
||
}
|
||
|
||
async fn handle_incoming(&self, msg: Value) {
|
||
let method = msg.get("method").and_then(|m| m.as_str()).map(String::from);
|
||
let id = msg.get("id").cloned();
|
||
match (method, id) {
|
||
// Response to one of our requests (stdio).
|
||
(None, Some(id)) => {
|
||
if let (Transport::Stdio { pending, .. }, Some(id)) = (&self.transport, id.as_u64()) {
|
||
if let Some(tx) = pending.lock().unwrap().remove(&id) {
|
||
let _ = tx.send(Self::result_of(msg));
|
||
}
|
||
}
|
||
}
|
||
// Server → client request.
|
||
(Some(method), Some(id)) => {
|
||
let result = match method.as_str() {
|
||
"ping" => Ok(json!({})),
|
||
"roots/list" => {
|
||
let root = self.app.upgrade().map(|a| a.workspace()).unwrap_or_default();
|
||
Ok(json!({"roots": [{"uri": format!("file://{}", root.display()), "name": "workspace"}]}))
|
||
}
|
||
"elicitation/create" => self.elicit(&msg["params"]).await,
|
||
other => Err(anyhow!("method not supported: {other}")),
|
||
};
|
||
let reply = match result {
|
||
Ok(r) => json!({"jsonrpc": "2.0", "id": id, "result": r}),
|
||
Err(e) => json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32601, "message": e.to_string()}}),
|
||
};
|
||
if let Err(e) = self.send_raw(&reply).await {
|
||
tracing::warn!("mcp {}: failed to reply: {e:#}", self.name);
|
||
}
|
||
}
|
||
(Some(method), None) => {
|
||
if method == "notifications/tools/list_changed" {
|
||
if let Err(e) = self.refresh_tools().await {
|
||
tracing::warn!("mcp {}: tool refresh failed: {e:#}", self.name);
|
||
}
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
/// Turn an elicitation form into UI questions and the answers back into content.
|
||
async fn elicit(&self, params: &Value) -> Result<Value> {
|
||
let app = self.app.upgrade().ok_or_else(|| anyhow!("app gone"))?;
|
||
let session = self
|
||
.current_session
|
||
.lock()
|
||
.unwrap()
|
||
.clone()
|
||
.ok_or_else(|| anyhow!("no active session for elicitation"))?;
|
||
let message = params["message"].as_str().unwrap_or("").to_string();
|
||
let props = params["requestedSchema"]["properties"].as_object().cloned().unwrap_or_default();
|
||
let mut questions = Vec::new();
|
||
let mut keys = Vec::new();
|
||
for (key, schema) in &props {
|
||
let title = schema["title"].as_str().unwrap_or(key);
|
||
let desc = schema["description"].as_str().unwrap_or("");
|
||
let mut options: Vec<QuestionOption> = Vec::new();
|
||
if let Some(e) = schema["enum"].as_array() {
|
||
let names = schema["enumNames"].as_array();
|
||
for (i, v) in e.iter().enumerate() {
|
||
let label = v.as_str().map(String::from).unwrap_or_else(|| v.to_string());
|
||
let description = names.and_then(|n| n.get(i)).and_then(|n| n.as_str()).unwrap_or("").to_string();
|
||
options.push(QuestionOption { label, description });
|
||
}
|
||
} else if schema["type"] == "boolean" {
|
||
options = vec![
|
||
QuestionOption { label: "true".into(), description: "Yes".into() },
|
||
QuestionOption { label: "false".into(), description: "No".into() },
|
||
];
|
||
}
|
||
let question = if props.len() == 1 && !message.is_empty() {
|
||
format!("{message}\n\n{title}{}", if desc.is_empty() { String::new() } else { format!(" – {desc}") })
|
||
} else {
|
||
format!("{title}{}", if desc.is_empty() { String::new() } else { format!(" – {desc}") })
|
||
};
|
||
questions.push(Question { question, header: key.chars().take(12).collect(), options, multi_select: false });
|
||
keys.push((key.clone(), schema.clone()));
|
||
}
|
||
if questions.is_empty() {
|
||
questions.push(Question {
|
||
question: message.clone(),
|
||
header: "Confirm".into(),
|
||
options: vec![
|
||
QuestionOption { label: "accept".into(), description: String::new() },
|
||
QuestionOption { label: "decline".into(), description: String::new() },
|
||
],
|
||
multi_select: false,
|
||
});
|
||
}
|
||
let source = format!("mcp:{}", self.name);
|
||
let answers = app.ask_user(&session, &source, questions.clone(), CancellationToken::new()).await?;
|
||
let Some(answers) = answers else { return Ok(json!({"action": "cancel"})) };
|
||
if keys.is_empty() {
|
||
let a = answers.get(&questions[0].question).and_then(|v| v.as_str()).unwrap_or("decline");
|
||
return Ok(json!({"action": if a == "accept" { "accept" } else { "decline" }}));
|
||
}
|
||
let mut content = Map::new();
|
||
for ((key, schema), q) in keys.iter().zip(&questions) {
|
||
let Some(raw) = answers.get(&q.question).and_then(|v| v.as_str()) else { continue };
|
||
let v = match schema["type"].as_str() {
|
||
Some("boolean") => json!(raw == "true" || raw.eq_ignore_ascii_case("yes")),
|
||
Some("integer") => raw.trim().parse::<i64>().map(|n| json!(n)).unwrap_or(json!(raw)),
|
||
Some("number") => raw.trim().parse::<f64>().map(|n| json!(n)).unwrap_or(json!(raw)),
|
||
_ => json!(raw),
|
||
};
|
||
content.insert(key.clone(), v);
|
||
}
|
||
Ok(json!({"action": "accept", "content": content}))
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// All configured MCP servers with their connection state.
|
||
#[derive(Default)]
|
||
pub struct McpManager {
|
||
servers: Mutex<BTreeMap<String, Arc<McpServer>>>,
|
||
states: Mutex<BTreeMap<String, McpState>>,
|
||
}
|
||
|
||
impl McpManager {
|
||
/// Connection state per server (for the status overlay).
|
||
pub fn states(&self) -> BTreeMap<String, McpState> {
|
||
self.states.lock().unwrap().clone()
|
||
}
|
||
|
||
pub fn status_json(&self) -> Value {
|
||
let states = self.states.lock().unwrap();
|
||
let servers = self.servers.lock().unwrap();
|
||
let mut out = Map::new();
|
||
for (name, st) in states.iter() {
|
||
let tools: Vec<Value> = servers
|
||
.get(name)
|
||
.map(|s| s.tools.lock().unwrap().iter().map(|t| json!({"name": t.name, "description": t.description})).collect())
|
||
.unwrap_or_default();
|
||
out.insert(name.clone(), json!({"status": st, "tools": tools}));
|
||
}
|
||
Value::Object(out)
|
||
}
|
||
|
||
/// (Re)connect all servers from config. Runs connections concurrently.
|
||
pub async fn start_all(self: &Arc<Self>, app: Weak<App>, configs: BTreeMap<String, McpConfig>) {
|
||
self.servers.lock().unwrap().clear();
|
||
self.states.lock().unwrap().clear();
|
||
let mut tasks = Vec::new();
|
||
for (name, cfg) in configs {
|
||
if !cfg.enabled {
|
||
self.states.lock().unwrap().insert(name, McpState::Disabled);
|
||
continue;
|
||
}
|
||
self.states.lock().unwrap().insert(name.clone(), McpState::Connecting);
|
||
let me = self.clone();
|
||
let app = app.clone();
|
||
tasks.push(tokio::spawn(async move {
|
||
let res = tokio::time::timeout(Duration::from_secs(60), McpServer::connect(&name, &cfg, app.clone())).await;
|
||
let state = match res {
|
||
Ok(Ok(server)) => {
|
||
let n = server.tools.lock().unwrap().len();
|
||
me.servers.lock().unwrap().insert(name.clone(), server);
|
||
McpState::Connected { tools: n }
|
||
}
|
||
Ok(Err(e)) => McpState::Failed { error: format!("{e:#}") },
|
||
Err(_) => McpState::Failed { error: "connection timed out".into() },
|
||
};
|
||
tracing::info!("mcp {name}: {state:?}");
|
||
me.states.lock().unwrap().insert(name, state);
|
||
if let Some(app) = app.upgrade() {
|
||
app.emit(crate::types::Event::McpStatus { servers: me.status_json() });
|
||
}
|
||
}));
|
||
}
|
||
for t in tasks {
|
||
let _ = t.await;
|
||
}
|
||
}
|
||
|
||
/// All MCP tools, split into tools given to the model directly and tools
|
||
/// that are only reachable through mcp_search / mcp_call.
|
||
pub fn tools(&self, cfg: &crate::config::Config, user: Option<&crate::store::User>) -> (Vec<Arc<dyn Tool>>, Vec<Arc<McpTool>>) {
|
||
let servers = self.servers.lock().unwrap();
|
||
let mut direct: Vec<Arc<dyn Tool>> = Vec::new();
|
||
let mut deferred = Vec::new();
|
||
for (sname, server) in servers.iter() {
|
||
if !cfg.mcp_allowed(sname, user) {
|
||
continue;
|
||
}
|
||
let tools = server.tools.lock().unwrap();
|
||
let load = cfg.mcp.get(sname).and_then(|m| m.load.clone()).unwrap_or_else(|| "auto".into());
|
||
let on_demand = match load.as_str() {
|
||
"always" => false,
|
||
"on_demand" => true,
|
||
_ => tools.len() > cfg.server.mcp_direct_limit,
|
||
};
|
||
for t in tools.iter() {
|
||
let tool = Arc::new(McpTool { full_name: tool_name(sname, &t.name), server: server.clone(), info: t.clone() });
|
||
if on_demand {
|
||
deferred.push(tool);
|
||
} else {
|
||
direct.push(tool);
|
||
}
|
||
}
|
||
}
|
||
(direct, deferred)
|
||
}
|
||
|
||
/// Short overview of on-demand servers for the system prompt.
|
||
pub fn deferred_summary(&self, cfg: &crate::config::Config, user: Option<&crate::store::User>) -> Vec<(String, usize, Vec<String>)> {
|
||
let (_, deferred) = self.tools(cfg, user);
|
||
let mut by_server: BTreeMap<String, Vec<String>> = BTreeMap::new();
|
||
for t in deferred {
|
||
by_server.entry(t.server.name.clone()).or_default().push(t.info.name.clone());
|
||
}
|
||
by_server
|
||
.into_iter()
|
||
.map(|(s, names)| {
|
||
let n = names.len();
|
||
(s, n, names.into_iter().take(12).collect())
|
||
})
|
||
.collect()
|
||
}
|
||
}
|
||
|
||
/// Search the on-demand MCP tools by keywords.
|
||
pub struct McpSearch;
|
||
|
||
#[async_trait]
|
||
impl Tool for McpSearch {
|
||
fn name(&self) -> &str {
|
||
"mcp_search"
|
||
}
|
||
fn description(&self) -> String {
|
||
"Find tools of the on-demand MCP servers listed in the system prompt. Returns matching tools \
|
||
with their input schema; call them with mcp_call. Search with a few keywords (e.g. \
|
||
\"container restart\", \"light living room\")."
|
||
.into()
|
||
}
|
||
fn schema(&self) -> Value {
|
||
json!({
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string", "description": "Keywords"},
|
||
"server": {"type": "string", "description": "Only search this server (optional)"}
|
||
},
|
||
"required": ["query"]
|
||
})
|
||
}
|
||
fn read_only(&self) -> bool {
|
||
true
|
||
}
|
||
fn title(&self, input: &Value) -> String {
|
||
format!("Search MCP tools: {}", input["query"].as_str().unwrap_or(""))
|
||
}
|
||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||
let query = input["query"].as_str().unwrap_or("").to_lowercase();
|
||
let server = input["server"].as_str();
|
||
let words: Vec<&str> = query.split_whitespace().collect();
|
||
let (_, deferred) = ctx.app.mcp.tools(&ctx.app.config(), owner(ctx).as_ref());
|
||
let mut scored: Vec<(usize, &Arc<McpTool>)> = deferred
|
||
.iter()
|
||
.filter(|t| server.is_none_or(|s| t.server.name == s))
|
||
.map(|t| {
|
||
let name = t.info.name.to_lowercase();
|
||
let desc = t.info.description.to_lowercase();
|
||
let score: usize = words.iter().map(|w| if name.contains(w) { 3 } else if desc.contains(w) { 1 } else { 0 }).sum();
|
||
(score, t)
|
||
})
|
||
.filter(|(s, _)| *s > 0 || words.is_empty())
|
||
.collect();
|
||
scored.sort_by(|a, b| b.0.cmp(&a.0));
|
||
if scored.is_empty() {
|
||
return Ok(ToolOutput::text("No matching tools. Try other keywords or another server."));
|
||
}
|
||
let mut out = String::new();
|
||
for (_, t) in scored.iter().take(8) {
|
||
let desc: String = t.info.description.chars().take(300).collect();
|
||
out.push_str(&format!("## {}\n{}\nInput schema: {}\n\n", t.full_name, desc, t.info.input_schema));
|
||
}
|
||
if scored.len() > 8 {
|
||
out.push_str(&format!("({} more matches – refine the query)\n", scored.len() - 8));
|
||
}
|
||
Ok(ToolOutput::text(out))
|
||
}
|
||
}
|
||
|
||
/// Call an on-demand MCP tool found with mcp_search. Permissions are checked
|
||
/// against the target tool by the runner.
|
||
pub struct McpCall;
|
||
|
||
#[async_trait]
|
||
impl Tool for McpCall {
|
||
fn name(&self) -> &str {
|
||
"mcp_call"
|
||
}
|
||
fn description(&self) -> String {
|
||
"Call an MCP tool found with mcp_search, by its full name (mcp__<server>__<tool>).".into()
|
||
}
|
||
fn schema(&self) -> Value {
|
||
json!({
|
||
"type": "object",
|
||
"properties": {
|
||
"tool": {"type": "string", "description": "Full tool name, e.g. mcp__dockhand__list_containers"},
|
||
"arguments": {"type": "object", "description": "Arguments matching the tool's input schema"}
|
||
},
|
||
"required": ["tool"]
|
||
})
|
||
}
|
||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||
let target = resolve_call_target(ctx, &input)?;
|
||
target.run(input.get("arguments").cloned().unwrap_or(json!({})), ctx).await
|
||
}
|
||
}
|
||
|
||
/// The session owner (MCP access is decided per user).
|
||
fn owner(ctx: &ToolCtx) -> Option<crate::store::User> {
|
||
if let Some(id) = &ctx.owner_id {
|
||
return ctx.app.store.user(id).ok().flatten();
|
||
}
|
||
let s = ctx.app.session(&ctx.session_id).ok()?;
|
||
ctx.app.store.user(s.owner_id.as_deref()?).ok().flatten()
|
||
}
|
||
|
||
/// Find the MCP tool an mcp_call refers to.
|
||
pub fn resolve_call_target(ctx: &ToolCtx, input: &Value) -> Result<Arc<McpTool>> {
|
||
let name = input["tool"].as_str().unwrap_or("");
|
||
let (_, deferred) = ctx.app.mcp.tools(&ctx.app.config(), owner(ctx).as_ref());
|
||
deferred
|
||
.into_iter()
|
||
.find(|t| t.full_name == name)
|
||
.ok_or_else(|| anyhow!("unknown MCP tool '{name}' – use mcp_search to find tool names"))
|
||
}
|
||
|
||
/// Tool names must match ^[a-zA-Z0-9_-]{1,64}$ for most providers.
|
||
pub fn tool_name(server: &str, tool: &str) -> String {
|
||
let clean = |s: &str| s.chars().map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' }).collect::<String>();
|
||
let mut n = format!("mcp__{}__{}", clean(server), clean(tool));
|
||
n.truncate(64);
|
||
n
|
||
}
|
||
|
||
pub struct McpTool {
|
||
pub full_name: String,
|
||
server: Arc<McpServer>,
|
||
info: McpToolInfo,
|
||
}
|
||
|
||
#[async_trait]
|
||
impl Tool for McpTool {
|
||
fn name(&self) -> &str {
|
||
&self.full_name
|
||
}
|
||
fn description(&self) -> String {
|
||
format!("[MCP {}] {}", self.server.name, self.info.description)
|
||
}
|
||
fn schema(&self) -> Value {
|
||
let mut s = self.info.input_schema.clone();
|
||
if s.get("type").is_none() {
|
||
s["type"] = json!("object");
|
||
}
|
||
s
|
||
}
|
||
fn title(&self, _input: &Value) -> String {
|
||
format!("{} → {}", self.server.name, self.info.name)
|
||
}
|
||
fn read_only(&self) -> bool {
|
||
self.info.read_only
|
||
}
|
||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||
self.server.call_tool(&ctx.session_id, &self.info.name, input, ctx.cancel.clone()).await
|
||
}
|
||
}
|