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

mcp_streamable_proxy/proxy_handler/
lifecycle.rs
