Update: 将子项目从 submodule 转为完整内容
- 移除 GovAI, nomifun-tauri, 算力盒子 的 submodule 引用 - 添加所有子项目的完整源代码 - 保留原始 .git 为 .git.bak 备份
This commit is contained in:
freedak
@@ -0,0 +1,2 @@
|
||||
// Re-export MCP configuration types from nomi-config to avoid duplication.
|
||||
pub use nomi_config::config::{McpConfig, McpServerConfig, TransportType};
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod config;
|
||||
pub mod manager;
|
||||
pub mod protocol;
|
||||
pub mod tool_proxy;
|
||||
pub mod transport;
|
||||
@@ -0,0 +1,767 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::config::{McpServerConfig, TransportType};
|
||||
use super::protocol::{
|
||||
ClientCapabilities, ClientInfo, InitializeParams, InitializeResult, JsonRpcRequest,
|
||||
McpResource, McpToolDef, McpToolResult, ResourcesListResult, ResourcesReadResult,
|
||||
ToolsListResult,
|
||||
};
|
||||
use super::transport::sse::SseTransport;
|
||||
use super::transport::stdio::StdioTransport;
|
||||
use super::transport::streamable_http::StreamableHttpTransport;
|
||||
use super::transport::{McpError, McpTransport};
|
||||
|
||||
/// Structured result of an MCP tool call: text and images kept separate so
|
||||
/// `McpToolProxy` can feed images back into the multimodal context instead of
|
||||
/// flattening them into a `[image: mime]` literal (which discarded the base64).
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct McpCallOutput {
|
||||
/// All text content joined with `\n` (preserves the pre-existing behaviour).
|
||||
pub text: String,
|
||||
/// Image content `(base64 data, mime_type)` in order of appearance.
|
||||
pub images: Vec<McpImageOut>,
|
||||
}
|
||||
|
||||
/// A single image returned by an MCP tool call.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct McpImageOut {
|
||||
/// Raw base64 (straight from the server — not re-encoded).
|
||||
pub data: String,
|
||||
/// MIME type, e.g. "image/png".
|
||||
pub mime_type: String,
|
||||
}
|
||||
|
||||
/// A connected MCP server with its discovered tools and capabilities
|
||||
struct McpServer {
|
||||
#[allow(dead_code)]
|
||||
name: String,
|
||||
transport: Box<dyn McpTransport>,
|
||||
tools: Vec<McpToolDef>,
|
||||
/// Whether the server declared resources capability in its initialize response
|
||||
supports_resources: bool,
|
||||
}
|
||||
|
||||
/// Manages connections to multiple MCP servers
|
||||
pub struct McpManager {
|
||||
servers: HashMap<String, McpServer>,
|
||||
/// Monotonically increasing request ID counter for all JSON-RPC calls
|
||||
next_id: AtomicU64,
|
||||
}
|
||||
|
||||
/// Timeout for connecting + initializing a single MCP server (transport spawn,
|
||||
/// `initialize` handshake, `tools/list`). Without it, a server that starts but
|
||||
/// never answers the handshake would hang the entire agent bootstrap — and thus
|
||||
/// any solo nomi conversation that injects the guide MCP — indefinitely, with no
|
||||
/// error surfaced to the user.
|
||||
const MCP_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
impl McpManager {
|
||||
/// Connect to all configured MCP servers
|
||||
pub async fn connect_all(configs: &HashMap<String, McpServerConfig>) -> Result<Self, McpError> {
|
||||
let mut servers = HashMap::new();
|
||||
|
||||
for (name, config) in configs {
|
||||
match tokio::time::timeout(MCP_CONNECT_TIMEOUT, Self::connect_server(name, config)).await {
|
||||
Ok(Ok(server)) => {
|
||||
tracing::info!(target: "nomi_mcp", server = %name, tools = server.tools.len(), resources = server.supports_resources, "mcp server connected");
|
||||
servers.insert(name.clone(), server);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
// Non-fatal: continue with other servers
|
||||
tracing::warn!(target: "nomi_mcp", server = %name, error = %e, "mcp server connection failed");
|
||||
}
|
||||
Err(_) => {
|
||||
// Non-fatal: a hung handshake must not block the other
|
||||
// servers or the agent bootstrap. Skip this server.
|
||||
tracing::warn!(target: "nomi_mcp", server = %name, timeout_secs = MCP_CONNECT_TIMEOUT.as_secs(), "mcp server connection timed out");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
servers,
|
||||
next_id: AtomicU64::new(10),
|
||||
})
|
||||
}
|
||||
|
||||
/// Connect a single additional MCP server after initial setup.
|
||||
/// Returns the list of tool names exposed by the server.
|
||||
pub async fn connect_one(
|
||||
&mut self,
|
||||
name: String,
|
||||
config: &McpServerConfig,
|
||||
) -> Result<Vec<String>, McpError> {
|
||||
let server = match tokio::time::timeout(MCP_CONNECT_TIMEOUT, Self::connect_server(&name, config)).await {
|
||||
Ok(result) => result?,
|
||||
Err(_) => {
|
||||
return Err(McpError::InitFailed(format!(
|
||||
"MCP server '{name}' connection timed out after {}s",
|
||||
MCP_CONNECT_TIMEOUT.as_secs()
|
||||
)));
|
||||
}
|
||||
};
|
||||
let tool_names: Vec<String> = server.tools.iter().map(|t| t.name.clone()).collect();
|
||||
tracing::info!(target: "nomi_mcp", server = %name, tools = server.tools.len(), resources = server.supports_resources, "mcp server connected");
|
||||
self.servers.insert(name, server);
|
||||
Ok(tool_names)
|
||||
}
|
||||
|
||||
/// Connect to a single MCP server: create transport, initialize, discover tools
|
||||
async fn connect_server(name: &str, config: &McpServerConfig) -> Result<McpServer, McpError> {
|
||||
let empty_map = HashMap::new();
|
||||
|
||||
// 1. Create transport
|
||||
let transport: Box<dyn McpTransport> = match config.transport {
|
||||
TransportType::Stdio => {
|
||||
let command = config.command.as_deref().ok_or_else(|| {
|
||||
McpError::InitFailed("stdio transport requires 'command'".into())
|
||||
})?;
|
||||
let args = config.args.as_deref().unwrap_or(&[]);
|
||||
let env = config.env.as_ref().unwrap_or(&empty_map);
|
||||
Box::new(StdioTransport::spawn(command, args, env).await?)
|
||||
}
|
||||
TransportType::Sse => {
|
||||
let url = config
|
||||
.url
|
||||
.as_deref()
|
||||
.ok_or_else(|| McpError::InitFailed("SSE transport requires 'url'".into()))?;
|
||||
let headers = config.headers.as_ref().unwrap_or(&empty_map);
|
||||
Box::new(SseTransport::connect(url, headers).await?)
|
||||
}
|
||||
TransportType::StreamableHttp => {
|
||||
let url = config.url.as_deref().ok_or_else(|| {
|
||||
McpError::InitFailed("streamable-http transport requires 'url'".into())
|
||||
})?;
|
||||
let headers = config.headers.as_ref().unwrap_or(&empty_map);
|
||||
Box::new(StreamableHttpTransport::connect(url, headers).await?)
|
||||
}
|
||||
};
|
||||
|
||||
// 2. Initialize handshake
|
||||
let init_params = InitializeParams {
|
||||
protocol_version: "2025-03-26".to_string(),
|
||||
capabilities: ClientCapabilities {
|
||||
tools: Some(json!({})),
|
||||
},
|
||||
client_info: ClientInfo {
|
||||
name: "nomi".to_string(),
|
||||
version: "0.3.0".to_string(),
|
||||
},
|
||||
};
|
||||
|
||||
let init_req = JsonRpcRequest::new(
|
||||
1,
|
||||
"initialize",
|
||||
Some(serde_json::to_value(&init_params).map_err(|e| {
|
||||
McpError::InitFailed(format!("Failed to serialize init params: {}", e))
|
||||
})?),
|
||||
);
|
||||
|
||||
let init_response = transport.request(&init_req).await?;
|
||||
let init_result: InitializeResult = serde_json::from_value(
|
||||
init_response
|
||||
.result
|
||||
.ok_or_else(|| McpError::InitFailed("No result in initialize response".into()))?,
|
||||
)
|
||||
.map_err(|e| McpError::InitFailed(format!("Failed to parse init result: {}", e)))?;
|
||||
|
||||
// Check whether server declared resources capability
|
||||
let supports_resources = init_result
|
||||
.capabilities
|
||||
.get("resources")
|
||||
.map(|v| !v.is_null())
|
||||
.unwrap_or(false);
|
||||
|
||||
// 3. Send initialized notification
|
||||
let initialized_notification =
|
||||
JsonRpcRequest::notification("notifications/initialized", None);
|
||||
transport.notify(&initialized_notification).await?;
|
||||
|
||||
// 4. List tools
|
||||
let list_req = JsonRpcRequest::new(2, "tools/list", None);
|
||||
let list_response = transport.request(&list_req).await?;
|
||||
let tools_result: ToolsListResult = serde_json::from_value(
|
||||
list_response
|
||||
.result
|
||||
.ok_or_else(|| McpError::InitFailed("No result in tools/list response".into()))?,
|
||||
)
|
||||
.map_err(|e| McpError::InitFailed(format!("Failed to parse tools list: {}", e)))?;
|
||||
|
||||
Ok(McpServer {
|
||||
name: name.to_string(),
|
||||
transport,
|
||||
tools: tools_result.tools,
|
||||
supports_resources,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get all discovered tools with their server names
|
||||
pub fn all_tools(&self) -> Vec<(&str, &McpToolDef)> {
|
||||
let mut result = Vec::new();
|
||||
for (server_name, server) in &self.servers {
|
||||
for tool in &server.tools {
|
||||
result.push((server_name.as_str(), tool));
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Check if a tool name exists across any server
|
||||
pub fn has_tool_name(&self, name: &str) -> bool {
|
||||
self.servers
|
||||
.values()
|
||||
.any(|s| s.tools.iter().any(|t| t.name == name))
|
||||
}
|
||||
|
||||
/// Count how many servers have a tool with the given name
|
||||
pub fn tool_name_count(&self, name: &str) -> usize {
|
||||
self.servers
|
||||
.values()
|
||||
.filter(|s| s.tools.iter().any(|t| t.name == name))
|
||||
.count()
|
||||
}
|
||||
|
||||
/// Execute a tool on a specific server.
|
||||
///
|
||||
/// Returns a structured [`McpCallOutput`] keeping text and image content
|
||||
/// separate so the proxy can feed screenshots into the multimodal context.
|
||||
/// Image base64 `data` is preserved verbatim (previously it was discarded
|
||||
/// and replaced with a `[image: mime]` text placeholder).
|
||||
pub async fn call_tool(
|
||||
&self,
|
||||
server_name: &str,
|
||||
tool_name: &str,
|
||||
arguments: serde_json::Value,
|
||||
) -> Result<McpCallOutput, McpError> {
|
||||
let server = self
|
||||
.servers
|
||||
.get(server_name)
|
||||
.ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;
|
||||
|
||||
let request = JsonRpcRequest::new(
|
||||
0, // id doesn't matter for stdio, will be used for SSE/HTTP
|
||||
"tools/call",
|
||||
Some(json!({
|
||||
"name": tool_name,
|
||||
"arguments": arguments
|
||||
})),
|
||||
);
|
||||
|
||||
let response = server.transport.request(&request).await?;
|
||||
|
||||
let result_value = response
|
||||
.result
|
||||
.ok_or_else(|| McpError::Transport("No result in tool call response".into()))?;
|
||||
|
||||
// Parse result, keeping text and image content separate.
|
||||
let tool_result: McpToolResult = serde_json::from_value(result_value)
|
||||
.map_err(|e| McpError::Transport(format!("Failed to parse tool result: {}", e)))?;
|
||||
|
||||
let mut out = McpCallOutput::default();
|
||||
let mut text_parts: Vec<String> = Vec::new();
|
||||
for content in &tool_result.content {
|
||||
match content {
|
||||
super::protocol::McpContent::Text { text } => text_parts.push(text.clone()),
|
||||
super::protocol::McpContent::Image { data, mime_type } => {
|
||||
// Preserve the base64 so the proxy can route it into
|
||||
// ToolResult.images; no longer flattened to a text literal.
|
||||
out.images.push(McpImageOut {
|
||||
data: data.clone(),
|
||||
mime_type: mime_type.clone(),
|
||||
});
|
||||
}
|
||||
super::protocol::McpContent::Resource { .. } => {
|
||||
text_parts.push("[resource]".to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
out.text = text_parts.join("\n");
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Get names of all connected servers.
|
||||
pub fn server_names(&self) -> Vec<String> {
|
||||
self.servers.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Check if a connected server declared the resources capability.
|
||||
pub fn server_supports_resources(&self, server_name: &str) -> bool {
|
||||
self.servers
|
||||
.get(server_name)
|
||||
.map(|s| s.supports_resources)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// List all resources from a server.
|
||||
pub async fn list_resources(&self, server_name: &str) -> Result<Vec<McpResource>, McpError> {
|
||||
let server = self
|
||||
.servers
|
||||
.get(server_name)
|
||||
.ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;
|
||||
|
||||
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
||||
let request = JsonRpcRequest::new(id, "resources/list", None);
|
||||
let response = server.transport.request(&request).await?;
|
||||
|
||||
let result_value = response
|
||||
.result
|
||||
.ok_or_else(|| McpError::Transport("No result in resources/list response".into()))?;
|
||||
|
||||
let list_result: ResourcesListResult = serde_json::from_value(result_value)
|
||||
.map_err(|e| McpError::Transport(format!("Failed to parse resources/list: {}", e)))?;
|
||||
|
||||
Ok(list_result.resources)
|
||||
}
|
||||
|
||||
/// Read a single resource by URI from a server. Returns the text content.
|
||||
pub async fn read_resource(&self, server_name: &str, uri: &str) -> Result<String, McpError> {
|
||||
let server = self
|
||||
.servers
|
||||
.get(server_name)
|
||||
.ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;
|
||||
|
||||
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
||||
let request = JsonRpcRequest::new(id, "resources/read", Some(json!({ "uri": uri })));
|
||||
let response = server.transport.request(&request).await?;
|
||||
|
||||
let result_value = response
|
||||
.result
|
||||
.ok_or_else(|| McpError::Transport("No result in resources/read response".into()))?;
|
||||
|
||||
let read_result: ResourcesReadResult = serde_json::from_value(result_value)
|
||||
.map_err(|e| McpError::Transport(format!("Failed to parse resources/read: {}", e)))?;
|
||||
|
||||
// Return the first text content found
|
||||
read_result
|
||||
.contents
|
||||
.into_iter()
|
||||
.find_map(|c| c.text)
|
||||
.ok_or_else(|| McpError::Transport(format!("No text content in resource '{}'", uri)))
|
||||
}
|
||||
|
||||
/// Gracefully shutdown all servers
|
||||
pub async fn shutdown(&self) {
|
||||
for (name, server) in &self.servers {
|
||||
if let Err(e) = server.transport.close().await {
|
||||
tracing::warn!(target: "nomi_mcp", server = %name, error = %e, "error closing mcp server");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Test-only constructor: build a manager from pre-configured servers.
|
||||
#[cfg(any(test, feature = "test-utils"))]
|
||||
pub fn new_for_test(
|
||||
entries: Vec<(&str, bool, Box<dyn super::transport::McpTransport>)>,
|
||||
) -> Self {
|
||||
let mut servers = HashMap::new();
|
||||
for (name, supports_resources, transport) in entries {
|
||||
servers.insert(
|
||||
name.to_string(),
|
||||
McpServer {
|
||||
name: name.to_string(),
|
||||
transport,
|
||||
tools: vec![],
|
||||
supports_resources,
|
||||
},
|
||||
);
|
||||
}
|
||||
Self {
|
||||
servers,
|
||||
next_id: AtomicU64::new(10),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::protocol::JsonRpcResponse;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Mutex;
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// MockTransport: returns pre-configured JSON-RPC responses
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
struct MockTransport {
|
||||
/// Responses returned in order for each request call
|
||||
responses: Mutex<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl MockTransport {
|
||||
fn new(responses: Vec<serde_json::Value>) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(responses),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for MockTransport {
|
||||
async fn request(&self, _req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError> {
|
||||
let mut guard = self.responses.lock().unwrap();
|
||||
let value = if guard.is_empty() {
|
||||
json!(null)
|
||||
} else {
|
||||
guard.remove(0)
|
||||
};
|
||||
Ok(JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: Some(1),
|
||||
result: Some(value),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn notify(&self, _req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
struct ErrorTransport;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for ErrorTransport {
|
||||
async fn request(&self, _req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError> {
|
||||
Err(McpError::Transport("mock transport error".into()))
|
||||
}
|
||||
|
||||
async fn notify(&self, _req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Test helpers: build McpManager with pre-configured servers
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
fn make_manager_with_servers(entries: Vec<(&str, bool, Box<dyn McpTransport>)>) -> McpManager {
|
||||
McpManager::new_for_test(entries)
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-2.x: server_supports_resources [黑盒 + 白盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn tc_2_1_server_supports_resources_true() {
|
||||
// [黑盒] TC-2.1: server with resources capability returns true
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"test-server",
|
||||
true,
|
||||
Box::new(MockTransport::new(vec![])),
|
||||
)]);
|
||||
|
||||
assert!(manager.server_supports_resources("test-server"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_2_2_server_supports_resources_false() {
|
||||
// [黑盒] TC-2.2: server without resources capability returns false
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"no-resources-server",
|
||||
false,
|
||||
Box::new(MockTransport::new(vec![])),
|
||||
)]);
|
||||
|
||||
assert!(!manager.server_supports_resources("no-resources-server"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_2_3_server_supports_resources_unknown_server() {
|
||||
// [黑盒] TC-2.3: unknown server name returns false (not error)
|
||||
let manager = make_manager_with_servers(vec![]);
|
||||
|
||||
assert!(!manager.server_supports_resources("unknown-server"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_2_wb_supports_resources_from_capabilities_null_value() {
|
||||
// [白盒] capabilities.get("resources") = null → supports_resources = false
|
||||
// This is tested via the parsed field; we verify via make_manager helper
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"server",
|
||||
false, // null resources → false per impl: !v.is_null() = false
|
||||
Box::new(MockTransport::new(vec![])),
|
||||
)]);
|
||||
|
||||
assert!(!manager.server_supports_resources("server"));
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-2.10/2.11: server_names [黑盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn tc_2_10_server_names_returns_all() {
|
||||
// [黑盒] TC-2.10: server_names returns all connected server names
|
||||
let manager = make_manager_with_servers(vec![
|
||||
("server-a", false, Box::new(MockTransport::new(vec![]))),
|
||||
("server-b", true, Box::new(MockTransport::new(vec![]))),
|
||||
]);
|
||||
|
||||
let mut names = manager.server_names();
|
||||
names.sort();
|
||||
assert_eq!(names, vec!["server-a", "server-b"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_2_11_server_names_empty_manager() {
|
||||
// [黑盒] TC-2.11: no connected servers → empty vec
|
||||
let manager = make_manager_with_servers(vec![]);
|
||||
|
||||
assert!(manager.server_names().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_2_wb_server_names_returns_owned_strings() {
|
||||
// [白盒] Decision 1: server_names() returns Vec<String> not Vec<&str>
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"my-server",
|
||||
false,
|
||||
Box::new(MockTransport::new(vec![])),
|
||||
)]);
|
||||
|
||||
let names: Vec<String> = manager.server_names();
|
||||
assert_eq!(names, vec!["my-server"]);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-2.4/2.5: list_resources [黑盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_4_list_resources_normal() {
|
||||
// [黑盒] TC-2.4: list_resources returns resources from server
|
||||
let resources_response = json!({
|
||||
"resources": [
|
||||
{"uri": "skill://skill-a"},
|
||||
{"uri": "skill://skill-b", "name": "Skill B"}
|
||||
]
|
||||
});
|
||||
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"test-server",
|
||||
true,
|
||||
Box::new(MockTransport::new(vec![resources_response])),
|
||||
)]);
|
||||
|
||||
let result = manager.list_resources("test-server").await.unwrap();
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0].uri, "skill://skill-a");
|
||||
assert_eq!(result[1].uri, "skill://skill-b");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_5_list_resources_empty() {
|
||||
// [黑盒] TC-2.5: list_resources returns empty list when server has no resources
|
||||
let resources_response = json!({"resources": []});
|
||||
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"test-server",
|
||||
true,
|
||||
Box::new(MockTransport::new(vec![resources_response])),
|
||||
)]);
|
||||
|
||||
let result = manager.list_resources("test-server").await.unwrap();
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_6_list_resources_server_not_found() {
|
||||
// [黑盒] TC-2.6: list_resources returns error when server does not exist
|
||||
let manager = make_manager_with_servers(vec![]);
|
||||
|
||||
let result = manager.list_resources("nonexistent").await;
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
McpError::ServerNotFound(name) => assert_eq!(name, "nonexistent"),
|
||||
e => panic!("expected ServerNotFound, got {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-2.7/2.8/2.9: read_resource [黑盒 + 白盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_7_read_resource_returns_text() {
|
||||
// [黑盒] TC-2.7: read_resource returns text content
|
||||
let read_response = json!({
|
||||
"contents": [{"uri": "skill://my-skill", "mimeType": "text/plain", "text": "---\ndescription: A skill\n---\n# My Skill\n"}]
|
||||
});
|
||||
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"test-server",
|
||||
true,
|
||||
Box::new(MockTransport::new(vec![read_response])),
|
||||
)]);
|
||||
|
||||
let result = manager
|
||||
.read_resource("test-server", "skill://my-skill")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.contains("description: A skill"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_8_read_resource_transport_error() {
|
||||
// [黑盒] TC-2.8: read_resource returns error when server returns transport error
|
||||
let manager =
|
||||
make_manager_with_servers(vec![("test-server", true, Box::new(ErrorTransport))]);
|
||||
|
||||
let result = manager
|
||||
.read_resource("test-server", "skill://nonexistent")
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_9_read_resource_server_not_found() {
|
||||
// [黑盒] TC-2.9: read_resource returns error when server does not exist
|
||||
let manager = make_manager_with_servers(vec![]);
|
||||
|
||||
let result = manager
|
||||
.read_resource("nonexistent", "skill://my-skill")
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
McpError::ServerNotFound(name) => assert_eq!(name, "nonexistent"),
|
||||
e => panic!("expected ServerNotFound, got {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_wb_read_resource_no_text_content_returns_error() {
|
||||
// [白盒] Decision 3: find_map returns None when all contents have text=None → error
|
||||
let read_response = json!({
|
||||
"contents": [{"uri": "skill://binary", "mimeType": "application/octet-stream"}]
|
||||
});
|
||||
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"test-server",
|
||||
true,
|
||||
Box::new(MockTransport::new(vec![read_response])),
|
||||
)]);
|
||||
|
||||
let result = manager.read_resource("test-server", "skill://binary").await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tc_2_wb_read_resource_find_map_first_text() {
|
||||
// [白盒] Decision 3: find_map returns first content with non-None text
|
||||
let read_response = json!({
|
||||
"contents": [
|
||||
{"uri": "skill://x"},
|
||||
{"uri": "skill://x", "text": "actual content"}
|
||||
]
|
||||
});
|
||||
|
||||
let manager = make_manager_with_servers(vec![(
|
||||
"test-server",
|
||||
true,
|
||||
Box::new(MockTransport::new(vec![read_response])),
|
||||
)]);
|
||||
|
||||
let result = manager
|
||||
.read_resource("test-server", "skill://x")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result, "actual content");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_2_wb_next_id_starts_at_10() {
|
||||
// [白盒] Decision 4: AtomicU64 counter starts at 10 to avoid conflict with connect_server IDs 1/2
|
||||
let manager = make_manager_with_servers(vec![]);
|
||||
// next_id is private — we verify by doing two fetch_adds and checking values are 10 and 11
|
||||
let id1 = manager
|
||||
.next_id
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
let id2 = manager
|
||||
.next_id
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
assert_eq!(id1, 10, "first ID should be 10");
|
||||
assert_eq!(id2, 11, "second ID should be 11");
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// call_tool: image content passthrough into McpCallOutput
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_tool_preserves_image_data_and_mime() {
|
||||
// tools/call response: one text block + one png image
|
||||
let resp = json!({ "content": [
|
||||
{"type":"text","text":"done"},
|
||||
{"type":"image","data":"AAAAbase64==","mimeType":"image/png"}
|
||||
]});
|
||||
let mgr = make_manager_with_servers(vec![(
|
||||
"srv",
|
||||
false,
|
||||
Box::new(MockTransport::new(vec![resp])),
|
||||
)]);
|
||||
let out = mgr.call_tool("srv", "shot", json!({})).await.unwrap();
|
||||
assert_eq!(out.text, "done");
|
||||
assert_eq!(out.images.len(), 1);
|
||||
assert_eq!(out.images[0].data, "AAAAbase64==");
|
||||
assert_eq!(out.images[0].mime_type, "image/png");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_tool_text_only_has_no_images() {
|
||||
// Pure-text result must not regress: text preserved, no images.
|
||||
let resp = json!({ "content": [{"type":"text","text":"hello"}] });
|
||||
let mgr = make_manager_with_servers(vec![(
|
||||
"srv",
|
||||
false,
|
||||
Box::new(MockTransport::new(vec![resp])),
|
||||
)]);
|
||||
let out = mgr.call_tool("srv", "echo", json!({})).await.unwrap();
|
||||
assert_eq!(out.text, "hello");
|
||||
assert!(out.images.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_tool_multi_image_and_resource() {
|
||||
// Multiple images interleaved with text + a resource placeholder.
|
||||
let resp = json!({ "content": [
|
||||
{"type":"image","data":"img1","mimeType":"image/png"},
|
||||
{"type":"text","text":"between"},
|
||||
{"type":"image","data":"img2","mimeType":"image/jpeg"},
|
||||
{"type":"resource","resource":{"uri":"x://y"}}
|
||||
]});
|
||||
let mgr = make_manager_with_servers(vec![(
|
||||
"srv",
|
||||
false,
|
||||
Box::new(MockTransport::new(vec![resp])),
|
||||
)]);
|
||||
let out = mgr.call_tool("srv", "t", json!({})).await.unwrap();
|
||||
assert_eq!(out.images.len(), 2);
|
||||
assert_eq!(out.images[0].mime_type, "image/png");
|
||||
assert_eq!(out.images[1].mime_type, "image/jpeg");
|
||||
// text holds the text block and the [resource] placeholder, no [image:..]
|
||||
assert!(out.text.contains("between"));
|
||||
assert!(out.text.contains("[resource]"));
|
||||
assert!(!out.text.contains("[image"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,521 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// JSON-RPC 2.0 request
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct JsonRpcRequest {
|
||||
pub jsonrpc: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub id: Option<u64>,
|
||||
pub method: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub params: Option<Value>,
|
||||
}
|
||||
|
||||
impl JsonRpcRequest {
|
||||
pub fn new(id: u64, method: &str, params: Option<Value>) -> Self {
|
||||
Self {
|
||||
jsonrpc: "2.0",
|
||||
id: Some(id),
|
||||
method: method.to_string(),
|
||||
params,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn notification(method: &str, params: Option<Value>) -> Self {
|
||||
Self {
|
||||
jsonrpc: "2.0",
|
||||
id: None,
|
||||
method: method.to_string(),
|
||||
params,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 response
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct JsonRpcResponse {
|
||||
#[allow(dead_code)]
|
||||
pub jsonrpc: String,
|
||||
pub id: Option<u64>,
|
||||
pub result: Option<Value>,
|
||||
pub error: Option<JsonRpcError>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct JsonRpcError {
|
||||
pub code: i64,
|
||||
pub message: String,
|
||||
#[allow(dead_code)]
|
||||
pub data: Option<Value>,
|
||||
}
|
||||
|
||||
/// MCP tool definition returned by tools/list
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct McpToolDef {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
#[serde(rename = "inputSchema")]
|
||||
pub input_schema: Value,
|
||||
/// Optional behaviour hints declared by the server (MCP `annotations`).
|
||||
/// Drives approval classification (see `McpToolProxy::category`). Absent on
|
||||
/// servers that predate the annotations field — `None` is then treated as
|
||||
/// "no hints", i.e. the from-strict default (approval required).
|
||||
#[serde(default)]
|
||||
pub annotations: Option<ToolAnnotations>,
|
||||
}
|
||||
|
||||
/// MCP `ToolAnnotations` — behaviour hints a server may attach to each tool.
|
||||
///
|
||||
/// All fields are advisory `Option<bool>` per the MCP spec (camelCase on the
|
||||
/// wire). We only act on `read_only_hint` / `destructive_hint` for approval
|
||||
/// gating today, but parse and retain the full set so future policy (and the
|
||||
/// human-readable `title`) is available without another protocol change.
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
pub struct ToolAnnotations {
|
||||
/// Human-readable title for the tool (display only).
|
||||
#[serde(default)]
|
||||
pub title: Option<String>,
|
||||
/// If true, the tool does not modify its environment.
|
||||
#[serde(rename = "readOnlyHint", default)]
|
||||
pub read_only_hint: Option<bool>,
|
||||
/// If true, the tool may perform destructive updates (only meaningful when
|
||||
/// not read-only).
|
||||
#[serde(rename = "destructiveHint", default)]
|
||||
pub destructive_hint: Option<bool>,
|
||||
/// If true, repeated calls with the same args have no additional effect.
|
||||
#[serde(rename = "idempotentHint", default)]
|
||||
pub idempotent_hint: Option<bool>,
|
||||
/// If true, the tool may interact with an "open world" of external entities.
|
||||
#[serde(rename = "openWorldHint", default)]
|
||||
pub open_world_hint: Option<bool>,
|
||||
}
|
||||
|
||||
/// MCP tool call result
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct McpToolResult {
|
||||
pub content: Vec<McpContent>,
|
||||
}
|
||||
|
||||
/// Content types returned by MCP tool calls
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum McpContent {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "image")]
|
||||
Image {
|
||||
data: String,
|
||||
#[serde(rename = "mimeType")]
|
||||
mime_type: String,
|
||||
},
|
||||
#[serde(rename = "resource")]
|
||||
Resource {
|
||||
#[allow(dead_code)]
|
||||
resource: Value,
|
||||
},
|
||||
}
|
||||
|
||||
/// Initialize request params
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct InitializeParams {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
pub protocol_version: String,
|
||||
pub capabilities: ClientCapabilities,
|
||||
#[serde(rename = "clientInfo")]
|
||||
pub client_info: ClientInfo,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ClientCapabilities {
|
||||
pub tools: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ClientInfo {
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
}
|
||||
|
||||
/// Initialize response result
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct InitializeResult {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
#[allow(dead_code)]
|
||||
pub protocol_version: String,
|
||||
#[allow(dead_code)]
|
||||
pub capabilities: Value,
|
||||
#[serde(rename = "serverInfo")]
|
||||
#[allow(dead_code)]
|
||||
pub server_info: Option<Value>,
|
||||
}
|
||||
|
||||
/// Tools list response
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ToolsListResult {
|
||||
pub tools: Vec<McpToolDef>,
|
||||
}
|
||||
|
||||
/// MCP resource definition returned by resources/list
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct McpResource {
|
||||
pub uri: String,
|
||||
pub name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
#[serde(rename = "mimeType", default)]
|
||||
pub mime_type: Option<String>,
|
||||
}
|
||||
|
||||
/// resources/list response
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ResourcesListResult {
|
||||
pub resources: Vec<McpResource>,
|
||||
}
|
||||
|
||||
/// resources/read response
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ResourcesReadResult {
|
||||
pub contents: Vec<ResourceContent>,
|
||||
}
|
||||
|
||||
/// Content of a single resource from resources/read
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ResourceContent {
|
||||
#[allow(dead_code)]
|
||||
pub uri: String,
|
||||
#[serde(rename = "mimeType", default)]
|
||||
pub mime_type: Option<String>,
|
||||
/// Text content — None for blob resources (binary); skill resources are always text
|
||||
#[serde(default)]
|
||||
pub text: Option<String>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_jsonrpc_request_serialization() {
|
||||
// Verify that a regular request serializes with jsonrpc, id, method and params
|
||||
let req = JsonRpcRequest::new(1, "tools/list", Some(json!({"cursor": null})));
|
||||
let value = serde_json::to_value(&req).unwrap();
|
||||
|
||||
assert_eq!(value["jsonrpc"], "2.0");
|
||||
assert_eq!(value["id"], 1u64);
|
||||
assert_eq!(value["method"], "tools/list");
|
||||
assert!(value.get("params").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jsonrpc_request_notification() {
|
||||
// Notifications must not include the "id" field when serialized
|
||||
let req = JsonRpcRequest::notification("notifications/initialized", None);
|
||||
let value = serde_json::to_value(&req).unwrap();
|
||||
|
||||
assert_eq!(value["jsonrpc"], "2.0");
|
||||
assert_eq!(value["method"], "notifications/initialized");
|
||||
// id should be absent because it is None and marked skip_serializing_if
|
||||
assert!(value.get("id").is_none() || value["id"].is_null());
|
||||
// When skip_serializing_if fires the key is absent entirely
|
||||
assert!(!value.as_object().unwrap().contains_key("id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jsonrpc_response_deserialization_success() {
|
||||
// Deserialize a successful JSON-RPC response and check result field
|
||||
let json_str = r#"{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}"#;
|
||||
let resp: JsonRpcResponse = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(resp.id, Some(1));
|
||||
assert!(resp.result.is_some());
|
||||
assert!(resp.error.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jsonrpc_response_deserialization_error() {
|
||||
// Deserialize an error JSON-RPC response and check error fields
|
||||
let json_str =
|
||||
r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32601,"message":"Method not found"}}"#;
|
||||
let resp: JsonRpcResponse = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(resp.id, Some(2));
|
||||
assert!(resp.result.is_none());
|
||||
let err = resp.error.expect("error field should be present");
|
||||
assert_eq!(err.code, -32601);
|
||||
assert_eq!(err.message, "Method not found");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_tool_def_deserialization() {
|
||||
// Deserialize a McpToolDef including the camelCase inputSchema rename
|
||||
let json_str = r#"{
|
||||
"name": "read_file",
|
||||
"description": "Read a file from disk",
|
||||
"inputSchema": {"type": "object", "properties": {}}
|
||||
}"#;
|
||||
let tool: McpToolDef = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(tool.name, "read_file");
|
||||
assert_eq!(tool.description.as_deref(), Some("Read a file from disk"));
|
||||
assert_eq!(tool.input_schema["type"], "object");
|
||||
// No annotations field → None (old-server compatible).
|
||||
assert!(tool.annotations.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_tool_def_with_annotations() {
|
||||
// A tool advertising read-only behaviour plus a title. Verify every
|
||||
// camelCase hint maps onto the snake_case Rust field.
|
||||
let json_str = r#"{
|
||||
"name": "browser_snapshot",
|
||||
"description": "Capture an accessibility snapshot",
|
||||
"inputSchema": {"type": "object"},
|
||||
"annotations": {
|
||||
"title": "Snapshot",
|
||||
"readOnlyHint": true,
|
||||
"destructiveHint": false,
|
||||
"idempotentHint": true,
|
||||
"openWorldHint": true
|
||||
}
|
||||
}"#;
|
||||
let tool: McpToolDef = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
let ann = tool.annotations.expect("annotations should be parsed");
|
||||
assert_eq!(ann.title.as_deref(), Some("Snapshot"));
|
||||
assert_eq!(ann.read_only_hint, Some(true));
|
||||
assert_eq!(ann.destructive_hint, Some(false));
|
||||
assert_eq!(ann.idempotent_hint, Some(true));
|
||||
assert_eq!(ann.open_world_hint, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_tool_def_partial_annotations() {
|
||||
// Only destructiveHint declared; the rest stay None (advisory + absent).
|
||||
let json_str = r#"{
|
||||
"name": "delete_all",
|
||||
"inputSchema": {"type": "object"},
|
||||
"annotations": {"destructiveHint": true}
|
||||
}"#;
|
||||
let tool: McpToolDef = serde_json::from_str(json_str).unwrap();
|
||||
let ann = tool.annotations.expect("annotations should be parsed");
|
||||
assert_eq!(ann.destructive_hint, Some(true));
|
||||
assert!(ann.read_only_hint.is_none());
|
||||
assert!(ann.title.is_none());
|
||||
assert!(ann.idempotent_hint.is_none());
|
||||
assert!(ann.open_world_hint.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tools_list_result_preserves_annotations() {
|
||||
// tools/list must carry annotations through into each McpToolDef.
|
||||
let json_str = r#"{
|
||||
"tools": [
|
||||
{"name": "ro", "inputSchema": {}, "annotations": {"readOnlyHint": true}},
|
||||
{"name": "plain", "inputSchema": {}}
|
||||
]
|
||||
}"#;
|
||||
let result: ToolsListResult = serde_json::from_str(json_str).unwrap();
|
||||
assert_eq!(result.tools.len(), 2);
|
||||
assert_eq!(
|
||||
result.tools[0]
|
||||
.annotations
|
||||
.as_ref()
|
||||
.and_then(|a| a.read_only_hint),
|
||||
Some(true)
|
||||
);
|
||||
assert!(result.tools[1].annotations.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_content_text() {
|
||||
// Deserialize McpContent::Text using the internally-tagged "type" field
|
||||
let json_str = r#"{"type":"text","text":"hello world"}"#;
|
||||
let content: McpContent = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
match content {
|
||||
McpContent::Text { text } => assert_eq!(text, "hello world"),
|
||||
other => panic!("expected McpContent::Text, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_content_image() {
|
||||
// Deserialize McpContent::Image including the camelCase mimeType rename
|
||||
let json_str = r#"{"type":"image","data":"base64data==","mimeType":"image/png"}"#;
|
||||
let content: McpContent = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
match content {
|
||||
McpContent::Image { data, mime_type } => {
|
||||
assert_eq!(data, "base64data==");
|
||||
assert_eq!(mime_type, "image/png");
|
||||
}
|
||||
other => panic!("expected McpContent::Image, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-1.x: McpResource deserialization [黑盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn tc_1_1_mcp_resource_all_fields() {
|
||||
// [黑盒] TC-1.1: McpResource complete deserialization with all optional fields
|
||||
let json_str = r#"{
|
||||
"uri": "skill://my-skill",
|
||||
"name": "My Skill",
|
||||
"description": "A test skill",
|
||||
"mimeType": "text/plain"
|
||||
}"#;
|
||||
let resource: McpResource = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(resource.uri, "skill://my-skill");
|
||||
assert_eq!(resource.name.as_deref(), Some("My Skill"));
|
||||
assert_eq!(resource.description.as_deref(), Some("A test skill"));
|
||||
assert_eq!(resource.mime_type.as_deref(), Some("text/plain"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_2_mcp_resource_uri_only() {
|
||||
// [黑盒] TC-1.2: McpResource with only the required uri field — all options are None
|
||||
let json_str = r#"{"uri": "skill://minimal"}"#;
|
||||
let resource: McpResource = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(resource.uri, "skill://minimal");
|
||||
assert!(resource.name.is_none());
|
||||
assert!(resource.description.is_none());
|
||||
assert!(resource.mime_type.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_3_mcp_resource_mime_type_camel_case_mapping() {
|
||||
// [白盒] TC-1.3: JSON field "mimeType" (camelCase) maps to Rust field mime_type via serde rename
|
||||
let json_str = r#"{"uri": "skill://x", "mimeType": "text/markdown"}"#;
|
||||
let resource: McpResource = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(resource.mime_type.as_deref(), Some("text/markdown"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_3b_mcp_resource_mime_type_snake_case_absent() {
|
||||
// [白盒] TC-1.3b: snake_case "mime_type" key is not accepted — mime_type stays None
|
||||
let json_str = r#"{"uri": "skill://x", "mime_type": "text/markdown"}"#;
|
||||
let resource: McpResource = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
// The snake_case key is unknown and ignored; mime_type should be None (default)
|
||||
assert!(resource.mime_type.is_none());
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-1.4: ResourcesListResult deserialization [黑盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn tc_1_4_resources_list_result_multiple() {
|
||||
// [黑盒] TC-1.4: ResourcesListResult with multiple resources
|
||||
let json_str = r#"{
|
||||
"resources": [
|
||||
{"uri": "skill://skill-a"},
|
||||
{"uri": "skill://skill-b", "name": "Skill B"}
|
||||
]
|
||||
}"#;
|
||||
let result: ResourcesListResult = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(result.resources.len(), 2);
|
||||
assert_eq!(result.resources[0].uri, "skill://skill-a");
|
||||
assert_eq!(result.resources[1].uri, "skill://skill-b");
|
||||
assert_eq!(result.resources[1].name.as_deref(), Some("Skill B"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_5_resources_list_result_empty() {
|
||||
// [黑盒] TC-1.5: ResourcesListResult with empty resources array
|
||||
let json_str = r#"{"resources": []}"#;
|
||||
let result: ResourcesListResult = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert!(result.resources.is_empty());
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// TC-1.6/1.7/1.8: ResourcesReadResult and ResourceContent [黑盒]
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn tc_1_6_resources_read_result_with_text() {
|
||||
// [黑盒] TC-1.6: ResourcesReadResult with text content
|
||||
let json_str = r#"{
|
||||
"contents": [
|
||||
{
|
||||
"uri": "skill://my-skill",
|
||||
"mimeType": "text/plain",
|
||||
"text": "---\ndescription: My skill\n---\n# My Skill"
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
let result: ResourcesReadResult = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(result.contents.len(), 1);
|
||||
let content = &result.contents[0];
|
||||
assert_eq!(content.uri, "skill://my-skill");
|
||||
assert_eq!(content.mime_type.as_deref(), Some("text/plain"));
|
||||
assert!(
|
||||
content
|
||||
.text
|
||||
.as_deref()
|
||||
.unwrap()
|
||||
.contains("description: My skill")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_7_resource_content_no_text_field() {
|
||||
// [黑盒] TC-1.7: ResourceContent without text (blob resource) — text is None
|
||||
let json_str = r#"{"uri": "skill://binary", "mimeType": "application/octet-stream"}"#;
|
||||
let content: ResourceContent = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(content.uri, "skill://binary");
|
||||
assert_eq!(
|
||||
content.mime_type.as_deref(),
|
||||
Some("application/octet-stream")
|
||||
);
|
||||
assert!(content.text.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_8_resource_content_no_mime_type() {
|
||||
// [黑盒] TC-1.8: ResourceContent without mimeType — mime_type is None
|
||||
let json_str = r#"{"uri": "skill://no-mime", "text": "content"}"#;
|
||||
let content: ResourceContent = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(content.uri, "skill://no-mime");
|
||||
assert!(content.mime_type.is_none());
|
||||
assert_eq!(content.text.as_deref(), Some("content"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_wb_resource_content_mime_type_camel_case() {
|
||||
// [白盒] ResourceContent.mimeType uses same serde rename as McpResource
|
||||
let json_str = r#"{"uri": "skill://x", "mimeType": "text/markdown", "text": "hello"}"#;
|
||||
let content: ResourceContent = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(content.mime_type.as_deref(), Some("text/markdown"));
|
||||
assert_eq!(content.text.as_deref(), Some("hello"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tc_1_wb_resources_read_result_multiple_contents() {
|
||||
// [白盒] TC: read_resource uses find_map — multiple contents, only first text is used
|
||||
// Here we verify the protocol type itself can hold multiple contents
|
||||
let json_str = r#"{
|
||||
"contents": [
|
||||
{"uri": "skill://x", "text": null},
|
||||
{"uri": "skill://x", "text": "actual content"}
|
||||
]
|
||||
}"#;
|
||||
let result: ResourcesReadResult = serde_json::from_str(json_str).unwrap();
|
||||
|
||||
assert_eq!(result.contents.len(), 2);
|
||||
assert!(result.contents[0].text.is_none());
|
||||
assert_eq!(result.contents[1].text.as_deref(), Some("actual content"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,642 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::config::McpServerConfig;
|
||||
use super::manager::McpManager;
|
||||
use super::protocol::ToolAnnotations;
|
||||
use nomi_protocol::events::ToolCategory;
|
||||
use nomi_tools::Tool;
|
||||
use nomi_types::tool::{JsonSchema, ToolImage, ToolResult};
|
||||
|
||||
/// Upper bound on a single MCP image's decoded byte size before it is dropped.
|
||||
///
|
||||
/// Browser/Playwright screenshots routinely run 1–5 MB; routing one verbatim
|
||||
/// into the message history would balloon the context. Aligned with
|
||||
/// `nomi-tools` `read.rs` `MAX_IMAGE_BYTES` (5 MiB). We estimate the decoded
|
||||
/// size from the base64 length (decoded ≈ len * 3 / 4) to avoid decoding the
|
||||
/// whole payload just to measure it.
|
||||
const MCP_MAX_IMAGE_BYTES: usize = 5 * 1024 * 1024;
|
||||
|
||||
/// Estimate the decoded byte length of a (possibly padded) base64 string
|
||||
/// without allocating/decoding it. Standard base64 encodes every 3 bytes as 4
|
||||
/// characters; trailing `=` padding marks 1–2 missing bytes. Whitespace and a
|
||||
/// `data:...;base64,` prefix (if any) are ignored so we measure the payload.
|
||||
fn decoded_base64_len(data: &str) -> usize {
|
||||
// Strip a data-URL prefix if a server happened to send one.
|
||||
let payload = match data.split_once(";base64,") {
|
||||
Some((_, b64)) => b64,
|
||||
None => data,
|
||||
};
|
||||
let mut chars = 0usize;
|
||||
let mut padding = 0usize;
|
||||
for b in payload.bytes() {
|
||||
match b {
|
||||
b'=' => padding += 1,
|
||||
b if b.is_ascii_whitespace() => {}
|
||||
_ => chars += 1,
|
||||
}
|
||||
}
|
||||
let total_units = chars + padding;
|
||||
// Each 4-char group decodes to 3 bytes; subtract for padding.
|
||||
let bytes = total_units / 4 * 3;
|
||||
bytes.saturating_sub(padding.min(2))
|
||||
}
|
||||
|
||||
/// Wraps an MCP server tool as a local Tool trait implementation.
|
||||
/// Uses naming convention "mcp__{server}__{tool}" when collisions exist,
|
||||
/// otherwise uses the tool's original name.
|
||||
pub struct McpToolProxy {
|
||||
/// Display name used for registration (may be prefixed)
|
||||
display_name: String,
|
||||
/// Original tool name on the MCP server
|
||||
tool_name: String,
|
||||
/// Server this tool belongs to
|
||||
server_name: String,
|
||||
description: String,
|
||||
input_schema: JsonSchema,
|
||||
manager: Arc<McpManager>,
|
||||
/// Whether this tool's schema should be deferred (sent as name-only stub).
|
||||
deferred: bool,
|
||||
/// MCP behaviour hints used to derive the approval category. `None` means
|
||||
/// the server declared no annotations → safe default (`Exec`, needs approval).
|
||||
annotations: Option<ToolAnnotations>,
|
||||
}
|
||||
|
||||
impl McpToolProxy {
|
||||
// One positional arg per proxy field; a builder would add ceremony without
|
||||
// value for two internal call sites. The `annotations` param (added for
|
||||
// approval classification) pushes this past clippy's 7-arg threshold.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
display_name: String,
|
||||
tool_name: String,
|
||||
server_name: String,
|
||||
description: String,
|
||||
input_schema: JsonSchema,
|
||||
manager: Arc<McpManager>,
|
||||
deferred: bool,
|
||||
annotations: Option<ToolAnnotations>,
|
||||
) -> Self {
|
||||
Self {
|
||||
display_name,
|
||||
tool_name,
|
||||
server_name,
|
||||
description,
|
||||
input_schema,
|
||||
manager,
|
||||
deferred,
|
||||
annotations,
|
||||
}
|
||||
}
|
||||
|
||||
/// Map MCP annotations to an approval [`ToolCategory`].
|
||||
///
|
||||
/// Rule (mirrors codex `requires_mcp_tool_approval`, collapsed onto nomi's
|
||||
/// Info/Exec axis):
|
||||
/// - `readOnlyHint == Some(true)` → [`ToolCategory::Info`] (approval-free).
|
||||
/// - everything else — `destructiveHint`, no hints, or an old server with no
|
||||
/// `annotations` block at all — → [`ToolCategory::Exec`] (needs approval).
|
||||
///
|
||||
/// The from-strict default is deliberate: an unannotated tool could mutate
|
||||
/// the world, so we never silently auto-approve it.
|
||||
fn category_from_annotations(&self) -> ToolCategory {
|
||||
if self.is_read_only() {
|
||||
ToolCategory::Info
|
||||
} else {
|
||||
ToolCategory::Exec
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the tool declared `readOnlyHint == true` (no side effects).
|
||||
/// Drives both the approval category and concurrency-safety.
|
||||
fn is_read_only(&self) -> bool {
|
||||
self.annotations
|
||||
.as_ref()
|
||||
.and_then(|a| a.read_only_hint)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for McpToolProxy {
|
||||
fn name(&self) -> &str {
|
||||
&self.display_name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.description
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> JsonSchema {
|
||||
self.input_schema.clone()
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
// Read-only MCP tools have no side effects → safe to run in parallel with
|
||||
// other read-only calls (mirrors built-in Read/Grep/Glob). Mutating or
|
||||
// unannotated tools stay serial.
|
||||
self.is_read_only()
|
||||
}
|
||||
|
||||
fn is_deferred(&self) -> bool {
|
||||
self.deferred
|
||||
}
|
||||
|
||||
async fn execute(&self, input: Value) -> ToolResult {
|
||||
match self
|
||||
.manager
|
||||
.call_tool(&self.server_name, &self.tool_name, input)
|
||||
.await
|
||||
{
|
||||
Ok(out) => {
|
||||
let mut text = out.text;
|
||||
let mut images: Vec<ToolImage> = Vec::with_capacity(out.images.len());
|
||||
|
||||
for img in out.images {
|
||||
// Estimate decoded byte size from base64 length to gate
|
||||
// oversized screenshots before they reach the context.
|
||||
let decoded_len = decoded_base64_len(&img.data);
|
||||
if decoded_len > MCP_MAX_IMAGE_BYTES {
|
||||
tracing::warn!(
|
||||
target: "nomi_mcp",
|
||||
server = %self.server_name,
|
||||
tool = %self.tool_name,
|
||||
bytes = decoded_len,
|
||||
limit = MCP_MAX_IMAGE_BYTES,
|
||||
"dropping oversized MCP image"
|
||||
);
|
||||
let placeholder = format!("[image too large: {} bytes, dropped]", decoded_len);
|
||||
if text.is_empty() {
|
||||
text = placeholder;
|
||||
} else {
|
||||
text.push('\n');
|
||||
text.push_str(&placeholder);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
images.push(ToolImage {
|
||||
media_type: img.mime_type, // mime_type → media_type field name
|
||||
data: img.data, // raw base64, passed through
|
||||
});
|
||||
}
|
||||
|
||||
if images.is_empty() {
|
||||
// Pure-text MCP tool: behaviour identical to before this change.
|
||||
ToolResult::text(text)
|
||||
} else {
|
||||
// Multimodal: text → content, images → ToolResult.images so the
|
||||
// downstream provider adapters feed them back to the model.
|
||||
ToolResult::text(text).with_images(images)
|
||||
}
|
||||
}
|
||||
Err(e) => ToolResult::error(format!("MCP tool error: {}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
fn category(&self) -> ToolCategory {
|
||||
// Annotation-driven: readOnly tools are approval-free Info, everything
|
||||
// else (destructive or unannotated) is Exec → needs approval. We no
|
||||
// longer collapse every MCP tool into the single `Mcp` bucket, which
|
||||
// forced even read-only snapshots through the approval gate.
|
||||
self.category_from_annotations()
|
||||
}
|
||||
|
||||
fn describe(&self, input: &Value) -> String {
|
||||
format!(
|
||||
"MCP {}/{}: {}",
|
||||
self.server_name,
|
||||
self.tool_name,
|
||||
serde_json::to_string(input).unwrap_or_default()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Register all MCP tools into the tool registry, handling name collisions.
|
||||
///
|
||||
/// Strategy:
|
||||
/// - If tool name doesn't collide with built-in or other MCP tools → use as-is
|
||||
/// - If collision detected → prefix with "mcp__{server_name}__"
|
||||
///
|
||||
/// Each tool's deferred flag is read from the server's config:
|
||||
/// `McpServerConfig::deferred` — defaults to `true` when absent.
|
||||
pub fn register_mcp_tools(
|
||||
registry: &mut nomi_tools::registry::ToolRegistry,
|
||||
manager: &Arc<McpManager>,
|
||||
builtin_names: &[String],
|
||||
server_configs: &HashMap<String, McpServerConfig>,
|
||||
) {
|
||||
let all_tools = manager.all_tools();
|
||||
|
||||
// Determine which names need prefixing
|
||||
for (server_name, tool_def) in &all_tools {
|
||||
let original_name = &tool_def.name;
|
||||
|
||||
// Check collision with built-in tools
|
||||
let collides_builtin = builtin_names.iter().any(|n| n == original_name);
|
||||
|
||||
// Check collision with other MCP servers' tools
|
||||
let cross_server_collision = manager.tool_name_count(original_name) > 1;
|
||||
|
||||
let display_name = if collides_builtin || cross_server_collision {
|
||||
format!("mcp__{}_{}", server_name, original_name)
|
||||
} else {
|
||||
original_name.clone()
|
||||
};
|
||||
|
||||
// MCP tools are deferred by default; server config can override.
|
||||
let deferred = server_configs
|
||||
.get(*server_name)
|
||||
.and_then(|c| c.deferred)
|
||||
.unwrap_or(true);
|
||||
|
||||
let proxy = McpToolProxy::new(
|
||||
display_name,
|
||||
original_name.clone(),
|
||||
server_name.to_string(),
|
||||
tool_def.description.clone().unwrap_or_default(),
|
||||
tool_def.input_schema.clone(),
|
||||
Arc::clone(manager),
|
||||
deferred,
|
||||
tool_def.annotations.clone(),
|
||||
);
|
||||
|
||||
registry.register(Box::new(proxy));
|
||||
}
|
||||
}
|
||||
|
||||
/// Register tools from a single newly-connected MCP server.
|
||||
/// Uses the same collision-detection logic as `register_mcp_tools`.
|
||||
pub fn register_single_server_tools(
|
||||
registry: &mut nomi_tools::registry::ToolRegistry,
|
||||
manager: &Arc<McpManager>,
|
||||
server_name: &str,
|
||||
builtin_names: &[String],
|
||||
deferred: bool,
|
||||
) {
|
||||
let all_tools = manager.all_tools();
|
||||
let server_tools: Vec<_> = all_tools
|
||||
.iter()
|
||||
.filter(|(sn, _)| *sn == server_name)
|
||||
.collect();
|
||||
|
||||
for (_, tool_def) in &server_tools {
|
||||
let original_name = &tool_def.name;
|
||||
let collides_builtin = builtin_names.iter().any(|n| n == original_name);
|
||||
let cross_server_collision = manager.tool_name_count(original_name) > 1;
|
||||
|
||||
let display_name = if collides_builtin || cross_server_collision {
|
||||
format!("mcp__{}_{}", server_name, original_name)
|
||||
} else {
|
||||
original_name.clone()
|
||||
};
|
||||
|
||||
let proxy = McpToolProxy::new(
|
||||
display_name,
|
||||
original_name.clone(),
|
||||
server_name.to_string(),
|
||||
tool_def.description.clone().unwrap_or_default(),
|
||||
tool_def.input_schema.clone(),
|
||||
Arc::clone(manager),
|
||||
deferred,
|
||||
tool_def.annotations.clone(),
|
||||
);
|
||||
|
||||
registry.register(Box::new(proxy));
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::protocol::ToolAnnotations;
|
||||
use nomi_config::config::TransportType;
|
||||
use serde_json::json;
|
||||
|
||||
// Minimal MockTransport local to this test module (the one in manager.rs's
|
||||
// test mod is not cross-module visible). McpTransport is crate-visible with
|
||||
// 3 methods, so duplicating it keeps the modules decoupled.
|
||||
use crate::protocol::{JsonRpcRequest, JsonRpcResponse};
|
||||
use crate::transport::{McpError, McpTransport};
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Mutex;
|
||||
|
||||
struct MockTransport {
|
||||
responses: Mutex<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl MockTransport {
|
||||
fn new(responses: Vec<serde_json::Value>) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(responses),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for MockTransport {
|
||||
async fn request(&self, _req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError> {
|
||||
let mut guard = self.responses.lock().unwrap();
|
||||
let value = if guard.is_empty() {
|
||||
json!(null)
|
||||
} else {
|
||||
guard.remove(0)
|
||||
};
|
||||
Ok(JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: Some(1),
|
||||
result: Some(value),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn notify(&self, _req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn make_proxy(deferred: bool) -> McpToolProxy {
|
||||
// manager is only used during execute(), which we don't call in these
|
||||
// tests, so we can construct one with no servers.
|
||||
let manager = Arc::new(McpManager::new_for_test(vec![]));
|
||||
McpToolProxy::new(
|
||||
"test_tool".into(),
|
||||
"test_tool".into(),
|
||||
"test_server".into(),
|
||||
"A test tool".into(),
|
||||
json!({"type": "object"}),
|
||||
manager,
|
||||
deferred,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
/// Build a proxy with the given annotations (no real transport needed —
|
||||
/// category() does not touch the manager).
|
||||
fn make_proxy_with_annotations(annotations: Option<ToolAnnotations>) -> McpToolProxy {
|
||||
let manager = Arc::new(McpManager::new_for_test(vec![]));
|
||||
McpToolProxy::new(
|
||||
"test_tool".into(),
|
||||
"test_tool".into(),
|
||||
"test_server".into(),
|
||||
"A test tool".into(),
|
||||
json!({"type": "object"}),
|
||||
manager,
|
||||
true,
|
||||
annotations,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_deferred_true_returns_true() {
|
||||
let proxy = make_proxy(true);
|
||||
assert!(proxy.is_deferred());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_deferred_false_returns_false() {
|
||||
let proxy = make_proxy(false);
|
||||
assert!(!proxy.is_deferred());
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// category(): annotations → approval class (readOnly→Info, else→Exec)
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn category_read_only_hint_true_is_info() {
|
||||
let proxy = make_proxy_with_annotations(Some(ToolAnnotations {
|
||||
read_only_hint: Some(true),
|
||||
..Default::default()
|
||||
}));
|
||||
assert_eq!(proxy.category(), ToolCategory::Info);
|
||||
// category_for delegates to category() for tools that don't override it.
|
||||
assert_eq!(proxy.category_for(&json!({})), ToolCategory::Info);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn category_destructive_hint_is_exec() {
|
||||
// destructive (and not read-only) must require approval.
|
||||
let proxy = make_proxy_with_annotations(Some(ToolAnnotations {
|
||||
destructive_hint: Some(true),
|
||||
..Default::default()
|
||||
}));
|
||||
assert_eq!(proxy.category(), ToolCategory::Exec);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn category_read_only_false_is_exec() {
|
||||
// Explicit readOnlyHint=false → still needs approval.
|
||||
let proxy = make_proxy_with_annotations(Some(ToolAnnotations {
|
||||
read_only_hint: Some(false),
|
||||
..Default::default()
|
||||
}));
|
||||
assert_eq!(proxy.category(), ToolCategory::Exec);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_hint_true_is_concurrency_safe() {
|
||||
// A read-only MCP tool has no side effects, so concurrent execution with
|
||||
// other read-only calls is safe (mirrors built-in Read/Grep/Glob).
|
||||
let proxy = make_proxy_with_annotations(Some(ToolAnnotations {
|
||||
read_only_hint: Some(true),
|
||||
..Default::default()
|
||||
}));
|
||||
assert!(proxy.is_concurrency_safe(&json!({})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_read_only_is_not_concurrency_safe() {
|
||||
let explicit_false = make_proxy_with_annotations(Some(ToolAnnotations {
|
||||
read_only_hint: Some(false),
|
||||
..Default::default()
|
||||
}));
|
||||
assert!(!explicit_false.is_concurrency_safe(&json!({})));
|
||||
// No annotations → safe default: assume side effects, run serially.
|
||||
let none = make_proxy_with_annotations(None);
|
||||
assert!(!none.is_concurrency_safe(&json!({})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn category_no_annotations_defaults_to_exec() {
|
||||
// Old server with no annotations block → safe default: Exec.
|
||||
let proxy = make_proxy_with_annotations(None);
|
||||
assert_eq!(proxy.category(), ToolCategory::Exec);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn category_empty_annotations_defaults_to_exec() {
|
||||
// annotations present but no readOnlyHint → from-strict default: Exec.
|
||||
let proxy = make_proxy_with_annotations(Some(ToolAnnotations::default()));
|
||||
assert_eq!(proxy.category(), ToolCategory::Exec);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn category_read_only_wins_when_no_destructive() {
|
||||
// openWorld + idempotent set but readOnly true → still Info.
|
||||
let proxy = make_proxy_with_annotations(Some(ToolAnnotations {
|
||||
read_only_hint: Some(true),
|
||||
open_world_hint: Some(true),
|
||||
idempotent_hint: Some(true),
|
||||
..Default::default()
|
||||
}));
|
||||
assert_eq!(proxy.category(), ToolCategory::Info);
|
||||
}
|
||||
|
||||
fn make_server_config(deferred: Option<bool>) -> McpServerConfig {
|
||||
McpServerConfig {
|
||||
transport: TransportType::Stdio,
|
||||
command: Some("echo".into()),
|
||||
args: None,
|
||||
env: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
deferred,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn register_defaults_to_deferred_when_config_omits_field() {
|
||||
let manager = Arc::new(McpManager::new_for_test(vec![]));
|
||||
let mut registry = nomi_tools::registry::ToolRegistry::new();
|
||||
// Empty server configs — deferred field absent
|
||||
let configs = HashMap::new();
|
||||
|
||||
register_mcp_tools(&mut registry, &manager, &[], &configs);
|
||||
|
||||
// No tools registered because manager has no tools, but the logic
|
||||
// is tested via the deferred default path. Test with a real config below.
|
||||
assert!(registry.tool_names().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_config_deferred_none_defaults_true() {
|
||||
let config = make_server_config(None);
|
||||
let deferred = config.deferred.unwrap_or(true);
|
||||
assert!(deferred, "deferred should default to true when None");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_config_deferred_explicit_false() {
|
||||
let config = make_server_config(Some(false));
|
||||
let deferred = config.deferred.unwrap_or(true);
|
||||
assert!(!deferred, "deferred should be false when explicitly set");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_config_deferred_explicit_true() {
|
||||
let config = make_server_config(Some(true));
|
||||
let deferred = config.deferred.unwrap_or(true);
|
||||
assert!(deferred, "deferred should be true when explicitly set");
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// execute: image content → ToolResult.images (end-to-end mapping)
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
fn proxy_with_response(
|
||||
tool: &str,
|
||||
resp: serde_json::Value,
|
||||
) -> McpToolProxy {
|
||||
let mgr = Arc::new(McpManager::new_for_test(vec![(
|
||||
"srv",
|
||||
false,
|
||||
Box::new(MockTransport::new(vec![resp])),
|
||||
)]));
|
||||
McpToolProxy::new(
|
||||
tool.into(),
|
||||
tool.into(),
|
||||
"srv".into(),
|
||||
"desc".into(),
|
||||
json!({"type":"object"}),
|
||||
mgr,
|
||||
true,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_execute_maps_image_to_tool_result_images() {
|
||||
let resp = json!({ "content": [
|
||||
{"type":"text","text":"ok"},
|
||||
{"type":"image","data":"ZZZ","mimeType":"image/png"}
|
||||
]});
|
||||
let proxy = proxy_with_response("shot", resp);
|
||||
let r = proxy.execute(json!({})).await;
|
||||
assert!(!r.is_error);
|
||||
assert_eq!(r.content, "ok");
|
||||
assert_eq!(r.images.len(), 1);
|
||||
assert_eq!(r.images[0].media_type, "image/png");
|
||||
assert_eq!(r.images[0].data, "ZZZ");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_execute_text_only_no_images() {
|
||||
let resp = json!({ "content": [{"type":"text","text":"plain"}] });
|
||||
let proxy = proxy_with_response("echo", resp);
|
||||
let r = proxy.execute(json!({})).await;
|
||||
assert!(!r.is_error);
|
||||
assert_eq!(r.content, "plain");
|
||||
assert!(r.images.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_execute_drops_oversized_image_with_placeholder() {
|
||||
// Build a base64 string whose decoded size exceeds MCP_MAX_IMAGE_BYTES.
|
||||
// decoded ≈ len * 3/4, so we need > 5 MiB * 4/3 base64 chars.
|
||||
let huge = "A".repeat((MCP_MAX_IMAGE_BYTES + 1024) * 4 / 3 + 8);
|
||||
let resp = json!({ "content": [
|
||||
{"type":"text","text":"shot taken"},
|
||||
{"type":"image","data": huge,"mimeType":"image/png"}
|
||||
]});
|
||||
let proxy = proxy_with_response("big_shot", resp);
|
||||
let r = proxy.execute(json!({})).await;
|
||||
assert!(!r.is_error);
|
||||
// Oversized image dropped → no images survive.
|
||||
assert!(r.images.is_empty());
|
||||
// Original text preserved + placeholder appended.
|
||||
assert!(r.content.contains("shot taken"));
|
||||
assert!(
|
||||
r.content.contains("image too large"),
|
||||
"expected placeholder text, got: {}",
|
||||
r.content
|
||||
);
|
||||
assert!(r.content.contains("dropped"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_execute_keeps_image_just_under_limit() {
|
||||
// A small image must survive the size guard untouched.
|
||||
let resp = json!({ "content": [
|
||||
{"type":"image","data":"c21hbGw=","mimeType":"image/jpeg"}
|
||||
]});
|
||||
let proxy = proxy_with_response("small_shot", resp);
|
||||
let r = proxy.execute(json!({})).await;
|
||||
assert_eq!(r.images.len(), 1);
|
||||
assert_eq!(r.images[0].media_type, "image/jpeg");
|
||||
assert_eq!(r.images[0].data, "c21hbGw=");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoded_base64_len_estimates_size() {
|
||||
// "QQ==" decodes to 1 byte; "QUI=" → 2 bytes; "QUJD" → 3 bytes.
|
||||
assert_eq!(decoded_base64_len("QQ=="), 1);
|
||||
assert_eq!(decoded_base64_len("QUI="), 2);
|
||||
assert_eq!(decoded_base64_len("QUJD"), 3);
|
||||
// Whitespace is ignored.
|
||||
assert_eq!(decoded_base64_len("QU\nJD"), 3);
|
||||
// data-URL prefix is stripped before measuring.
|
||||
assert_eq!(decoded_base64_len("data:image/png;base64,QUJD"), 3);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
pub mod sse;
|
||||
pub mod stdio;
|
||||
pub mod streamable_http;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::protocol::{JsonRpcRequest, JsonRpcResponse};
|
||||
|
||||
/// Find the next SSE event boundary (blank line) in `buf`, returning
|
||||
/// `(offset, delimiter_len)` for the earliest match.
|
||||
///
|
||||
/// Per the SSE spec an event is terminated by a blank line, which may be framed
|
||||
/// with LF (`\n\n`), CRLF (`\r\n\r\n`), or bare CR (`\r\r`). The MCP spec and
|
||||
/// most servers use `\n\n`, but some MCP servers behind new-api / one-api style
|
||||
/// proxies emit `\r\n\r\n`; matching only `\n\n` there finds no event boundary,
|
||||
/// parses zero events, and yields a silent connection failure ("No endpoint
|
||||
/// event received" / "SSE stream ended without JSON-RPC response").
|
||||
///
|
||||
/// Returns the boundary with the smallest offset; if only a partial delimiter
|
||||
/// sits at the end of the buffer (e.g. a chunk split mid-`\r\n\r\n`), none match
|
||||
/// yet and the caller waits for more bytes — same as the original `\n\n` logic.
|
||||
/// Mirrors `find_sse_event_boundary` in `nomi-providers/src/anthropic_shared.rs`.
|
||||
pub(crate) fn find_sse_event_boundary(buf: &str) -> Option<(usize, usize)> {
|
||||
[
|
||||
buf.find("\r\n\r\n").map(|i| (i, 4)),
|
||||
buf.find("\n\n").map(|i| (i, 2)),
|
||||
buf.find("\r\r").map(|i| (i, 2)),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.min_by_key(|&(offset, _)| offset)
|
||||
}
|
||||
|
||||
/// Transport abstraction for MCP communication
|
||||
#[async_trait]
|
||||
pub trait McpTransport: Send + Sync {
|
||||
/// Send a JSON-RPC request and receive the response
|
||||
async fn request(&self, req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError>;
|
||||
|
||||
/// Send a notification (no response expected)
|
||||
async fn notify(&self, req: &JsonRpcRequest) -> Result<(), McpError>;
|
||||
|
||||
/// Close the transport
|
||||
async fn close(&self) -> Result<(), McpError>;
|
||||
}
|
||||
|
||||
/// Errors from MCP transport and protocol
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum McpError {
|
||||
#[error("Transport error: {0}")]
|
||||
Transport(String),
|
||||
|
||||
#[error("JSON-RPC error {code}: {message}")]
|
||||
JsonRpc { code: i64, message: String },
|
||||
|
||||
#[error("Server not found: {0}")]
|
||||
ServerNotFound(String),
|
||||
|
||||
#[error("Tool not found: {server}/{tool}")]
|
||||
ToolNotFound { server: String, tool: String },
|
||||
|
||||
#[error("Initialization failed: {0}")]
|
||||
InitFailed(String),
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::find_sse_event_boundary;
|
||||
|
||||
#[test]
|
||||
fn lf_framing() {
|
||||
assert_eq!(find_sse_event_boundary("a\n\nb"), Some((1, 2)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crlf_framing() {
|
||||
// new-api / one-api proxies frame SSE events with CRLF.
|
||||
assert_eq!(find_sse_event_boundary("a\r\n\r\nb"), Some((1, 4)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bare_cr_framing() {
|
||||
assert_eq!(find_sse_event_boundary("a\r\rb"), Some((1, 2)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn earliest_boundary_wins() {
|
||||
// A CRLF boundary at offset 1 must beat an LF boundary later in the buffer.
|
||||
assert_eq!(find_sse_event_boundary("a\r\n\r\nb\n\nc"), Some((1, 4)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_delimiter_waits() {
|
||||
// A chunk split mid-CRLF must not match yet.
|
||||
assert_eq!(find_sse_event_boundary("data: {}\r\n\r"), None);
|
||||
assert_eq!(find_sse_event_boundary("data: {}\n"), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use reqwest::header::{HeaderMap, HeaderValue};
|
||||
use tokio::sync::{Mutex, oneshot};
|
||||
|
||||
use super::{McpError, McpTransport, find_sse_event_boundary};
|
||||
use crate::protocol::{JsonRpcRequest, JsonRpcResponse};
|
||||
|
||||
/// SSE transport: connects to an SSE endpoint for server→client events,
|
||||
/// sends requests via POST to the endpoint URL received from the SSE stream
|
||||
pub struct SseTransport {
|
||||
client: reqwest::Client,
|
||||
/// The POST endpoint URL (received from the SSE stream's "endpoint" event)
|
||||
post_url: String,
|
||||
headers: HeaderMap,
|
||||
/// Pending request-response channels, keyed by JSON-RPC id
|
||||
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<JsonRpcResponse>>>>,
|
||||
next_id: AtomicU64,
|
||||
/// Handle to the background SSE listener task
|
||||
_listener: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl SseTransport {
|
||||
/// Connect to an SSE MCP server
|
||||
pub async fn connect(url: &str, headers: &HashMap<String, String>) -> Result<Self, McpError> {
|
||||
let mut header_map = HeaderMap::new();
|
||||
for (k, v) in headers {
|
||||
let name = reqwest::header::HeaderName::from_bytes(k.as_bytes())
|
||||
.map_err(|e| McpError::Transport(format!("Invalid header name '{}': {}", k, e)))?;
|
||||
let value = HeaderValue::from_str(v)
|
||||
.map_err(|e| McpError::Transport(format!("Invalid header value '{}': {}", v, e)))?;
|
||||
header_map.insert(name, value);
|
||||
}
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// GET the SSE endpoint to establish the event stream
|
||||
let response = client
|
||||
.get(url)
|
||||
.headers(header_map.clone())
|
||||
.header("Accept", "text/event-stream")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("SSE connection failed: {}", e)))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(McpError::Transport(format!(
|
||||
"SSE connection returned status: {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
let pending: Arc<Mutex<HashMap<u64, oneshot::Sender<JsonRpcResponse>>>> =
|
||||
Arc::new(Mutex::new(HashMap::new()));
|
||||
|
||||
// Parse the SSE stream to find the endpoint URL
|
||||
// The server sends an "endpoint" event with the POST URL
|
||||
let base_url = extract_base_url(url);
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
let mut post_url: Option<String> = None;
|
||||
|
||||
use futures::StreamExt;
|
||||
// Read initial events to get the endpoint URL
|
||||
while let Some(chunk) = bytes_stream.next().await {
|
||||
let chunk = chunk.map_err(|e| McpError::Transport(format!("SSE read error: {}", e)))?;
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// Parse SSE events from buffer. Events may be framed with LF, CRLF
|
||||
// (new-api / one-api proxies), or bare CR — see find_sse_event_boundary.
|
||||
while let Some((event_end, delim_len)) = find_sse_event_boundary(&buffer) {
|
||||
let event_block = buffer[..event_end].to_string();
|
||||
buffer = buffer[event_end + delim_len..].to_string();
|
||||
|
||||
let (event_type, event_data) = parse_sse_event(&event_block);
|
||||
|
||||
if event_type == "endpoint" {
|
||||
// The endpoint might be relative or absolute
|
||||
let endpoint = if event_data.starts_with("http") {
|
||||
event_data.clone()
|
||||
} else {
|
||||
format!("{}{}", base_url, event_data)
|
||||
};
|
||||
post_url = Some(endpoint);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if post_url.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let post_url = post_url
|
||||
.ok_or_else(|| McpError::Transport("No endpoint event received from SSE".into()))?;
|
||||
|
||||
// Spawn background task to listen for SSE responses
|
||||
let pending_clone = pending.clone();
|
||||
let listener = tokio::spawn(async move {
|
||||
let mut buf = buffer; // carry over remaining buffer
|
||||
while let Some(chunk) = bytes_stream.next().await {
|
||||
let Ok(chunk) = chunk else { break };
|
||||
buf.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
while let Some((event_end, delim_len)) = find_sse_event_boundary(&buf) {
|
||||
let event_block = buf[..event_end].to_string();
|
||||
buf = buf[event_end + delim_len..].to_string();
|
||||
|
||||
let (event_type, event_data) = parse_sse_event(&event_block);
|
||||
|
||||
if (event_type == "message" || event_type.is_empty())
|
||||
&& let Ok(response) = serde_json::from_str::<JsonRpcResponse>(&event_data)
|
||||
&& let Some(id) = response.id
|
||||
{
|
||||
let mut map: tokio::sync::MutexGuard<
|
||||
'_,
|
||||
HashMap<u64, oneshot::Sender<JsonRpcResponse>>,
|
||||
> = pending_clone.lock().await;
|
||||
if let Some(sender) = map.remove(&id) {
|
||||
let _ = sender.send(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
client,
|
||||
post_url,
|
||||
headers: header_map,
|
||||
pending,
|
||||
next_id: AtomicU64::new(1),
|
||||
_listener: listener,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn next_id(&self) -> u64 {
|
||||
self.next_id.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for SseTransport {
|
||||
async fn request(&self, req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError> {
|
||||
let req_id = req
|
||||
.id
|
||||
.ok_or_else(|| McpError::Transport("Request must have an id".into()))?;
|
||||
|
||||
// Set up response channel before sending
|
||||
let (tx, rx) = oneshot::channel::<JsonRpcResponse>();
|
||||
{
|
||||
let mut map: tokio::sync::MutexGuard<
|
||||
'_,
|
||||
HashMap<u64, oneshot::Sender<JsonRpcResponse>>,
|
||||
> = self.pending.lock().await;
|
||||
map.insert(req_id, tx);
|
||||
}
|
||||
|
||||
// POST the request
|
||||
let body = serde_json::to_string(req)
|
||||
.map_err(|e| McpError::Transport(format!("JSON serialize error: {}", e)))?;
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&self.post_url)
|
||||
.headers(self.headers.clone())
|
||||
.header("Content-Type", "application/json")
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("POST request failed: {}", e)))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
// Clean up pending
|
||||
self.pending.lock().await.remove(&req_id);
|
||||
return Err(McpError::Transport(format!(
|
||||
"POST returned status: {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
// Wait for response from SSE stream
|
||||
let rpc_response = rx
|
||||
.await
|
||||
.map_err(|_| McpError::Transport("Response channel closed unexpectedly".into()))?;
|
||||
|
||||
if let Some(err) = &rpc_response.error {
|
||||
return Err(McpError::JsonRpc {
|
||||
code: err.code,
|
||||
message: err.message.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(rpc_response)
|
||||
}
|
||||
|
||||
async fn notify(&self, req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
let body = serde_json::to_string(req)
|
||||
.map_err(|e| McpError::Transport(format!("JSON serialize error: {}", e)))?;
|
||||
|
||||
self.client
|
||||
.post(&self.post_url)
|
||||
.headers(self.headers.clone())
|
||||
.header("Content-Type", "application/json")
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Notification POST failed: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
self._listener.abort();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a single SSE event block into (event_type, data)
|
||||
fn parse_sse_event(block: &str) -> (String, String) {
|
||||
let mut event_type = String::new();
|
||||
let mut data_lines = Vec::new();
|
||||
|
||||
for line in block.lines() {
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_type = value.trim().to_string();
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
(event_type, data_lines.join("\n"))
|
||||
}
|
||||
|
||||
/// Extract base URL (scheme + host + port) from a full URL
|
||||
fn extract_base_url(url: &str) -> String {
|
||||
// Find the position after "://"
|
||||
if let Some(scheme_end) = url.find("://") {
|
||||
let rest = &url[scheme_end + 3..];
|
||||
if let Some(path_start) = rest.find('/') {
|
||||
return url[..scheme_end + 3 + path_start].to_string();
|
||||
}
|
||||
}
|
||||
url.to_string()
|
||||
}
|
||||
@@ -0,0 +1,625 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
|
||||
use tokio::process::{Child, ChildStdin, ChildStdout};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::{McpError, McpTransport};
|
||||
use crate::protocol::{
|
||||
ClientCapabilities, ClientInfo, InitializeParams, JsonRpcRequest, JsonRpcResponse,
|
||||
};
|
||||
|
||||
/// Maximum number of automatic respawns within [`RESPAWN_WINDOW`] before the
|
||||
/// transport gives up and surfaces a hard error. Without a ceiling a server
|
||||
/// that crashes on every `initialize` would spin forever (crashloop). Codex's
|
||||
/// rmcp client makes the same trade-off — bounded restarts, then fail loud.
|
||||
const MAX_RESPAWNS: u32 = 3;
|
||||
|
||||
/// Sliding window over which [`MAX_RESPAWNS`] is counted. A server that is
|
||||
/// healthy for this long resets its respawn budget, so a single crash months
|
||||
/// into a long session does not consume the lifetime quota.
|
||||
const RESPAWN_WINDOW: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Base backoff before the first respawn; doubled per consecutive attempt
|
||||
/// (200ms, 400ms, 800ms…) and capped at [`MAX_BACKOFF`]. Gives a flapping
|
||||
/// child a moment to settle without stalling the caller for long.
|
||||
const BASE_BACKOFF: Duration = Duration::from_millis(200);
|
||||
const MAX_BACKOFF: Duration = Duration::from_secs(2);
|
||||
|
||||
/// Immutable parameters needed to (re)spawn the child and redo the MCP
|
||||
/// handshake. Captured once at construction so respawn never needs the caller.
|
||||
struct SpawnSpec {
|
||||
command: String,
|
||||
args: Vec<String>,
|
||||
env: HashMap<String, String>,
|
||||
init_params: InitializeParams,
|
||||
}
|
||||
|
||||
/// The live child process and its piped stdio. Replaced wholesale on respawn so
|
||||
/// a half-dead connection (e.g. stdin alive but stdout at EOF) is never reused.
|
||||
struct Connection {
|
||||
stdin: BufWriter<ChildStdin>,
|
||||
stdout: BufReader<ChildStdout>,
|
||||
child: Child,
|
||||
}
|
||||
|
||||
/// Stdio transport: communicates with an MCP server via a child process's
|
||||
/// stdin/stdout. On a detected pipe failure (EOF / broken pipe), it transparently
|
||||
/// respawns the child and re-runs the `initialize` handshake, with a bounded
|
||||
/// retry budget to avoid crashlooping.
|
||||
pub struct StdioTransport {
|
||||
conn: Mutex<Connection>,
|
||||
spec: SpawnSpec,
|
||||
next_id: AtomicU64,
|
||||
/// Respawn bookkeeping: count within the current window + window start.
|
||||
respawn_state: Mutex<RespawnState>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RespawnState {
|
||||
/// Respawns inside the current window.
|
||||
count: u32,
|
||||
/// When the current window began (monotonic). `None` until the first respawn.
|
||||
window_start: Option<std::time::Instant>,
|
||||
}
|
||||
|
||||
impl StdioTransport {
|
||||
/// Spawn a child process and return the transport.
|
||||
///
|
||||
/// `init_params` are retained so a respawn can replay the `initialize`
|
||||
/// handshake without the manager's involvement.
|
||||
pub async fn spawn(
|
||||
command: &str,
|
||||
args: &[String],
|
||||
env: &HashMap<String, String>,
|
||||
) -> Result<Self, McpError> {
|
||||
// Default handshake params; the manager normally drives `initialize`
|
||||
// itself on first connect, but a respawn must be self-contained.
|
||||
let init_params = InitializeParams {
|
||||
protocol_version: "2025-03-26".to_string(),
|
||||
capabilities: ClientCapabilities {
|
||||
tools: Some(serde_json::json!({})),
|
||||
},
|
||||
client_info: ClientInfo {
|
||||
name: "nomi".to_string(),
|
||||
version: "0.3.0".to_string(),
|
||||
},
|
||||
};
|
||||
Self::spawn_with_init(command, args, env, init_params).await
|
||||
}
|
||||
|
||||
/// Spawn with explicit handshake params (kept for the respawn path and for
|
||||
/// callers that want to control the `initialize` payload).
|
||||
pub async fn spawn_with_init(
|
||||
command: &str,
|
||||
args: &[String],
|
||||
env: &HashMap<String, String>,
|
||||
init_params: InitializeParams,
|
||||
) -> Result<Self, McpError> {
|
||||
let spec = SpawnSpec {
|
||||
command: command.to_string(),
|
||||
args: args.to_vec(),
|
||||
env: env.clone(),
|
||||
init_params,
|
||||
};
|
||||
let conn = Self::spawn_child(&spec)?;
|
||||
Ok(Self {
|
||||
conn: Mutex::new(conn),
|
||||
spec,
|
||||
next_id: AtomicU64::new(1),
|
||||
respawn_state: Mutex::new(RespawnState::default()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Launch the child process and capture its piped stdio.
|
||||
fn spawn_child(spec: &SpawnSpec) -> Result<Connection, McpError> {
|
||||
let mut cmd = tokio::process::Command::new(&spec.command);
|
||||
cmd.args(&spec.args)
|
||||
.stdin(std::process::Stdio::piped())
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::inherit())
|
||||
.envs(&spec.env)
|
||||
// Reap the child when the transport is dropped so a respawned-away
|
||||
// or session-ending process never leaks. Mirrors codex rmcp-client.
|
||||
.kill_on_drop(true);
|
||||
// Put the child in its own process group so killing it takes down any
|
||||
// grandchildren (npx → node, etc.) instead of orphaning them.
|
||||
#[cfg(unix)]
|
||||
cmd.process_group(0);
|
||||
// CREATE_NO_WINDOW: MCP stdio servers (npx/node/bun/python) must not
|
||||
// flash a console window under a GUI host.
|
||||
#[cfg(windows)]
|
||||
cmd.creation_flags(0x0800_0000);
|
||||
|
||||
let mut child = cmd.spawn().map_err(|e| {
|
||||
McpError::Transport(format!("Failed to spawn '{}': {}", spec.command, e))
|
||||
})?;
|
||||
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| McpError::Transport("Failed to capture child stdin".into()))?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| McpError::Transport("Failed to capture child stdout".into()))?;
|
||||
|
||||
Ok(Connection {
|
||||
stdin: BufWriter::new(stdin),
|
||||
stdout: BufReader::new(stdout),
|
||||
child,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the next request ID
|
||||
pub fn next_id(&self) -> u64 {
|
||||
self.next_id.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Serialize and write a JSON-RPC message to the child's stdin (one line +
|
||||
/// newline + flush). Errors here mean the write pipe is broken.
|
||||
async fn send_on(conn: &mut Connection, req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
let json = serde_json::to_string(req)
|
||||
.map_err(|e| McpError::Transport(format!("JSON serialize error: {}", e)))?;
|
||||
|
||||
conn.stdin
|
||||
.write_all(json.as_bytes())
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Write to stdin failed: {}", e)))?;
|
||||
conn.stdin
|
||||
.write_all(b"\n")
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Write newline failed: {}", e)))?;
|
||||
conn.stdin
|
||||
.flush()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Flush stdin failed: {}", e)))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read a single JSON-RPC response line from the child's stdout, skipping
|
||||
/// blank lines. A zero-byte read means the child closed stdout (EOF).
|
||||
async fn read_response_on(conn: &mut Connection) -> Result<JsonRpcResponse, McpError> {
|
||||
let mut line = String::new();
|
||||
loop {
|
||||
line.clear();
|
||||
let bytes_read = conn
|
||||
.stdout
|
||||
.read_line(&mut line)
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Read from stdout failed: {}", e)))?;
|
||||
|
||||
if bytes_read == 0 {
|
||||
return Err(McpError::Transport("Child process stdout closed".into()));
|
||||
}
|
||||
|
||||
let trimmed = line.trim();
|
||||
if !trimmed.is_empty() {
|
||||
let response: JsonRpcResponse = serde_json::from_str(trimmed).map_err(|e| {
|
||||
McpError::Transport(format!(
|
||||
"Failed to parse JSON-RPC response: {} — raw: {}",
|
||||
e, trimmed
|
||||
))
|
||||
})?;
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One round-trip on the given connection: write request, read response,
|
||||
/// surface any JSON-RPC error. Used both directly and during the re-handshake.
|
||||
async fn roundtrip_on(
|
||||
conn: &mut Connection,
|
||||
req: &JsonRpcRequest,
|
||||
) -> Result<JsonRpcResponse, McpError> {
|
||||
Self::send_on(conn, req).await?;
|
||||
let response = Self::read_response_on(conn).await?;
|
||||
if let Some(err) = &response.error {
|
||||
return Err(McpError::JsonRpc {
|
||||
code: err.code,
|
||||
message: err.message.clone(),
|
||||
});
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// True for failures that indicate the child/pipe is gone and a respawn is
|
||||
/// warranted. JSON-RPC application errors (the server answered, just with an
|
||||
/// error) and serialize/parse failures are NOT respawn-worthy — respawning
|
||||
/// would not change the outcome.
|
||||
fn is_pipe_failure(err: &McpError) -> bool {
|
||||
match err {
|
||||
McpError::Transport(msg) => {
|
||||
msg.contains("stdout closed")
|
||||
|| msg.contains("Write to stdin failed")
|
||||
|| msg.contains("Write newline failed")
|
||||
|| msg.contains("Flush stdin failed")
|
||||
|| msg.contains("Read from stdout failed")
|
||||
}
|
||||
McpError::Io(_) => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Respawn the child and replay the `initialize` + `notifications/initialized`
|
||||
/// handshake, honouring the bounded retry budget. On success the live
|
||||
/// connection is swapped in place. Returns an error (without panicking) when
|
||||
/// the budget is exhausted or the new child fails to handshake.
|
||||
async fn respawn(&self) -> Result<(), McpError> {
|
||||
// Enforce the crashloop ceiling within a sliding window.
|
||||
{
|
||||
let mut state = self.respawn_state.lock().await;
|
||||
let now = std::time::Instant::now();
|
||||
match state.window_start {
|
||||
Some(start) if now.duration_since(start) <= RESPAWN_WINDOW => {
|
||||
if state.count >= MAX_RESPAWNS {
|
||||
return Err(McpError::Transport(format!(
|
||||
"MCP stdio server '{}' exceeded {} respawns within {}s; giving up",
|
||||
self.spec.command,
|
||||
MAX_RESPAWNS,
|
||||
RESPAWN_WINDOW.as_secs()
|
||||
)));
|
||||
}
|
||||
state.count += 1;
|
||||
}
|
||||
_ => {
|
||||
// First respawn, or the previous window has elapsed → reset.
|
||||
state.window_start = Some(now);
|
||||
state.count = 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Backoff before respawning (exponential, capped). Read attempt count
|
||||
// again under the lock-free local; `count` was just incremented above.
|
||||
let attempt = {
|
||||
let state = self.respawn_state.lock().await;
|
||||
state.count
|
||||
};
|
||||
let backoff = BASE_BACKOFF
|
||||
.saturating_mul(1u32 << attempt.saturating_sub(1).min(5))
|
||||
.min(MAX_BACKOFF);
|
||||
tokio::time::sleep(backoff).await;
|
||||
|
||||
tracing::warn!(
|
||||
target: "nomi_mcp",
|
||||
command = %self.spec.command,
|
||||
attempt,
|
||||
backoff_ms = backoff.as_millis() as u64,
|
||||
"respawning crashed MCP stdio server"
|
||||
);
|
||||
|
||||
// Spawn a fresh child and run the handshake on it before publishing it,
|
||||
// so a half-initialized child never becomes the live connection.
|
||||
let mut new_conn = Self::spawn_child(&self.spec)?;
|
||||
|
||||
let init_req = JsonRpcRequest::new(
|
||||
1,
|
||||
"initialize",
|
||||
Some(serde_json::to_value(&self.spec.init_params).map_err(|e| {
|
||||
McpError::InitFailed(format!("Failed to serialize init params: {}", e))
|
||||
})?),
|
||||
);
|
||||
Self::roundtrip_on(&mut new_conn, &init_req)
|
||||
.await
|
||||
.map_err(|e| McpError::InitFailed(format!("respawn initialize failed: {}", e)))?;
|
||||
|
||||
let initialized = JsonRpcRequest::notification("notifications/initialized", None);
|
||||
Self::send_on(&mut new_conn, &initialized).await?;
|
||||
|
||||
// Swap in the healthy connection. The old `Connection` is dropped here;
|
||||
// `kill_on_drop(true)` reaps the dead child's process group.
|
||||
{
|
||||
let mut conn = self.conn.lock().await;
|
||||
*conn = new_conn;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
target: "nomi_mcp",
|
||||
command = %self.spec.command,
|
||||
"MCP stdio server respawned and re-initialized"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for StdioTransport {
|
||||
async fn request(&self, req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError> {
|
||||
// First attempt on the current connection.
|
||||
let first = {
|
||||
let mut conn = self.conn.lock().await;
|
||||
Self::roundtrip_on(&mut conn, req).await
|
||||
};
|
||||
|
||||
match first {
|
||||
Ok(resp) => Ok(resp),
|
||||
// Do NOT auto-respawn while the handshake itself is in flight: the
|
||||
// respawn path *runs* `initialize`, so retrying an `initialize`
|
||||
// request afterwards would double-initialize the fresh child and
|
||||
// desync the protocol. First-connect handshake failures are already
|
||||
// handled non-fatally by the manager.
|
||||
Err(err) if Self::is_pipe_failure(&err) && !is_handshake_method(&req.method) => {
|
||||
// The child/pipe died. Respawn + re-handshake, then retry once.
|
||||
self.respawn().await?;
|
||||
let mut conn = self.conn.lock().await;
|
||||
Self::roundtrip_on(&mut conn, req).await
|
||||
}
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
async fn notify(&self, req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
let first = {
|
||||
let mut conn = self.conn.lock().await;
|
||||
Self::send_on(&mut conn, req).await
|
||||
};
|
||||
match first {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) if Self::is_pipe_failure(&err) && !is_handshake_method(&req.method) => {
|
||||
self.respawn().await?;
|
||||
let mut conn = self.conn.lock().await;
|
||||
Self::send_on(&mut conn, req).await
|
||||
}
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
// Kill the child gracefully; `kill_on_drop` is the backstop.
|
||||
let mut conn = self.conn.lock().await;
|
||||
let _ = conn.child.kill().await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Handshake methods must not trigger the auto-respawn retry (respawn already
|
||||
/// replays the handshake; retrying would double-initialize the new child).
|
||||
fn is_handshake_method(method: &str) -> bool {
|
||||
matches!(method, "initialize" | "notifications/initialized")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
/// Pure-unit checks on the failure classifier — these need no child process.
|
||||
#[test]
|
||||
fn pipe_failure_classifies_transport_eof_and_io() {
|
||||
assert!(StdioTransport::is_pipe_failure(&McpError::Transport(
|
||||
"Child process stdout closed".into()
|
||||
)));
|
||||
assert!(StdioTransport::is_pipe_failure(&McpError::Transport(
|
||||
"Write to stdin failed: broken pipe".into()
|
||||
)));
|
||||
assert!(StdioTransport::is_pipe_failure(&McpError::Transport(
|
||||
"Read from stdout failed: x".into()
|
||||
)));
|
||||
assert!(StdioTransport::is_pipe_failure(&McpError::Io(
|
||||
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "boom")
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pipe_failure_excludes_jsonrpc_and_parse_errors() {
|
||||
// A JSON-RPC application error means the server answered — not a dead
|
||||
// pipe; respawning would not help, so it must NOT be classified as one.
|
||||
assert!(!StdioTransport::is_pipe_failure(&McpError::JsonRpc {
|
||||
code: -32601,
|
||||
message: "method not found".into(),
|
||||
}));
|
||||
// A parse failure is a protocol/serialize issue, not a broken pipe.
|
||||
assert!(!StdioTransport::is_pipe_failure(&McpError::Transport(
|
||||
"Failed to parse JSON-RPC response: x — raw: {".into()
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handshake_methods_are_excluded_from_respawn() {
|
||||
assert!(is_handshake_method("initialize"));
|
||||
assert!(is_handshake_method("notifications/initialized"));
|
||||
assert!(!is_handshake_method("tools/call"));
|
||||
assert!(!is_handshake_method("tools/list"));
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Respawn integration test: a mock stdio MCP server that crashes once.
|
||||
// Uses /bin/sh, so it is gated to unix.
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Write a mock MCP stdio server shell script that speaks line-delimited
|
||||
/// JSON-RPC. The script tracks how many times it has been *launched* via a
|
||||
/// shared counter file: launch #1 answers `initialize` + exactly one
|
||||
/// `tools/call`, then exits (EOF) to simulate a crash. Launch #2+ answers
|
||||
/// `initialize` and then every `tools/call` indefinitely.
|
||||
#[cfg(unix)]
|
||||
fn write_mock_server(dir: &std::path::Path) -> std::path::PathBuf {
|
||||
use std::io::Write;
|
||||
let launch_counter = dir.join("launches");
|
||||
let script_path = dir.join("mock_server.sh");
|
||||
// The script reads JSON-RPC lines from stdin and replies on stdout.
|
||||
// `initialize` → result; `tools/call` → a text content result; the
|
||||
// `notifications/initialized` notification gets no reply.
|
||||
let script = format!(
|
||||
r#"#!/bin/sh
|
||||
COUNTER="{counter}"
|
||||
# Record this launch (atomic-enough for a single-writer test).
|
||||
n=$(cat "$COUNTER" 2>/dev/null || echo 0)
|
||||
n=$((n + 1))
|
||||
echo "$n" > "$COUNTER"
|
||||
calls=0
|
||||
while IFS= read -r line; do
|
||||
case "$line" in
|
||||
*'"method":"initialize"'*)
|
||||
printf '{{"jsonrpc":"2.0","id":1,"result":{{"protocolVersion":"2025-03-26","capabilities":{{}},"serverInfo":{{"name":"mock","version":"0"}}}}}}\n'
|
||||
;;
|
||||
*'notifications/initialized'*)
|
||||
: # notification, no response
|
||||
;;
|
||||
*'"method":"tools/call"'*)
|
||||
calls=$((calls + 1))
|
||||
# On the very first launch, crash right after answering one call.
|
||||
if [ "$n" -eq 1 ] && [ "$calls" -ge 1 ]; then
|
||||
printf '{{"jsonrpc":"2.0","id":0,"result":{{"content":[{{"type":"text","text":"before-crash"}}]}}}}\n'
|
||||
exit 0
|
||||
fi
|
||||
printf '{{"jsonrpc":"2.0","id":0,"result":{{"content":[{{"type":"text","text":"after-respawn"}}]}}}}\n'
|
||||
;;
|
||||
*)
|
||||
: # ignore anything else
|
||||
;;
|
||||
esac
|
||||
done
|
||||
"#,
|
||||
counter = launch_counter.display()
|
||||
);
|
||||
let mut f = std::fs::File::create(&script_path).unwrap();
|
||||
f.write_all(script.as_bytes()).unwrap();
|
||||
f.flush().unwrap();
|
||||
drop(f);
|
||||
// Make it executable.
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let mut perms = std::fs::metadata(&script_path).unwrap().permissions();
|
||||
perms.set_mode(0o755);
|
||||
std::fs::set_permissions(&script_path, perms).unwrap();
|
||||
script_path
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn handshake(transport: &StdioTransport) {
|
||||
// Drive the same handshake the manager would, so the first connection
|
||||
// is fully initialized before we exercise tools/call.
|
||||
let init = JsonRpcRequest::new(1, "initialize", Some(json!({})));
|
||||
transport.request(&init).await.expect("initialize");
|
||||
let initialized = JsonRpcRequest::notification("notifications/initialized", None);
|
||||
transport.notify(&initialized).await.expect("initialized");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn respawn_recovers_after_child_crash() {
|
||||
let tmp = std::env::temp_dir().join(format!("nomi_mcp_respawn_{}", std::process::id()));
|
||||
std::fs::create_dir_all(&tmp).unwrap();
|
||||
let script = write_mock_server(&tmp);
|
||||
|
||||
let transport =
|
||||
StdioTransport::spawn("/bin/sh", &[script.to_string_lossy().into_owned()], &HashMap::new())
|
||||
.await
|
||||
.expect("spawn mock server");
|
||||
|
||||
handshake(&transport).await;
|
||||
|
||||
// First tools/call: the child answers "before-crash" then exits (EOF).
|
||||
// The next call detects the dead pipe, respawns, re-handshakes, retries.
|
||||
let call = JsonRpcRequest::new(2, "tools/call", Some(json!({"name": "t", "arguments": {}})));
|
||||
let r1 = transport.request(&call).await.expect("first call ok");
|
||||
assert_eq!(
|
||||
r1.result.unwrap()["content"][0]["text"],
|
||||
"before-crash",
|
||||
"first call should be served by the original child"
|
||||
);
|
||||
|
||||
// The second call lands after the child has exited → triggers respawn.
|
||||
// It must succeed against the freshly respawned (stable) child.
|
||||
let r2 = transport
|
||||
.request(&call)
|
||||
.await
|
||||
.expect("second call must recover via respawn");
|
||||
assert_eq!(
|
||||
r2.result.unwrap()["content"][0]["text"],
|
||||
"after-respawn",
|
||||
"second call should be served by the respawned child"
|
||||
);
|
||||
|
||||
// The respawn counter must show exactly one respawn (launch #2).
|
||||
let launches: u32 = std::fs::read_to_string(tmp.join("launches"))
|
||||
.unwrap()
|
||||
.trim()
|
||||
.parse()
|
||||
.unwrap();
|
||||
assert!(
|
||||
launches >= 2,
|
||||
"child should have been launched at least twice (got {launches})"
|
||||
);
|
||||
|
||||
let _ = transport.close().await;
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn respawn_budget_is_bounded() {
|
||||
// A server that exits immediately on every launch must not respawn
|
||||
// forever: after MAX_RESPAWNS the transport surfaces a hard error.
|
||||
use std::io::Write;
|
||||
let tmp = std::env::temp_dir().join(format!("nomi_mcp_crashloop_{}", std::process::id()));
|
||||
std::fs::create_dir_all(&tmp).unwrap();
|
||||
let script_path = tmp.join("always_crash.sh");
|
||||
// Answers initialize once, then exits the moment a tools/call arrives —
|
||||
// and the respawn's own re-handshake initialize also gets answered, but
|
||||
// the subsequent retried tools/call again hits EOF → respawn → ...
|
||||
let script = r#"#!/bin/sh
|
||||
while IFS= read -r line; do
|
||||
case "$line" in
|
||||
*'"method":"initialize"'*)
|
||||
printf '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-03-26","capabilities":{}}}\n'
|
||||
;;
|
||||
*'"method":"tools/call"'*)
|
||||
exit 0
|
||||
;;
|
||||
*) : ;;
|
||||
esac
|
||||
done
|
||||
"#;
|
||||
let mut f = std::fs::File::create(&script_path).unwrap();
|
||||
f.write_all(script.as_bytes()).unwrap();
|
||||
f.flush().unwrap();
|
||||
drop(f);
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let mut perms = std::fs::metadata(&script_path).unwrap().permissions();
|
||||
perms.set_mode(0o755);
|
||||
std::fs::set_permissions(&script_path, perms).unwrap();
|
||||
|
||||
let transport = StdioTransport::spawn(
|
||||
"/bin/sh",
|
||||
&[script_path.to_string_lossy().into_owned()],
|
||||
&HashMap::new(),
|
||||
)
|
||||
.await
|
||||
.expect("spawn crashloop server");
|
||||
handshake(&transport).await;
|
||||
|
||||
let call = JsonRpcRequest::new(2, "tools/call", Some(json!({"name": "t", "arguments": {}})));
|
||||
// Each request EOFs and respawns once, then the retried call EOFs again
|
||||
// → that request returns Err. Repeated requests keep respawning, but the
|
||||
// per-window budget (MAX_RESPAWNS) must stop the bleeding: once exhausted,
|
||||
// respawn() itself errors instead of forking yet another doomed child.
|
||||
// Every attempt must therefore return Err (never hang, never panic).
|
||||
for i in 0..(MAX_RESPAWNS as usize + 3) {
|
||||
let result = transport.request(&call).await;
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"attempt {i}: a server that crashes every call must surface an error, not hang"
|
||||
);
|
||||
}
|
||||
|
||||
// After the budget is spent, respawn() must report the crashloop ceiling
|
||||
// rather than silently keep trying.
|
||||
let final_err = transport.request(&call).await.unwrap_err();
|
||||
let msg = final_err.to_string();
|
||||
assert!(
|
||||
msg.contains("exceeded") && msg.contains("respawns"),
|
||||
"expected a crashloop-ceiling error once the budget is spent, got: {msg}"
|
||||
);
|
||||
|
||||
let _ = transport.close().await;
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use reqwest::header::{HeaderMap, HeaderValue};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::{McpError, McpTransport, find_sse_event_boundary};
|
||||
use crate::protocol::{JsonRpcRequest, JsonRpcResponse};
|
||||
|
||||
/// Streamable HTTP transport: uses HTTP POST for both requests and responses
|
||||
/// Supports optional SSE streaming for server responses
|
||||
pub struct StreamableHttpTransport {
|
||||
client: reqwest::Client,
|
||||
url: String,
|
||||
headers: HeaderMap,
|
||||
session_id: Mutex<Option<String>>,
|
||||
next_id: AtomicU64,
|
||||
}
|
||||
|
||||
impl StreamableHttpTransport {
|
||||
/// Create a new Streamable HTTP transport
|
||||
pub async fn connect(url: &str, headers: &HashMap<String, String>) -> Result<Self, McpError> {
|
||||
let mut header_map = HeaderMap::new();
|
||||
for (k, v) in headers {
|
||||
let name = reqwest::header::HeaderName::from_bytes(k.as_bytes())
|
||||
.map_err(|e| McpError::Transport(format!("Invalid header name '{}': {}", k, e)))?;
|
||||
let value = HeaderValue::from_str(v)
|
||||
.map_err(|e| McpError::Transport(format!("Invalid header value '{}': {}", v, e)))?;
|
||||
header_map.insert(name, value);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
client: reqwest::Client::new(),
|
||||
url: url.to_string(),
|
||||
headers: header_map,
|
||||
session_id: Mutex::new(None),
|
||||
next_id: AtomicU64::new(1),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn next_id(&self) -> u64 {
|
||||
self.next_id.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Build request with session ID header if available
|
||||
async fn build_request(&self, body: &str) -> reqwest::RequestBuilder {
|
||||
let mut req = self
|
||||
.client
|
||||
.post(&self.url)
|
||||
.headers(self.headers.clone())
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream");
|
||||
|
||||
if let Some(sid) = self.session_id.lock().await.as_ref() {
|
||||
req = req.header("Mcp-Session-Id", sid.as_str());
|
||||
}
|
||||
|
||||
req.body(body.to_string())
|
||||
}
|
||||
|
||||
/// Parse response based on content type
|
||||
async fn parse_response(
|
||||
&self,
|
||||
response: reqwest::Response,
|
||||
) -> Result<JsonRpcResponse, McpError> {
|
||||
// Capture session ID from response headers
|
||||
if let Some(sid) = response.headers().get("mcp-session-id")
|
||||
&& let Ok(sid_str) = sid.to_str()
|
||||
{
|
||||
*self.session_id.lock().await = Some(sid_str.to_string());
|
||||
}
|
||||
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if content_type.contains("text/event-stream") {
|
||||
// SSE response: parse events to find the JSON-RPC response
|
||||
self.parse_sse_response(response).await
|
||||
} else {
|
||||
// Direct JSON response
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Read response body failed: {}", e)))?;
|
||||
serde_json::from_str(&text).map_err(|e| {
|
||||
McpError::Transport(format!("Parse JSON response failed: {} — raw: {}", e, text))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse an SSE stream response to extract JSON-RPC response
|
||||
async fn parse_sse_response(
|
||||
&self,
|
||||
response: reqwest::Response,
|
||||
) -> Result<JsonRpcResponse, McpError> {
|
||||
use futures::StreamExt;
|
||||
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| McpError::Transport(format!("SSE read error: {}", e)))?;
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// Parse SSE events. Events may be framed with LF, CRLF (new-api /
|
||||
// one-api proxies), or bare CR — see find_sse_event_boundary.
|
||||
while let Some((event_end, delim_len)) = find_sse_event_boundary(&buffer) {
|
||||
let event_block = buffer[..event_end].to_string();
|
||||
buffer = buffer[event_end + delim_len..].to_string();
|
||||
|
||||
// Extract data lines
|
||||
let mut data_lines = Vec::new();
|
||||
for line in event_block.lines() {
|
||||
if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let data = data_lines.join("\n");
|
||||
if !data.is_empty()
|
||||
&& let Ok(rpc_response) = serde_json::from_str::<JsonRpcResponse>(&data)
|
||||
{
|
||||
return Ok(rpc_response);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(McpError::Transport(
|
||||
"SSE stream ended without JSON-RPC response".into(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for StreamableHttpTransport {
|
||||
async fn request(&self, req: &JsonRpcRequest) -> Result<JsonRpcResponse, McpError> {
|
||||
let body = serde_json::to_string(req)
|
||||
.map_err(|e| McpError::Transport(format!("JSON serialize error: {}", e)))?;
|
||||
|
||||
let http_req = self.build_request(&body).await;
|
||||
let response = http_req
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("HTTP request failed: {}", e)))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(McpError::Transport(format!(
|
||||
"HTTP request returned status: {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
let rpc_response = self.parse_response(response).await?;
|
||||
|
||||
if let Some(err) = &rpc_response.error {
|
||||
return Err(McpError::JsonRpc {
|
||||
code: err.code,
|
||||
message: err.message.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(rpc_response)
|
||||
}
|
||||
|
||||
async fn notify(&self, req: &JsonRpcRequest) -> Result<(), McpError> {
|
||||
let body = serde_json::to_string(req)
|
||||
.map_err(|e| McpError::Transport(format!("JSON serialize error: {}", e)))?;
|
||||
|
||||
let http_req = self.build_request(&body).await;
|
||||
http_req
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| McpError::Transport(format!("Notification request failed: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
// No persistent connection to close for HTTP
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user