提交qiming-mcp-proxy

This commit is contained in:
Codex
2026-06-01 13:03:20 +08:00
parent 9e9486b7c2
commit afb3d9f4e6
394 changed files with 124494 additions and 0 deletions

View File

@@ -0,0 +1,353 @@
//! SSE Client Connection Module
//!
//! Provides a high-level API for connecting to MCP servers via SSE protocol.
//! This module encapsulates the rmcp transport details and exposes a simple interface.
use anyhow::{Context, Result};
use mcp_common::McpClientConfig;
use rmcp::{
RoleClient, ServiceExt,
model::{ClientCapabilities, ClientInfo, Implementation, ProtocolVersion},
service::RunningService,
transport::{
SseClientTransport, common::client_side_sse::SseRetryPolicy, sse_client::SseClientConfig,
},
};
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;
use tracing::{debug, info};
use crate::sse_handler::SseHandler;
use mcp_common::ToolFilter;
/// 自定义的指数退避重试策略,支持最大间隔限制
///
/// 重试间隔按照指数增长,但不会超过 max_interval
/// - 第 1 次重试base_duration × 2^0
/// - 第 2 次重试base_duration × 2^1
/// - ...
/// - 第 n 次重试min(base_duration × 2^(n-1), max_interval)
#[derive(Debug, Clone)]
pub struct CappedExponentialBackoff {
/// 最大重试次数None 表示无限制
pub max_times: Option<usize>,
/// 基础延迟时间(第一次重试前的等待时间)
pub base_duration: Duration,
/// 最大延迟间隔(重试间隔不会超过这个值)
pub max_interval: Duration,
}
impl CappedExponentialBackoff {
/// 创建一个新的带上限的指数退避策略
///
/// # Arguments
/// * `max_times` - 最大重试次数None 表示无限制
/// * `base_duration` - 基础延迟时间
/// * `max_interval` - 最大延迟间隔
pub fn new(max_times: Option<usize>, base_duration: Duration, max_interval: Duration) -> Self {
Self {
max_times,
base_duration,
max_interval,
}
}
}
impl Default for CappedExponentialBackoff {
fn default() -> Self {
Self {
max_times: None,
base_duration: Duration::from_secs(1),
max_interval: Duration::from_secs(60),
}
}
}
impl SseRetryPolicy for CappedExponentialBackoff {
fn retry(&self, current_times: usize) -> Option<Duration> {
// 检查是否超过最大重试次数
if let Some(max_times) = self.max_times
&& current_times >= max_times
{
return None;
}
// 计算指数退避时间
let exponential_delay = self.base_duration * (2u32.pow(current_times as u32));
// 限制最大间隔
Some(exponential_delay.min(self.max_interval))
}
}
/// Opaque wrapper for SSE client connection
///
/// This type encapsulates an active connection to an MCP server via SSE protocol.
/// It hides the internal `RunningService` type and provides only the methods
/// needed by consuming code.
///
/// Note: This type is not Clone because the underlying RunningService
/// is designed for single-owner use. Use `into_handler()` or `into_running_service()`
/// to consume the connection.
///
/// # Example
///
/// ```rust,ignore
/// use mcp_sse_proxy::{SseClientConnection, McpClientConfig};
///
/// let config = McpClientConfig::new("http://localhost:8080/sse")
/// .with_header("Authorization", "Bearer token");
///
/// let conn = SseClientConnection::connect(config).await?;
/// let tools = conn.list_tools().await?;
/// println!("Available tools: {:?}", tools);
/// ```
pub struct SseClientConnection {
inner: RunningService<RoleClient, ClientInfo>,
}
impl SseClientConnection {
/// Connect to an SSE MCP server
///
/// # Arguments
/// * `config` - Client configuration including URL and headers
///
/// # Returns
/// * `Ok(SseClientConnection)` - Successfully connected client
/// * `Err` - Connection failed
pub async fn connect(config: McpClientConfig) -> Result<Self> {
let start = Instant::now();
info!("🔗 Start establishing SSE connection: {}", config.url);
debug!("Building the HTTP client configuration...");
let http_client = build_http_client(&config)?;
// 配置指数退避重试策略,最大间隔 1 分钟,不限制重试次数
let retry_policy = CappedExponentialBackoff::new(
None, // 不限制重试次数
Duration::from_secs(1), // 基础延迟 1 秒
Duration::from_secs(60), // 最大间隔 60 秒
);
let sse_config = SseClientConfig {
sse_endpoint: config.url.clone().into(),
retry_policy: Arc::new(retry_policy),
..Default::default()
};
debug!("Starting the SSE transport layer...");
let transport: SseClientTransport<reqwest::Client> =
SseClientTransport::start_with_client(http_client, sse_config)
.await
.context("Failed to start SSE transport")?;
let transport_elapsed = start.elapsed();
debug!(
"SSE transport layer startup completed, time taken: {:?}",
transport_elapsed
);
debug!("Initializing MCP client handshake...");
let client_info = create_default_client_info();
let running = client_info
.serve(transport)
.await
.context("Failed to initialize MCP client")?;
let total_elapsed = start.elapsed();
// 记录详细的连接状态信息
{
use std::ops::Deref;
let transport_closed = running.deref().is_transport_closed();
let peer_info = running.peer_info();
info!(
"✅ SSE connection established successfully - total time taken: {:?}, transport_closed: {}, peer_info: {:?}",
total_elapsed, transport_closed, peer_info
);
if let Some(info) = peer_info {
info!(
"Server information: name={}, version={}, capabilities={:?}",
info.server_info.name, info.server_info.version, info.capabilities
);
}
}
Ok(Self { inner: running })
}
/// List available tools from the MCP server
pub async fn list_tools(&self) -> Result<Vec<ToolInfo>> {
let result = self.inner.list_tools(None).await?;
Ok(result
.tools
.into_iter()
.map(|t| ToolInfo {
name: t.name.to_string(),
description: t.description.map(|d| d.to_string()),
})
.collect())
}
/// Check if the connection is closed
pub fn is_closed(&self) -> bool {
use std::ops::Deref;
self.inner.deref().is_transport_closed()
}
/// Get the peer info from the server
pub fn peer_info(&self) -> Option<&rmcp::model::ServerInfo> {
self.inner.peer_info()
}
/// Convert this connection into an SseHandler for serving
///
/// This consumes the connection and creates an SseHandler that can
/// proxy requests to the backend MCP server.
///
/// # Arguments
/// * `mcp_id` - Identifier for logging purposes
/// * `tool_filter` - Tool filtering configuration
pub fn into_handler(self, mcp_id: String, tool_filter: ToolFilter) -> SseHandler {
SseHandler::with_tool_filter(self.inner, mcp_id, tool_filter)
}
/// Extract the internal RunningService for use with swap_backend
///
/// This is used internally to support backend hot-swapping.
pub fn into_running_service(self) -> RunningService<RoleClient, ClientInfo> {
self.inner
}
}
/// Simplified tool information
#[derive(Clone, Debug)]
pub struct ToolInfo {
/// Tool name
pub name: String,
/// Tool description (optional)
pub description: Option<String>,
}
/// Build an HTTP client with the given configuration
fn build_http_client(config: &McpClientConfig) -> Result<reqwest::Client> {
let mut headers = reqwest::header::HeaderMap::new();
for (key, value) in &config.headers {
let header_name = key
.parse::<reqwest::header::HeaderName>()
.with_context(|| format!("Invalid header name: {}", key))?;
let header_value = value
.parse()
.with_context(|| format!("Invalid header value for {}: {}", key, value))?;
headers.insert(header_name, header_value);
}
let mut builder = reqwest::Client::builder().default_headers(headers);
if let Some(timeout) = config.connect_timeout {
builder = builder.connect_timeout(timeout);
}
if let Some(timeout) = config.read_timeout {
builder = builder.timeout(timeout);
}
builder.build().context("Failed to build HTTP client")
}
/// Create default client info for MCP handshake
fn create_default_client_info() -> ClientInfo {
ClientInfo {
protocol_version: ProtocolVersion::V_2024_11_05,
capabilities: ClientCapabilities::builder()
.enable_experimental()
.enable_roots()
.enable_roots_list_changed()
.enable_sampling()
.build(),
client_info: Implementation {
name: "mcp-sse-proxy-client".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
title: None,
website_url: None,
icons: None,
},
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tool_info() {
let info = ToolInfo {
name: "test_tool".to_string(),
description: Some("A test tool".to_string()),
};
assert_eq!(info.name, "test_tool");
assert_eq!(info.description, Some("A test tool".to_string()));
}
#[test]
fn test_capped_exponential_backoff() {
// 测试带上限的指数退避策略
let policy = CappedExponentialBackoff::new(
None, // 不限制重试次数
Duration::from_secs(1), // 基础延迟 1 秒
Duration::from_secs(60), // 最大间隔 60 秒
);
// 验证第 1 次重试1 秒
assert_eq!(policy.retry(0), Some(Duration::from_secs(1)));
// 验证第 2 次重试2 秒
assert_eq!(policy.retry(1), Some(Duration::from_secs(2)));
// 验证第 3 次重试4 秒
assert_eq!(policy.retry(2), Some(Duration::from_secs(4)));
// 验证第 4 次重试8 秒
assert_eq!(policy.retry(3), Some(Duration::from_secs(8)));
// 验证第 5 次重试16 秒
assert_eq!(policy.retry(4), Some(Duration::from_secs(16)));
// 验证第 6 次重试32 秒
assert_eq!(policy.retry(5), Some(Duration::from_secs(32)));
// 验证第 7 次重试64 秒 -> 会被限制为 60 秒
assert_eq!(policy.retry(6), Some(Duration::from_secs(60)));
// 验证第 8 次重试128 秒 -> 会被限制为 60 秒
assert_eq!(policy.retry(7), Some(Duration::from_secs(60)));
}
#[test]
fn test_capped_exponential_backoff_with_max_times() {
// 测试带最大重试次数的限制
let policy = CappedExponentialBackoff::new(
Some(3), // 最多重试 3 次
Duration::from_secs(1), // 基础延迟 1 秒
Duration::from_secs(60), // 最大间隔 60 秒
);
// 验证前 3 次重试都有延迟时间
assert_eq!(policy.retry(0), Some(Duration::from_secs(1)));
assert_eq!(policy.retry(1), Some(Duration::from_secs(2)));
assert_eq!(policy.retry(2), Some(Duration::from_secs(4)));
// 验证第 4 次重试(重试次数已达到上限)
assert_eq!(policy.retry(3), None);
}
#[test]
fn test_capped_exponential_backoff_default() {
// 测试默认配置
let policy = CappedExponentialBackoff::default();
// 验证默认配置
assert_eq!(policy.max_times, None);
assert_eq!(policy.base_duration, Duration::from_secs(1));
assert_eq!(policy.max_interval, Duration::from_secs(60));
// 验证重试行为
assert_eq!(policy.retry(0), Some(Duration::from_secs(1)));
assert_eq!(policy.retry(5), Some(Duration::from_secs(32)));
assert_eq!(policy.retry(10), Some(Duration::from_secs(60)));
}
}

View File

@@ -0,0 +1,28 @@
//! Configuration types for SSE proxy
use serde::{Deserialize, Serialize};
/// SSE server configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SseConfig {
/// Bind address for the HTTP server
#[serde(default = "default_bind_addr")]
pub bind_addr: String,
/// Quiet mode (suppress startup messages)
#[serde(default)]
pub quiet: bool,
}
impl Default for SseConfig {
fn default() -> Self {
Self {
bind_addr: default_bind_addr(),
quiet: false,
}
}
}
fn default_bind_addr() -> String {
"127.0.0.1:3001".to_string()
}

View File

@@ -0,0 +1,200 @@
use futures::StreamExt;
use reqwest::header::{ACCEPT, CONTENT_TYPE, HeaderMap, HeaderValue};
use std::collections::HashMap;
use std::time::Duration;
use tracing::debug;
/// Detect if a URL supports the SSE (Server-Sent Events) MCP protocol
///
/// Convenience wrapper around [`is_sse_with_headers`] that passes no custom headers.
pub async fn is_sse(url: &str) -> bool {
is_sse_with_headers(url, None).await
}
/// Detect if a URL supports the MCP SSE protocol, with optional custom headers
///
/// MCP SSE protocol has a unique characteristic: upon GET connection to the SSE
/// endpoint, the server sends an `event: endpoint` event containing the URL for
/// POSTing messages. This `endpoint` event is exclusive to MCP SSE and never appears
/// in Streamable HTTP, making it the definitive distinguishing feature.
///
/// # Detection logic
///
/// 1. Send GET request with `Accept: text/event-stream`
/// 2. Verify response Content-Type is `text/event-stream`
/// 3. Read the first few events from the SSE stream
/// 4. If an `event: endpoint` is found → confirmed MCP SSE
///
/// # Candidate URLs
///
/// - If URL ends with `/sse`, try it as-is
/// - Otherwise, try `{url}/sse` first (MCP SSE convention), then the original URL
///
/// # Arguments
///
/// * `url` - The URL to test
/// * `custom_headers` - Optional custom headers (e.g., Authorization)
///
/// # Returns
///
/// Returns `true` if the URL supports MCP SSE protocol, `false` otherwise.
pub async fn is_sse_with_headers(
url: &str,
custom_headers: Option<&HashMap<String, String>>,
) -> bool {
let client = match reqwest::Client::builder()
.connect_timeout(Duration::from_secs(5))
.build()
{
Ok(c) => c,
Err(_) => return false,
};
let mut headers = HeaderMap::new();
headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream"));
// Merge custom headers
if let Some(custom) = custom_headers {
for (key, value) in custom {
if let (Ok(name), Ok(val)) = (
reqwest::header::HeaderName::try_from(key.as_str()),
HeaderValue::from_str(value),
) {
headers.insert(name, val);
}
}
}
// Build candidate URLs
let trimmed = url.trim_end_matches('/');
let candidates: Vec<String> = if trimmed.ends_with("/sse") {
vec![url.to_string()]
} else {
// Try /sse suffix first (MCP SSE convention), then original URL
vec![format!("{}/sse", trimmed), url.to_string()]
};
for probe_url in &candidates {
debug!("SSE probe: try {}", probe_url);
match tokio::time::timeout(
Duration::from_secs(5),
probe_sse_endpoint(&client, probe_url, &headers),
)
.await
{
Ok(true) => {
debug!(
"SSE probe: Confirm {} is MCP SSE protocol (discover endpoint event)",
probe_url
);
return true;
}
Ok(false) => {
debug!("SSE probe: {} is not MCP SSE protocol", probe_url);
}
Err(_) => {
debug!("SSE probe: {} timeout", probe_url);
}
}
}
false
}
/// Probe a single URL for MCP SSE protocol
///
/// Returns `true` only if the response is `text/event-stream` AND contains
/// an `event: endpoint` SSE event (the MCP SSE distinguishing feature).
async fn probe_sse_endpoint(client: &reqwest::Client, url: &str, headers: &HeaderMap) -> bool {
let response = match client.get(url).headers(headers.clone()).send().await {
Ok(r) => r,
Err(_) => return false,
};
if !response.status().is_success() {
return false;
}
// Verify Content-Type is text/event-stream
let is_event_stream = response
.headers()
.get(CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.is_some_and(|ct| ct.contains("text/event-stream"));
if !is_event_stream {
return false;
}
// Read the SSE stream looking for "event: endpoint" — the MCP SSE signature.
// The endpoint event is sent immediately upon connection, so we only need
// to read the first few chunks (up to 4KB).
read_sse_for_endpoint_event(response).await
}
/// Read SSE stream and check for the `endpoint` event
///
/// MCP SSE servers send `event: endpoint\ndata: <message_url>\n\n` as the
/// first event after connection. This event is never sent by Streamable HTTP.
async fn read_sse_for_endpoint_event(response: reqwest::Response) -> bool {
let mut stream = response.bytes_stream();
let mut buffer = String::new();
const MAX_BYTES: usize = 4096;
while let Some(chunk) = stream.next().await {
match chunk {
Ok(bytes) => {
if let Ok(text) = std::str::from_utf8(&bytes) {
buffer.push_str(text);
}
// SSE event format: "event: endpoint\n" or "event:endpoint\n"
if buffer.contains("event: endpoint") || buffer.contains("event:endpoint") {
return true;
}
// Don't read too much — if endpoint event hasn't appeared
// in the first 4KB, it's not MCP SSE
if buffer.len() > MAX_BYTES {
return false;
}
}
Err(_) => return false,
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_is_sse_nonexistent_server() {
let result = is_sse("http://localhost:99999/mcp").await;
assert!(!result);
}
#[tokio::test]
async fn test_is_sse_with_headers_no_panic() {
let mut headers = HashMap::new();
headers.insert("Authorization".to_string(), "Bearer test-token".to_string());
let result = is_sse_with_headers("http://localhost:99999/mcp", Some(&headers)).await;
assert!(!result);
}
#[tokio::test]
async fn test_candidate_urls_with_sse_suffix() {
// URL already ending with /sse should not get /sse appended
let result = is_sse("http://localhost:99999/sse").await;
assert!(!result);
}
#[tokio::test]
async fn test_candidate_urls_with_trailing_slash() {
// Trailing slash should be trimmed before appending /sse
let result = is_sse("http://localhost:99999/mcp/").await;
assert!(!result);
}
}

View File

@@ -0,0 +1,66 @@
//! MCP SSE Proxy Module
//!
//! This module provides a proxy implementation for MCP (Model Context Protocol)
//! using SSE (Server-Sent Events) transport.
//!
//! # Features
//!
//! - **SSE Support**: Uses rmcp 0.10 with SSE transport (removed in 0.12+)
//! - **Stable Protocol**: Production-ready SSE implementation
//! - **Hot Swap**: Supports backend connection replacement
//! - **Fallback Option**: Used when Streamable HTTP is not supported
//! - **High-level Client API**: Simple connection interface hiding transport details
//!
//! # Architecture
//!
//! ```text
//! Client → SSE → SseHandler → Backend MCP Service
//! ```
//!
//! # Example
//!
//! ```rust,ignore
//! use mcp_sse_proxy::{SseClientConnection, McpClientConfig};
//!
//! // Connect to an MCP server
//! let config = McpClientConfig::new("http://localhost:8080/sse");
//! let conn = SseClientConnection::connect(config).await?;
//!
//! // List available tools
//! let tools = conn.list_tools().await?;
//! ```
pub mod client;
pub mod config;
pub mod detector;
pub mod server;
pub mod server_builder;
pub mod sse_handler;
// Re-export SSE protocol detection function
pub use detector::{is_sse, is_sse_with_headers};
// Re-export main types
pub use mcp_common::McpServiceConfig;
pub use server::{run_sse_server, run_sse_server_from_config};
pub use sse_handler::{BackendSessionHandler, SseHandler, SseServerHandler, ToolFilter};
// Re-export server builder API
pub use server_builder::{BackendConfig, SseServerBuilder, SseServerBuilderConfig};
// Re-export client connection types
pub use client::{SseClientConnection, ToolInfo};
pub use mcp_common::McpClientConfig;
// Re-export commonly used rmcp types
pub use rmcp::{
RoleClient, RoleServer, ServerHandler, ServiceExt,
model::{CallToolRequestParam, ClientCapabilities, ClientInfo, Implementation, ServerInfo},
service::{Peer, RunningService},
};
// Re-export transport types for SSE protocol (rmcp 0.10)
pub use rmcp::transport::{
SseClientTransport, SseServer, child_process::TokioChildProcess, sse_client::SseClientConfig,
sse_server::SseServerConfig, stdio,
};

View File

@@ -0,0 +1,329 @@
//! SSE server implementation
//!
//! This module provides the SSE server using rmcp 0.10's stable SSE transport.
use anyhow::{Result, bail};
use mcp_common::{McpServiceConfig, check_windows_command, wrap_process_v8};
use rmcp::{
ServiceExt,
model::{ClientCapabilities, ClientInfo, ProtocolVersion},
transport::{
TokioChildProcess,
sse_server::{SseServer, SseServerConfig},
},
};
use std::process::Stdio;
use tokio_util::sync::CancellationToken;
use tracing::{error, info};
// 进程组管理(跨平台子进程清理)
use process_wrap::tokio::{KillOnDrop, TokioCommandWrap};
use crate::SseHandler;
/// 从配置启动 SSE 服务器
///
/// # Features
///
/// - **SSE Protocol**: 使用稳定的 rmcp 0.10 SSE 实现
/// - **Hot Swap**: 支持后端连接热替换
/// - **Full Lifecycle**: 自动创建子进程、连接、handler、服务器
///
/// # Arguments
///
/// * `config` - MCP 服务配置
/// * `std_listener` - 预先绑定的 TCP 监听器(端口在重试循环前绑定,保证端口占用)
/// * `quiet` - 静默模式,不输出启动信息
pub async fn run_sse_server_from_config(
config: McpServiceConfig,
std_listener: &std::net::TcpListener,
quiet: bool,
) -> Result<()> {
// 1. 使用 process-wrap 创建子进程命令(跨平台进程清理)
// process-wrap 会自动处理进程组Unix或 Job ObjectWindows
// 并且在 Drop 时自动清理子进程树
// 子进程默认继承父进程的所有环境变量
// 🔧 Windows 特殊处理:检测并转换 .cmd/.bat 文件避免弹窗
// 如果用户配置了 npm 全局安装的 MCP 服务(如 npx some-server 或 some-server.cmd
// 直接运行会弹 CMD 窗口。这里尝试转换
check_windows_command(&config.command);
info!(
"[Subprocess][{}] Command: {} {:?}",
config.name,
config.command,
config.args.as_ref().unwrap_or(&vec![])
);
let mut wrapped_cmd = TokioCommandWrap::with_new(&config.command, |command| {
if let Some(ref cmd_args) = config.args {
command.args(cmd_args);
}
// 子进程默认继承父进程的所有环境变量
// 设置 MCP JSON 配置中的环境变量(会覆盖继承的同名变量)
if let Some(ref env_vars) = config.env {
for (k, v) in env_vars {
command.env(k, v);
}
}
});
// 应用平台特定的进程包装Unix: ProcessGroup, Windows: CREATE_NO_WINDOW + JobObject
wrap_process_v8!(wrapped_cmd);
// 所有平台: Drop 时自动清理进程
wrapped_cmd.wrap(KillOnDrop);
// 2. 启动子进程rmcp 的 TokioChildProcess 已经支持 process-wrap
// 使用 builder 模式捕获 stderr便于诊断子 MCP 服务初始化失败
let (tokio_process, child_stderr) = TokioChildProcess::builder(wrapped_cmd)
.stderr(Stdio::piped())
.spawn()?;
// 启动 stderr 日志读取任务
if let Some(stderr_pipe) = child_stderr {
mcp_common::spawn_stderr_reader(stderr_pipe, config.name.clone());
}
// 3. 创建客户端信息
let client_info = ClientInfo {
protocol_version: ProtocolVersion::V_2024_11_05,
capabilities: ClientCapabilities::builder()
.enable_experimental()
.enable_roots()
.enable_roots_list_changed()
.enable_sampling()
.build(),
..Default::default()
};
// 4. 连接到子进程
let client = client_info.serve(tokio_process).await?;
// 记录子进程启动到日志文件
info!(
"[Subprocess startup] SSE - Service name: {}, Command: {} {:?}",
config.name,
config.command,
config.args.as_ref().unwrap_or(&vec![])
);
if !quiet {
eprintln!("✅ The child process has been started");
// 获取并打印工具列表
match client.list_tools(None).await {
Ok(tools_result) => {
let tools = &tools_result.tools;
if tools.is_empty() {
info!(
"[Tool list] Tool list is empty - Service name: {}",
config.name
);
eprintln!("Tool list is empty");
} else {
info!(
"[Tool list] Service name: {}, Number of tools: {}",
config.name,
tools.len()
);
eprintln!("🔧 Available tools ({}):", tools.len());
for tool in tools.iter().take(10) {
let desc = tool.description.as_deref().unwrap_or("无描述");
let desc_short = if desc.len() > 50 {
format!("{}...", &desc[..50])
} else {
desc.to_string()
};
eprintln!(" - {} : {}", tool.name, desc_short);
}
if tools.len() > 10 {
eprintln!("... and {} other tools", tools.len() - 10);
}
}
}
Err(e) => {
error!(
"[Tool List] Failed to obtain tool list - Service name: {}, Error: {}",
config.name, e
);
eprintln!("⚠️ Failed to obtain tool list: {}", e);
}
}
} else {
// 即使静默模式也记录日志
match client.list_tools(None).await {
Ok(tools_result) => {
info!(
"[Tool list] Service name: {}, Number of tools: {}",
config.name,
tools_result.tools.len()
);
}
Err(e) => {
error!(
"[Tool List] Failed to obtain tool list - Service name: {}, Error: {}",
config.name, e
);
}
}
}
// 5. 创建 SseHandler
let sse_handler = if let Some(tool_filter) = config.tool_filter {
SseHandler::with_tool_filter(client, config.name.clone(), tool_filter)
} else {
SseHandler::with_mcp_id(client, config.name.clone())
};
// 6. 启动服务器(使用预绑定的 listener
let listener = tokio::net::TcpListener::from_std(std_listener.try_clone()?)?;
run_sse_server(sse_handler, listener, quiet).await
}
/// Run SSE server with rmcp 0.10
///
/// # Features
///
/// - **SSE Protocol**: 使用稳定的 rmcp 0.10 SSE 实现
/// - **Hot Swap**: 支持后端连接热替换(通过 SseHandler
/// - **Keep-Alive**: 15秒心跳防止连接被中间代理关闭
///
/// # Arguments
///
/// * `sse_handler` - SseHandler 实例(包含热替换逻辑)
/// * `listener` - 已绑定的 tokio TcpListener
/// * `quiet` - 静默模式,不输出启动信息
pub async fn run_sse_server(
sse_handler: SseHandler,
listener: tokio::net::TcpListener,
quiet: bool,
) -> Result<()> {
// 从 listener 获取绑定地址
let bind_addr = listener.local_addr()?;
let bind_addr_str = bind_addr.to_string();
// 默认的 SSE 和消息路径
let sse_path = "/sse".to_string();
let message_path = "/message".to_string();
let mcp_id = sse_handler.mcp_id().to_string();
// 记录服务启动到日志文件
info!(
"[HTTP service startup] SSE service startup - Address: {}, MCP ID: {}, SSE endpoint: {}, Message endpoint: {}",
bind_addr_str, mcp_id, sse_path, message_path
);
if !quiet {
eprintln!("📡 SSE service startup: http://{}", bind_addr_str);
eprintln!("SSE endpoint: http://{}{}", bind_addr_str, sse_path);
eprintln!("Message endpoint: http://{}{}", bind_addr_str, message_path);
eprintln!(
"💡 MCP client can be used directly: http://{} (automatic redirection)",
bind_addr_str
);
eprintln!("🔄 Backend hot replacement: enabled");
eprintln!("💡 Press Ctrl+C to stop the service");
}
// 配置 SSE 服务器
let config = SseServerConfig {
bind: bind_addr,
sse_path: sse_path.clone(),
post_path: message_path.clone(),
ct: CancellationToken::new(),
sse_keep_alive: Some(std::time::Duration::from_secs(15)),
};
// 创建 SSE 服务器
let (sse_server, sse_router) = SseServer::new(config);
let ct = sse_server.with_service(move || sse_handler.clone());
// 根路径兼容处理器 - 自动重定向到正确的端点
let sse_path_for_fallback = sse_path.clone();
let message_path_for_fallback = message_path.clone();
let fallback_handler = move |method: axum::http::Method, headers: axum::http::HeaderMap| {
let sse_path = sse_path_for_fallback.clone();
let message_path = message_path_for_fallback.clone();
async move {
match method {
axum::http::Method::GET => {
// 检查 Accept 头
let accept = headers
.get("accept")
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if accept.contains("text/event-stream") {
// SSE 请求,重定向到 /sse
(
axum::http::StatusCode::TEMPORARY_REDIRECT,
[("Location", sse_path)],
"Redirecting to SSE endpoint".to_string(),
)
} else {
// 普通 GET 请求,返回服务信息
(
axum::http::StatusCode::OK,
[("Content-Type", "application/json".to_string())],
serde_json::json!({
"status": "running",
"protocol": "SSE",
"endpoints": {
"sse": sse_path,
"message": message_path
},
"usage": "Connect your MCP client to this URL or the SSE endpoint directly"
}).to_string(),
)
}
}
axum::http::Method::POST => {
// POST 请求,重定向到 /message
(
axum::http::StatusCode::TEMPORARY_REDIRECT,
[("Location", message_path)],
"Redirecting to message endpoint".to_string(),
)
}
_ => (
axum::http::StatusCode::METHOD_NOT_ALLOWED,
[("Allow", "GET, POST".to_string())],
"Method not allowed".to_string(),
),
}
}
};
// 合并路由SSE 路由 + 根路径兼容处理
let router = sse_router.fallback(fallback_handler);
// 使用传入的 listener 启动 HTTP 服务器
// 使用 select 处理 Ctrl+C 和服务器
tokio::select! {
result = axum::serve(listener, router) => {
if let Err(e) = result {
error!(
"[HTTP Service Error] SSE Server Error - MCP ID: {}, Error: {}",
mcp_id, e
);
bail!("服务器错误: {}", e);
}
}
_ = tokio::signal::ctrl_c() => {
info!(
"[HTTP service shutdown] Received exit signal, closing SSE service - MCP ID: {}",
mcp_id
);
if !quiet {
eprintln!("\\n🛑 Received exit signal, closing...");
}
ct.cancel();
}
}
Ok(())
}

View File

@@ -0,0 +1,653 @@
//! SSE Server Builder
//!
//! This module provides a high-level Builder API for creating SSE MCP servers.
//! It encapsulates all rmcp-specific types and provides a simple interface for mcp-proxy.
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use anyhow::Result;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
// 进程组管理(跨平台子进程清理)
use process_wrap::tokio::{KillOnDrop, TokioCommandWrap};
#[cfg(unix)]
use process_wrap::tokio::ProcessGroup;
#[cfg(windows)]
use process_wrap::tokio::JobObject;
use rmcp::{
ServiceExt,
model::{ClientCapabilities, ClientInfo, ProtocolVersion},
transport::{
SseClientTransport, TokioChildProcess,
sse_client::SseClientConfig,
sse_server::{SseServer, SseServerConfig},
},
};
use crate::{
sse_handler::{BackendSessionHandler, SseServerHandler},
SseHandler, ToolFilter,
};
/// Performance warning threshold for stdio (child process) backend connections
const STDIO_SLOW_THRESHOLD_SECS: u64 = 30;
/// Performance warning threshold for HTTP-based backend connections (SSE/Stream)
const HTTP_SLOW_THRESHOLD_SECS: u64 = 10;
/// Backend configuration for the MCP server
///
/// Defines how the proxy connects to the upstream MCP service.
#[derive(Clone)]
pub enum BackendConfig {
/// Connect to a local command via stdio
Stdio {
/// Command to execute (e.g., "npx", "python", etc.)
command: String,
/// Arguments for the command
args: Option<Vec<String>>,
/// Environment variables
env: Option<HashMap<String, String>>,
},
/// Connect to a remote URL using SSE protocol
SseUrl {
/// URL of the MCP SSE service
url: String,
/// Custom HTTP headers (including Authorization)
headers: Option<HashMap<String, String>>,
},
/// Use a pre-connected backend via BackendBridge trait
///
/// The caller is responsible for establishing the backend connection and
/// passing it as an `Arc<dyn BackendBridge>`. This enables protocol conversion
/// (e.g., Streamable HTTP backend → SSE frontend) without this module
/// depending on the specific backend implementation.
BackendBridge(Arc<dyn mcp_common::BackendBridge>),
}
impl std::fmt::Debug for BackendConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BackendConfig::Stdio { command, args, .. } => f
.debug_struct("Stdio")
.field("command", command)
.field("args", args)
.finish(),
BackendConfig::SseUrl { url, .. } => {
f.debug_struct("SseUrl").field("url", url).finish()
}
BackendConfig::BackendBridge(bridge) => f
.debug_struct("BackendBridge")
.field("mcp_id", &bridge.mcp_id())
.finish(),
}
}
}
/// Configuration for the SSE server
#[derive(Debug, Clone)]
pub struct SseServerBuilderConfig {
/// SSE endpoint path (default: "/sse")
pub sse_path: String,
/// Message endpoint path (default: "/message")
pub post_path: String,
/// MCP service identifier for logging
pub mcp_id: Option<String>,
/// Tool filter configuration
pub tool_filter: Option<ToolFilter>,
/// Keep-alive interval in seconds (default: 15)
pub keep_alive_secs: u64,
/// Enable stateful mode with full MCP initialization (default: true)
/// When false, uses `with_service_directly` which skips initialization for faster responses
pub stateful: bool,
}
impl Default for SseServerBuilderConfig {
fn default() -> Self {
Self {
sse_path: "/sse".into(),
post_path: "/message".into(),
mcp_id: None,
tool_filter: None,
keep_alive_secs: 15,
stateful: true,
}
}
}
/// Log connection timing with optional performance warning
///
/// # Arguments
///
/// * `mcp_id` - MCP service identifier
/// * `backend_type` - Type of backend (e.g., "stdio", "SSE", "Streamable HTTP")
/// * `total_duration` - Total connection time
/// * `breakdown` - Optional breakdown of timing components
/// * `warn_threshold_secs` - Threshold for performance warning
/// * `warn_message` - Message to show if threshold exceeded
fn log_connection_timing(
mcp_id: &str,
backend_type: &str,
total_duration: Duration,
breakdown: &[(&str, Duration)],
warn_threshold_secs: u64,
warn_message: &str,
) {
let breakdown_str: Vec<String> = breakdown
.iter()
.map(|(name, dur)| format!("{}: {:?}", name, dur))
.collect();
info!(
"[SseServerBuilder] {} backend connected successfully - MCP ID: {}, total: {:?} ({})",
backend_type,
mcp_id,
total_duration,
breakdown_str.join(", ")
);
if total_duration.as_secs() >= warn_threshold_secs {
warn!(
"[SseServerBuilder] {} backend connection takes a long time - MCP ID: {}, time: {:?}, {}",
backend_type, mcp_id, total_duration, warn_message
);
}
}
/// Builder for creating SSE MCP servers
///
/// Provides a fluent API for configuring and building MCP proxy servers.
///
/// # Example
///
/// ```rust,ignore
/// use mcp_sse_proxy::server_builder::{SseServerBuilder, BackendConfig};
///
/// // Create a server with stdio backend
/// let (router, ct) = SseServerBuilder::new(BackendConfig::Stdio {
/// command: "npx".into(),
/// args: Some(vec!["-y".into(), "@modelcontextprotocol/server-filesystem".into()]),
/// env: None,
/// })
/// .mcp_id("my-server")
/// .sse_path("/custom/sse")
/// .post_path("/custom/message")
/// .stateful(false) // Disable stateful mode for OneShot services (faster responses)
/// .build()
/// .await?;
/// ```
pub struct SseServerBuilder {
backend_config: BackendConfig,
server_config: SseServerBuilderConfig,
}
impl SseServerBuilder {
/// Create a new builder with the given backend configuration
pub fn new(backend: BackendConfig) -> Self {
Self {
backend_config: backend,
server_config: SseServerBuilderConfig::default(),
}
}
/// Set the SSE endpoint path
pub fn sse_path(mut self, path: impl Into<String>) -> Self {
self.server_config.sse_path = path.into();
self
}
/// Set the message endpoint path
pub fn post_path(mut self, path: impl Into<String>) -> Self {
self.server_config.post_path = path.into();
self
}
/// Set the MCP service identifier
///
/// Used for logging and service identification.
pub fn mcp_id(mut self, id: impl Into<String>) -> Self {
self.server_config.mcp_id = Some(id.into());
self
}
/// Set the tool filter configuration
pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
self.server_config.tool_filter = Some(filter);
self
}
/// Set the keep-alive interval in seconds
pub fn keep_alive(mut self, secs: u64) -> Self {
self.server_config.keep_alive_secs = secs;
self
}
/// Set stateful mode (default: true)
///
/// When false, uses `with_service_directly` which skips MCP initialization
/// for faster responses. This is recommended for OneShot services.
pub fn stateful(mut self, stateful: bool) -> Self {
self.server_config.stateful = stateful;
self
}
/// Build the server and return an axum Router, CancellationToken, and handler
///
/// The router can be merged with other axum routers or served directly.
/// The CancellationToken can be used to gracefully shut down the service.
/// The handler is a unified type that can be either SseHandler or BackendSessionHandler.
pub async fn build(self) -> Result<(axum::Router, CancellationToken, SseServerHandler)> {
let mcp_id = self
.server_config
.mcp_id
.clone()
.unwrap_or_else(|| "sse-proxy".into());
// Create client info for connecting to backend
let client_info = ClientInfo {
protocol_version: ProtocolVersion::V_2024_11_05,
capabilities: ClientCapabilities::builder()
.enable_experimental()
.enable_roots()
.enable_roots_list_changed()
.enable_sampling()
.build(),
..Default::default()
};
// Connect to backend based on configuration
let client = match &self.backend_config {
BackendConfig::Stdio { command, args, env } => {
self.connect_stdio(command, args, env, &client_info).await?
}
BackendConfig::SseUrl { url, headers } => {
self.connect_sse_url(url, headers, &client_info).await?
}
BackendConfig::BackendBridge(bridge) => {
let handler = BackendSessionHandler::new(bridge.clone(), mcp_id.clone());
let (router, ct, handler_for_return) =
self.create_server_with_backend_session(handler).await?;
info!(
"[SseServerBuilder] Server created with backend bridge \
- mcp_id: {}, sse_path: {}, post_path: {}",
mcp_id, self.server_config.sse_path, self.server_config.post_path
);
return Ok((
router,
ct,
SseServerHandler::BackendSession(handler_for_return),
));
}
};
// Create SSE handler
let sse_handler = if let Some(ref tool_filter) = self.server_config.tool_filter {
SseHandler::with_tool_filter(client, mcp_id.clone(), tool_filter.clone())
} else {
SseHandler::with_mcp_id(client, mcp_id.clone())
};
// Clone handler before creating server (create_server uses sse_handler.clone() internally)
let handler_for_return = sse_handler.clone();
// Create SSE server
let (router, ct) = self.create_server(sse_handler)?;
info!(
"[SseServerBuilder] Server created - mcp_id: {}, sse_path: {}, post_path: {}",
mcp_id, self.server_config.sse_path, self.server_config.post_path
);
Ok((router, ct, SseServerHandler::Sse(handler_for_return)))
}
/// Connect to a stdio backend (child process)
async fn connect_stdio(
&self,
command: &str,
args: &Option<Vec<String>>,
env: &Option<HashMap<String, String>>,
client_info: &ClientInfo,
) -> Result<rmcp::service::RunningService<rmcp::RoleClient, ClientInfo>> {
use std::time::Instant;
let start_time = Instant::now();
let mcp_id = self
.server_config
.mcp_id
.clone()
.unwrap_or_else(|| "unknown".into());
// 使用 process-wrap 创建子进程命令(跨平台进程清理)
// process-wrap 会自动处理进程组Unix或 Job ObjectWindows
// 并且在 Drop 时自动清理子进程树
// 子进程默认继承父进程的所有环境变量
let mut wrapped_cmd = TokioCommandWrap::with_new(command, |cmd| {
if let Some(cmd_args) = args {
cmd.args(cmd_args);
}
// 设置 MCP JSON 配置中的环境变量(会覆盖继承的同名变量)
if let Some(env_vars) = env {
for (k, v) in env_vars {
cmd.env(k, v);
}
}
});
// Unix: 创建进程组,支持 killpg 清理整个进程树
#[cfg(unix)]
wrapped_cmd.wrap(ProcessGroup::leader());
// Windows: 使用 CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP 隐藏控制台窗口
#[cfg(windows)]
{
use process_wrap::tokio::CreationFlags;
use windows::Win32::System::Threading::{CREATE_NEW_PROCESS_GROUP, CREATE_NO_WINDOW};
wrapped_cmd.wrap(CreationFlags(CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP));
wrapped_cmd.wrap(JobObject);
}
// 所有平台: Drop 时自动清理进程
wrapped_cmd.wrap(KillOnDrop);
info!(
"[SseServerBuilder] Starting child process - MCP ID: {}, command: {}, args: {:?}",
mcp_id,
command,
args.as_ref().unwrap_or(&vec![])
);
// 诊断日志:子进程关键环境变量
mcp_common::diagnostic::log_stdio_spawn_context("SseServerBuilder", &mcp_id, env);
let process_start = Instant::now();
// MCP 服务通过 stdin/stdout 进行 JSON-RPC 通信,必须使用 piped默认行为
// 使用 builder 模式捕获 stderr便于诊断子 MCP 服务初始化失败
let (tokio_process, child_stderr) = TokioChildProcess::builder(wrapped_cmd)
.stderr(std::process::Stdio::piped())
.spawn()
.map_err(|e| {
anyhow::anyhow!(
"{}",
mcp_common::diagnostic::format_spawn_error(&mcp_id, command, args, e)
)
})?;
// 启动 stderr 日志读取任务
if let Some(stderr_pipe) = child_stderr {
mcp_common::spawn_stderr_reader(stderr_pipe, mcp_id.clone());
}
let process_duration = process_start.elapsed();
debug!(
"[SseServerBuilder] Child process spawned - MCP ID: {}, spawn time: {:?}",
mcp_id, process_duration
);
let serve_start = Instant::now();
let client = client_info.clone().serve(tokio_process).await?;
let serve_duration = serve_start.elapsed();
let total_duration = start_time.elapsed();
let warn_msg = "建议的优化方案: \
1) 检查网络连接速度 (npm 包下载) \
2) 配置国内 npm 镜像 (如淘宝镜像: npm config set registry https://registry.npmmirror.com) \
3) 预热服务 (启动 mcp-proxy 时预先加载常用服务) \
4) 检查命令参数是否正确";
log_connection_timing(
&mcp_id,
"Stdio",
total_duration,
&[("spawn", process_duration), ("serve", serve_duration)],
STDIO_SLOW_THRESHOLD_SECS,
warn_msg,
);
Ok(client)
}
/// Connect to an SSE URL backend
async fn connect_sse_url(
&self,
url: &str,
headers: &Option<HashMap<String, String>>,
client_info: &ClientInfo,
) -> Result<rmcp::service::RunningService<rmcp::RoleClient, ClientInfo>> {
use std::time::Instant;
let start_time = Instant::now();
let mcp_id = self
.server_config
.mcp_id
.clone()
.unwrap_or_else(|| "unknown".into());
info!(
"[SseServerBuilder] Connecting to SSE URL backend - MCP ID: {}, URL: {}",
mcp_id, url
);
// Build HTTP client with custom headers
let mut req_headers = reqwest::header::HeaderMap::new();
if let Some(config_headers) = headers {
for (key, value) in config_headers {
req_headers.insert(
reqwest::header::HeaderName::try_from(key)
.map_err(|e| anyhow::anyhow!("Invalid header name '{}': {}", key, e))?,
value.parse().map_err(|e| {
anyhow::anyhow!("Invalid header value for '{}': {}", key, e)
})?,
);
}
}
let http_client = reqwest::Client::builder()
.default_headers(req_headers)
.build()
.map_err(|e| anyhow::anyhow!("Failed to create HTTP client: {}", e))?;
// Create SSE client configuration
let sse_config = SseClientConfig {
sse_endpoint: url.to_string().into(),
..Default::default()
};
let transport_start = Instant::now();
let sse_transport = SseClientTransport::start_with_client(http_client, sse_config).await?;
let transport_duration = transport_start.elapsed();
let serve_start = Instant::now();
let client = client_info.clone().serve(sse_transport).await?;
let serve_duration = serve_start.elapsed();
let total_duration = start_time.elapsed();
log_connection_timing(
&mcp_id,
"SSE",
total_duration,
&[("transport", transport_duration), ("serve", serve_duration)],
HTTP_SLOW_THRESHOLD_SECS,
"建议: 检查网络连接和后端服务状态",
);
Ok(client)
}
/// Create the SSE server
fn create_server(&self, sse_handler: SseHandler) -> Result<(axum::Router, CancellationToken)> {
// SSE server uses bind address 0.0.0.0:0 since we're returning a router
// The actual binding will be done by the caller
let config = SseServerConfig {
bind: "0.0.0.0:0".parse()?,
sse_path: self.server_config.sse_path.clone(),
post_path: self.server_config.post_path.clone(),
ct: CancellationToken::new(),
sse_keep_alive: Some(std::time::Duration::from_secs(
self.server_config.keep_alive_secs,
)),
};
let (sse_server, router) = SseServer::new(config);
// Use with_service_directly for non-stateful mode (OneShot services)
// This skips MCP initialization for faster responses
let ct = if self.server_config.stateful {
sse_server.with_service(move || sse_handler.clone())
} else {
sse_server.with_service_directly(move || sse_handler.clone())
};
Ok((router, ct))
}
/// Create SSE server with BackendSessionHandler (for StreamUrl backend)
async fn create_server_with_backend_session(
&self,
handler: BackendSessionHandler,
) -> Result<(axum::Router, CancellationToken, BackendSessionHandler)> {
let config = SseServerConfig {
bind: "0.0.0.0:0".parse()?,
sse_path: self.server_config.sse_path.clone(),
post_path: self.server_config.post_path.clone(),
ct: CancellationToken::new(),
sse_keep_alive: Some(std::time::Duration::from_secs(
self.server_config.keep_alive_secs,
)),
};
let (sse_server, router) = SseServer::new(config);
// Clone handler before passing to closure since we need to return it
let handler_for_return = handler.clone();
let ct = if self.server_config.stateful {
sse_server.with_service(move || handler.clone())
} else {
sse_server.with_service_directly(move || handler.clone())
};
Ok((router, ct, handler_for_return))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_builder_creation() {
let builder = SseServerBuilder::new(BackendConfig::Stdio {
command: "echo".into(),
args: Some(vec!["hello".into()]),
env: None,
})
.mcp_id("test")
.sse_path("/custom/sse")
.post_path("/custom/message");
assert!(builder.server_config.mcp_id.is_some());
assert_eq!(builder.server_config.mcp_id.as_deref(), Some("test"));
assert_eq!(builder.server_config.sse_path, "/custom/sse");
assert_eq!(builder.server_config.post_path, "/custom/message");
}
#[test]
fn test_default_config() {
let config = SseServerBuilderConfig::default();
assert_eq!(config.sse_path, "/sse");
assert_eq!(config.post_path, "/message");
assert_eq!(config.keep_alive_secs, 15);
assert!(
config.stateful,
"default stateful should be true for backward compatibility"
);
}
#[test]
fn test_stateful_flag_default() {
let builder = SseServerBuilder::new(BackendConfig::Stdio {
command: "echo".into(),
args: None,
env: None,
});
assert!(
builder.server_config.stateful,
"stateful should default to true"
);
}
#[test]
fn test_stateful_flag_disabled() {
let builder = SseServerBuilder::new(BackendConfig::Stdio {
command: "echo".into(),
args: None,
env: None,
})
.stateful(false);
assert!(
!builder.server_config.stateful,
"stateful should be false when set"
);
}
#[test]
fn test_stateful_flag_enabled() {
let builder = SseServerBuilder::new(BackendConfig::Stdio {
command: "echo".into(),
args: None,
env: None,
})
.stateful(true);
assert!(
builder.server_config.stateful,
"stateful should be true when set"
);
}
#[test]
fn test_timing_constants() {
assert_eq!(STDIO_SLOW_THRESHOLD_SECS, 30);
assert_eq!(HTTP_SLOW_THRESHOLD_SECS, 10);
}
#[test]
fn test_log_connection_timing_format() {
use std::time::Duration;
// Test that the function doesn't panic and formats correctly
log_connection_timing(
"test-mcp",
"TestBackend",
Duration::from_millis(1500),
&[
("step1", Duration::from_millis(500)),
("step2", Duration::from_millis(1000)),
],
10,
"Test warning message",
);
// If we get here, the function works correctly
}
#[test]
fn test_log_connection_timing_no_breakdown() {
use std::time::Duration;
// Test with empty breakdown
log_connection_timing(
"test-mcp",
"TestBackend",
Duration::from_millis(500),
&[],
10,
"Test warning message",
);
}
}

File diff suppressed because it is too large Load Diff