mcp_streamable_proxy/proxy_handler/
mod.rs

1/**
2 * Create a local SSE server that proxies requests to a stdio MCP server.
3 */
4use crate::backend_client::{
5    BackendNotificationBridge, ProgressRouteGuard, UpstreamPeerRegistry, UpstreamPeerSlot,
6};
7use crate::backend_connector::{BackendConnector, BackendIsolation};
8use arc_swap::ArcSwapOption;
9pub use mcp_common::ToolFilter;
10use rmcp::{
11    ErrorData, RoleClient, RoleServer, ServerHandler, ServiceError,
12    model::{
13        CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ClientRequest,
14        ContentBlock, GetPromptResponse, GetTaskParams, GetTaskRequest, GetTaskResult,
15        Implementation, InitializeRequestParams, InitializeResult, ListToolsResult,
16        PaginatedRequestParams, ProtocolVersion, ReadResourceResponse, RequestMetaObject,
17        RequestParamsMeta, ServerConfig, ServerResult, SetLevelRequestMethod,
18        SubscribeRequestMethod, SubscribeRequestParams, UnsubscribeRequestMethod,
19        UnsubscribeRequestParams,
20    },
21    service::{NotificationContext, Peer, RequestContext, RunningService},
22};
23use std::future::Future;
24use std::sync::Arc;
25use std::sync::atomic::{AtomicU64, Ordering};
26use std::time::Instant;
27use tracing::{debug, error, info, warn};
28
29/// 全局请求计数器，用于生成唯一的请求 ID
30static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(1);
31
32/// Merge `_meta` from the inbound request context into outbound params.
33///
34/// rmcp lifts wire `params._meta` into [`RequestContext::meta`], leaving
35/// `params.meta` often empty. Without this merge, `call_tool_once` / similar
36/// helpers send requests to the backend without `progressToken`.
37fn merge_context_meta_into_params<P: RequestParamsMeta>(
38    params: &mut P,
39    context_meta: &RequestMetaObject,
40) {
41    if context_meta.is_empty() {
42        return;
43    }
44    let dest = params.meta_or_default();
45    for (key, value) in context_meta.iter() {
46        dest.entry(key.clone()).or_insert_with(|| value.clone());
47    }
48}
49
50/// Echo client-requested protocol version when known (mirrors rmcp negotiate).
51fn negotiate_protocol_version(
52    client_requested: &ProtocolVersion,
53    server_fallback: ProtocolVersion,
54) -> ProtocolVersion {
55    if ProtocolVersion::KNOWN_VERSIONS.contains(client_requested) {
56        client_requested.clone()
57    } else {
58        warn!(
59            client_requested = %client_requested,
60            server_fallback = %server_fallback,
61            "client requested unsupported protocol version; falling back to server default"
62        );
63        server_fallback
64    }
65}
66
67/// Convert a backend's [`ServerPeerInfo`] (rmcp 3.1.0 `peer_info()`) into the
68/// [`ServerConfig`] (= `InitializeResult`) shape stored in the discovery cache.
69///
70/// rmcp 3.1.0 narrowed the peer-info type from `InitializeResult` to
71/// `ServerPeerInfo`; the proxy caches the full `ServerConfig`, so we rebuild it.
72/// `InitializeResult` is `#[non_exhaustive]`, so we use its builder instead of a
73/// struct literal. `_meta` has no setter and is dropped (consistent with
74/// `default_server_info`, which also omits it).
75pub(crate) fn peer_info_to_server_info(peer: rmcp::model::ServerPeerInfo) -> ServerConfig {
76    let mut info =
77        InitializeResult::new(peer.capabilities).with_protocol_version(peer.protocol_version);
78    if let Some(server_info) = peer.server_info {
79        info = info.with_server_info(server_info);
80    }
81    if let Some(instructions) = peer.instructions {
82        info = info.with_instructions(instructions);
83    }
84    info
85}
86
87/// Running backend service with notification bridge.
88pub type BackendRunningService = RunningService<RoleClient, BackendNotificationBridge>;
89
90/// 包装后端连接和运行服务
91/// 用于 ArcSwap 热替换
92#[derive(Debug)]
93struct PeerInner {
94    /// Peer 用于发送请求
95    peer: Peer<RoleClient>,
96    /// 保持 RunningService 的所有权，确保服务生命周期
97    #[allow(dead_code)]
98    _running: Arc<BackendRunningService>,
99}
100
101/// Lazy-connect state for [`BackendIsolation::PerSession`].
102#[derive(Clone)]
103struct PerSessionState {
104    connector: Arc<dyn BackendConnector>,
105    notify_slot: UpstreamPeerSlot,
106    connect_lock: Arc<tokio::sync::Mutex<()>>,
107}
108
109impl std::fmt::Debug for PerSessionState {
110    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111        f.debug_struct("PerSessionState")
112            .field("notify_slot", &self.notify_slot)
113            .finish_non_exhaustive()
114    }
115}
116
117/// A proxy handler that forwards requests to a client based on the server's capabilities
118///
119/// **Isolation**:
120/// - [`BackendIsolation::Shared`]: one backend, clones share Arcs (stdio).
121/// - [`BackendIsolation::PerSession`]: factory creates a disconnected handler; backend
122///   is opened during `initialize` (so capabilities are real) and dropped with the session.
123#[derive(Clone, Debug)]
124pub struct ProxyHandler {
125    /// 后端连接（ArcSwap 支持无锁原子替换）
126    /// None 表示后端断开/重连中
127    peer: Arc<ArcSwapOption<PeerInner>>,
128    /// 所有 clone 共享、原子更新的发现快照
129    discovery: DiscoveryCache,
130    /// MCP ID 用于日志记录
131    mcp_id: String,
132    /// 工具过滤配置
133    tool_filter: ToolFilter,
134    /// 后端版本号（每次 swap_backend 递增）
135    backend_version: Arc<AtomicU64>,
136    /// 上游 Streamable HTTP session peers（shared 模式通知桥）
137    upstream_peers: Arc<UpstreamPeerRegistry>,
138    /// Backend isolation mode
139    isolation: BackendIsolation,
140    /// Per-session lazy connect (URL); `None` in shared mode
141    per_session: Option<PerSessionState>,
142}
143
144mod backend_bridge;
145mod connection_swap;
146mod discovery;
147mod forwarding;
148mod lifecycle;
149mod raw_dispatch;
150mod server_handler;
151
152use discovery::DiscoveryCache;
mcp_streamable_proxy/proxy_handler/mod.rs

mcp_streamable_proxy/proxy_handler/
mod.rs
