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

mcp_streamable_proxy/proxy_handler/
mod.rs
