Skip to main content

mcp_streamable_proxy/
server.rs

1//! Streamable HTTP server implementation
2//!
3//! This module provides the HTTP server that uses ProxyAwareSessionManager
4//! for stateful session management with backend version control.
5
6use anyhow::{Result, bail};
7use mcp_common::{McpServiceConfig, check_windows_command, wrap_process_v9};
8use rmcp::{
9    ServiceExt,
10    transport::{
11        TokioChildProcess,
12        streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService},
13    },
14};
15use std::process::Stdio;
16use std::sync::Arc;
17use tracing::{error, info, warn};
18
19// 进程组管理(跨平台子进程清理)
20// process-wrap 9.0 使用 CommandWrap 而不是 TokioCommandWrap
21use process_wrap::tokio::{CommandWrap, KillOnDrop};
22
23use crate::backend_client::{BackendNotificationBridge, UpstreamPeerRegistry};
24use crate::{ProxyAwareSessionManager, ProxyHandler};
25
26/// 从配置启动 Streamable HTTP 服务器
27///
28/// # Features
29///
30/// - **Stateful Mode**: `stateful_mode: true` 支持 session 管理和服务端推送
31/// - **Version Control**: 自动检测后端重连,使旧 session 失效
32/// - **Full Lifecycle**: 自动创建子进程、连接、handler、服务器
33///
34/// # Arguments
35///
36/// * `config` - MCP 服务配置
37/// * `std_listener` - 预先绑定的 TCP 监听器(端口在重试循环前绑定,保证端口占用)
38/// * `quiet` - 静默模式,不输出启动信息
39pub async fn run_stream_server_from_config(
40    config: McpServiceConfig,
41    std_listener: &std::net::TcpListener,
42    quiet: bool,
43) -> Result<()> {
44    // 1. 使用 process-wrap 创建子进程命令(跨平台进程清理)
45    // process-wrap 会自动处理进程组(Unix)或 Job Object(Windows)
46    // 并且在 Drop 时自动清理子进程树
47    // 子进程默认继承父进程的所有环境变量
48
49    // 🔧 Windows 特殊处理:检测并转换 .cmd/.bat 文件避免弹窗
50    // 如果用户配置了 npm 全局安装的 MCP 服务(如 npx some-server 或 some-server.cmd),
51    // 直接运行会弹 CMD 窗口。这里尝试转换
52    check_windows_command(&config.command);
53
54    info!(
55        "[Subprocess][{}] Command: {} {:?}",
56        config.name,
57        config.command,
58        config.args.as_ref().unwrap_or(&vec![])
59    );
60
61    let mut wrapped_cmd = CommandWrap::with_new(&config.command, |command| {
62        if let Some(ref cmd_args) = config.args {
63            command.args(cmd_args);
64        }
65        // 子进程默认继承父进程的所有环境变量
66        // 设置 MCP JSON 配置中的环境变量(会覆盖继承的同名变量)
67        if let Some(ref env_vars) = config.env {
68            for (k, v) in env_vars {
69                command.env(k, v);
70            }
71        }
72    });
73
74    // 应用平台特定的进程包装(Unix: ProcessGroup, Windows: CREATE_NO_WINDOW + JobObject)
75    wrap_process_v9!(wrapped_cmd);
76
77    // 所有平台: Drop 时自动清理进程
78    wrapped_cmd.wrap(KillOnDrop);
79
80    // 2. 启动子进程(rmcp 的 TokioChildProcess 已经支持 process-wrap)
81    //    使用 builder 模式捕获 stderr,便于诊断子 MCP 服务初始化失败
82    let (tokio_process, child_stderr) = TokioChildProcess::builder(wrapped_cmd)
83        .stderr(Stdio::piped())
84        .spawn()?;
85
86    // 启动 stderr 日志读取任务
87    if let Some(stderr_pipe) = child_stderr {
88        mcp_common::spawn_stderr_reader(stderr_pipe, config.name.clone());
89    }
90
91    // 3. 创建带通知桥的后端客户端并连接子进程
92    let upstream_peers = Arc::new(UpstreamPeerRegistry::new());
93    let bridge = BackendNotificationBridge::with_default_info(upstream_peers);
94    let client = bridge.serve(tokio_process).await?;
95
96    // 记录子进程启动到日志文件
97    info!(
98        "[Subprocess startup] Streamable HTTP - Service name: {}, Command: {} {:?}",
99        config.name,
100        config.command,
101        config.args.as_ref().unwrap_or(&vec![])
102    );
103
104    if !quiet {
105        eprintln!("✅ The child process has been started");
106
107        // 获取并打印工具列表
108        match client.list_tools(None).await {
109            Ok(tools_result) => {
110                let tools = &tools_result.tools;
111                if tools.is_empty() {
112                    warn!(
113                        "[Tool list] Tool list is empty - Service name: {}",
114                        config.name
115                    );
116                    eprintln!("⚠️Tool list is empty");
117                } else {
118                    info!(
119                        "[Tool list] Service name: {}, Number of tools: {}",
120                        config.name,
121                        tools.len()
122                    );
123                    eprintln!("🔧 Available tools ({}):", tools.len());
124                    for tool in tools.iter().take(10) {
125                        let desc = tool.description.as_deref().unwrap_or("无描述");
126                        let desc_short = if desc.len() > 50 {
127                            format!("{}...", &desc[..50])
128                        } else {
129                            desc.to_string()
130                        };
131                        eprintln!("   - {} : {}", tool.name, desc_short);
132                    }
133                    if tools.len() > 10 {
134                        eprintln!("... and {} other tools", tools.len() - 10);
135                    }
136                }
137            }
138            Err(e) => {
139                error!(
140                    "[Tool List] Failed to obtain tool list - Service name: {}, Error: {}",
141                    config.name, e
142                );
143                eprintln!("⚠️ Failed to obtain tool list: {}", e);
144            }
145        }
146    } else {
147        // 即使静默模式也记录日志
148        match client.list_tools(None).await {
149            Ok(tools_result) => {
150                info!(
151                    "[Tool list] Service name: {}, Number of tools: {}",
152                    config.name,
153                    tools_result.tools.len()
154                );
155            }
156            Err(e) => {
157                error!(
158                    "[Tool List] Failed to obtain tool list - Service name: {}, Error: {}",
159                    config.name, e
160                );
161            }
162        }
163    }
164
165    // 5. 创建 ProxyHandler
166    let proxy_handler = if let Some(tool_filter) = config.tool_filter {
167        ProxyHandler::with_tool_filter(client, config.name.clone(), tool_filter)
168    } else {
169        ProxyHandler::with_mcp_id(client, config.name.clone())
170    };
171
172    // 6. 启动服务器(使用预绑定的 listener)
173    let listener = tokio::net::TcpListener::from_std(std_listener.try_clone()?)?;
174    run_stream_server(proxy_handler, listener, quiet).await
175}
176
177/// Run Streamable HTTP server with ProxyAwareSessionManager
178///
179/// # Features
180///
181/// - **Stateful Mode**: `stateful_mode: true` 支持 session 管理和服务端推送
182/// - **Version Control**: 自动检测后端重连,使旧 session 失效
183/// - **Hot Swap**: 支持后端连接热替换
184///
185/// # Arguments
186///
187/// * `proxy_handler` - ProxyHandler 实例(包含后端版本控制)
188/// * `listener` - 已绑定的 tokio TcpListener
189/// * `quiet` - 静默模式,不输出启动信息
190pub async fn run_stream_server(
191    proxy_handler: ProxyHandler,
192    listener: tokio::net::TcpListener,
193    quiet: bool,
194) -> Result<()> {
195    let bind_addr = listener
196        .local_addr()
197        .map(|a| a.to_string())
198        .unwrap_or_else(|_| "<unknown>".to_string());
199    let mcp_id = proxy_handler.mcp_id().to_string();
200
201    // 记录服务启动到日志文件
202    info!(
203        "[HTTP service startup] Streamable HTTP service startup - Address: {}, MCP ID: {}",
204        bind_addr, mcp_id
205    );
206
207    if !quiet {
208        eprintln!("📡 Streamable HTTP service startup: http://{}", bind_addr);
209        eprintln!("💡 MCP client can be used directly: http://{}", bind_addr);
210        eprintln!("✨ Feature: stateful_mode (session management + server push)");
211        eprintln!("🔄 Backend version control: Enable (automatically handles reconnections)");
212        eprintln!("💡 Press Ctrl+C to stop the service");
213    }
214
215    // 包装 handler 为 Arc,供 SessionManager 和 service factory 共享
216    let handler = Arc::new(proxy_handler);
217
218    // 创建自定义 SessionManager(带版本控制)
219    let session_manager = ProxyAwareSessionManager::new(handler.clone());
220
221    // 创建 Streamable HTTP 服务
222    // service factory 每次请求都会调用,返回 handler 的克隆
223    //
224    // Host 校验:`StreamableHttpServerConfig::default()` 的 `allowed_hosts` 仅放行
225    // loopback(localhost / 127.0.0.1 / ::1),用于防 DNS rebinding。
226    // 非本机部署若需公网/内网域名,应调用 `.with_allowed_hosts([...])`
227    // 或 `.disable_allowed_hosts()`(后者不推荐用于公网)。
228    let handler_for_service = handler.clone();
229    let mut server_config = StreamableHttpServerConfig::default();
230    server_config.legacy_session_mode = true; // 关键:启用有状态(legacy session)模式
231    let service = StreamableHttpService::new(
232        move || Ok((*handler_for_service).clone()),
233        session_manager.into(), // 转换为 Arc<dyn SessionManager>
234        server_config,
235    );
236
237    // Streamable HTTP 直接在根路径提供服务
238    let router = axum::Router::new().fallback_service(service);
239
240    // 使用传入的 listener 启动 HTTP 服务器
241
242    // 使用 select 处理 Ctrl+C 和服务器
243    tokio::select! {
244        result = axum::serve(listener, router) => {
245            if let Err(e) = result {
246                error!(
247                    "[HTTP Service Error] Streamable HTTP Server Error - MCP ID: {}, Error: {}",
248                    mcp_id, e
249                );
250                bail!("服务器错误: {}", e);
251            }
252        }
253        _ = tokio::signal::ctrl_c() => {
254            info!(
255                "[HTTP service shutdown] Received exit signal, closing Streamable HTTP service - MCP ID: {}",
256                mcp_id
257            );
258            if !quiet {
259                eprintln!("\\n🛑 Received exit signal, closing...");
260            }
261        }
262    }
263
264    Ok(())
265}