mcp_streamable_proxy/
client.rs

1//! Streamable HTTP Client Connection Module
2//!
3//! Provides a high-level API for connecting to MCP servers via Streamable HTTP protocol.
4//! This module encapsulates the rmcp transport details and exposes a simple interface.
5
6use anyhow::{Context, Result};
7use mcp_common::McpClientConfig;
8use rmcp::{
9    ServiceExt,
10    transport::{
11        common::client_side_sse::SseRetryPolicy,
12        streamable_http_client::{
13            StreamableHttpClientTransport, StreamableHttpClientTransportConfig,
14        },
15    },
16};
17use std::sync::Arc;
18use std::time::Duration;
19
20use crate::backend_client::{BackendNotificationBridge, UpstreamPeerRegistry};
21use crate::proxy_handler::{BackendRunningService, ProxyHandler};
22use mcp_common::ToolFilter;
23
24/// 自定义的指数退避重试策略，支持最大间隔限制
25///
26/// 重试间隔按照指数增长，但不会超过 max_interval
27/// - 第 1 次重试：base_duration × 2^0
28/// - 第 2 次重试：base_duration × 2^1
29/// - ...
30/// - 第 n 次重试：min(base_duration × 2^(n-1), max_interval)
31#[derive(Debug, Clone)]
32pub struct CappedExponentialBackoff {
33    /// 最大重试次数，None 表示无限制
34    pub max_times: Option<usize>,
35    /// 基础延迟时间（第一次重试前的等待时间）
36    pub base_duration: Duration,
37    /// 最大延迟间隔（重试间隔不会超过这个值）
38    pub max_interval: Duration,
39}
40
41impl CappedExponentialBackoff {
42    /// 创建一个新的带上限的指数退避策略
43    ///
44    /// # Arguments
45    /// * `max_times` - 最大重试次数，None 表示无限制
46    /// * `base_duration` - 基础延迟时间
47    /// * `max_interval` - 最大延迟间隔
48    pub fn new(max_times: Option<usize>, base_duration: Duration, max_interval: Duration) -> Self {
49        Self {
50            max_times,
51            base_duration,
52            max_interval,
53        }
54    }
55}
56
57impl Default for CappedExponentialBackoff {
58    fn default() -> Self {
59        Self {
60            max_times: None,
61            base_duration: Duration::from_secs(1),
62            max_interval: Duration::from_secs(60),
63        }
64    }
65}
66
67impl SseRetryPolicy for CappedExponentialBackoff {
68    fn retry(&self, current_times: usize) -> Option<Duration> {
69        // 检查是否超过最大重试次数
70        if let Some(max_times) = self.max_times
71            && current_times >= max_times
72        {
73            return None;
74        }
75
76        // 计算指数退避时间
77        let exponential_delay = self.base_duration * (2u32.pow(current_times as u32));
78
79        // 限制最大间隔
80        Some(exponential_delay.min(self.max_interval))
81    }
82}
83
84/// Opaque wrapper for Streamable HTTP client connection
85///
86/// This type encapsulates an active connection to an MCP server via Streamable HTTP protocol.
87/// It hides the internal `RunningService` type and provides only the methods
88/// needed by consuming code.
89///
90/// Note: This type is not Clone because the underlying RunningService
91/// is designed for single-owner use. Use `into_handler()` or `into_running_service()`
92/// to consume the connection.
93///
94/// # Example
95///
96/// ```rust,ignore
97/// use mcp_streamable_proxy::{StreamClientConnection, McpClientConfig};
98///
99/// let config = McpClientConfig::new("http://localhost:8080/mcp")
100///     .with_header("Authorization", "Bearer token");
101///
102/// let conn = StreamClientConnection::connect(config).await?;
103/// let tools = conn.list_tools().await?;
104/// println!("Available tools: {:?}", tools);
105/// ```
106pub struct StreamClientConnection {
107    inner: BackendRunningService,
108}
109
110impl StreamClientConnection {
111    /// Connect to a Streamable HTTP MCP server (new upstream peer registry).
112    pub async fn connect(config: McpClientConfig) -> Result<Self> {
113        Self::connect_with_peers(config, Arc::new(UpstreamPeerRegistry::new())).await
114    }
115
116    /// Connect using a shared [`UpstreamPeerRegistry`] (for proxy reconnect / hot-swap).
117    pub async fn connect_with_peers(
118        config: McpClientConfig,
119        upstream_peers: Arc<UpstreamPeerRegistry>,
120    ) -> Result<Self> {
121        let http_client = build_http_client(&config)?;
122
123        // 配置指数退避重试策略，最大间隔 1 分钟，不限制重试次数
124        let retry_policy = CappedExponentialBackoff::new(
125            None,                    // 不限制重试次数
126            Duration::from_secs(1),  // 基础延迟 1 秒
127            Duration::from_secs(60), // 最大间隔 60 秒
128        );
129
130        let mut transport_config =
131            StreamableHttpClientTransportConfig::with_uri(config.url.clone());
132        transport_config.retry_config = Arc::new(retry_policy);
133
134        let transport = StreamableHttpClientTransport::with_client(http_client, transport_config);
135
136        let bridge = BackendNotificationBridge::with_default_info(upstream_peers);
137        let running = bridge
138            .serve(transport)
139            .await
140            .context("Failed to initialize MCP client")?;
141
142        Ok(Self { inner: running })
143    }
144
145    /// List available tools from the MCP server
146    pub async fn list_tools(&self) -> Result<Vec<ToolInfo>> {
147        let result = self.inner.list_tools(None).await?;
148        Ok(result
149            .tools
150            .into_iter()
151            .map(|t| ToolInfo {
152                name: t.name.to_string(),
153                description: t.description.map(|d| d.to_string()),
154            })
155            .collect())
156    }
157
158    /// Check if the connection is closed
159    pub fn is_closed(&self) -> bool {
160        use std::ops::Deref;
161        self.inner.deref().is_transport_closed()
162    }
163
164    /// Get the peer info from the server
165    pub fn peer_info(&self) -> Option<Arc<rmcp::model::ServerInfo>> {
166        self.inner.peer_info()
167    }
168
169    /// Convert this connection into a ProxyHandler for serving
170    ///
171    /// This consumes the connection and creates a ProxyHandler that can
172    /// proxy requests to the backend MCP server.
173    ///
174    /// # Arguments
175    /// * `mcp_id` - Identifier for logging purposes
176    /// * `tool_filter` - Tool filtering configuration
177    pub fn into_handler(self, mcp_id: String, tool_filter: ToolFilter) -> ProxyHandler {
178        ProxyHandler::with_tool_filter(self.inner, mcp_id, tool_filter)
179    }
180
181    /// Extract the internal RunningService for use with swap_backend
182    ///
183    /// This is used internally to support backend hot-swapping.
184    pub fn into_running_service(self) -> BackendRunningService {
185        self.inner
186    }
187}
188
189/// Simplified tool information
190#[derive(Clone, Debug)]
191pub struct ToolInfo {
192    /// Tool name
193    pub name: String,
194    /// Tool description (optional)
195    pub description: Option<String>,
196}
197
198/// Build an HTTP client with the given configuration
199fn build_http_client(config: &McpClientConfig) -> Result<reqwest::Client> {
200    let mut headers = reqwest::header::HeaderMap::new();
201    for (key, value) in &config.headers {
202        let header_name = key
203            .parse::<reqwest::header::HeaderName>()
204            .with_context(|| format!("Invalid header name: {}", key))?;
205        let header_value = value
206            .parse()
207            .with_context(|| format!("Invalid header value for {}: {}", key, value))?;
208        headers.insert(header_name, header_value);
209    }
210
211    let mut builder = reqwest::Client::builder().default_headers(headers);
212
213    if let Some(timeout) = config.connect_timeout {
214        builder = builder.connect_timeout(timeout);
215    }
216
217    if let Some(timeout) = config.read_timeout {
218        builder = builder.timeout(timeout);
219    }
220
221    builder.build().context("Failed to build HTTP client")
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    #[test]
229    fn test_tool_info() {
230        let info = ToolInfo {
231            name: "test_tool".to_string(),
232            description: Some("A test tool".to_string()),
233        };
234        assert_eq!(info.name, "test_tool");
235        assert_eq!(info.description, Some("A test tool".to_string()));
236    }
237
238    #[test]
239    fn test_capped_exponential_backoff() {
240        // 测试带上限的指数退避策略
241        let policy = CappedExponentialBackoff::new(
242            None,                    // 不限制重试次数
243            Duration::from_secs(1),  // 基础延迟 1 秒
244            Duration::from_secs(60), // 最大间隔 60 秒
245        );
246
247        // 验证第 1 次重试：1 秒
248        assert_eq!(policy.retry(0), Some(Duration::from_secs(1)));
249        // 验证第 2 次重试：2 秒
250        assert_eq!(policy.retry(1), Some(Duration::from_secs(2)));
251        // 验证第 3 次重试：4 秒
252        assert_eq!(policy.retry(2), Some(Duration::from_secs(4)));
253        // 验证第 7 次重试：64 秒，但被限制为 60 秒
254        assert_eq!(policy.retry(6), Some(Duration::from_secs(60)));
255        // 验证第 10 次重试：仍然是 60 秒
256        assert_eq!(policy.retry(9), Some(Duration::from_secs(60)));
257    }
258
259    #[test]
260    fn test_capped_exponential_backoff_with_max_times() {
261        let policy =
262            CappedExponentialBackoff::new(Some(3), Duration::from_secs(1), Duration::from_secs(60));
263
264        assert_eq!(policy.retry(0), Some(Duration::from_secs(1)));
265        assert_eq!(policy.retry(1), Some(Duration::from_secs(2)));
266        assert_eq!(policy.retry(2), Some(Duration::from_secs(4)));
267        // 超过最大次数
268        assert_eq!(policy.retry(3), None);
269    }
270}
mcp_streamable_proxy/client.rs

mcp_streamable_proxy/
client.rs
