First Version
This commit is contained in:
@@ -0,0 +1,14 @@
|
||||
[package]
|
||||
name = "mcp-client"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1", features = ["sync", "process", "io-util", "time", "rt", "macros"] }
|
||||
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
anyhow = "1"
|
||||
tracing = "0.1"
|
||||
async-trait = "0.1"
|
||||
base64 = "0.22"
|
||||
@@ -0,0 +1,29 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct McpServerConfig {
|
||||
pub name: String,
|
||||
pub transport: McpTransport,
|
||||
/// stdio only: the executable to spawn.
|
||||
pub command: Option<String>,
|
||||
/// stdio only: arguments passed to the command.
|
||||
pub args: Option<Vec<String>>,
|
||||
/// stdio only: extra environment variables (values support `${VAR}` interpolation).
|
||||
pub env: Option<HashMap<String, String>>,
|
||||
/// http only: base URL of the MCP server.
|
||||
pub url: Option<String>,
|
||||
/// http only: API key sent as `Authorization: Bearer <key>` (supports `${VAR}` interpolation).
|
||||
pub api_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum McpTransport {
|
||||
Stdio,
|
||||
/// Streamable HTTP (remote MCP servers — previously called SSE).
|
||||
Http,
|
||||
/// Alias kept for backwards compatibility.
|
||||
Sse,
|
||||
}
|
||||
@@ -0,0 +1,423 @@
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use reqwest::header::{ACCEPT, AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue};
|
||||
use serde_json::{Value, json};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::{McpCallResult, McpServerClient, McpTool, extract_text, interpolate_env};
|
||||
use crate::config::McpServerConfig;
|
||||
|
||||
const CALL_TIMEOUT_SECS: u64 = 120;
|
||||
|
||||
/// Best-effort cancellation for an in-flight HTTP `tools/call`: if dropped while
|
||||
/// armed (a `/stop` drops the request future, or the request timed out), it POSTs
|
||||
/// `notifications/cancelled` so the server can stop. Correlation is weaker than on
|
||||
/// stdio — the server must map `requestId` to the abandoned POST, which not every
|
||||
/// server does — hence best-effort. Disarmed once the server responds (or when a
|
||||
/// non-timeout send error proves the server never received the request).
|
||||
struct HttpCancelOnDrop {
|
||||
id: u64,
|
||||
client: reqwest::Client,
|
||||
url: String,
|
||||
headers: HeaderMap,
|
||||
name: String,
|
||||
reason: &'static str,
|
||||
armed: bool,
|
||||
}
|
||||
|
||||
impl HttpCancelOnDrop {
|
||||
fn disarm(&mut self) {
|
||||
self.armed = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for HttpCancelOnDrop {
|
||||
fn drop(&mut self) {
|
||||
if !self.armed {
|
||||
return;
|
||||
}
|
||||
let (id, client, url, headers, name, reason) =
|
||||
(self.id, self.client.clone(), self.url.clone(), self.headers.clone(), self.name.clone(), self.reason);
|
||||
tokio::spawn(async move {
|
||||
debug!("MCP http '{name}': notifications/cancelled for request {id} ({reason})");
|
||||
let _ = client.post(&url).headers(headers)
|
||||
.json(&crate::cancelled_notification(id, reason))
|
||||
.send().await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Cooperative `tasks/cancel` for a block-and-poll `poll_task` over HTTP: if the
|
||||
/// poll future is dropped while still polling (a `/stop`) or hits its deadline,
|
||||
/// POST `tasks/cancel` best-effort. Disarmed once the task reaches a terminal state.
|
||||
struct HttpTaskCancelOnDrop {
|
||||
request_id: u64,
|
||||
task_id: String,
|
||||
client: reqwest::Client,
|
||||
url: String,
|
||||
headers: HeaderMap,
|
||||
name: String,
|
||||
armed: bool,
|
||||
}
|
||||
|
||||
impl HttpTaskCancelOnDrop {
|
||||
fn disarm(&mut self) {
|
||||
self.armed = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for HttpTaskCancelOnDrop {
|
||||
fn drop(&mut self) {
|
||||
if !self.armed {
|
||||
return;
|
||||
}
|
||||
let (request_id, task_id, client, url, headers, name) =
|
||||
(self.request_id, self.task_id.clone(), self.client.clone(), self.url.clone(), self.headers.clone(), self.name.clone());
|
||||
tokio::spawn(async move {
|
||||
debug!("MCP http '{name}': tasks/cancel for task {task_id}");
|
||||
let _ = client.post(&url).headers(headers)
|
||||
.json(&crate::tasks_cancel_request(request_id, &task_id))
|
||||
.send().await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub struct McpHttpServer {
|
||||
name: String,
|
||||
url: String,
|
||||
client: reqwest::Client,
|
||||
headers: HeaderMap,
|
||||
/// Set after the `initialize` response — required by stateful servers like Tavily.
|
||||
session_id: Mutex<Option<String>>,
|
||||
/// Protocol version negotiated in the `initialize` response (falls back to
|
||||
/// [`crate::PROTOCOL_VERSION`]). Once set, echoed in the `MCP-Protocol-Version`
|
||||
/// header on every post-initialize request, per the Streamable HTTP spec.
|
||||
protocol_version: Mutex<Option<String>>,
|
||||
next_id: AtomicU64,
|
||||
tools: Vec<McpTool>,
|
||||
/// Capabilities the server advertised in its `InitializeResult`. Captured so a
|
||||
/// future Tasks polling loop can gate on `tasks` support; unused for now.
|
||||
server_capabilities: Value,
|
||||
}
|
||||
|
||||
impl McpHttpServer {
|
||||
pub async fn start(cfg: &McpServerConfig) -> Result<Self> {
|
||||
let url = cfg.url.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("http server '{}' requires 'url'", cfg.name))?
|
||||
.trim_end_matches('/')
|
||||
.to_string();
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
headers.insert(ACCEPT, HeaderValue::from_static("application/json, text/event-stream"));
|
||||
|
||||
if let Some(key) = &cfg.api_key {
|
||||
let val = interpolate_env(key);
|
||||
let bearer = format!("Bearer {val}");
|
||||
headers.insert(AUTHORIZATION, bearer.parse()
|
||||
.map_err(|_| anyhow::anyhow!("invalid api_key for '{}'", cfg.name))?);
|
||||
}
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(CALL_TIMEOUT_SECS))
|
||||
.build()?;
|
||||
|
||||
let server = McpHttpServer {
|
||||
name: cfg.name.clone(),
|
||||
url,
|
||||
client,
|
||||
headers,
|
||||
session_id: Mutex::new(None),
|
||||
protocol_version: Mutex::new(None),
|
||||
next_id: AtomicU64::new(1),
|
||||
tools: Vec::new(),
|
||||
server_capabilities: json!({}),
|
||||
};
|
||||
|
||||
let init = server.request("initialize", json!({
|
||||
// The HTTP transport doesn't service the ElicitationHandler (stdio-only),
|
||||
// so it must NOT advertise the elicitation capability.
|
||||
"protocolVersion": crate::PROTOCOL_VERSION,
|
||||
// Experimental Tasks marker only (recognise-but-don't-poll); see the
|
||||
// stdio transport for the rationale behind keeping it under `experimental`.
|
||||
"capabilities": { "experimental": { "tasks": {} } },
|
||||
"clientInfo": { "name": "skald", "version": env!("CARGO_PKG_VERSION") },
|
||||
})).await?;
|
||||
// Capture the negotiated version (fall back to our own) so post-initialize
|
||||
// requests can echo it in the MCP-Protocol-Version header; tolerate a
|
||||
// downgrade with a warning rather than disconnecting.
|
||||
let negotiated = init["protocolVersion"].as_str().unwrap_or(crate::PROTOCOL_VERSION);
|
||||
if negotiated != crate::PROTOCOL_VERSION {
|
||||
warn!("MCP http '{}': server negotiated protocol {negotiated} (we requested {}); proceeding",
|
||||
server.name, crate::PROTOCOL_VERSION);
|
||||
}
|
||||
*server.protocol_version.lock().unwrap() = Some(negotiated.to_string());
|
||||
// Capture the server's advertised capabilities for a future Tasks poller.
|
||||
let server_capabilities = init.get("capabilities").cloned().unwrap_or_else(|| json!({}));
|
||||
|
||||
if let Err(e) = server.notify("notifications/initialized", json!({})).await {
|
||||
warn!("MCP http '{}': initialized notification failed (ignoring): {e}", server.name);
|
||||
}
|
||||
|
||||
// Follow `nextCursor` across pages so large tool lists aren't silently
|
||||
// truncated; capped at `MAX_TOOL_PAGES` against a stuck cursor.
|
||||
let mut tools: Vec<McpTool> = Vec::new();
|
||||
let mut cursor: Option<String> = None;
|
||||
for page_n in 0..crate::MAX_TOOL_PAGES {
|
||||
let params = match &cursor {
|
||||
Some(c) => json!({ "cursor": c }),
|
||||
None => json!({}),
|
||||
};
|
||||
let page = server.request("tools/list", params).await?;
|
||||
if let Some(arr) = page["tools"].as_array() {
|
||||
tools.extend(arr.iter().map(|t| McpTool::from_json(&cfg.name, t)));
|
||||
}
|
||||
cursor = page["nextCursor"].as_str().filter(|s| !s.is_empty()).map(str::to_string);
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
if page_n + 1 == crate::MAX_TOOL_PAGES {
|
||||
warn!("MCP http '{}': tools/list hit {}-page cap; some tools may be omitted",
|
||||
server.name, crate::MAX_TOOL_PAGES);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(McpHttpServer { tools, server_capabilities, ..server })
|
||||
}
|
||||
|
||||
pub fn tools(&self) -> &[McpTool] {
|
||||
&self.tools
|
||||
}
|
||||
|
||||
/// Capabilities the server advertised at `initialize`. Exposed for a future
|
||||
/// Tasks polling loop to gate on `tasks` support.
|
||||
pub fn server_capabilities(&self) -> &Value {
|
||||
&self.server_capabilities
|
||||
}
|
||||
|
||||
pub async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> {
|
||||
let mut params = json!({ "name": name, "arguments": args });
|
||||
if self.wants_task(name) {
|
||||
// Opt into deferred execution for a task-capable tool (experimental Tasks).
|
||||
params["task"] = json!({});
|
||||
}
|
||||
let result = self.request("tools/call", params).await?;
|
||||
|
||||
if result["isError"].as_bool().unwrap_or(false) {
|
||||
anyhow::bail!("MCP tool error: {}", extract_text(&result));
|
||||
}
|
||||
match crate::extract_call_result(&result) {
|
||||
McpCallResult::Task(task) => self.poll_task(task).await,
|
||||
other => Ok(other),
|
||||
}
|
||||
}
|
||||
|
||||
/// True when tool `name` advertises `execution.taskSupport` as `required`/
|
||||
/// `optional`, so we opt into deferred (Task) execution.
|
||||
fn wants_task(&self, name: &str) -> bool {
|
||||
self.tools.iter()
|
||||
.find(|t| t.name == name)
|
||||
.and_then(|t| t.task_support.as_deref())
|
||||
.is_some_and(|s| s == "required" || s == "optional")
|
||||
}
|
||||
|
||||
/// Drives a deferred Task to completion (experimental Tasks, block-and-poll):
|
||||
/// polls `tasks/get` until a terminal status, then fetches the real result via
|
||||
/// `tasks/result`. A [`HttpTaskCancelOnDrop`] guard POSTs `tasks/cancel` if this
|
||||
/// future is dropped (a `/stop`) or the deadline is hit. The overall wait is
|
||||
/// bounded only by the task's `ttl`, so long tasks no longer hit the 120s wall.
|
||||
async fn poll_task(&self, task: crate::CreateTaskResult) -> Result<McpCallResult> {
|
||||
let task_id = task.task_id.as_str();
|
||||
let cancel_id = self.next_id.fetch_add(1, Ordering::SeqCst);
|
||||
let mut guard = HttpTaskCancelOnDrop {
|
||||
request_id: cancel_id,
|
||||
task_id: task.task_id.clone(),
|
||||
client: self.client.clone(),
|
||||
url: self.url.clone(),
|
||||
headers: self.request_headers(),
|
||||
name: self.name.clone(),
|
||||
armed: true,
|
||||
};
|
||||
|
||||
let deadline = crate::poll_deadline(task.ttl_ms);
|
||||
let mut interval = task.poll_interval_ms;
|
||||
|
||||
loop {
|
||||
tokio::time::sleep(crate::clamp_poll_interval(interval)).await;
|
||||
if std::time::Instant::now() >= deadline {
|
||||
anyhow::bail!("MCP http '{}' task '{task_id}' exceeded max wait", self.name);
|
||||
}
|
||||
let get = self.request("tasks/get", json!({ "taskId": task_id })).await?;
|
||||
let Some(state) = crate::CreateTaskResult::parse(&get) else {
|
||||
anyhow::bail!("MCP http '{}' task '{task_id}': malformed tasks/get response", self.name);
|
||||
};
|
||||
interval = state.poll_interval_ms.or(interval);
|
||||
match state.status {
|
||||
crate::TaskStatus::Working => continue,
|
||||
crate::TaskStatus::Completed => break,
|
||||
crate::TaskStatus::Failed => {
|
||||
guard.disarm();
|
||||
anyhow::bail!("MCP http '{}' task '{task_id}' failed: {}", self.name, extract_text(&get));
|
||||
}
|
||||
crate::TaskStatus::Cancelled => {
|
||||
guard.disarm();
|
||||
anyhow::bail!("MCP http '{}' task '{task_id}' was cancelled by the server", self.name);
|
||||
}
|
||||
crate::TaskStatus::InputRequired =>
|
||||
anyhow::bail!("MCP http '{}' task '{task_id}' requires input mid-task, which isn't supported yet", self.name),
|
||||
}
|
||||
}
|
||||
|
||||
// Task is terminal (completed) — nothing left to cancel.
|
||||
guard.disarm();
|
||||
let result = self.request("tasks/result", json!({ "taskId": task_id })).await?;
|
||||
|
||||
if result["isError"].as_bool().unwrap_or(false) {
|
||||
anyhow::bail!("MCP tool error: {}", extract_text(&result));
|
||||
}
|
||||
Ok(crate::extract_call_result(&result))
|
||||
}
|
||||
|
||||
/// Builds per-request headers: the static base plus the captured
|
||||
/// `Mcp-Session-Id` and `MCP-Protocol-Version`. Both are set only after the
|
||||
/// `initialize` response, so they're naturally absent on the initialize call
|
||||
/// itself (the spec scopes the version header to post-initialize requests).
|
||||
fn request_headers(&self) -> HeaderMap {
|
||||
let mut headers = self.headers.clone();
|
||||
if let Some(sid) = self.session_id.lock().unwrap().as_deref() {
|
||||
if let Ok(val) = HeaderValue::from_str(sid) {
|
||||
headers.insert("mcp-session-id", val);
|
||||
}
|
||||
}
|
||||
if let Some(ver) = self.protocol_version.lock().unwrap().as_deref() {
|
||||
if let Ok(val) = HeaderValue::from_str(ver) {
|
||||
headers.insert("mcp-protocol-version", val);
|
||||
}
|
||||
}
|
||||
headers
|
||||
}
|
||||
|
||||
async fn request(&self, method: &str, params: Value) -> Result<Value> {
|
||||
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
|
||||
let body = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"method": method,
|
||||
"params": params,
|
||||
});
|
||||
|
||||
let req_headers = self.request_headers();
|
||||
|
||||
// Arm a best-effort cancellation guard for cancellable operations only
|
||||
// (`tools/call`): a `/stop` that drops this future, or a request timeout,
|
||||
// then POSTs `notifications/cancelled`. Disarmed once the server responds.
|
||||
let mut cancel_guard = (method == "tools/call").then(|| HttpCancelOnDrop {
|
||||
id,
|
||||
client: self.client.clone(),
|
||||
url: self.url.clone(),
|
||||
headers: req_headers.clone(),
|
||||
name: self.name.clone(),
|
||||
reason: "cancelled by client",
|
||||
armed: true,
|
||||
});
|
||||
|
||||
let resp = match self.client
|
||||
.post(&self.url)
|
||||
.headers(req_headers)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
// A timeout may leave the server still working → cancel it. Any
|
||||
// other send failure means the server never got the request, so
|
||||
// there is nothing to cancel.
|
||||
if let Some(g) = cancel_guard.as_mut() {
|
||||
if e.is_timeout() { g.reason = "timeout"; } else { g.disarm(); }
|
||||
}
|
||||
anyhow::bail!("MCP http '{}' request failed: {e}", self.name);
|
||||
}
|
||||
};
|
||||
|
||||
// The server responded — the request completed on its side; disarm.
|
||||
if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
|
||||
|
||||
if let Some(sid) = resp.headers().get("mcp-session-id") {
|
||||
if let Ok(sid_str) = sid.to_str() {
|
||||
debug!("MCP http '{}': captured session id", self.name);
|
||||
*self.session_id.lock().unwrap() = Some(sid_str.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let status = resp.status();
|
||||
let content_type = resp.headers()
|
||||
.get(CONTENT_TYPE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
let msg: Value = if content_type.contains("text/event-stream") {
|
||||
parse_sse_response(resp).await
|
||||
.map_err(|e| anyhow::anyhow!("MCP http '{}' SSE parse error: {e}", self.name))?
|
||||
} else {
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("MCP http '{}' HTTP {status}: {body}", self.name);
|
||||
}
|
||||
resp.json::<Value>().await
|
||||
.map_err(|e| anyhow::anyhow!("MCP http '{}' JSON decode error: {e}", self.name))?
|
||||
};
|
||||
|
||||
if let Some(error) = msg.get("error") {
|
||||
anyhow::bail!("MCP http '{}' protocol error: {error}", self.name);
|
||||
}
|
||||
Ok(msg["result"].clone())
|
||||
}
|
||||
|
||||
async fn notify(&self, method: &str, params: Value) -> Result<()> {
|
||||
let body = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": method,
|
||||
"params": params,
|
||||
});
|
||||
|
||||
let req_headers = self.request_headers();
|
||||
|
||||
self.client
|
||||
.post(&self.url)
|
||||
.headers(req_headers)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("MCP http notify failed: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpServerClient for McpHttpServer {
|
||||
fn tools(&self) -> &[McpTool] { self.tools() }
|
||||
async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> { self.call_tool(name, args).await }
|
||||
}
|
||||
|
||||
async fn parse_sse_response(resp: reqwest::Response) -> Result<Value> {
|
||||
let text = resp.text().await?;
|
||||
for line in text.lines() {
|
||||
let data = match line.strip_prefix("data:") {
|
||||
Some(d) => d.trim(),
|
||||
None => continue,
|
||||
};
|
||||
if data == "[DONE]" { break; }
|
||||
if let Ok(msg) = serde_json::from_str::<Value>(data) {
|
||||
if msg.get("result").is_some() || msg.get("error").is_some() {
|
||||
return Ok(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::bail!("no JSON-RPC result found in SSE response")
|
||||
}
|
||||
@@ -0,0 +1,565 @@
|
||||
pub mod config;
|
||||
pub mod http_server;
|
||||
pub mod log;
|
||||
pub mod server;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub use config::{McpServerConfig, McpTransport};
|
||||
pub use log::{McpLogLine, McpLogTx};
|
||||
pub use server::{
|
||||
ElicitationAction, ElicitationHandler, ElicitationReply, ElicitationRequest, McpNotification,
|
||||
};
|
||||
|
||||
/// MCP protocol version this client advertises in `initialize` and (over HTTP)
|
||||
/// echoes in the `MCP-Protocol-Version` header on every post-initialize request.
|
||||
/// Shared by both transports so they can never drift apart. The host tolerates a
|
||||
/// server negotiating a different (older) version — see the HTTP transport, which
|
||||
/// captures the server's reply and warns rather than disconnecting.
|
||||
pub const PROTOCOL_VERSION: &str = "2025-11-25";
|
||||
|
||||
/// Safety cap on `tools/list` pagination: stop following `nextCursor` after this
|
||||
/// many pages so a buggy or hostile server that never clears the cursor can't
|
||||
/// loop the client forever.
|
||||
pub(crate) const MAX_TOOL_PAGES: usize = 50;
|
||||
|
||||
/// Builds a `notifications/cancelled` message for an in-flight request. Shared by
|
||||
/// both transports (like [`PROTOCOL_VERSION`]) so the wire shape can't drift. Per
|
||||
/// the MCP spec the client MUST NOT cancel the `initialize` request; callers only
|
||||
/// arm this for cancellable operations (`tools/call`).
|
||||
pub(crate) fn cancelled_notification(request_id: u64, reason: &str) -> Value {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/cancelled",
|
||||
"params": { "requestId": request_id, "reason": reason },
|
||||
})
|
||||
}
|
||||
|
||||
/// Builds a `tasks/cancel` request for a task the client is abandoning
|
||||
/// (experimental Tasks). Sent fire-and-forget from the poll cancel-guard — the
|
||||
/// response is ignored — so it carries a pre-allocated `request_id`. Shared by both
|
||||
/// transports.
|
||||
pub(crate) fn tasks_cancel_request(request_id: u64, task_id: &str) -> Value {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"method": "tasks/cancel",
|
||||
"params": { "taskId": task_id },
|
||||
})
|
||||
}
|
||||
|
||||
// ── Public types ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct McpTool {
|
||||
pub server_name: String,
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub input_schema: Value,
|
||||
/// Optional human-readable title (MCP 2025-06-18+).
|
||||
pub title: Option<String>,
|
||||
/// Optional JSON Schema describing the tool's structured output
|
||||
/// (`outputSchema`, MCP 2025-06-18+). Captured for future validation and to
|
||||
/// know a tool can return `structuredContent`.
|
||||
pub output_schema: Option<Value>,
|
||||
/// Optional tool annotations (`readOnlyHint`, `destructiveHint`, …), treated
|
||||
/// as untrusted hints per the spec.
|
||||
pub annotations: Option<Value>,
|
||||
/// Optional per-tool Tasks negotiation (`execution.taskSupport`:
|
||||
/// `required`/`optional`/`forbidden`, MCP 2025-11-25 experimental). Captured
|
||||
/// as an untrusted hint so a future polling implementation knows which tools
|
||||
/// may return a [`CreateTaskResult`]. See [`McpCallResult::Task`].
|
||||
pub task_support: Option<String>,
|
||||
}
|
||||
|
||||
impl McpTool {
|
||||
/// Builds an [`McpTool`] from one entry of a `tools/list` `tools[]` array.
|
||||
/// Shared by both transports so the field mapping (incl. the 2025-06-18+
|
||||
/// `title`/`outputSchema`/`annotations`) stays in one place.
|
||||
pub fn from_json(server_name: &str, t: &Value) -> McpTool {
|
||||
McpTool {
|
||||
server_name: server_name.to_string(),
|
||||
name: t["name"].as_str().unwrap_or("").to_string(),
|
||||
description: t["description"].as_str().unwrap_or("").to_string(),
|
||||
input_schema: t.get("inputSchema").cloned().unwrap_or_else(|| json!({
|
||||
"type": "object", "properties": {},
|
||||
})),
|
||||
title: t.get("title").and_then(Value::as_str).map(str::to_string),
|
||||
output_schema: t.get("outputSchema").cloned(),
|
||||
annotations: t.get("annotations").cloned(),
|
||||
task_support: t.get("execution")
|
||||
.and_then(|e| e.get("taskSupport"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tool_id(&self) -> String {
|
||||
format!("mcp__{}__{}", self.server_name, self.name)
|
||||
}
|
||||
|
||||
pub fn to_openai_definition(&self) -> Value {
|
||||
let params = if self.input_schema.is_object() {
|
||||
self.input_schema.clone()
|
||||
} else {
|
||||
json!({ "type": "object", "properties": {} })
|
||||
};
|
||||
json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.tool_id(),
|
||||
"description": format!("[{}] {}", self.server_name, self.description),
|
||||
"parameters": params,
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Status of a configured MCP server.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[serde(tag = "status", rename_all = "snake_case")]
|
||||
pub enum McpServerStatus {
|
||||
Running { tools: Vec<String> },
|
||||
Error { message: String },
|
||||
Disabled,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct McpServerInfo {
|
||||
pub name: String,
|
||||
pub transport: String,
|
||||
pub description: Option<String>,
|
||||
pub friendly_name: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub status: McpServerStatus,
|
||||
}
|
||||
|
||||
// ── Tool-result media ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Kind of a non-text content block carried in an MCP `tools/call` result.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum McpMediaKind {
|
||||
Image,
|
||||
Audio,
|
||||
/// Embedded or linked resource (any other binary/file payload).
|
||||
Resource,
|
||||
}
|
||||
|
||||
/// Payload of one media content block: either decoded inline bytes (from an
|
||||
/// `image`/`audio`/embedded-`resource` base64 field) or a link to a remote
|
||||
/// resource (`resource_link`), which the crate does not download.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum McpMediaData {
|
||||
/// Decoded bytes plus the server-declared MIME type.
|
||||
Inline { bytes: Vec<u8>, mime: String },
|
||||
/// A `resource_link` URI, passed through untouched (no fetch here).
|
||||
Link { uri: String, mime: Option<String> },
|
||||
}
|
||||
|
||||
/// One non-text content block extracted from a `tools/call` result. The crate
|
||||
/// stays a generic transport: it decodes base64 to bytes but never touches the
|
||||
/// disk — persistence is the host's job (`McpManager`).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct McpMedia {
|
||||
pub kind: McpMediaKind,
|
||||
pub data: McpMediaData,
|
||||
}
|
||||
|
||||
// ── Tasks (MCP 2025-11-25, experimental) ────────────────────────────────────────
|
||||
|
||||
/// Lifecycle state of a Task (`CreateTaskResult.status`). `Working` transitions to
|
||||
/// `Completed`/`Failed`/`Cancelled`; `InputRequired` means the receiver needs more
|
||||
/// input (e.g. an elicitation) before it can continue.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskStatus {
|
||||
Working,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
InputRequired,
|
||||
}
|
||||
|
||||
impl TaskStatus {
|
||||
/// A terminal status: the task will not change further, so polling stops.
|
||||
/// `InputRequired` is **not** terminal (the task resumes once input is given),
|
||||
/// but the current block-and-poll client can't fulfil input mid-task, so it
|
||||
/// treats it as an error rather than looping forever — see `poll_task`.
|
||||
pub fn is_terminal(self) -> bool {
|
||||
matches!(self, TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled)
|
||||
}
|
||||
}
|
||||
|
||||
/// Clamps a server-suggested `pollInterval` (ms) into a sane range: default 1s when
|
||||
/// absent, floored at 500ms (don't hammer the server) and capped at 30s (stay
|
||||
/// responsive). Shared by both transports so their poll cadence can't drift.
|
||||
pub(crate) fn clamp_poll_interval(ms: Option<u64>) -> std::time::Duration {
|
||||
std::time::Duration::from_millis(ms.unwrap_or(1000).clamp(500, 30_000))
|
||||
}
|
||||
|
||||
/// Deadline after which the block-and-poll loop gives up on a task and cancels it:
|
||||
/// the server's `ttl` when given, otherwise a generous 1-hour cap so a stuck task
|
||||
/// can't pin a session forever.
|
||||
pub(crate) fn poll_deadline(ttl_ms: Option<u64>) -> std::time::Instant {
|
||||
let ttl = std::time::Duration::from_millis(ttl_ms.unwrap_or(3_600_000));
|
||||
std::time::Instant::now() + ttl
|
||||
}
|
||||
|
||||
/// A durable task handle returned by a receiver when a request is executed as a
|
||||
/// deferred, pollable operation instead of blocking (MCP 2025-11-25 *experimental*
|
||||
/// Tasks). The client polls `tasks/get` until a terminal status, then fetches the
|
||||
/// real result via `tasks/result` (see `poll_task` on each transport). Also used to
|
||||
/// parse each `tasks/get` response (same shape).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateTaskResult {
|
||||
pub task_id: String,
|
||||
pub status: TaskStatus,
|
||||
/// Suggested delay between `tasks/get` polls, in milliseconds.
|
||||
pub poll_interval_ms: Option<u64>,
|
||||
/// Time-to-live of the task handle, in milliseconds.
|
||||
pub ttl_ms: Option<u64>,
|
||||
}
|
||||
|
||||
impl CreateTaskResult {
|
||||
/// Parses a `tools/call` result that carries a task handle. Returns `None` if
|
||||
/// `v` is not a task-shaped object (no `taskId`). The handle may live at the
|
||||
/// top level or under a `task` field, depending on the server.
|
||||
pub(crate) fn parse(v: &Value) -> Option<CreateTaskResult> {
|
||||
let obj = v.get("task").filter(|t| t.is_object()).unwrap_or(v);
|
||||
let task_id = obj.get("taskId").and_then(Value::as_str)?.to_string();
|
||||
let status = obj.get("status")
|
||||
.and_then(|s| serde_json::from_value::<TaskStatus>(s.clone()).ok())
|
||||
.unwrap_or(TaskStatus::Working);
|
||||
Some(CreateTaskResult {
|
||||
task_id,
|
||||
status,
|
||||
poll_interval_ms: obj.get("pollInterval").and_then(Value::as_u64),
|
||||
ttl_ms: obj.get("ttl").and_then(Value::as_u64),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ── Transport trait ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Typed result of an MCP `tools/call`. Mirrors the host's tool-result shape
|
||||
/// without coupling this crate to it; the host (`McpManager`) maps it. Per the
|
||||
/// MCP spec, `structuredContent` is canonical when present (servers SHOULD also
|
||||
/// mirror it in a `TextContent` block for backwards compatibility); we prefer it
|
||||
/// and fall back to the joined `text` items, which fixes the silent empty-result
|
||||
/// case for servers that omit the text mirror. Non-text content blocks
|
||||
/// (`image`/`audio`/`resource`/`resource_link`) are surfaced via [`McpCallResult::Media`]
|
||||
/// so the host can persist them instead of dropping them.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum McpCallResult {
|
||||
/// Joined `text` content items.
|
||||
Text(String),
|
||||
/// Canonical structured payload from `structuredContent`.
|
||||
Json(Value),
|
||||
/// At least one non-text media block was present. `text`/`structured` carry
|
||||
/// any accompanying textual/structured payload from the same result.
|
||||
Media {
|
||||
text: Option<String>,
|
||||
structured: Option<Value>,
|
||||
items: Vec<McpMedia>,
|
||||
},
|
||||
/// The server deferred the call as a Task (experimental Tasks, 2025-11-25) and
|
||||
/// returned a durable handle instead of a result. Skald surfaces the handle;
|
||||
/// polling for the real result is a follow-up.
|
||||
Task(CreateTaskResult),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait McpServerClient: Send + Sync {
|
||||
fn tools(&self) -> &[McpTool];
|
||||
async fn call_tool(&self, name: &str, args: Value) -> anyhow::Result<McpCallResult>;
|
||||
}
|
||||
|
||||
// ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Parses `mcp__<server>__<tool>` → `(server, tool)`.
|
||||
pub fn parse_mcp_tool_name(name: &str) -> Option<(&str, &str)> {
|
||||
let rest = name.strip_prefix("mcp__")?;
|
||||
let sep = rest.find("__")?;
|
||||
Some((&rest[..sep], &rest[sep + 2..]))
|
||||
}
|
||||
|
||||
/// Extracts text content from an MCP tool result value (the `text` items of
|
||||
/// `content[]`, plus the inline `text` of embedded text resources). Used for the
|
||||
/// `isError` path and as the text component of a successful result.
|
||||
pub(crate) fn extract_text(result: &Value) -> String {
|
||||
result["content"]
|
||||
.as_array()
|
||||
.map(|arr| classify_content(arr).0)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Decodes a standard-base64 string to bytes, returning `None` on malformed input.
|
||||
fn decode_b64(s: &str) -> Option<Vec<u8>> {
|
||||
use base64::Engine;
|
||||
base64::engine::general_purpose::STANDARD.decode(s).ok()
|
||||
}
|
||||
|
||||
/// Walks `content[]`, returning the joined text and any non-text media blocks
|
||||
/// (`image`/`audio`/embedded-`resource`/`resource_link`). Base64 payloads are
|
||||
/// decoded to bytes here; persistence is left to the host.
|
||||
fn classify_content(content: &[Value]) -> (String, Vec<McpMedia>) {
|
||||
let mut texts: Vec<String> = Vec::new();
|
||||
let mut media: Vec<McpMedia> = Vec::new();
|
||||
|
||||
for item in content {
|
||||
match item.get("type").and_then(Value::as_str) {
|
||||
Some("text") => {
|
||||
if let Some(t) = item.get("text").and_then(Value::as_str) {
|
||||
texts.push(t.to_string());
|
||||
}
|
||||
}
|
||||
Some(kind @ ("image" | "audio")) => {
|
||||
let mime = item.get("mimeType").and_then(Value::as_str)
|
||||
.unwrap_or("application/octet-stream").to_string();
|
||||
if let Some(bytes) = item.get("data").and_then(Value::as_str).and_then(decode_b64) {
|
||||
let kind = if kind == "image" { McpMediaKind::Image } else { McpMediaKind::Audio };
|
||||
media.push(McpMedia { kind, data: McpMediaData::Inline { bytes, mime } });
|
||||
}
|
||||
}
|
||||
Some("resource") => {
|
||||
// Embedded resource: a base64 `blob` (binary) or inline `text`.
|
||||
let res = item.get("resource");
|
||||
let mime = res.and_then(|r| r.get("mimeType")).and_then(Value::as_str)
|
||||
.unwrap_or("application/octet-stream").to_string();
|
||||
if let Some(bytes) = res.and_then(|r| r.get("blob")).and_then(Value::as_str).and_then(decode_b64) {
|
||||
media.push(McpMedia { kind: McpMediaKind::Resource, data: McpMediaData::Inline { bytes, mime } });
|
||||
} else if let Some(t) = res.and_then(|r| r.get("text")).and_then(Value::as_str) {
|
||||
texts.push(t.to_string());
|
||||
}
|
||||
}
|
||||
Some("resource_link") => {
|
||||
if let Some(uri) = item.get("uri").and_then(Value::as_str) {
|
||||
let mime = item.get("mimeType").and_then(Value::as_str).map(str::to_string);
|
||||
media.push(McpMedia {
|
||||
kind: McpMediaKind::Resource,
|
||||
data: McpMediaData::Link { uri: uri.to_string(), mime },
|
||||
});
|
||||
}
|
||||
}
|
||||
// Unknown/typeless block: capture a bare `text` field if present.
|
||||
_ => {
|
||||
if let Some(t) = item.get("text").and_then(Value::as_str) {
|
||||
texts.push(t.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
(texts.join("\n"), media)
|
||||
}
|
||||
|
||||
/// Builds the typed result of an MCP `tools/call`. When any non-text media block
|
||||
/// is present it returns [`McpCallResult::Media`] (so the host can persist the
|
||||
/// bytes instead of dropping them). Otherwise it preserves the original
|
||||
/// precedence: `structuredContent` is canonical when present, else the joined
|
||||
/// `text` items — which also fixes the silent empty-result case for servers that
|
||||
/// return only `structuredContent` without the recommended text mirror.
|
||||
pub(crate) fn extract_call_result(result: &Value) -> McpCallResult {
|
||||
// A Task handle (deferred execution) has no `content[]`; recognise it first so
|
||||
// it isn't mistaken for an empty result.
|
||||
if result.get("content").is_none() {
|
||||
if let Some(task) = CreateTaskResult::parse(result) {
|
||||
return McpCallResult::Task(task);
|
||||
}
|
||||
}
|
||||
|
||||
let content = result["content"].as_array().cloned().unwrap_or_default();
|
||||
let (text, media) = classify_content(&content);
|
||||
let structured = result.get("structuredContent").filter(|v| !v.is_null()).cloned();
|
||||
|
||||
if !media.is_empty() {
|
||||
return McpCallResult::Media {
|
||||
text: (!text.is_empty()).then_some(text),
|
||||
structured,
|
||||
items: media,
|
||||
};
|
||||
}
|
||||
if let Some(sc) = structured {
|
||||
return McpCallResult::Json(sc);
|
||||
}
|
||||
McpCallResult::Text(text)
|
||||
}
|
||||
|
||||
/// Interpolates `${VAR}` references in a string from the process environment.
|
||||
pub(crate) fn interpolate_env(s: &str) -> String {
|
||||
let mut result = s.to_string();
|
||||
loop {
|
||||
let Some(start) = result.find("${") else { break };
|
||||
let Some(rel_end) = result[start..].find('}') else { break };
|
||||
let var_name = result[start + 2..start + rel_end].to_string();
|
||||
let value = std::env::var(&var_name).unwrap_or_else(|_| {
|
||||
tracing::warn!("MCP env var ${{{var_name}}} not set");
|
||||
String::new()
|
||||
});
|
||||
result = format!("{}{}{}", &result[..start], value, &result[start + rel_end + 1..]);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn b64(bytes: &[u8]) -> String {
|
||||
use base64::Engine;
|
||||
base64::engine::general_purpose::STANDARD.encode(bytes)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_block_becomes_media_with_decoded_bytes() {
|
||||
let png = [0x89u8, b'P', b'N', b'G'];
|
||||
let result = json!({
|
||||
"content": [
|
||||
{ "type": "text", "text": "here is your image" },
|
||||
{ "type": "image", "data": b64(&png), "mimeType": "image/png" }
|
||||
]
|
||||
});
|
||||
match extract_call_result(&result) {
|
||||
McpCallResult::Media { text, items, structured } => {
|
||||
assert_eq!(text.as_deref(), Some("here is your image"));
|
||||
assert!(structured.is_none());
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0].kind, McpMediaKind::Image);
|
||||
match &items[0].data {
|
||||
McpMediaData::Inline { bytes, mime } => {
|
||||
assert_eq!(bytes.as_slice(), &png);
|
||||
assert_eq!(mime, "image/png");
|
||||
}
|
||||
_ => panic!("expected inline media"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Media, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_only_stays_text() {
|
||||
let result = json!({ "content": [ { "type": "text", "text": "plain" } ] });
|
||||
match extract_call_result(&result) {
|
||||
McpCallResult::Text(t) => assert_eq!(t, "plain"),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_without_media_stays_json() {
|
||||
let result = json!({
|
||||
"content": [ { "type": "text", "text": "mirror" } ],
|
||||
"structuredContent": { "ok": true }
|
||||
});
|
||||
match extract_call_result(&result) {
|
||||
McpCallResult::Json(v) => assert_eq!(v, json!({ "ok": true })),
|
||||
other => panic!("expected Json, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancelled_notification_has_request_id_and_reason() {
|
||||
let msg = cancelled_notification(7, "timeout");
|
||||
assert_eq!(msg["jsonrpc"], "2.0");
|
||||
assert_eq!(msg["method"], "notifications/cancelled");
|
||||
assert!(msg.get("id").is_none(), "notifications MUST NOT carry an id");
|
||||
assert_eq!(msg["params"]["requestId"], 7);
|
||||
assert_eq!(msg["params"]["reason"], "timeout");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_json_captures_task_support() {
|
||||
let t = json!({
|
||||
"name": "gen_video",
|
||||
"execution": { "taskSupport": "required" }
|
||||
});
|
||||
let tool = McpTool::from_json("srv", &t);
|
||||
assert_eq!(tool.task_support.as_deref(), Some("required"));
|
||||
|
||||
// Absent execution → None.
|
||||
let plain = McpTool::from_json("srv", &json!({ "name": "echo" }));
|
||||
assert!(plain.task_support.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_task_result_parses_top_level_and_nested() {
|
||||
// Top-level handle.
|
||||
let top = json!({
|
||||
"taskId": "abc-123", "status": "working",
|
||||
"pollInterval": 500, "ttl": 60000
|
||||
});
|
||||
let r = CreateTaskResult::parse(&top).expect("top-level task");
|
||||
assert_eq!(r.task_id, "abc-123");
|
||||
assert_eq!(r.status, TaskStatus::Working);
|
||||
assert_eq!(r.poll_interval_ms, Some(500));
|
||||
assert_eq!(r.ttl_ms, Some(60000));
|
||||
|
||||
// Nested under `task`, unknown status → defaults to Working.
|
||||
let nested = json!({ "task": { "taskId": "x", "status": "input_required" } });
|
||||
let r = CreateTaskResult::parse(&nested).expect("nested task");
|
||||
assert_eq!(r.task_id, "x");
|
||||
assert_eq!(r.status, TaskStatus::InputRequired);
|
||||
|
||||
// Not task-shaped → None.
|
||||
assert!(CreateTaskResult::parse(&json!({ "content": [] })).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_status_is_terminal() {
|
||||
assert!(TaskStatus::Completed.is_terminal());
|
||||
assert!(TaskStatus::Failed.is_terminal());
|
||||
assert!(TaskStatus::Cancelled.is_terminal());
|
||||
assert!(!TaskStatus::Working.is_terminal());
|
||||
assert!(!TaskStatus::InputRequired.is_terminal());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clamp_poll_interval_defaults_and_bounds() {
|
||||
use std::time::Duration;
|
||||
assert_eq!(clamp_poll_interval(None), Duration::from_millis(1000)); // default
|
||||
assert_eq!(clamp_poll_interval(Some(10)), Duration::from_millis(500)); // floor
|
||||
assert_eq!(clamp_poll_interval(Some(99_999)), Duration::from_millis(30_000)); // cap
|
||||
assert_eq!(clamp_poll_interval(Some(2000)), Duration::from_millis(2000)); // passthrough
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tasks_cancel_request_shape() {
|
||||
let msg = tasks_cancel_request(42, "job-1");
|
||||
assert_eq!(msg["jsonrpc"], "2.0");
|
||||
assert_eq!(msg["id"], 42);
|
||||
assert_eq!(msg["method"], "tasks/cancel");
|
||||
assert_eq!(msg["params"]["taskId"], "job-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_call_result_recognises_task_handle() {
|
||||
let result = json!({ "taskId": "job-9", "status": "working", "ttl": 120000 });
|
||||
match extract_call_result(&result) {
|
||||
McpCallResult::Task(t) => {
|
||||
assert_eq!(t.task_id, "job-9");
|
||||
assert_eq!(t.status, TaskStatus::Working);
|
||||
assert_eq!(t.ttl_ms, Some(120000));
|
||||
}
|
||||
other => panic!("expected Task, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resource_link_passes_through_without_fetch() {
|
||||
let result = json!({
|
||||
"content": [ { "type": "resource_link", "uri": "https://x/y.mp4", "mimeType": "video/mp4" } ]
|
||||
});
|
||||
match extract_call_result(&result) {
|
||||
McpCallResult::Media { items, .. } => match &items[0].data {
|
||||
McpMediaData::Link { uri, mime } => {
|
||||
assert_eq!(uri, "https://x/y.mp4");
|
||||
assert_eq!(mime.as_deref(), Some("video/mp4"));
|
||||
}
|
||||
_ => panic!("expected link"),
|
||||
},
|
||||
other => panic!("expected Media, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
//! Per-server diagnostic log lines.
|
||||
//!
|
||||
//! An MCP server produces diagnostics from three places: its process `stderr`
|
||||
//! (stdio transport), MCP `notifications/message` log records, and connection
|
||||
//! lifecycle events (start failure / disconnect). This crate stays a generic
|
||||
//! transport — it *emits* [`McpLogLine`]s over a channel but never touches the
|
||||
//! disk. The host (`McpManager`) owns the channel and writes each line to a
|
||||
//! per-server file `logs/mcp/<name>.log`.
|
||||
//!
|
||||
//! Note: the MCP `logging` utility (`notifications/message` + `logging/setLevel`)
|
||||
//! is deprecated from the 2026-07-28 draft (SEP-2577) in favour of `stderr`, so
|
||||
//! `stderr` is the primary, future-proof source; `notifications/message` is
|
||||
//! captured for interop with servers that still emit it.
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// A single diagnostic line from an MCP server, routed by the host to that
|
||||
/// server's log file.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct McpLogLine {
|
||||
/// The MCP server that produced the line (selects the target file).
|
||||
pub server: String,
|
||||
/// Severity/kind tag, written inside `[...]`: `"stderr"` for raw child
|
||||
/// stderr, an MCP log level (`"debug"`..`"emergency"`) for
|
||||
/// `notifications/message`, or `"lifecycle"` for start/disconnect events.
|
||||
pub level: Cow<'static, str>,
|
||||
/// Human-readable text, already flattened to a single line (no trailing `\n`).
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
/// Channel the host installs to receive [`McpLogLine`]s from a running server.
|
||||
pub type McpLogTx = mpsc::UnboundedSender<McpLogLine>;
|
||||
|
||||
impl McpLogLine {
|
||||
/// A raw line drained from the child process's `stderr` (stdio transport).
|
||||
/// The level is unknown, so it's tagged `stderr`; the text itself usually
|
||||
/// carries the server's own `ERROR`/`WARNING` marker.
|
||||
pub fn stderr(server: impl Into<String>, text: impl Into<String>) -> Self {
|
||||
Self { server: server.into(), level: Cow::Borrowed("stderr"), text: text.into() }
|
||||
}
|
||||
|
||||
/// A connection lifecycle event (start failure, disconnect). Emitted for every
|
||||
/// transport so even HTTP servers — which have no `stderr` — get a connection
|
||||
/// history in their file.
|
||||
pub fn lifecycle(server: impl Into<String>, text: impl Into<String>) -> Self {
|
||||
Self { server: server.into(), level: Cow::Borrowed("lifecycle"), text: text.into() }
|
||||
}
|
||||
|
||||
/// Builds a line from the `params` of an MCP `notifications/message`
|
||||
/// (`{ level, logger?, data }`). Missing `level` falls back to `info`; `data`
|
||||
/// strings pass through, other JSON is compacted to one line so nothing is
|
||||
/// lost and the file stays greppable.
|
||||
pub fn from_message(server: impl Into<String>, params: &Value) -> Self {
|
||||
let level = params.get("level").and_then(Value::as_str).unwrap_or("info").to_string();
|
||||
let logger = params.get("logger").and_then(Value::as_str);
|
||||
let text = format_message_data(logger, params.get("data"));
|
||||
Self { server: server.into(), level: Cow::Owned(level), text }
|
||||
}
|
||||
}
|
||||
|
||||
/// Flattens the `data` of a `notifications/message` into one line, prefixing the
|
||||
/// optional `logger` name. Strings pass through; any other JSON is compacted.
|
||||
fn format_message_data(logger: Option<&str>, data: Option<&Value>) -> String {
|
||||
let body = match data {
|
||||
Some(Value::String(s)) => s.clone(),
|
||||
Some(v) => serde_json::to_string(v).unwrap_or_else(|_| v.to_string()),
|
||||
None => String::new(),
|
||||
};
|
||||
match logger {
|
||||
Some(l) if !l.is_empty() => format!("{l}: {body}"),
|
||||
_ => body,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn from_message_uses_level_and_string_data() {
|
||||
let line = McpLogLine::from_message("srv", &json!({ "level": "error", "data": "boom" }));
|
||||
assert_eq!(line.level, "error");
|
||||
assert_eq!(line.text, "boom");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_message_compacts_object_data_and_prefixes_logger() {
|
||||
// firecrawl-shaped payload.
|
||||
let params = json!({
|
||||
"level": "info",
|
||||
"logger": "scraper",
|
||||
"data": { "message": "Scraping URL", "context": { "url": "https://x/y" } }
|
||||
});
|
||||
let line = McpLogLine::from_message("firecrawl", ¶ms);
|
||||
assert_eq!(line.level, "info");
|
||||
assert!(line.text.starts_with("scraper: "));
|
||||
assert!(line.text.contains("Scraping URL"));
|
||||
assert!(line.text.contains("https://x/y"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_message_defaults_level_to_info() {
|
||||
let line = McpLogLine::from_message("srv", &json!({ "data": "hi" }));
|
||||
assert_eq!(line.level, "info");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stderr_and_lifecycle_tags() {
|
||||
assert_eq!(McpLogLine::stderr("s", "x").level, "stderr");
|
||||
assert_eq!(McpLogLine::lifecycle("s", "x").level, "lifecycle");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,608 @@
|
||||
use std::collections::HashMap;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::{ChildStdin, Command};
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::{McpCallResult, McpLogLine, McpLogTx, McpServerClient, McpTool, extract_text, interpolate_env};
|
||||
use crate::config::McpServerConfig;
|
||||
|
||||
/// A server-initiated notification from an MCP server: `(server_name, full JSON-RPC message)`.
|
||||
pub type McpNotification = (String, Value);
|
||||
|
||||
const CALL_TIMEOUT_SECS: u64 = 120;
|
||||
|
||||
// ── Elicitation (MCP spec 2025-06-18) ──────────────────────────────────────────
|
||||
|
||||
/// A server-initiated elicitation request: the server asks the user for input
|
||||
/// *during* a tool call (`elicitation/create`). The secret/value never reaches
|
||||
/// the LLM and is never persisted.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ElicitationRequest {
|
||||
/// Human-readable message to show the user.
|
||||
pub message: String,
|
||||
/// JSON Schema (flat object of typed fields) describing the requested input.
|
||||
pub requested_schema: Value,
|
||||
}
|
||||
|
||||
/// The user's decision on an elicitation request.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ElicitationAction {
|
||||
Accept,
|
||||
Decline,
|
||||
Cancel,
|
||||
}
|
||||
|
||||
impl ElicitationAction {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
ElicitationAction::Accept => "accept",
|
||||
ElicitationAction::Decline => "decline",
|
||||
ElicitationAction::Cancel => "cancel",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The reply sent back to the server for an `elicitation/create` request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ElicitationReply {
|
||||
pub action: ElicitationAction,
|
||||
/// Field values for `accept`; `None` for `decline`/`cancel`.
|
||||
pub content: Option<Value>,
|
||||
}
|
||||
|
||||
/// Bridges a server-initiated elicitation to whatever surfaces it to the user
|
||||
/// (in Skald: the Agent Inbox via `ElicitationManager`). The crate writes the
|
||||
/// JSON-RPC reply itself — the handler only produces the decision. The returned
|
||||
/// value (and any secret it carries) flows straight to the server's stdin and is
|
||||
/// never logged here.
|
||||
#[async_trait]
|
||||
pub trait ElicitationHandler: Send + Sync {
|
||||
async fn handle(&self, server_name: &str, request: ElicitationRequest) -> ElicitationReply;
|
||||
}
|
||||
|
||||
/// Serialises `msg` as a single JSON-RPC line and writes it to the child's stdin.
|
||||
async fn write_json_line(stdin: &Arc<Mutex<ChildStdin>>, msg: &Value) {
|
||||
if let Ok(mut line) = serde_json::to_string(msg) {
|
||||
line.push('\n');
|
||||
let _ = stdin.lock().await.write_all(line.as_bytes()).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Handles a server→client JSON-RPC request (e.g. `elicitation/create`) without
|
||||
/// blocking the read-loop: the user reply may take minutes, so the work is
|
||||
/// spawned and the JSON-RPC response is written back when it resolves.
|
||||
fn handle_server_request(
|
||||
server_name: &str,
|
||||
msg: Value,
|
||||
stdin: &Arc<Mutex<ChildStdin>>,
|
||||
handler: &Option<Arc<dyn ElicitationHandler>>,
|
||||
pending_elicitations: &Arc<AtomicUsize>,
|
||||
) {
|
||||
let id = msg.get("id").cloned().unwrap_or(Value::Null);
|
||||
let method = msg.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
|
||||
if method == "elicitation/create" {
|
||||
let params = msg.get("params").cloned().unwrap_or_else(|| json!({}));
|
||||
let request = ElicitationRequest {
|
||||
message: params.get("message").and_then(Value::as_str).unwrap_or("").to_string(),
|
||||
requested_schema: params.get("requestedSchema").cloned().unwrap_or_else(|| json!({})),
|
||||
};
|
||||
let stdin = Arc::clone(stdin);
|
||||
|
||||
let Some(handler) = handler.clone() else {
|
||||
// Capability declared but no handler wired: cancel so the server
|
||||
// doesn't hang waiting for input we can't collect.
|
||||
tokio::spawn(async move {
|
||||
write_json_line(&stdin, &json!({
|
||||
"jsonrpc": "2.0", "id": id, "result": { "action": "cancel" },
|
||||
})).await;
|
||||
});
|
||||
return;
|
||||
};
|
||||
|
||||
pending_elicitations.fetch_add(1, Ordering::SeqCst);
|
||||
let counter = Arc::clone(pending_elicitations);
|
||||
let server = server_name.to_string();
|
||||
tokio::spawn(async move {
|
||||
let reply = handler.handle(&server, request).await;
|
||||
let mut result = json!({ "action": reply.action.as_str() });
|
||||
if reply.action == ElicitationAction::Accept {
|
||||
if let Some(content) = reply.content {
|
||||
result["content"] = content;
|
||||
}
|
||||
}
|
||||
write_json_line(&stdin, &json!({
|
||||
"jsonrpc": "2.0", "id": id, "result": result,
|
||||
})).await;
|
||||
counter.fetch_sub(1, Ordering::SeqCst);
|
||||
});
|
||||
} else {
|
||||
// Unknown server→client request: reply method-not-found so the server
|
||||
// isn't left hanging.
|
||||
let stdin = Arc::clone(stdin);
|
||||
let method = method.to_string();
|
||||
tokio::spawn(async move {
|
||||
write_json_line(&stdin, &json!({
|
||||
"jsonrpc": "2.0", "id": id,
|
||||
"error": { "code": -32601, "message": format!("method not found: {method}") },
|
||||
})).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Guards an in-flight `tools/call`: if this is dropped while still armed — the
|
||||
/// caller aborted the tool (a `/stop` drops the work future) or the call timed out
|
||||
/// — it tells the server to stop via `notifications/cancelled` and drops the now
|
||||
/// orphaned `pending` entry (which would otherwise leak). Disarmed once the server
|
||||
/// replies or disconnects, where cancelling is pointless. The spec forbids
|
||||
/// cancelling `initialize`, so only `tools/call` arms a guard.
|
||||
struct CancelOnDrop {
|
||||
id: u64,
|
||||
stdin: Arc<Mutex<ChildStdin>>,
|
||||
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
|
||||
name: String,
|
||||
reason: &'static str,
|
||||
armed: bool,
|
||||
}
|
||||
|
||||
impl CancelOnDrop {
|
||||
fn disarm(&mut self) {
|
||||
self.armed = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CancelOnDrop {
|
||||
fn drop(&mut self) {
|
||||
if !self.armed {
|
||||
return;
|
||||
}
|
||||
let (id, stdin, pending, name, reason) =
|
||||
(self.id, Arc::clone(&self.stdin), Arc::clone(&self.pending), self.name.clone(), self.reason);
|
||||
tokio::spawn(async move {
|
||||
pending.lock().await.remove(&id);
|
||||
debug!(target: "mcp_client", "[{name}] notifications/cancelled for request {id} ({reason})");
|
||||
write_json_line(&stdin, &crate::cancelled_notification(id, reason)).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Cooperative `tasks/cancel` for a block-and-poll `poll_task`: if the poll future
|
||||
/// is dropped while still polling (a `/stop`) or hits its deadline, tell the server
|
||||
/// to abandon the task. Fire-and-forget — the response carries no `pending` entry,
|
||||
/// so the read-loop simply discards it. Disarmed once the task reaches a terminal
|
||||
/// state (or fails), where cancelling is pointless.
|
||||
struct TaskCancelOnDrop {
|
||||
request_id: u64,
|
||||
task_id: String,
|
||||
stdin: Arc<Mutex<ChildStdin>>,
|
||||
name: String,
|
||||
armed: bool,
|
||||
}
|
||||
|
||||
impl TaskCancelOnDrop {
|
||||
fn disarm(&mut self) {
|
||||
self.armed = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TaskCancelOnDrop {
|
||||
fn drop(&mut self) {
|
||||
if !self.armed {
|
||||
return;
|
||||
}
|
||||
let (request_id, task_id, stdin, name) =
|
||||
(self.request_id, self.task_id.clone(), Arc::clone(&self.stdin), self.name.clone());
|
||||
tokio::spawn(async move {
|
||||
debug!(target: "mcp_client", "[{name}] tasks/cancel for task {task_id}");
|
||||
write_json_line(&stdin, &crate::tasks_cancel_request(request_id, &task_id)).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub struct McpServer {
|
||||
name: String,
|
||||
stdin: Arc<Mutex<ChildStdin>>,
|
||||
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
|
||||
next_id: AtomicU64,
|
||||
tools: Vec<McpTool>,
|
||||
/// Number of in-flight server→client elicitations awaiting a user reply.
|
||||
/// While > 0, an in-flight `tools/call` on this server must not time out
|
||||
/// (the user may still be typing a password into the Inbox).
|
||||
pending_elicitations: Arc<AtomicUsize>,
|
||||
/// Capabilities the server advertised in its `InitializeResult`. Captured so a
|
||||
/// future Tasks polling loop can gate on `tasks` support; unused for now.
|
||||
server_capabilities: Value,
|
||||
}
|
||||
|
||||
impl McpServer {
|
||||
pub async fn start(
|
||||
cfg: &McpServerConfig,
|
||||
notification_tx: Option<mpsc::UnboundedSender<McpNotification>>,
|
||||
log_tx: Option<McpLogTx>,
|
||||
elicitation_handler: Option<Arc<dyn ElicitationHandler>>,
|
||||
) -> Result<Self> {
|
||||
let command = cfg.command.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("stdio server '{}' requires 'command'", cfg.name))?;
|
||||
|
||||
let mut cmd = Command::new(command);
|
||||
if let Some(args) = &cfg.args {
|
||||
cmd.args(args);
|
||||
}
|
||||
if let Some(env_map) = &cfg.env {
|
||||
for (k, v) in env_map {
|
||||
cmd.env(k, interpolate_env(v));
|
||||
}
|
||||
}
|
||||
cmd.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
// Capture the child's stderr instead of inheriting it: many MCP
|
||||
// servers (e.g. FastMCP) print a startup banner, deprecation
|
||||
// warnings and INFO logs there, which would otherwise spill raw
|
||||
// onto our console. We drain it into tracing at `debug` level so it
|
||||
// stays quiet by default but is still available for diagnostics.
|
||||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
|
||||
// Detach the child into its own process group so that a terminal
|
||||
// Ctrl+C (SIGINT delivered to the whole foreground process group)
|
||||
// does not reach it directly. Otherwise Python-based MCP servers
|
||||
// catch the SIGINT and dump a KeyboardInterrupt traceback to stderr.
|
||||
// Cleanup still happens via `kill_on_drop`: when the app shuts down
|
||||
// and the reader task is dropped, the child receives SIGKILL silently.
|
||||
#[cfg(unix)]
|
||||
cmd.process_group(0);
|
||||
|
||||
let mut child = cmd.spawn()
|
||||
.map_err(|e| anyhow::anyhow!("failed to spawn '{}': {e}", cfg.name))?;
|
||||
|
||||
let stdin = child.stdin.take()
|
||||
.ok_or_else(|| anyhow::anyhow!("could not capture stdin for '{}'", cfg.name))?;
|
||||
// Wrap stdin early so both the struct and the read-loop (which writes
|
||||
// elicitation replies back to the server) can share the same handle.
|
||||
let stdin = Arc::new(Mutex::new(stdin));
|
||||
let stdout = child.stdout.take()
|
||||
.ok_or_else(|| anyhow::anyhow!("could not capture stdout for '{}'", cfg.name))?;
|
||||
|
||||
// Drain the child's stderr into tracing at `debug` (so banners/warnings
|
||||
// don't pollute our console at the default level) *and* forward each line
|
||||
// to the host's per-server log file via `log_tx`. Per the MCP 2026 draft,
|
||||
// `stderr` is the primary place servers should log, so this is the main
|
||||
// diagnostic source for stdio transports.
|
||||
if let Some(stderr) = child.stderr.take() {
|
||||
let server_name_err = cfg.name.clone();
|
||||
let log_tx_err = log_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut lines = BufReader::new(stderr).lines();
|
||||
while let Ok(Some(line)) = lines.next_line().await {
|
||||
if !line.trim().is_empty() {
|
||||
debug!(target: "mcp_client", "[{server_name_err}] {line}");
|
||||
if let Some(tx) = &log_tx_err {
|
||||
let _ = tx.send(McpLogLine::stderr(server_name_err.clone(), &line));
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>> =
|
||||
Arc::new(Mutex::new(HashMap::new()));
|
||||
let pending_elicitations = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let pending_bg = pending.clone();
|
||||
let server_name_bg = cfg.name.clone();
|
||||
let notification_tx_bg = notification_tx;
|
||||
let log_tx_bg = log_tx;
|
||||
let stdin_bg = Arc::clone(&stdin);
|
||||
let elicitation_handler_bg = elicitation_handler;
|
||||
let pending_elicitations_bg = Arc::clone(&pending_elicitations);
|
||||
tokio::spawn(async move {
|
||||
let mut child = child;
|
||||
let mut lines = BufReader::new(stdout).lines();
|
||||
loop {
|
||||
match lines.next_line().await {
|
||||
Ok(Some(line)) if !line.trim().is_empty() => {
|
||||
if let Ok(msg) = serde_json::from_str::<Value>(&line) {
|
||||
let has_method = msg.get("method").is_some();
|
||||
let has_id = msg.get("id").map(|v| !v.is_null()).unwrap_or(false);
|
||||
if has_method && has_id {
|
||||
// Server → client request (e.g. elicitation/create):
|
||||
// has *both* `method` and `id`, so it must be checked
|
||||
// before the response branch below.
|
||||
handle_server_request(
|
||||
&server_name_bg, msg, &stdin_bg,
|
||||
&elicitation_handler_bg, &pending_elicitations_bg,
|
||||
);
|
||||
} else if let Some(id) = msg["id"].as_u64() {
|
||||
if let Some(tx) = pending_bg.lock().await.remove(&id) {
|
||||
let _ = tx.send(msg);
|
||||
}
|
||||
} else if has_method {
|
||||
// `notifications/message` is the MCP logging utility
|
||||
// (deprecated 2026-07-28): route it to the per-server
|
||||
// log file, not to the notification queue that feeds
|
||||
// TIC — otherwise log records masquerade as business
|
||||
// events. Every other notification (e.g. the custom
|
||||
// `event/*` methods) flows on to `notification_tx`.
|
||||
if msg.get("method").and_then(Value::as_str) == Some("notifications/message") {
|
||||
if let Some(tx) = &log_tx_bg {
|
||||
let params = &msg["params"];
|
||||
let _ = tx.send(McpLogLine::from_message(server_name_bg.clone(), params));
|
||||
}
|
||||
} else if let Some(tx) = ¬ification_tx_bg {
|
||||
let _ = tx.send((server_name_bg.clone(), msg));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Some(_)) => {}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
let exit_info = match child.wait().await {
|
||||
Ok(status) if !status.success() => format!(
|
||||
"process exited with {}",
|
||||
status.code().map(|c| c.to_string()).unwrap_or_else(|| "signal".into())
|
||||
),
|
||||
_ => "process exited unexpectedly".into(),
|
||||
};
|
||||
let error_msg = format!("MCP '{}' disconnected: {exit_info}", server_name_bg);
|
||||
if let Some(tx) = &log_tx_bg {
|
||||
let _ = tx.send(McpLogLine::lifecycle(server_name_bg.clone(), format!("disconnected: {exit_info}")));
|
||||
}
|
||||
for (_, tx) in pending_bg.lock().await.drain() {
|
||||
let _ = tx.send(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"error": { "code": -32000, "message": error_msg }
|
||||
}));
|
||||
}
|
||||
});
|
||||
|
||||
let server = McpServer {
|
||||
name: cfg.name.clone(),
|
||||
stdin,
|
||||
pending,
|
||||
next_id: AtomicU64::new(1),
|
||||
tools: Vec::new(),
|
||||
pending_elicitations,
|
||||
server_capabilities: json!({}),
|
||||
};
|
||||
|
||||
let init = server.request("initialize", json!({
|
||||
// Declare the elicitation capability (form mode) so servers know they
|
||||
// may request input mid-call.
|
||||
"protocolVersion": crate::PROTOCOL_VERSION,
|
||||
"capabilities": {
|
||||
"elicitation": {},
|
||||
// Experimental Tasks marker: we *recognise* a CreateTaskResult
|
||||
// (see McpCallResult::Task) but don't poll yet, so we deliberately
|
||||
// avoid claiming the full `tasks` capability — that could make a
|
||||
// server defer results we can't retrieve. Kept under `experimental`
|
||||
// until the polling loop lands.
|
||||
"experimental": { "tasks": {} },
|
||||
},
|
||||
"clientInfo": { "name": "skald", "version": env!("CARGO_PKG_VERSION") },
|
||||
})).await?;
|
||||
// Tolerate a server that negotiates a different (older) version — warn but
|
||||
// keep going rather than disconnecting.
|
||||
if let Some(v) = init["protocolVersion"].as_str() {
|
||||
if v != crate::PROTOCOL_VERSION {
|
||||
warn!("MCP '{}': server negotiated protocol {v} (we requested {}); proceeding",
|
||||
server.name, crate::PROTOCOL_VERSION);
|
||||
}
|
||||
}
|
||||
// Capture the server's advertised capabilities for a future Tasks poller.
|
||||
let server_capabilities = init.get("capabilities").cloned().unwrap_or_else(|| json!({}));
|
||||
|
||||
server.notify("notifications/initialized", json!({})).await?;
|
||||
|
||||
// Follow `nextCursor` across pages so servers with large tool lists aren't
|
||||
// silently truncated; capped at `MAX_TOOL_PAGES` against a stuck cursor.
|
||||
let mut tools: Vec<McpTool> = Vec::new();
|
||||
let mut cursor: Option<String> = None;
|
||||
for page_n in 0..crate::MAX_TOOL_PAGES {
|
||||
let params = match &cursor {
|
||||
Some(c) => json!({ "cursor": c }),
|
||||
None => json!({}),
|
||||
};
|
||||
let page = server.request("tools/list", params).await?;
|
||||
if let Some(arr) = page["tools"].as_array() {
|
||||
tools.extend(arr.iter().map(|t| McpTool::from_json(&cfg.name, t)));
|
||||
}
|
||||
cursor = page["nextCursor"].as_str().filter(|s| !s.is_empty()).map(str::to_string);
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
if page_n + 1 == crate::MAX_TOOL_PAGES {
|
||||
warn!("MCP '{}': tools/list hit {}-page cap; some tools may be omitted",
|
||||
server.name, crate::MAX_TOOL_PAGES);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(McpServer { tools, server_capabilities, ..server })
|
||||
}
|
||||
|
||||
pub fn tools(&self) -> &[McpTool] {
|
||||
&self.tools
|
||||
}
|
||||
|
||||
/// Capabilities the server advertised at `initialize`. Exposed for a future
|
||||
/// Tasks polling loop to gate on `tasks` support.
|
||||
pub fn server_capabilities(&self) -> &Value {
|
||||
&self.server_capabilities
|
||||
}
|
||||
|
||||
pub async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> {
|
||||
let mut params = json!({ "name": name, "arguments": args });
|
||||
if self.wants_task(name) {
|
||||
// Opt into deferred execution for a task-capable tool (experimental
|
||||
// Tasks). Per spec, adding the `task` field to the request is the
|
||||
// opt-in for `tools/call` (the client is the requestor).
|
||||
params["task"] = json!({});
|
||||
}
|
||||
let result = self.request("tools/call", params).await?;
|
||||
|
||||
if result["isError"].as_bool().unwrap_or(false) {
|
||||
anyhow::bail!("MCP tool error: {}", extract_text(&result));
|
||||
}
|
||||
match crate::extract_call_result(&result) {
|
||||
McpCallResult::Task(task) => self.poll_task(task).await,
|
||||
other => Ok(other),
|
||||
}
|
||||
}
|
||||
|
||||
/// True when tool `name` advertises `execution.taskSupport` as `required`/
|
||||
/// `optional`, so we opt into deferred (Task) execution.
|
||||
fn wants_task(&self, name: &str) -> bool {
|
||||
self.tools.iter()
|
||||
.find(|t| t.name == name)
|
||||
.and_then(|t| t.task_support.as_deref())
|
||||
.is_some_and(|s| s == "required" || s == "optional")
|
||||
}
|
||||
|
||||
/// Drives a deferred Task to completion (experimental Tasks, block-and-poll):
|
||||
/// polls `tasks/get` until a terminal status, then fetches the real result via
|
||||
/// `tasks/result`. A [`TaskCancelOnDrop`] guard sends `tasks/cancel` if this
|
||||
/// future is dropped (a `/stop`) or the deadline is hit. Each poll request is a
|
||||
/// normal short `request()` (subject to the 120s timeout); the *overall* wait is
|
||||
/// bounded only by the task's `ttl` — so long tasks no longer hit that wall.
|
||||
async fn poll_task(&self, task: crate::CreateTaskResult) -> Result<McpCallResult> {
|
||||
let task_id = task.task_id.as_str();
|
||||
let cancel_id = self.next_id.fetch_add(1, Ordering::SeqCst);
|
||||
let mut guard = TaskCancelOnDrop {
|
||||
request_id: cancel_id,
|
||||
task_id: task.task_id.clone(),
|
||||
stdin: Arc::clone(&self.stdin),
|
||||
name: self.name.clone(),
|
||||
armed: true,
|
||||
};
|
||||
|
||||
let deadline = crate::poll_deadline(task.ttl_ms);
|
||||
let mut interval = task.poll_interval_ms;
|
||||
|
||||
loop {
|
||||
tokio::time::sleep(crate::clamp_poll_interval(interval)).await;
|
||||
if std::time::Instant::now() >= deadline {
|
||||
anyhow::bail!("MCP '{}' task '{task_id}' exceeded max wait", self.name);
|
||||
}
|
||||
let get = self.request("tasks/get", json!({ "taskId": task_id })).await?;
|
||||
let Some(state) = crate::CreateTaskResult::parse(&get) else {
|
||||
anyhow::bail!("MCP '{}' task '{task_id}': malformed tasks/get response", self.name);
|
||||
};
|
||||
interval = state.poll_interval_ms.or(interval);
|
||||
match state.status {
|
||||
crate::TaskStatus::Working => continue,
|
||||
crate::TaskStatus::Completed => break,
|
||||
crate::TaskStatus::Failed => {
|
||||
guard.disarm();
|
||||
anyhow::bail!("MCP '{}' task '{task_id}' failed: {}", self.name, extract_text(&get));
|
||||
}
|
||||
crate::TaskStatus::Cancelled => {
|
||||
guard.disarm();
|
||||
anyhow::bail!("MCP '{}' task '{task_id}' was cancelled by the server", self.name);
|
||||
}
|
||||
// Still alive but blocked on input we can't supply mid-task: abandon
|
||||
// it (guard stays armed → tasks/cancel). See follow-up.
|
||||
crate::TaskStatus::InputRequired =>
|
||||
anyhow::bail!("MCP '{}' task '{task_id}' requires input mid-task, which isn't supported yet", self.name),
|
||||
}
|
||||
}
|
||||
|
||||
// Task is terminal (completed) — nothing left to cancel.
|
||||
guard.disarm();
|
||||
let result = self.request("tasks/result", json!({ "taskId": task_id })).await?;
|
||||
|
||||
if result["isError"].as_bool().unwrap_or(false) {
|
||||
anyhow::bail!("MCP tool error: {}", extract_text(&result));
|
||||
}
|
||||
Ok(crate::extract_call_result(&result))
|
||||
}
|
||||
|
||||
async fn request(&self, method: &str, params: Value) -> Result<Value> {
|
||||
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.pending.lock().await.insert(id, tx);
|
||||
|
||||
// Arm a cancellation guard for cancellable operations only (`tools/call`):
|
||||
// if this future is dropped by a `/stop` or times out before the server
|
||||
// replies, the guard notifies the server. Disarmed on a real reply.
|
||||
let mut cancel_guard = (method == "tools/call").then(|| CancelOnDrop {
|
||||
id,
|
||||
stdin: Arc::clone(&self.stdin),
|
||||
pending: Arc::clone(&self.pending),
|
||||
name: self.name.clone(),
|
||||
reason: "cancelled by client",
|
||||
armed: true,
|
||||
});
|
||||
|
||||
let msg = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"method": method,
|
||||
"params": params,
|
||||
});
|
||||
let mut line = serde_json::to_string(&msg)?;
|
||||
line.push('\n');
|
||||
self.stdin.lock().await.write_all(line.as_bytes()).await?;
|
||||
|
||||
// Wait for the response, but re-arm the timeout while an elicitation is
|
||||
// in flight on this server: the user may still be typing a secret into
|
||||
// the Inbox, and the server won't answer `tools/call` until then.
|
||||
tokio::pin!(rx);
|
||||
let response = loop {
|
||||
tokio::select! {
|
||||
res = &mut rx => match res {
|
||||
Ok(v) => break v,
|
||||
Err(_) => {
|
||||
// Server disconnected: nothing to cancel.
|
||||
if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
|
||||
anyhow::bail!("MCP '{}' disconnected", self.name);
|
||||
}
|
||||
},
|
||||
_ = tokio::time::sleep(Duration::from_secs(CALL_TIMEOUT_SECS)) => {
|
||||
if self.pending_elicitations.load(Ordering::SeqCst) == 0 {
|
||||
// Leave the guard armed so it fires `notifications/cancelled`;
|
||||
// tag the reason as a timeout.
|
||||
if let Some(g) = cancel_guard.as_mut() { g.reason = "timeout"; }
|
||||
anyhow::bail!("MCP '{}' timed out on '{method}'", self.name);
|
||||
}
|
||||
// Elicitation pending: keep waiting for the user.
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Server replied (even with an error): the request completed — disarm.
|
||||
if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
|
||||
|
||||
if let Some(error) = response.get("error") {
|
||||
anyhow::bail!("MCP '{}' protocol error: {error}", self.name);
|
||||
}
|
||||
Ok(response["result"].clone())
|
||||
}
|
||||
|
||||
async fn notify(&self, method: &str, params: Value) -> Result<()> {
|
||||
let msg = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": method,
|
||||
"params": params,
|
||||
});
|
||||
let mut line = serde_json::to_string(&msg)?;
|
||||
line.push('\n');
|
||||
self.stdin.lock().await.write_all(line.as_bytes()).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpServerClient for McpServer {
|
||||
fn tools(&self) -> &[McpTool] { self.tools() }
|
||||
async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> { self.call_tool(name, args).await }
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
//! End-to-end test of the elicitation path against a real stdio MCP subprocess.
|
||||
//!
|
||||
//! A tiny Python "server" exposes one tool that, when called, issues a
|
||||
//! server→client `elicitation/create` and echoes back whatever value it receives.
|
||||
//! This exercises the read-loop request branch, the spawned reply writer, the
|
||||
//! capability handshake, and `content` passthrough. Skipped if `python3` is absent.
|
||||
|
||||
use std::io::Write;
|
||||
use std::process::Command;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use mcp_client::config::{McpServerConfig, McpTransport};
|
||||
use mcp_client::server::{
|
||||
ElicitationAction, ElicitationHandler, ElicitationReply, ElicitationRequest, McpServer,
|
||||
};
|
||||
use mcp_client::McpCallResult;
|
||||
|
||||
const FAKE_SERVER: &str = r#"
|
||||
import sys, json
|
||||
|
||||
def send(obj):
|
||||
sys.stdout.write(json.dumps(obj) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def readline():
|
||||
line = sys.stdin.readline()
|
||||
if not line:
|
||||
return None
|
||||
line = line.strip()
|
||||
return line if line else readline()
|
||||
|
||||
while True:
|
||||
raw = readline()
|
||||
if raw is None:
|
||||
break
|
||||
msg = json.loads(raw)
|
||||
mid = msg.get("id")
|
||||
method = msg.get("method")
|
||||
if method == "initialize":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"protocolVersion": "2025-06-18", "capabilities": {},
|
||||
"serverInfo": {"name": "fake", "version": "0"}}})
|
||||
elif method == "notifications/initialized":
|
||||
pass
|
||||
elif method == "tools/list":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": [
|
||||
{"name": "need_secret", "description": "asks for a secret",
|
||||
"inputSchema": {"type": "object", "properties": {}}}]}})
|
||||
elif method == "tools/call":
|
||||
eid = "elicit-1"
|
||||
send({"jsonrpc": "2.0", "id": eid, "method": "elicitation/create", "params": {
|
||||
"message": "Enter password",
|
||||
"requestedSchema": {"type": "object", "properties": {
|
||||
"password": {"type": "string", "format": "password"}}}}})
|
||||
reply = None
|
||||
while reply is None:
|
||||
r = readline()
|
||||
if r is None:
|
||||
break
|
||||
rr = json.loads(r)
|
||||
if rr.get("id") == eid:
|
||||
reply = rr
|
||||
val = (reply or {}).get("result", {}).get("content", {}).get("password", "<none>")
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"content": [{"type": "text", "text": "got:" + val}], "isError": False}})
|
||||
elif mid is not None:
|
||||
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
|
||||
"#;
|
||||
|
||||
struct AcceptHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl ElicitationHandler for AcceptHandler {
|
||||
async fn handle(&self, _server: &str, req: ElicitationRequest) -> ElicitationReply {
|
||||
assert_eq!(req.message, "Enter password");
|
||||
assert!(req.requested_schema.get("properties").is_some());
|
||||
ElicitationReply {
|
||||
action: ElicitationAction::Accept,
|
||||
content: Some(json!({ "password": "hunter2" })),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn python3_available() -> bool {
|
||||
Command::new("python3").arg("--version").output().is_ok()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn elicitation_roundtrip_returns_secret_to_server() {
|
||||
if !python3_available() {
|
||||
eprintln!("python3 not found — skipping elicitation integration test");
|
||||
return;
|
||||
}
|
||||
|
||||
let script_path = std::env::temp_dir().join(format!("skald_elicit_{}.py", std::process::id()));
|
||||
std::fs::File::create(&script_path)
|
||||
.unwrap()
|
||||
.write_all(FAKE_SERVER.as_bytes())
|
||||
.unwrap();
|
||||
|
||||
let cfg = McpServerConfig {
|
||||
name: "fake".to_string(),
|
||||
transport: McpTransport::Stdio,
|
||||
command: Some("python3".to_string()),
|
||||
args: Some(vec![script_path.to_string_lossy().to_string()]),
|
||||
env: None,
|
||||
url: None,
|
||||
api_key: None,
|
||||
};
|
||||
|
||||
let server = McpServer::start(&cfg, None, None, Some(Arc::new(AcceptHandler)))
|
||||
.await
|
||||
.expect("server should start");
|
||||
|
||||
let result = server
|
||||
.call_tool("need_secret", json!({}))
|
||||
.await
|
||||
.expect("tool call should succeed");
|
||||
|
||||
// The fake server returns text content (no structuredContent) → Text variant.
|
||||
match result {
|
||||
McpCallResult::Text(s) => assert_eq!(s, "got:hunter2"),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
|
||||
let _ = std::fs::remove_file(&script_path);
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
//! End-to-end test of per-server diagnostic capture against a real stdio MCP
|
||||
//! subprocess.
|
||||
//!
|
||||
//! A tiny Python "server" prints a banner to **stderr** and, right after
|
||||
//! `notifications/initialized`, emits both a `notifications/message` log record and
|
||||
//! a business `event/ping` notification to stdout. The test asserts that:
|
||||
//! - the stderr banner arrives on `log_tx` tagged `stderr`,
|
||||
//! - the `notifications/message` arrives on `log_tx` with its MCP level, and is
|
||||
//! **not** delivered to `notification_tx` (it's diverted away from TIC),
|
||||
//! - the business `event/ping` still arrives on `notification_tx`.
|
||||
//! Skipped if `python3` is absent.
|
||||
|
||||
use std::io::Write;
|
||||
use std::process::Command;
|
||||
use std::time::Duration;
|
||||
|
||||
use mcp_client::config::{McpServerConfig, McpTransport};
|
||||
use mcp_client::server::{McpNotification, McpServer};
|
||||
use mcp_client::McpLogLine;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
const FAKE_SERVER: &str = r#"
|
||||
import sys, json
|
||||
|
||||
def send(obj):
|
||||
sys.stdout.write(json.dumps(obj) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def readline():
|
||||
line = sys.stdin.readline()
|
||||
if not line:
|
||||
return None
|
||||
line = line.strip()
|
||||
return line if line else readline()
|
||||
|
||||
while True:
|
||||
raw = readline()
|
||||
if raw is None:
|
||||
break
|
||||
msg = json.loads(raw)
|
||||
mid = msg.get("id")
|
||||
method = msg.get("method")
|
||||
if method == "initialize":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"protocolVersion": "2025-11-25", "capabilities": {},
|
||||
"serverInfo": {"name": "fake", "version": "0"}}})
|
||||
elif method == "notifications/initialized":
|
||||
# A diagnostic banner on stderr (the primary, future-proof log source).
|
||||
print("startup banner on stderr", file=sys.stderr, flush=True)
|
||||
# An MCP logging record (should be diverted to the log file, NOT TIC).
|
||||
send({"jsonrpc": "2.0", "method": "notifications/message",
|
||||
"params": {"level": "warning", "logger": "test", "data": "disk almost full"}})
|
||||
# A business event (should still reach the notification queue / TIC).
|
||||
send({"jsonrpc": "2.0", "method": "event/ping", "params": {"n": 1}})
|
||||
elif method == "tools/list":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": []}})
|
||||
elif mid is not None:
|
||||
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
|
||||
"#;
|
||||
|
||||
fn python3_available() -> bool {
|
||||
Command::new("python3").arg("--version").output().is_ok()
|
||||
}
|
||||
|
||||
fn write_script() -> std::path::PathBuf {
|
||||
let path = std::env::temp_dir().join(format!("skald_logging_{}.py", std::process::id()));
|
||||
std::fs::File::create(&path).unwrap().write_all(FAKE_SERVER.as_bytes()).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
/// Drains a channel for up to `budget`, returning everything received.
|
||||
async fn drain<T>(rx: &mut mpsc::UnboundedReceiver<T>, budget: Duration) -> Vec<T> {
|
||||
let mut out = Vec::new();
|
||||
let deadline = tokio::time::Instant::now() + budget;
|
||||
while let Ok(Some(item)) = tokio::time::timeout_at(deadline, rx.recv()).await {
|
||||
out.push(item);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stderr_and_log_records_are_captured_and_diverted() {
|
||||
if !python3_available() {
|
||||
eprintln!("python3 not found — skipping per-server logging integration test");
|
||||
return;
|
||||
}
|
||||
let script = write_script();
|
||||
let cfg = McpServerConfig {
|
||||
name: "fake".to_string(),
|
||||
transport: McpTransport::Stdio,
|
||||
command: Some("python3".to_string()),
|
||||
args: Some(vec![script.to_string_lossy().to_string()]),
|
||||
env: None,
|
||||
url: None,
|
||||
api_key: None,
|
||||
};
|
||||
|
||||
let (notif_tx, mut notif_rx) = mpsc::unbounded_channel::<McpNotification>();
|
||||
let (log_tx, mut log_rx) = mpsc::unbounded_channel::<McpLogLine>();
|
||||
|
||||
let _server = McpServer::start(&cfg, Some(notif_tx), Some(log_tx), None)
|
||||
.await
|
||||
.expect("server should start");
|
||||
|
||||
let logs: Vec<McpLogLine> = drain(&mut log_rx, Duration::from_millis(1500)).await;
|
||||
let notes: Vec<McpNotification> = drain(&mut notif_rx, Duration::from_millis(500)).await;
|
||||
|
||||
// stderr banner captured on the log channel.
|
||||
assert!(
|
||||
logs.iter().any(|l| l.level == "stderr" && l.text.contains("startup banner on stderr")),
|
||||
"expected a [stderr] log line, got: {logs:?}",
|
||||
);
|
||||
// notifications/message captured with its MCP level + logger prefix.
|
||||
assert!(
|
||||
logs.iter().any(|l| l.level == "warning" && l.text.contains("disk almost full") && l.text.contains("test")),
|
||||
"expected a [warning] log line from notifications/message, got: {logs:?}",
|
||||
);
|
||||
// The log record was diverted away from the notification queue…
|
||||
assert!(
|
||||
!notes.iter().any(|(_, msg)| msg.get("method").and_then(Value::as_str) == Some("notifications/message")),
|
||||
"notifications/message must NOT reach the notification queue, got: {notes:?}",
|
||||
);
|
||||
// …while the business event still flowed through it.
|
||||
assert!(
|
||||
notes.iter().any(|(_, msg)| msg.get("method").and_then(Value::as_str) == Some("event/ping")),
|
||||
"expected event/ping on the notification queue, got: {notes:?}",
|
||||
);
|
||||
|
||||
let _ = std::fs::remove_file(&script);
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
//! End-to-end test of `tools/list` cursor pagination against a real stdio MCP
|
||||
//! subprocess.
|
||||
//!
|
||||
//! A tiny Python "server" returns its tools across two pages: the first
|
||||
//! `tools/list` (no `cursor`) yields one tool plus a `nextCursor`; the follow-up
|
||||
//! (with that `cursor`) yields the rest and no `nextCursor`. This exercises the
|
||||
//! cursor loop in `McpServer::start`. Skipped if `python3` is absent.
|
||||
|
||||
use std::io::Write;
|
||||
use std::process::Command;
|
||||
|
||||
use mcp_client::config::{McpServerConfig, McpTransport};
|
||||
use mcp_client::server::McpServer;
|
||||
|
||||
const FAKE_SERVER: &str = r#"
|
||||
import sys, json
|
||||
|
||||
def send(obj):
|
||||
sys.stdout.write(json.dumps(obj) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def readline():
|
||||
line = sys.stdin.readline()
|
||||
if not line:
|
||||
return None
|
||||
line = line.strip()
|
||||
return line if line else readline()
|
||||
|
||||
while True:
|
||||
raw = readline()
|
||||
if raw is None:
|
||||
break
|
||||
msg = json.loads(raw)
|
||||
mid = msg.get("id")
|
||||
method = msg.get("method")
|
||||
if method == "initialize":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"protocolVersion": "2025-11-25", "capabilities": {},
|
||||
"serverInfo": {"name": "fake", "version": "0"}}})
|
||||
elif method == "notifications/initialized":
|
||||
pass
|
||||
elif method == "tools/list":
|
||||
cursor = (msg.get("params") or {}).get("cursor")
|
||||
if cursor is None:
|
||||
# Page 1: one tool + a cursor pointing at the next page.
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"tools": [{"name": "alpha", "description": "first",
|
||||
"inputSchema": {"type": "object", "properties": {}}}],
|
||||
"nextCursor": "page-2"}})
|
||||
elif cursor == "page-2":
|
||||
# Page 2: the rest, no further cursor → loop terminates.
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"tools": [{"name": "beta", "description": "second",
|
||||
"inputSchema": {"type": "object", "properties": {}}},
|
||||
{"name": "gamma", "description": "third",
|
||||
"inputSchema": {"type": "object", "properties": {}}}]}})
|
||||
else:
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": []}})
|
||||
elif mid is not None:
|
||||
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
|
||||
"#;
|
||||
|
||||
fn python3_available() -> bool {
|
||||
Command::new("python3").arg("--version").output().is_ok()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tools_list_follows_next_cursor_across_pages() {
|
||||
if !python3_available() {
|
||||
eprintln!("python3 not found — skipping pagination integration test");
|
||||
return;
|
||||
}
|
||||
|
||||
let script_path =
|
||||
std::env::temp_dir().join(format!("skald_paginate_{}.py", std::process::id()));
|
||||
std::fs::File::create(&script_path)
|
||||
.unwrap()
|
||||
.write_all(FAKE_SERVER.as_bytes())
|
||||
.unwrap();
|
||||
|
||||
let cfg = McpServerConfig {
|
||||
name: "fake".to_string(),
|
||||
transport: McpTransport::Stdio,
|
||||
command: Some("python3".to_string()),
|
||||
args: Some(vec![script_path.to_string_lossy().to_string()]),
|
||||
env: None,
|
||||
url: None,
|
||||
api_key: None,
|
||||
};
|
||||
|
||||
let server = McpServer::start(&cfg, None, None, None)
|
||||
.await
|
||||
.expect("server should start");
|
||||
|
||||
let names: Vec<&str> = server.tools().iter().map(|t| t.name.as_str()).collect();
|
||||
assert_eq!(
|
||||
names,
|
||||
vec!["alpha", "beta", "gamma"],
|
||||
"all pages should be collected"
|
||||
);
|
||||
|
||||
let _ = std::fs::remove_file(&script_path);
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
//! End-to-end test of the block-and-poll Tasks client against a real stdio MCP
|
||||
//! subprocess.
|
||||
//!
|
||||
//! A tiny Python "server" exposes one tool with `execution.taskSupport: "required"`.
|
||||
//! On `tools/call` it returns a `CreateTaskResult` (a durable handle, no `content`)
|
||||
//! instead of a result; the client then polls `tasks/get` until `completed` and
|
||||
//! fetches the real answer via `tasks/result`. A second mode never completes, so the
|
||||
//! test can drop the call mid-poll and assert the client sends `tasks/cancel`.
|
||||
//! Skipped if `python3` is absent.
|
||||
|
||||
use std::io::Write;
|
||||
use std::process::Command;
|
||||
use std::time::Duration;
|
||||
|
||||
use mcp_client::config::{McpServerConfig, McpTransport};
|
||||
use mcp_client::server::McpServer;
|
||||
use mcp_client::McpCallResult;
|
||||
|
||||
/// argv[1] = mode ("complete" | "cancel"); argv[2] = marker file (cancel mode).
|
||||
const FAKE_SERVER: &str = r#"
|
||||
import sys, json
|
||||
|
||||
mode = sys.argv[1] if len(sys.argv) > 1 else "complete"
|
||||
marker = sys.argv[2] if len(sys.argv) > 2 else None
|
||||
get_count = 0
|
||||
|
||||
def send(obj):
|
||||
sys.stdout.write(json.dumps(obj) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def readline():
|
||||
line = sys.stdin.readline()
|
||||
if not line:
|
||||
return None
|
||||
line = line.strip()
|
||||
return line if line else readline()
|
||||
|
||||
while True:
|
||||
raw = readline()
|
||||
if raw is None:
|
||||
break
|
||||
msg = json.loads(raw)
|
||||
mid = msg.get("id")
|
||||
method = msg.get("method")
|
||||
if method == "initialize":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"protocolVersion": "2025-11-25", "capabilities": {},
|
||||
"serverInfo": {"name": "fake", "version": "0"}}})
|
||||
elif method == "notifications/initialized":
|
||||
pass
|
||||
elif method == "tools/list":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": [
|
||||
{"name": "gen", "description": "deferred generator",
|
||||
"inputSchema": {"type": "object", "properties": {}},
|
||||
"execution": {"taskSupport": "required"}}]}})
|
||||
elif method == "tools/call":
|
||||
# Defer: return a task handle (no `content`).
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"taskId": "job-1", "status": "working", "pollInterval": 500}})
|
||||
elif method == "tasks/get":
|
||||
get_count += 1
|
||||
if mode == "complete" and get_count >= 2:
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"taskId": "job-1", "status": "completed"}})
|
||||
else:
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"taskId": "job-1", "status": "working", "pollInterval": 500}})
|
||||
elif method == "tasks/result":
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {
|
||||
"content": [{"type": "text", "text": "the real answer"}]}})
|
||||
elif method == "tasks/cancel":
|
||||
if marker:
|
||||
with open(marker, "w") as f:
|
||||
f.write((msg.get("params") or {}).get("taskId", ""))
|
||||
if mid is not None:
|
||||
send({"jsonrpc": "2.0", "id": mid, "result": {}})
|
||||
elif mid is not None:
|
||||
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
|
||||
"#;
|
||||
|
||||
fn python3_available() -> bool {
|
||||
Command::new("python3").arg("--version").output().is_ok()
|
||||
}
|
||||
|
||||
fn write_script() -> std::path::PathBuf {
|
||||
let path = std::env::temp_dir().join(format!("skald_tasks_{}.py", std::process::id()));
|
||||
std::fs::File::create(&path).unwrap().write_all(FAKE_SERVER.as_bytes()).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
fn cfg(script: &std::path::Path, mode: &str, marker: Option<&std::path::Path>) -> McpServerConfig {
|
||||
let mut args = vec![script.to_string_lossy().to_string(), mode.to_string()];
|
||||
if let Some(m) = marker {
|
||||
args.push(m.to_string_lossy().to_string());
|
||||
}
|
||||
McpServerConfig {
|
||||
name: "fake".to_string(),
|
||||
transport: McpTransport::Stdio,
|
||||
command: Some("python3".to_string()),
|
||||
args: Some(args),
|
||||
env: None,
|
||||
url: None,
|
||||
api_key: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_is_polled_to_completion_and_returns_real_result() {
|
||||
if !python3_available() {
|
||||
eprintln!("python3 not found — skipping tasks integration test");
|
||||
return;
|
||||
}
|
||||
let script = write_script();
|
||||
let server = McpServer::start(&cfg(&script, "complete", None), None, None, None)
|
||||
.await
|
||||
.expect("server should start");
|
||||
|
||||
// The tool opts into tasks (taskSupport: required); call_tool must add the
|
||||
// `task` field, poll to completion, and return the real result — not the handle.
|
||||
let result = server.call_tool("gen", serde_json::json!({}))
|
||||
.await
|
||||
.expect("task should complete");
|
||||
match result {
|
||||
McpCallResult::Text(t) => assert_eq!(t, "the real answer"),
|
||||
other => panic!("expected the polled Text result, got {other:?}"),
|
||||
}
|
||||
|
||||
let _ = std::fs::remove_file(&script);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropping_a_polling_call_sends_tasks_cancel() {
|
||||
if !python3_available() {
|
||||
eprintln!("python3 not found — skipping tasks cancel integration test");
|
||||
return;
|
||||
}
|
||||
let script = write_script();
|
||||
let marker = std::env::temp_dir().join(format!("skald_tasks_marker_{}", std::process::id()));
|
||||
let _ = std::fs::remove_file(&marker);
|
||||
|
||||
let server = McpServer::start(&cfg(&script, "cancel", Some(&marker)), None, None, None)
|
||||
.await
|
||||
.expect("server should start");
|
||||
|
||||
// In "cancel" mode tasks/get never completes, so the call polls forever; time out
|
||||
// to drop the future mid-poll, which must fire tasks/cancel via the drop-guard.
|
||||
let timed = tokio::time::timeout(
|
||||
Duration::from_millis(900),
|
||||
server.call_tool("gen", serde_json::json!({})),
|
||||
).await;
|
||||
assert!(timed.is_err(), "call should still be polling (timed out), not resolved");
|
||||
|
||||
// Give the guard's spawned tasks/cancel time to reach the server.
|
||||
tokio::time::sleep(Duration::from_millis(700)).await;
|
||||
let recorded = std::fs::read_to_string(&marker).unwrap_or_default();
|
||||
assert_eq!(recorded, "job-1", "server should have received tasks/cancel for job-1");
|
||||
|
||||
let _ = std::fs::remove_file(&script);
|
||||
let _ = std::fs::remove_file(&marker);
|
||||
}
|
||||
Reference in New Issue
Block a user