[router] move to mcp sdk instead (#10057)

This commit is contained in:
Simo Lin
2025-09-05 21:03:46 -04:00
committed by GitHub
parent ab62b135c1
commit db37422c92
10 changed files with 1234 additions and 1333 deletions

View File

@@ -2,18 +2,19 @@
// functionality required for SGLang responses API integration.
//
// Test Coverage:
// - Core MCP server functionality (Python tool_server.py parity)
// - Core MCP server functionality
// - Tool session management (individual and multi-tool)
// - Tool execution and error handling
// - Schema adaptation and validation
// - SSE parsing and protocol compliance
// - Mock server integration for reliable testing
mod common;
use common::mock_mcp_server::MockMCPServer;
use serde_json::json;
use sglang_router_rs::mcp::{parse_sse_event, MCPToolServer, MultiToolSessionManager, ToolSession};
use sglang_router_rs::mcp::{McpClientManager, McpConfig, McpError, McpServerConfig, McpTransport};
use std::collections::HashMap;
/// Create a new mock server for testing (each test gets its own)
async fn create_mock_server() -> MockMCPServer {
MockMCPServer::start()
@@ -21,49 +22,69 @@ async fn create_mock_server() -> MockMCPServer {
.expect("Failed to start mock MCP server")
}
// Core MCP Server Tests (Python parity)
// Core MCP Server Tests
#[tokio::test]
async fn test_mcp_server_initialization() {
let server = MCPToolServer::new();
// Test that we can create an empty configuration
let config = McpConfig { servers: vec![] };
assert!(!server.has_tool("any_tool"));
assert_eq!(server.list_tools().len(), 0);
assert_eq!(server.list_servers().len(), 0);
let stats = server.get_tool_stats();
assert_eq!(stats.total_tools, 0);
assert_eq!(stats.total_servers, 0);
// Should fail with no servers
let result = McpClientManager::new(config).await;
assert!(result.is_err(), "Should fail with no servers configured");
}
#[tokio::test]
async fn test_server_connection_with_mock() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
let result = mcp_server.add_tool_server(mock_server.url()).await;
let config = McpConfig {
servers: vec![McpServerConfig {
name: "mock_server".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
let result = McpClientManager::new(config).await;
assert!(result.is_ok(), "Should connect to mock server");
let stats = mcp_server.get_tool_stats();
assert_eq!(stats.total_tools, 2);
assert_eq!(stats.total_servers, 1);
let mut manager = result.unwrap();
assert!(mcp_server.has_tool("brave_web_search"));
assert!(mcp_server.has_tool("brave_local_search"));
let servers = manager.list_servers();
assert_eq!(servers.len(), 1);
assert!(servers.contains(&"mock_server".to_string()));
let tools = manager.list_tools();
assert_eq!(tools.len(), 2, "Should have 2 tools from mock server");
assert!(manager.has_tool("brave_web_search"));
assert!(manager.has_tool("brave_local_search"));
manager.shutdown().await;
}
#[tokio::test]
async fn test_tool_availability_checking() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
assert!(!mcp_server.has_tool("brave_web_search"));
let config = McpConfig {
servers: vec![McpServerConfig {
name: "mock_server".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
let mut manager = McpClientManager::new(config).await.unwrap();
let test_tools = vec!["brave_web_search", "brave_local_search", "calculator"];
for tool in test_tools {
let available = mcp_server.has_tool(tool);
let available = manager.has_tool(tool);
match tool {
"brave_web_search" | "brave_local_search" => {
assert!(
@@ -82,90 +103,77 @@ async fn test_tool_availability_checking() {
_ => {}
}
}
manager.shutdown().await;
}
#[tokio::test]
async fn test_multi_server_url_parsing() {
async fn test_multi_server_connection() {
let mock_server1 = create_mock_server().await;
let mock_server2 = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
let combined_urls = format!("{},{}", mock_server1.url(), mock_server2.url());
let result = mcp_server.add_tool_server(combined_urls).await;
assert!(result.is_ok(), "Should connect to multiple servers");
let config = McpConfig {
servers: vec![
McpServerConfig {
name: "mock_server_1".to_string(),
transport: McpTransport::Streamable {
url: mock_server1.url(),
token: None,
},
},
McpServerConfig {
name: "mock_server_2".to_string(),
transport: McpTransport::Streamable {
url: mock_server2.url(),
token: None,
},
},
],
};
let stats = mcp_server.get_tool_stats();
assert!(stats.total_servers >= 1);
assert!(stats.total_tools >= 2);
}
// Note: This will fail to connect to both servers in the current implementation
// since they return the same tools. The manager will connect to the first one.
let result = McpClientManager::new(config).await;
// Tool Session Management Tests
if let Ok(mut manager) = result {
let servers = manager.list_servers();
assert!(!servers.is_empty(), "Should have at least one server");
#[tokio::test]
async fn test_individual_tool_session_creation() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
let tools = manager.list_tools();
assert!(tools.len() >= 2, "Should have tools from servers");
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
let session_result = mcp_server.get_tool_session("brave_web_search").await;
assert!(session_result.is_ok(), "Should create tool session");
let session = session_result.unwrap();
assert!(session.is_ready(), "Session should be ready");
assert!(session.connection_info().contains("HTTP"));
}
#[tokio::test]
async fn test_multi_tool_session_manager() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
let available_tools = mcp_server.list_tools();
assert!(
!available_tools.is_empty(),
"Should have tools from mock server"
);
let session_manager_result = mcp_server
.create_multi_tool_session(available_tools.clone())
.await;
assert!(
session_manager_result.is_ok(),
"Should create session manager"
);
let session_manager = session_manager_result.unwrap();
for tool in &available_tools {
assert!(session_manager.has_tool(tool));
manager.shutdown().await;
}
let stats = session_manager.session_stats();
// After optimization: 1 session per server (not per tool)
assert_eq!(stats.total_sessions, 1); // One session for the mock server
assert_eq!(stats.ready_sessions, 1); // One ready session
assert_eq!(stats.unique_servers, 1); // One unique server
// But we still have all tools available
assert_eq!(session_manager.list_tools().len(), available_tools.len());
}
#[tokio::test]
async fn test_tool_execution_with_mock() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
let config = McpConfig {
servers: vec![McpServerConfig {
name: "mock_server".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
let result = mcp_server
let mut manager = McpClientManager::new(config).await.unwrap();
let result = manager
.call_tool(
"brave_web_search",
json!({
"query": "rust programming",
"count": 1
}),
Some(
json!({
"query": "rust programming",
"count": 1
})
.as_object()
.unwrap()
.clone(),
),
)
.await;
@@ -175,48 +183,53 @@ async fn test_tool_execution_with_mock() {
);
let response = result.unwrap();
assert!(
response.get("content").is_some(),
"Response should have content"
);
assert_eq!(response.get("isError").unwrap(), false);
assert!(!response.content.is_empty(), "Should have content");
let content = response.get("content").unwrap().as_array().unwrap();
let text = content[0].get("text").unwrap().as_str().unwrap();
assert!(text.contains("Mock search results for: rust programming"));
// Check the content
if let rmcp::model::RawContent::Text(text) = &response.content[0].raw {
assert!(text
.text
.contains("Mock search results for: rust programming"));
} else {
panic!("Expected text content");
}
manager.shutdown().await;
}
#[tokio::test]
async fn test_concurrent_tool_execution() {
let mock_server = create_mock_server().await;
let mut session_manager = MultiToolSessionManager::new();
session_manager
.add_tools_from_server(
mock_server.url(),
vec![
"brave_web_search".to_string(),
"brave_local_search".to_string(),
],
)
.await
.unwrap();
let config = McpConfig {
servers: vec![McpServerConfig {
name: "mock_server".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
let mut manager = McpClientManager::new(config).await.unwrap();
// Execute tools sequentially (true concurrent execution would require Arc<Mutex>)
let tool_calls = vec![
("brave_web_search".to_string(), json!({"query": "test1"})),
("brave_local_search".to_string(), json!({"query": "test2"})),
("brave_web_search", json!({"query": "test1"})),
("brave_local_search", json!({"query": "test2"})),
];
let results = session_manager.call_tools_concurrent(tool_calls).await;
assert_eq!(results.len(), 2, "Should return results for both tools");
for (tool_name, args) in tool_calls {
let result = manager
.call_tool(tool_name, Some(args.as_object().unwrap().clone()))
.await;
for (i, result) in results.iter().enumerate() {
assert!(result.is_ok(), "Tool {} should succeed with mock server", i);
let response = result.as_ref().unwrap();
assert!(response.get("content").is_some());
assert_eq!(response.get("isError").unwrap(), false);
assert!(result.is_ok(), "Tool {} should succeed", tool_name);
let response = result.unwrap();
assert!(!response.content.is_empty(), "Should have content");
}
manager.shutdown().await;
}
// Error Handling Tests
@@ -224,235 +237,221 @@ async fn test_concurrent_tool_execution() {
#[tokio::test]
async fn test_tool_execution_errors() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
let config = McpConfig {
servers: vec![McpServerConfig {
name: "mock_server".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
let result = mcp_server.call_tool("unknown_tool", json!({})).await;
let mut manager = McpClientManager::new(config).await.unwrap();
// Try to call unknown tool
let result = manager
.call_tool("unknown_tool", Some(serde_json::Map::new()))
.await;
assert!(result.is_err(), "Should fail for unknown tool");
let session = mcp_server
.get_tool_session("brave_web_search")
.await
.unwrap();
let session_result = session.call_tool("unknown_tool", json!({})).await;
assert!(
session_result.is_err(),
"Session should fail for unknown tool"
);
match result.unwrap_err() {
McpError::ToolNotFound(name) => {
assert_eq!(name, "unknown_tool");
}
_ => panic!("Expected ToolNotFound error"),
}
manager.shutdown().await;
}
#[tokio::test]
async fn test_connection_without_server() {
let mut server = MCPToolServer::new();
let config = McpConfig {
servers: vec![McpServerConfig {
name: "nonexistent".to_string(),
transport: McpTransport::Streamable {
url: "http://localhost:9999/mcp".to_string(),
token: None,
},
}],
};
let result = server
.add_tool_server("http://localhost:9999/mcp".to_string())
.await;
let result = McpClientManager::new(config).await;
assert!(result.is_err(), "Should fail when no server is running");
let error_msg = result.unwrap_err().to_string();
assert!(
error_msg.contains("Failed to connect") || error_msg.contains("Connection"),
"Error should be connection-related: {}",
error_msg
);
if let Err(e) = result {
let error_msg = e.to_string();
assert!(
error_msg.contains("Failed to connect") || error_msg.contains("Connection"),
"Error should be connection-related: {}",
error_msg
);
}
}
// Schema Adaptation Tests
// Schema Validation Tests
#[tokio::test]
async fn test_schema_validation() {
async fn test_tool_info_structure() {
let mock_server = create_mock_server().await;
let mut mcp_server = MCPToolServer::new();
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
let config = McpConfig {
servers: vec![McpServerConfig {
name: "mock_server".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
let description = mcp_server.get_tool_description("brave_web_search");
assert!(description.is_some(), "Should have tool description");
let manager = McpClientManager::new(config).await.unwrap();
let desc_value = description.unwrap();
assert!(desc_value.get("name").is_some());
assert!(desc_value.get("description").is_some());
let tools = manager.list_tools();
let brave_search = tools
.iter()
.find(|t| t.name == "brave_web_search")
.expect("Should have brave_web_search tool");
assert_eq!(brave_search.name, "brave_web_search");
assert!(brave_search.description.contains("Mock web search"));
assert_eq!(brave_search.server, "mock_server");
assert!(brave_search.parameters.is_some());
}
// SSE Parsing Tests
// SSE Parsing Tests (simplified since we don't expose parse_sse_event)
#[tokio::test]
async fn test_sse_event_parsing_success() {
let valid_event = "data: {\"jsonrpc\": \"2.0\", \"id\": \"1\", \"result\": {\"test\": \"success\", \"content\": [{\"type\": \"text\", \"text\": \"Hello\"}]}}";
async fn test_sse_connection() {
let mock_server = create_mock_server().await;
let result = parse_sse_event(valid_event);
assert!(result.is_ok(), "Valid SSE event should parse successfully");
// Test SSE transport configuration
let config = McpConfig {
servers: vec![McpServerConfig {
name: "sse_server".to_string(),
transport: McpTransport::Sse {
// Mock server doesn't support SSE, but we can test the config
url: format!("http://127.0.0.1:{}/sse", mock_server.port),
token: Some("test_token".to_string()),
},
}],
};
let parsed = result.unwrap();
assert!(parsed.is_some(), "Should return parsed data");
let response = parsed.unwrap();
assert_eq!(response["test"], "success");
assert!(response.get("content").is_some());
}
#[tokio::test]
async fn test_sse_event_parsing_error() {
let error_event = "data: {\"jsonrpc\": \"2.0\", \"id\": \"1\", \"error\": {\"code\": -1, \"message\": \"Rate limit exceeded\"}}";
let result = parse_sse_event(error_event);
assert!(result.is_err(), "Error SSE event should return error");
let error_msg = result.unwrap_err().to_string();
assert!(
error_msg.contains("Rate limit exceeded"),
"Should contain error message"
);
}
#[tokio::test]
async fn test_sse_event_parsing_empty() {
let empty_event = "";
let result = parse_sse_event(empty_event);
assert!(result.is_ok(), "Empty event should parse successfully");
assert!(result.unwrap().is_none(), "Empty event should return None");
let no_data_event = "event: ping\nid: 123";
let result2 = parse_sse_event(no_data_event);
assert!(result2.is_ok(), "Non-data event should parse successfully");
assert!(
result2.unwrap().is_none(),
"Non-data event should return None"
);
// This will fail to connect but tests the configuration
let result = McpClientManager::new(config).await;
assert!(result.is_err(), "Mock server doesn't support SSE");
}
// Connection Type Tests
#[tokio::test]
async fn test_connection_type_detection() {
let mock_server = create_mock_server().await;
async fn test_transport_types() {
// Test different transport configurations
let session_result = ToolSession::new(mock_server.url()).await;
assert!(session_result.is_ok(), "Should create HTTP session");
// HTTP/Streamable transport
let http_config = McpServerConfig {
name: "http_server".to_string(),
transport: McpTransport::Streamable {
url: "http://localhost:8080/mcp".to_string(),
token: Some("auth_token".to_string()),
},
};
assert_eq!(http_config.name, "http_server");
let session = session_result.unwrap();
assert!(session.connection_info().contains("HTTP"));
assert!(session.is_ready(), "HTTP session should be ready");
// SSE transport
let sse_config = McpServerConfig {
name: "sse_server".to_string(),
transport: McpTransport::Sse {
url: "http://localhost:8081/sse".to_string(),
token: None,
},
};
assert_eq!(sse_config.name, "sse_server");
// Stdio sessions are no longer supported - test invalid URL handling
let invalid_session = ToolSession::new("invalid-url".to_string()).await;
assert!(invalid_session.is_err(), "Should reject non-HTTP URLs");
// STDIO transport
let stdio_config = McpServerConfig {
name: "stdio_server".to_string(),
transport: McpTransport::Stdio {
command: "mcp-server".to_string(),
args: vec!["--port".to_string(), "8082".to_string()],
envs: HashMap::new(),
},
};
assert_eq!(stdio_config.name, "stdio_server");
}
// Integration Pattern Tests
#[tokio::test]
async fn test_responses_api_integration_patterns() {
async fn test_complete_workflow() {
let mock_server = create_mock_server().await;
// Server initialization
let mut mcp_server = MCPToolServer::new();
// 1. Initialize configuration
let config = McpConfig {
servers: vec![McpServerConfig {
name: "integration_test".to_string(),
transport: McpTransport::Streamable {
url: mock_server.url(),
token: None,
},
}],
};
// Tool server connection (like responses API startup)
match mcp_server.add_tool_server(mock_server.url()).await {
Ok(_) => {
let stats = mcp_server.get_tool_stats();
assert_eq!(stats.total_tools, 2);
assert_eq!(stats.total_servers, 1);
}
Err(e) => {
panic!("Should connect to mock server: {}", e);
}
}
// 2. Connect to server
let mut manager = McpClientManager::new(config)
.await
.expect("Should connect to mock server");
// Tool availability checking
let test_tools = vec!["brave_web_search", "brave_local_search", "calculator"];
for tool in &test_tools {
let _available = mcp_server.has_tool(tool);
}
// 3. Verify server connection
let servers = manager.list_servers();
assert_eq!(servers.len(), 1);
assert_eq!(servers[0], "integration_test");
// Tool session creation
if mcp_server.has_tool("brave_web_search") {
let session_result = mcp_server.get_tool_session("brave_web_search").await;
assert!(session_result.is_ok(), "Should create tool session");
}
// 4. Check available tools
let tools = manager.list_tools();
assert_eq!(tools.len(), 2);
// Multi-tool session creation
let available_tools = mcp_server.list_tools();
if !available_tools.is_empty() {
let session_manager_result = mcp_server.create_multi_tool_session(available_tools).await;
assert!(
session_manager_result.is_ok(),
"Should create multi-tool session"
);
}
// 5. Verify specific tools exist
assert!(manager.has_tool("brave_web_search"));
assert!(manager.has_tool("brave_local_search"));
assert!(!manager.has_tool("nonexistent_tool"));
// Tool execution
let result = mcp_server
// 6. Execute a tool
let result = manager
.call_tool(
"brave_web_search",
json!({
"query": "SGLang router MCP integration",
"count": 1
}),
Some(
json!({
"query": "SGLang router MCP integration",
"count": 1
})
.as_object()
.unwrap()
.clone(),
),
)
.await;
if result.is_err() {
// This might fail if called after another test that uses the same tool name
// Due to the shared mock server. That's OK, the main test covers this.
return;
}
assert!(result.is_ok(), "Should execute tool successfully");
}
// Complete Integration Test
assert!(result.is_ok(), "Tool execution should succeed");
let response = result.unwrap();
assert!(!response.content.is_empty(), "Should return content");
#[tokio::test]
async fn test_responses_api_integration() {
let mock_server = create_mock_server().await;
// Run through all functionality required for responses API integration
let mut mcp_server = MCPToolServer::new();
mcp_server.add_tool_server(mock_server.url()).await.unwrap();
// Test all core functionality
assert!(mcp_server.has_tool("brave_web_search"));
let session = mcp_server
.get_tool_session("brave_web_search")
.await
.unwrap();
assert!(session.is_ready());
let session_manager = mcp_server
.create_multi_tool_session(mcp_server.list_tools())
.await
.unwrap();
assert!(session_manager.session_stats().total_sessions > 0);
let result = mcp_server
.call_tool(
"brave_web_search",
json!({
"query": "test",
"count": 1
}),
)
.await
.unwrap();
assert!(result.get("content").is_some());
// 7. Clean shutdown
manager.shutdown().await;
// Verify all required capabilities for responses API integration
let capabilities = [
"MCP server initialization",
"Tool server connection and discovery",
"Tool availability checking",
"Individual tool session management",
"Multi-tool session manager (Python tool_session_ctxs pattern)",
"Concurrent tool execution",
"Direct tool execution",
"Tool execution",
"Error handling and robustness",
"Protocol compliance (SSE parsing)",
"Schema adaptation (Python parity)",
"Multi-server support",
"Schema adaptation",
"Mock server integration (no external dependencies)",
];
assert_eq!(capabilities.len(), 11);
assert_eq!(capabilities.len(), 8);
}