mcp_sse_proxy/sse_handler/
lifecycle.rs

1use super::direct_core::PeerInner;
2use super::*;
3
4impl SseHandler {
5    pub(super) fn backend_peer(&self) -> Result<Peer<RoleClient>, ErrorData> {
6        self.load_backend_peer()
7            .map_err(ForwardError::into_error_data)
8    }
9
10    /// 获取 capabilities 快照
11    #[inline]
12    pub(super) fn capabilities(&self) -> rmcp::model::ServerCapabilities {
13        self.discovery.capabilities()
14    }
15
16    pub(super) fn filter_tools(&self, mut result: ListToolsResult) -> ListToolsResult {
17        if self.tool_filter.is_enabled() {
18            result
19                .tools
20                .retain(|tool| self.tool_filter.is_allowed(&tool.name));
21        }
22        result
23    }
24
25    pub(super) fn cached_tools_or_error(
26        &self,
27        message: &str,
28    ) -> Result<ListToolsResult, ErrorData> {
29        if let Some(cached) = self.discovery.tools() {
30            warn!("{message}; returning cached tools snapshot");
31            return Ok(self.filter_tools(cached));
32        }
33        Err(ErrorData::internal_error(message.to_string(), None))
34    }
35
36    pub(super) fn update_cached_tools(&self, tools: ListToolsResult) {
37        self.discovery.update_tools(tools);
38    }
39
40    fn update_discovery(&self, info: ServerInfo, tools: Option<ListToolsResult>) {
41        self.discovery.update(info, tools);
42    }
43
44    /// 创建一个默认的 ServerInfo（用于断开状态）
45    fn default_server_info(mcp_id: &str) -> ServerInfo {
46        warn!(
47            "[SseHandler] Create default ServerInfo - MCP ID: {}",
48            mcp_id
49        );
50        ServerInfo {
51            protocol_version: ProtocolVersion::V_2024_11_05,
52            server_info: Implementation {
53                name: "MCP Proxy".to_string(),
54                version: "0.1.0".to_string(),
55                title: None,
56                website_url: None,
57                icons: None,
58            },
59            instructions: None,
60            capabilities: Default::default(),
61        }
62    }
63
64    /// 从 RunningService 提取 ServerInfo
65    fn extract_server_info(
66        client: &RunningService<RoleClient, ClientInfo>,
67        mcp_id: &str,
68    ) -> ServerInfo {
69        client
70            .peer_info()
71            .cloned()
72            .unwrap_or_else(|| Self::default_server_info(mcp_id))
73    }
74
75    /// 创建断开状态的 handler（用于初始化）
76    /// 后续通过 swap_backend() 注入实际的后端连接
77    pub fn new_disconnected(
78        mcp_id: String,
79        tool_filter: ToolFilter,
80        default_info: ServerInfo,
81    ) -> Self {
82        info!(
83            "[SseHandler] Create a disconnected handler - MCP ID: {}",
84            mcp_id
85        );
86
87        // 记录过滤器配置
88        if tool_filter.is_enabled() {
89            if let Some(ref allow_list) = tool_filter.allow_tools {
90                info!(
91                    "[SseHandler] Tool whitelist enabled - MCP ID: {}, allowed tools: {:?}",
92                    mcp_id, allow_list
93                );
94            }
95            if let Some(ref deny_list) = tool_filter.deny_tools {
96                info!(
97                    "[SseHandler] Tool blacklist enabled - MCP ID: {}, excluded tools: {:?}",
98                    mcp_id, deny_list
99                );
100            }
101        }
102
103        Self::new_disconnected_with_fallback(mcp_id, tool_filter, default_info, None)
104    }
105
106    pub fn new_disconnected_with_fallback(
107        mcp_id: String,
108        tool_filter: ToolFilter,
109        default_info: ServerInfo,
110        fallback_tools: Option<ListToolsResult>,
111    ) -> Self {
112        Self {
113            peer: Arc::new(ArcSwapOption::empty()),
114            discovery: DiscoveryCache::new(default_info, fallback_tools),
115            mcp_id,
116            tool_filter,
117        }
118    }
119
120    pub fn new(client: RunningService<RoleClient, ClientInfo>) -> Self {
121        Self::with_mcp_id(client, "unknown".to_string())
122    }
123
124    pub fn with_mcp_id(client: RunningService<RoleClient, ClientInfo>, mcp_id: String) -> Self {
125        Self::with_tool_filter(client, mcp_id, ToolFilter::default())
126    }
127
128    /// 创建带工具过滤器的 SseHandler（带初始后端连接）
129    pub fn with_tool_filter(
130        client: RunningService<RoleClient, ClientInfo>,
131        mcp_id: String,
132        tool_filter: ToolFilter,
133    ) -> Self {
134        use std::ops::Deref;
135
136        // 提取 ServerInfo
137        let cached_info = Self::extract_server_info(&client, &mcp_id);
138
139        // 克隆 Peer 用于并发请求（无需锁）
140        let peer = client.deref().clone();
141
142        // 记录过滤器配置
143        if tool_filter.is_enabled() {
144            if let Some(ref allow_list) = tool_filter.allow_tools {
145                info!(
146                    "[SseHandler] Tool whitelist enabled - MCP ID: {}, allowed tools: {:?}",
147                    mcp_id, allow_list
148                );
149            }
150            if let Some(ref deny_list) = tool_filter.deny_tools {
151                info!(
152                    "[SseHandler] Tool blacklist enabled - MCP ID: {}, excluded tools: {:?}",
153                    mcp_id, deny_list
154                );
155            }
156        }
157
158        // 创建 PeerInner
159        let inner = PeerInner {
160            peer,
161            _running: Arc::new(client),
162        };
163
164        Self {
165            peer: Arc::new(ArcSwapOption::from(Some(Arc::new(inner)))),
166            discovery: DiscoveryCache::new(cached_info, None),
167            mcp_id,
168            tool_filter,
169        }
170    }
171
172    /// 原子性替换后端连接
173    /// - Some(client): 设置新的后端连接
174    /// - None: 标记后端断开
175    pub fn swap_backend(&self, new_client: Option<RunningService<RoleClient, ClientInfo>>) {
176        self.swap_backend_with_discovery(new_client, None);
177    }
178
179    pub fn swap_backend_with_discovery(
180        &self,
181        new_client: Option<RunningService<RoleClient, ClientInfo>>,
182        complete_tools: Option<ListToolsResult>,
183    ) {
184        use std::ops::Deref;
185
186        match new_client {
187            Some(client) => {
188                let info = Self::extract_server_info(&client, &self.mcp_id);
189                self.update_discovery(info, complete_tools);
190                let peer = client.deref().clone();
191                let inner = PeerInner {
192                    peer,
193                    _running: Arc::new(client),
194                };
195                self.peer.store(Some(Arc::new(inner)));
196                info!(
197                    "[SseHandler] Backend connection updated - MCP ID: {}",
198                    self.mcp_id
199                );
200            }
201            None => {
202                self.peer.store(None);
203                info!(
204                    "[SseHandler] Backend connection disconnected - MCP ID: {}",
205                    self.mcp_id
206                );
207            }
208        }
209    }
210
211    /// 检查后端是否可用（快速检查，不发送请求）
212    pub fn is_backend_available(&self) -> bool {
213        match self.peer.load_full() {
214            Some(inner) => !inner.peer.is_transport_closed(),
215            None => false,
216        }
217    }
218
219    /// 检查 mcp 服务是否正常（异步版本，会发送验证请求）
220    pub async fn is_mcp_server_ready(&self) -> bool {
221        !self.is_terminated_async().await
222    }
223
224    /// 检查后端连接是否已关闭（同步版本，仅检查 transport 状态）
225    pub fn is_terminated(&self) -> bool {
226        !self.is_backend_available()
227    }
228
229    /// 异步检查后端连接是否已断开（会发送验证请求）
230    pub async fn is_terminated_async(&self) -> bool {
231        // 原子加载后端连接
232        let inner = match self.peer.load_full() {
233            Some(inner) => inner,
234            None => return true,
235        };
236
237        // 快速检查 transport 状态
238        if inner.peer.is_transport_closed() {
239            return true;
240        }
241
242        // 通过发送轻量级请求来验证连接
243        match inner.peer.list_tools(None).await {
244            Ok(_) => {
245                debug!("Backend connection status check: OK");
246                false
247            }
248            Err(e) => {
249                info!("Backend connection status check: Disconnected, reason: {e}");
250                true
251            }
252        }
253    }
254
255    /// 获取 MCP ID
256    pub fn mcp_id(&self) -> &str {
257        &self.mcp_id
258    }
259
260    /// Update backend from an SseClientConnection
261    ///
262    /// This method allows updating the backend connection using the high-level
263    /// `SseClientConnection` type, which is more convenient than the raw
264    /// `RunningService` type.
265    ///
266    /// # Arguments
267    /// * `conn` - Some(connection) to set new backend, None to mark disconnected
268    pub fn swap_backend_from_connection(&self, conn: Option<crate::client::SseClientConnection>) {
269        match conn {
270            Some(c) => {
271                let running = c.into_running_service();
272                self.swap_backend(Some(running));
273            }
274            None => {
275                self.swap_backend(None);
276            }
277        }
278    }
279}
mcp_sse_proxy/sse_handler/lifecycle.rs

mcp_sse_proxy/sse_handler/
lifecycle.rs
