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 =
76        InitializeResult::new(peer.capabilities).with_protocol_version(peer.protocol_version);
77    if let Some(server_info) = peer.server_info {
78        info = info.with_server_info(server_info);
79    }
80    if let Some(instructions) = peer.instructions {
81        info = info.with_instructions(instructions);
82    }
83    info
84}
85
86/// Running backend service with notification bridge.
87pub type BackendRunningService = RunningService<RoleClient, BackendNotificationBridge>;
88
89/// 包装后端连接和运行服务
90/// 用于 ArcSwap 热替换
91#[derive(Debug)]
92struct PeerInner {
93    /// Peer 用于发送请求
94    peer: Peer<RoleClient>,
95    /// 保持 RunningService 的所有权，确保服务生命周期
96    #[allow(dead_code)]
97    _running: Arc<BackendRunningService>,
98}
99
100/// Lazy-connect state for [`BackendIsolation::PerSession`].
101#[derive(Clone)]
102struct PerSessionState {
103    connector: Arc<dyn BackendConnector>,
104    notify_slot: UpstreamPeerSlot,
105    connect_lock: Arc<tokio::sync::Mutex<()>>,
106}
107
108impl std::fmt::Debug for PerSessionState {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        f.debug_struct("PerSessionState")
111            .field("notify_slot", &self.notify_slot)
112            .finish_non_exhaustive()
113    }
114}
115
116/// A proxy handler that forwards requests to a client based on the server's capabilities
117///
118/// **Isolation**:
119/// - [`BackendIsolation::Shared`]: one backend, clones share Arcs (stdio).
120/// - [`BackendIsolation::PerSession`]: factory creates a disconnected handler; backend
121///   is opened during `initialize` (so capabilities are real) and dropped with the session.
122#[derive(Clone, Debug)]
123pub struct ProxyHandler {
124    /// 后端连接（ArcSwap 支持无锁原子替换）
125    /// None 表示后端断开/重连中
126    peer: Arc<ArcSwapOption<PeerInner>>,
127    /// 所有 clone 共享、原子更新的发现快照
128    discovery: DiscoveryCache,
129    /// MCP ID 用于日志记录
130    mcp_id: String,
131    /// 工具过滤配置
132    tool_filter: ToolFilter,
133    /// 后端版本号（每次 swap_backend 递增）
134    backend_version: Arc<AtomicU64>,
135    /// 上游 Streamable HTTP session peers（shared 模式通知桥）
136    upstream_peers: Arc<UpstreamPeerRegistry>,
137    /// Backend isolation mode
138    isolation: BackendIsolation,
139    /// Per-session lazy connect (URL); `None` in shared mode
140    per_session: Option<PerSessionState>,
141}
142
143mod backend_bridge;
144mod connection_swap;
145mod discovery;
146mod forwarding;
147mod lifecycle;
148mod raw_dispatch;
149mod server_handler;
150
151use discovery::DiscoveryCache;
mcp_streamable_proxy/proxy_handler/mod.rs

mcp_streamable_proxy/proxy_handler/
mod.rs
