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

mcp_streamable_proxy/
client.rs
