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

mcp_streamable_proxy/proxy_handler/
mod.rs
