mcp_streamable_proxy/proxy_handler/
lifecycle.rs

1use super::*;
2
3impl ProxyHandler {
4    /// Shared upstream peer registry used across reconnects.
5    pub fn upstream_peers(&self) -> Arc<UpstreamPeerRegistry> {
6        self.upstream_peers.clone()
7    }
8
9    /// 创建一个默认的 ServerConfig（用于断开状态）
10    fn default_server_info(mcp_id: &str) -> ServerConfig {
11        warn!(
12            "[ProxyHandler] Create default ServerConfig - MCP ID: {}",
13            mcp_id
14        );
15        ServerConfig::new(rmcp::model::ServerCapabilities::default())
16            .with_server_info(Implementation::new("MCP Proxy", "0.1.0"))
17    }
18
19    /// 从 RunningService 提取 ServerConfig
20    pub(super) fn extract_server_info(
21        client: &BackendRunningService,
22        mcp_id: &str,
23    ) -> ServerConfig {
24        client
25            .peer_info()
26            .map(|peer_info| super::peer_info_to_server_info((*peer_info).clone()))
27            .unwrap_or_else(|| Self::default_server_info(mcp_id))
28    }
29
30    /// 创建断开状态的 handler（用于初始化）
31    /// 后续通过 swap_backend() 注入实际的后端连接
32    pub fn new_disconnected(
33        mcp_id: String,
34        tool_filter: ToolFilter,
35        default_info: ServerConfig,
36    ) -> Self {
37        info!(
38            "[ProxyHandler] Create a disconnected handler - MCP ID: {}",
39            mcp_id
40        );
41
42        // 记录过滤器配置
43        if tool_filter.is_enabled() {
44            if let Some(ref allow_list) = tool_filter.allow_tools {
45                info!(
46                    "[ProxyHandler] Tool whitelist enabled - MCP ID: {}, allowed tools: {:?}",
47                    mcp_id, allow_list
48                );
49            }
50            if let Some(ref deny_list) = tool_filter.deny_tools {
51                info!(
52                    "[ProxyHandler] Tool blacklist enabled - MCP ID: {}, excluded tools: {:?}",
53                    mcp_id, deny_list
54                );
55            }
56        }
57
58        Self::new_disconnected_with_fallback(mcp_id, tool_filter, default_info, None)
59    }
60
61    pub fn new_disconnected_with_fallback(
62        mcp_id: String,
63        tool_filter: ToolFilter,
64        default_info: ServerConfig,
65        fallback_tools: Option<ListToolsResult>,
66    ) -> Self {
67        Self {
68            peer: Arc::new(ArcSwapOption::empty()),
69            discovery: DiscoveryCache::new(default_info, fallback_tools),
70            mcp_id,
71            tool_filter,
72            backend_version: Arc::new(AtomicU64::new(0)), // 断开状态版本为 0
73            upstream_peers: Arc::new(UpstreamPeerRegistry::new()),
74            isolation: BackendIsolation::Shared,
75            per_session: None,
76        }
77    }
78
79    /// Create a disconnected per-session handler (URL isolation).
80    ///
81    /// Backend is opened during [`ServerHandler::initialize`] so the initialize
82    /// response carries real backend capabilities.
83    pub fn new_per_session(
84        connector: Arc<dyn BackendConnector>,
85        mcp_id: String,
86        tool_filter: ToolFilter,
87    ) -> Self {
88        info!(
89            "[ProxyHandler] Create per-session disconnected handler - MCP ID: {}",
90            mcp_id
91        );
92        let default_info = Self::default_server_info(&mcp_id);
93        Self {
94            peer: Arc::new(ArcSwapOption::empty()),
95            discovery: DiscoveryCache::new(default_info, None),
96            mcp_id,
97            tool_filter,
98            backend_version: Arc::new(AtomicU64::new(0)),
99            upstream_peers: Arc::new(UpstreamPeerRegistry::new()),
100            isolation: BackendIsolation::PerSession,
101            per_session: Some(PerSessionState {
102                connector,
103                notify_slot: UpstreamPeerSlot::new(),
104                connect_lock: Arc::new(tokio::sync::Mutex::new(())),
105            }),
106        }
107    }
108
109    /// Isolation mode for this handler.
110    pub fn isolation(&self) -> BackendIsolation {
111        self.isolation
112    }
113
114    pub fn new(client: BackendRunningService) -> Self {
115        Self::with_mcp_id(client, "unknown".to_string())
116    }
117
118    pub fn with_mcp_id(client: BackendRunningService, mcp_id: String) -> Self {
119        Self::with_tool_filter(client, mcp_id, ToolFilter::default())
120    }
121
122    /// 创建带工具过滤器的 ProxyHandler（带初始后端连接，shared 隔离）
123    pub fn with_tool_filter(
124        client: BackendRunningService,
125        mcp_id: String,
126        tool_filter: ToolFilter,
127    ) -> Self {
128        use std::ops::Deref;
129
130        // 提取 ServerConfig 与共享 upstream peer registry
131        let cached_info = Self::extract_server_info(&client, &mcp_id);
132        let upstream_peers = client
133            .service()
134            .upstream_peers()
135            .cloned()
136            .unwrap_or_else(|| Arc::new(UpstreamPeerRegistry::new()));
137
138        // 克隆 Peer 用于并发请求（无需锁）
139        let peer = client.deref().clone();
140
141        // 记录过滤器配置
142        if tool_filter.is_enabled() {
143            if let Some(ref allow_list) = tool_filter.allow_tools {
144                info!(
145                    "[ProxyHandler] Tool whitelist enabled - MCP ID: {}, allowed tools: {:?}",
146                    mcp_id, allow_list
147                );
148            }
149            if let Some(ref deny_list) = tool_filter.deny_tools {
150                info!(
151                    "[ProxyHandler] Tool blacklist enabled - MCP ID: {}, excluded tools: {:?}",
152                    mcp_id, deny_list
153                );
154            }
155        }
156
157        // 创建 PeerInner
158        let inner = PeerInner {
159            peer,
160            _running: Arc::new(client),
161        };
162
163        Self {
164            peer: Arc::new(ArcSwapOption::from(Some(Arc::new(inner)))),
165            discovery: DiscoveryCache::new(cached_info, None),
166            mcp_id,
167            tool_filter,
168            backend_version: Arc::new(AtomicU64::new(1)), // 初始版本为 1
169            upstream_peers,
170            isolation: BackendIsolation::Shared,
171            per_session: None,
172        }
173    }
174
175    /// 原子性替换后端连接
176    /// - Some(client): 设置新的后端连接
177    /// - None: 标记后端断开
178    ///
179    /// **版本控制**：每次调用都会递增 backend_version，使旧 session 失效
180    ///
181    /// **注意**：
182    /// - 新 client 应使用与本 handler 相同的 [`UpstreamPeerRegistry`]
183    ///   （见 `StreamClientConnection::connect_with_peers`）。
184    /// - 对 [`BackendIsolation::PerSession`] 的 **management stub**（builder 返回值）
185    ///   调用此方法没有运维意义：真实后端由各 session handler 持有。
186    pub fn swap_backend(&self, new_client: Option<BackendRunningService>) {
187        self.swap_backend_with_discovery(new_client, None);
188    }
189
190    pub fn swap_backend_with_discovery(
191        &self,
192        new_client: Option<BackendRunningService>,
193        complete_tools: Option<ListToolsResult>,
194    ) {
195        use std::ops::Deref;
196
197        if self.isolation == BackendIsolation::PerSession && self.peer.load().is_none() {
198            warn!(
199                "[ProxyHandler] swap_backend on disconnected per-session handler \
200                 (likely management stub) - MCP ID: {}",
201                self.mcp_id
202            );
203        }
204
205        match new_client {
206            Some(client) => {
207                if let Some(reg) = client.service().upstream_peers()
208                    && !Arc::ptr_eq(reg, &self.upstream_peers)
209                {
210                    warn!(
211                        "[ProxyHandler] swap_backend registry Arc differs from handler - MCP ID: {}",
212                        self.mcp_id
213                    );
214                }
215                // 刷新 cached capabilities / server info
216                let info = Self::extract_server_info(&client, &self.mcp_id);
217                self.update_discovery(info, complete_tools);
218
219                let peer = client.deref().clone();
220                let inner = PeerInner {
221                    peer,
222                    _running: Arc::new(client),
223                };
224                self.peer.store(Some(Arc::new(inner)));
225                info!(
226                    "[ProxyHandler] Backend connection updated - MCP ID: {}",
227                    self.mcp_id
228                );
229            }
230            None => {
231                self.peer.store(None);
232                info!(
233                    "[ProxyHandler] Backend connection disconnected - MCP ID: {}",
234                    self.mcp_id
235                );
236            }
237        }
238
239        // 关键：递增版本号，使所有旧 session 失效
240        let new_version = self.backend_version.fetch_add(1, Ordering::SeqCst) + 1;
241        info!(
242            "[ProxyHandler] Backend version update: {} - MCP ID: {}",
243            new_version, self.mcp_id
244        );
245    }
246
247    /// 检查后端是否可用（快速检查，不发送请求）
248    pub fn is_backend_available(&self) -> bool {
249        let inner_guard = self.peer.load();
250        match inner_guard.as_ref() {
251            Some(inner) => !inner.peer.is_transport_closed(),
252            None => false,
253        }
254    }
255
256    /// 检查 mcp 服务是否正常（异步版本，会发送验证请求）
257    pub async fn is_mcp_server_ready(&self) -> bool {
258        !self.is_terminated_async().await
259    }
260
261    /// 检查后端连接是否已关闭（同步版本，仅检查 transport 状态）
262    pub fn is_terminated(&self) -> bool {
263        !self.is_backend_available()
264    }
265
266    /// 异步检查后端连接是否已断开（会发送验证请求）
267    pub async fn is_terminated_async(&self) -> bool {
268        let peer = {
269            let inner_guard = self.peer.load();
270            let inner = match inner_guard.as_ref() {
271                Some(inner) => inner,
272                None => return true,
273            };
274            if inner.peer.is_transport_closed() {
275                return true;
276            }
277            inner.peer.clone()
278        };
279
280        match peer.list_tools(None).await {
281            Ok(_) => {
282                debug!("Backend connection status check: OK");
283                false
284            }
285            Err(e) => {
286                info!("Backend connection status check: Disconnected, reason: {e}");
287                true
288            }
289        }
290    }
291
292    /// 获取 MCP ID
293    pub fn mcp_id(&self) -> &str {
294        &self.mcp_id
295    }
296
297    /// 获取后端 ServerConfig 的 JSON 表示
298    ///
299    /// 用于跨 rmcp 版本桥接：将 rmcp 1.4.0 的 ServerConfig 序列化为 JSON，
300    /// 供 rmcp 0.10 侧反序列化使用。
301    pub fn get_server_info_json(&self) -> serde_json::Value {
302        match serde_json::to_value(self.discovery.info()) {
303            Ok(value) => value,
304            Err(error) => {
305                tracing::error!(%error, "failed to serialize cached server info");
306                serde_json::Value::Null
307            }
308        }
309    }
310
311    /// 获取当前后端版本号
312    ///
313    /// 版本号用于跟踪后端连接变化：
314    /// - 0: 断开状态
315    /// - 1+: 已连接，每次 swap_backend 递增
316    ///
317    /// **用途**：配合 ProxyAwareSessionManager 实现 session 版本控制
318    pub fn get_backend_version(&self) -> u64 {
319        self.backend_version.load(Ordering::SeqCst)
320    }
321}
mcp_streamable_proxy/proxy_handler/lifecycle.rs

mcp_streamable_proxy/proxy_handler/
lifecycle.rs
