提交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,311 @@
//! Streamable HTTP Client Connection Module
//!
//! Provides a high-level API for connecting to MCP servers via Streamable HTTP protocol.
//! This module encapsulates the rmcp 0.12 transport details and exposes a simple interface.
use anyhow::{Context, Result};
use mcp_common::McpClientConfig;
use rmcp::{
RoleClient, ServiceExt,
model::{ClientCapabilities, ClientInfo, Implementation},
service::RunningService,
transport::{
common::client_side_sse::SseRetryPolicy,
streamable_http_client::{
StreamableHttpClientTransport, StreamableHttpClientTransportConfig,
},
},
};
use std::sync::Arc;
use std::time::Duration;
use crate::proxy_handler::ProxyHandler;
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 Streamable HTTP client connection
///
/// This type encapsulates an active connection to an MCP server via Streamable HTTP 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_streamable_proxy::{StreamClientConnection, McpClientConfig};
///
/// let config = McpClientConfig::new("http://localhost:8080/mcp")
/// .with_header("Authorization", "Bearer token");
///
/// let conn = StreamClientConnection::connect(config).await?;
/// let tools = conn.list_tools().await?;
/// println!("Available tools: {:?}", tools);
/// ```
pub struct StreamClientConnection {
inner: RunningService<RoleClient, ClientInfo>,
}
impl StreamClientConnection {
/// Connect to a Streamable HTTP MCP server
///
/// # Arguments
/// * `config` - Client configuration including URL and headers
///
/// # Returns
/// * `Ok(StreamClientConnection)` - Successfully connected client
/// * `Err` - Connection failed
pub async fn connect(config: McpClientConfig) -> Result<Self> {
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 mut transport_config = StreamableHttpClientTransportConfig::with_uri(config.url.clone());
transport_config.retry_config = Arc::new(retry_policy);
let transport = StreamableHttpClientTransport::with_client(http_client, transport_config);
let client_info = create_default_client_info();
let running = client_info
.serve(transport)
.await
.context("Failed to initialize MCP client")?;
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 a ProxyHandler for serving
///
/// This consumes the connection and creates a ProxyHandler 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) -> ProxyHandler {
ProxyHandler::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 {
let capabilities = ClientCapabilities::builder()
.enable_experimental()
.enable_roots()
.enable_roots_list_changed()
.enable_sampling()
.build();
ClientInfo::new(
capabilities,
Implementation::new("mcp-streamable-proxy-client", env!("CARGO_PKG_VERSION")),
)
}
#[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)));
}
}