提交qiming-mcp-proxy
This commit is contained in:
Codex
252
qiming-mcp-proxy/mcp-sse-proxy/examples/test_mcp_server.rs
Normal file
252
qiming-mcp-proxy/mcp-sse-proxy/examples/test_mcp_server.rs
Normal file
@@ -0,0 +1,252 @@
|
||||
//! 用于集成测试的最小 MCP 服务器
|
||||
//!
|
||||
//! 支持 echo 和 counter 两个简单工具，用于测试连接、重连和工具调用功能。
|
||||
//!
|
||||
//! 运行方式:
|
||||
//! ```bash
|
||||
//! cargo run --example test_mcp_server -p mcp-sse-proxy
|
||||
//! ```
|
||||
|
||||
use rmcp::{
|
||||
ErrorData, RoleServer, ServerHandler, ServiceExt,
|
||||
model::{
|
||||
CallToolRequestParam, CallToolResult, Content, Implementation, JsonObject, ListToolsResult,
|
||||
PaginatedRequestParam, ProtocolVersion, ServerCapabilities, ServerInfo, Tool,
|
||||
},
|
||||
service::RequestContext,
|
||||
transport::stdio,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// 测试用 MCP 服务器
|
||||
///
|
||||
/// 提供简单工具:
|
||||
/// - `echo`: 回显输入消息
|
||||
/// - `increment`: 递增计数器并返回当前值
|
||||
/// - `reset`: 重置计数器
|
||||
/// - `get_counter`: 获取当前计数器值
|
||||
#[derive(Clone)]
|
||||
pub struct TestMcpServer {
|
||||
counter: Arc<Mutex<i32>>,
|
||||
}
|
||||
|
||||
impl TestMcpServer {
|
||||
/// 创建新的测试服务器实例
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
counter: Arc::new(Mutex::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建一个简单的空 schema
|
||||
fn empty_schema() -> Arc<JsonObject> {
|
||||
let mut schema = JsonObject::new();
|
||||
schema.insert("type".to_string(), serde_json::json!("object"));
|
||||
schema.insert("properties".to_string(), serde_json::json!({}));
|
||||
Arc::new(schema)
|
||||
}
|
||||
|
||||
/// 创建 echo 工具的 schema
|
||||
fn echo_schema() -> Arc<JsonObject> {
|
||||
let mut schema = JsonObject::new();
|
||||
schema.insert("type".to_string(), serde_json::json!("object"));
|
||||
schema.insert(
|
||||
"properties".to_string(),
|
||||
serde_json::json!({
|
||||
"message": {
|
||||
"type": "string",
|
||||
"description": "Message to echo back"
|
||||
}
|
||||
}),
|
||||
);
|
||||
schema.insert("required".to_string(), serde_json::json!(["message"]));
|
||||
Arc::new(schema)
|
||||
}
|
||||
|
||||
/// 获取工具定义列表
|
||||
fn get_tools() -> Vec<Tool> {
|
||||
vec![
|
||||
Tool::new("echo", "Echo back the input message", Self::echo_schema()),
|
||||
Tool::new(
|
||||
"increment",
|
||||
"Increment the counter by 1 and return new value",
|
||||
Self::empty_schema(),
|
||||
),
|
||||
Tool::new("reset", "Reset the counter to 0", Self::empty_schema()),
|
||||
Tool::new(
|
||||
"get_counter",
|
||||
"Get current counter value without changing it",
|
||||
Self::empty_schema(),
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
/// 处理 echo 工具调用
|
||||
async fn handle_echo(&self, args: &serde_json::Value) -> CallToolResult {
|
||||
let message = args
|
||||
.get("message")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("(no message)");
|
||||
CallToolResult::success(vec![Content::text(format!("Echo: {}", message))])
|
||||
}
|
||||
|
||||
/// 处理 increment 工具调用
|
||||
async fn handle_increment(&self) -> CallToolResult {
|
||||
let mut counter = self.counter.lock().await;
|
||||
*counter += 1;
|
||||
CallToolResult::success(vec![Content::text(format!("Counter: {}", *counter))])
|
||||
}
|
||||
|
||||
/// 处理 reset 工具调用
|
||||
async fn handle_reset(&self) -> CallToolResult {
|
||||
let mut counter = self.counter.lock().await;
|
||||
*counter = 0;
|
||||
CallToolResult::success(vec![Content::text("Counter reset to 0".to_string())])
|
||||
}
|
||||
|
||||
/// 处理 get_counter 工具调用
|
||||
async fn handle_get_counter(&self) -> CallToolResult {
|
||||
let counter = self.counter.lock().await;
|
||||
CallToolResult::success(vec![Content::text(format!("Counter: {}", *counter))])
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TestMcpServer {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ServerHandler for TestMcpServer {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo {
|
||||
protocol_version: ProtocolVersion::V_2024_11_05,
|
||||
capabilities: ServerCapabilities::builder().enable_tools().build(),
|
||||
server_info: Implementation {
|
||||
name: "test-mcp-server".to_string(),
|
||||
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||
title: Some("Test MCP Server".to_string()),
|
||||
website_url: None,
|
||||
icons: None,
|
||||
},
|
||||
instructions: Some(
|
||||
"A minimal MCP server for integration testing. \
|
||||
Provides echo, increment, reset, and get_counter tools."
|
||||
.to_string(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_tools(
|
||||
&self,
|
||||
_request: Option<PaginatedRequestParam>,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<ListToolsResult, ErrorData> {
|
||||
Ok(ListToolsResult {
|
||||
tools: Self::get_tools(),
|
||||
next_cursor: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
request: CallToolRequestParam,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<CallToolResult, ErrorData> {
|
||||
let args = request
|
||||
.arguments
|
||||
.as_ref()
|
||||
.map(|v| serde_json::Value::Object(v.clone()))
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
|
||||
// 使用 &str 来进行匹配
|
||||
let tool_name: &str = &request.name;
|
||||
let result = match tool_name {
|
||||
"echo" => self.handle_echo(&args).await,
|
||||
"increment" => self.handle_increment().await,
|
||||
"reset" => self.handle_reset().await,
|
||||
"get_counter" => self.handle_get_counter().await,
|
||||
_ => CallToolResult::error(vec![Content::text(format!(
|
||||
"Unknown tool: {}",
|
||||
request.name
|
||||
))]),
|
||||
};
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
/// 独立运行时作为 stdio MCP 服务器
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// 初始化日志（输出到 stderr，避免干扰 stdio 通信）
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::from_default_env()
|
||||
.add_directive(tracing::Level::INFO.into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
tracing::info!("Test MCP Server starting...");
|
||||
|
||||
let server = TestMcpServer::new();
|
||||
let transport = stdio();
|
||||
|
||||
tracing::info!("Serving on stdio...");
|
||||
let running = server.serve(transport).await?;
|
||||
running.waiting().await?;
|
||||
|
||||
tracing::info!("Test MCP Server stopped.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_counter_increment() {
|
||||
let server = TestMcpServer::new();
|
||||
|
||||
// 第一次递增
|
||||
let result = server.handle_increment().await;
|
||||
assert!(!result.is_error.unwrap_or(false));
|
||||
|
||||
// 第二次递增
|
||||
let result = server.handle_increment().await;
|
||||
assert!(!result.is_error.unwrap_or(false));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_echo() {
|
||||
let server = TestMcpServer::new();
|
||||
let args = serde_json::json!({"message": "hello"});
|
||||
let result = server.handle_echo(&args).await;
|
||||
assert!(!result.is_error.unwrap_or(false));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reset() {
|
||||
let server = TestMcpServer::new();
|
||||
|
||||
// 先递增几次
|
||||
server.handle_increment().await;
|
||||
server.handle_increment().await;
|
||||
|
||||
// 重置
|
||||
let result = server.handle_reset().await;
|
||||
assert!(!result.is_error.unwrap_or(false));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_tools() {
|
||||
let tools = TestMcpServer::get_tools();
|
||||
assert_eq!(tools.len(), 4);
|
||||
assert!(tools.iter().any(|t| &*t.name == "echo"));
|
||||
assert!(tools.iter().any(|t| &*t.name == "increment"));
|
||||
assert!(tools.iter().any(|t| &*t.name == "reset"));
|
||||
assert!(tools.iter().any(|t| &*t.name == "get_counter"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user