chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[router] refactor mcp to use LRU and fix pooling bug (#12346)

Browse Source
This commit is contained in:
Chang Su
2025-10-29 15:18:54 -04:00
committed by GitHub
parent a18161875c
commit d717e73ea8
11 changed files with 277 additions and 662 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-router/Cargo.toml
View File

@@ -38,6 +38,7 @@ futures-util = "0.3"
futures = "0.3"
pyo3 = { version = "0.27.1", features = ["extension-module", "abi3-py38"] }
dashmap = "6.1.0"
lru = "0.16.2"
blake3 = "1.5"
http = "1.1.0"
tokio = { version = "1.42.0", features = ["full"] }

9
sgl-router/src/mcp/config.rs
View File

@@ -1,13 +1,12 @@
//! MCP configuration types and utilities.
//!
//! Defines configuration structures for MCP servers, transports, proxies, and inventory.
use std::collections::HashMap;
// Re-export rmcp types for convenient access
pub use rmcp::model::{Prompt, RawResource, Tool};
use serde::{Deserialize, Serialize};
// ============================================================================
// Configuration Structures
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpConfig {
/// Static MCP servers (loaded at startup)

322
sgl-router/src/mcp/connection_pool.rs
View File

@@ -1,19 +1,10 @@
// MCP Connection Pool
//
// This module provides connection pooling for dynamic MCP servers (per-request).
// Connections are cached and reused to avoid 70-650ms connection overhead on every request.
//
// Performance target:
// - First request: 70-650ms (connection establishment)
// - Subsequent requests: <1ms (cache hit)
// - 90%+ reduction in per-request overhead
/// MCP Connection Pool
///
/// This module provides connection pooling for dynamic MCP servers (per-request).
use std::sync::Arc;
use std::{
sync::Arc,
time::{Duration, Instant},
};
use dashmap::DashMap;
use lru::LruCache;
use parking_lot::Mutex;
use rmcp::{service::RunningService, RoleClient};
use crate::mcp::{
@@ -24,13 +15,14 @@ use crate::mcp::{
/// Type alias for MCP client
type McpClient = RunningService<RoleClient, ()>;
/// Type alias for eviction callback
type EvictionCallback = Arc<dyn Fn(&str) + Send + Sync>;
/// Cached MCP connection with metadata
#[derive(Clone)]
pub struct CachedConnection {
/// The MCP client instance
pub client: Arc<McpClient>,
/// Last time this connection was accessed
pub last_used: Instant,
/// Server configuration used to create this connection
pub config: McpServerConfig,
}
@@ -38,89 +30,88 @@ pub struct CachedConnection {
impl CachedConnection {
/// Create a new cached connection
pub fn new(client: Arc<McpClient>, config: McpServerConfig) -> Self {
Self {
client,
last_used: Instant::now(),
config,
}
}
/// Update last_used timestamp
pub fn touch(&mut self) {
self.last_used = Instant::now();
}
/// Check if connection has been idle for longer than TTL
pub fn is_idle(&self, idle_ttl: Duration) -> bool {
self.last_used.elapsed() > idle_ttl
Self { client, config }
}
}
/// Connection pool for dynamic MCP servers
///
/// Provides thread-safe connection pooling with automatic cleanup of idle connections.
/// Provides thread-safe connection pooling with LRU eviction.
/// Connections are keyed by server URL and reused across requests.
pub struct McpConnectionPool {
/// Map of server_url -> cached connection
connections: DashMap<String, CachedConnection>,
/// LRU cache of server_url -> cached connection
connections: Arc<Mutex<LruCache<String, CachedConnection>>>,
/// Idle connection TTL (connections unused for this duration are cleaned up)
idle_ttl: Duration,
/// Maximum number of cached connections (prevents unbounded growth)
/// Maximum number of cached connections (LRU capacity)
max_connections: usize,
/// Global proxy configuration (applied to all dynamic servers)
/// Can be overridden per-server via McpServerConfig.proxy
global_proxy: Option<McpProxyConfig>,
/// Optional eviction callback (called when LRU evicts a connection)
/// Used to clean up tools from inventory
eviction_callback: Option<EvictionCallback>,
}
impl McpConnectionPool {
/// Default max connections for pool
const DEFAULT_MAX_CONNECTIONS: usize = 200;
/// Create a new connection pool with default settings
///
/// Default settings:
/// - idle_ttl: 300 seconds (5 minutes)
/// - max_connections: 100
/// - max_connections: 200
/// - global_proxy: Loaded from environment variables (MCP_HTTP_PROXY, etc.)
pub fn new() -> Self {
Self {
connections: DashMap::new(),
idle_ttl: Duration::from_secs(300),
max_connections: 100,
connections: Arc::new(Mutex::new(LruCache::new(
std::num::NonZeroUsize::new(Self::DEFAULT_MAX_CONNECTIONS).unwrap(),
))),
max_connections: Self::DEFAULT_MAX_CONNECTIONS,
global_proxy: McpProxyConfig::from_env(),
eviction_callback: None,
}
}
/// Create a new connection pool with custom settings
pub fn with_config(idle_ttl: Duration, max_connections: usize) -> Self {
/// Create a new connection pool with custom capacity
pub fn with_capacity(max_connections: usize) -> Self {
Self {
connections: DashMap::new(),
idle_ttl,
connections: Arc::new(Mutex::new(LruCache::new(
std::num::NonZeroUsize::new(max_connections).unwrap(),
))),
max_connections,
global_proxy: McpProxyConfig::from_env(),
eviction_callback: None,
}
}
/// Create a new connection pool with full custom configuration
pub fn with_full_config(
idle_ttl: Duration,
max_connections: usize,
global_proxy: Option<McpProxyConfig>,
) -> Self {
pub fn with_full_config(max_connections: usize, global_proxy: Option<McpProxyConfig>) -> Self {
Self {
connections: DashMap::new(),
idle_ttl,
connections: Arc::new(Mutex::new(LruCache::new(
std::num::NonZeroUsize::new(max_connections).unwrap(),
))),
max_connections,
global_proxy,
eviction_callback: None,
}
}
/// Set the eviction callback (called when LRU evicts a connection)
pub fn set_eviction_callback<F>(&mut self, callback: F)
where
F: Fn(&str) + Send + Sync + 'static,
{
self.eviction_callback = Some(Arc::new(callback));
}
/// Get an existing connection or create a new one
///
/// This method:
/// 1. Checks if a connection exists for the given URL
/// 2. If exists and fresh, updates last_used and returns it (fast path <1ms)
/// 3. If not exists or stale, creates new connection (slow path 70-650ms)
/// 1. Checks if a connection exists for the given URL (fast path <1ms)
/// 2. If exists, promotes it in LRU and returns it
/// 3. If not exists, creates new connection (slow path 70-650ms)
///
/// # Arguments
/// * `server_url` - The MCP server URL (used as cache key)
@@ -139,97 +130,78 @@ impl McpConnectionPool {
F: FnOnce(McpServerConfig, Option<McpProxyConfig>) -> Fut,
Fut: std::future::Future<Output = McpResult<McpClient>>,
{
// Fast path: Check if connection exists and is still fresh
if let Some(mut entry) = self.connections.get_mut(server_url) {
let cached = entry.value_mut();
// Check if connection is still within TTL
if !cached.is_idle(self.idle_ttl) {
// Update last_used and return cached connection
cached.touch();
// Fast path: Check if connection exists in LRU cache
{
let mut connections = self.connections.lock();
if let Some(cached) = connections.get(server_url) {
// LRU get() promotes the entry
return Ok(Arc::clone(&cached.client));
}
// Connection is stale, drop it and create new one
drop(entry);
self.connections.remove(server_url);
}
// Slow path: Create new connection
// Enforce max_connections limit
if self.connections.len() >= self.max_connections {
self.cleanup_idle_connections();
let client = connect_fn(server_config.clone(), self.global_proxy.clone()).await?;
let client_arc = Arc::new(client);
// If still at limit after cleanup, remove oldest connection
if self.connections.len() >= self.max_connections {
if let Some(oldest_key) = self.find_oldest_connection() {
self.connections.remove(&oldest_key);
// Cache the new connection (LRU will automatically evict oldest if at capacity)
let cached = CachedConnection::new(Arc::clone(&client_arc), server_config);
{
let mut connections = self.connections.lock();
if let Some((evicted_key, _evicted_conn)) =
connections.push(server_url.to_string(), cached)
{
// Call eviction callback if set
if let Some(callback) = &self.eviction_callback {
callback(&evicted_key);
}
}
}
// Create new MCP client using the provided connect function
let client = connect_fn(server_config.clone(), self.global_proxy.clone()).await?;
let client_arc = Arc::new(client);
// Cache the new connection
let cached = CachedConnection::new(Arc::clone(&client_arc), server_config);
self.connections.insert(server_url.to_string(), cached);
Ok(client_arc)
}
/// Remove all idle connections that have exceeded the TTL
///
/// This method is called:
/// - Automatically when max_connections limit is reached
/// - Can be called manually by background cleanup task
pub fn cleanup_idle_connections(&self) {
let now = Instant::now();
self.connections
.retain(|_, cached| now.duration_since(cached.last_used) < self.idle_ttl);
}
/// Find the oldest connection (by last_used timestamp)
///
/// Used for eviction when max_connections is reached and cleanup didn't free space
fn find_oldest_connection(&self) -> Option<String> {
self.connections
.iter()
.min_by_key(|entry| entry.value().last_used)
.map(|entry| entry.key().clone())
}
/// Get current number of cached connections
pub fn len(&self) -> usize {
self.connections.len()
self.connections.lock().len()
}
/// Check if pool is empty
pub fn is_empty(&self) -> bool {
self.connections.is_empty()
self.connections.lock().is_empty()
}
/// Clear all connections (useful for tests)
/// Clear all connections
pub fn clear(&self) {
self.connections.clear();
self.connections.lock().clear();
}
/// Get connection statistics
pub fn stats(&self) -> PoolStats {
let total = self.connections.len();
let idle_count = self
.connections
.iter()
.filter(|entry| entry.value().is_idle(self.idle_ttl))
.count();
let total = self.connections.lock().len();
PoolStats {
total_connections: total,
active_connections: total - idle_count,
idle_connections: idle_count,
capacity: self.max_connections,
}
}
/// List all server keys in the pool
pub fn list_server_keys(&self) -> Vec<String> {
self.connections
.lock()
.iter()
.map(|(key, _)| key.clone())
.collect()
}
/// Get a connection by server key without creating it
/// Promotes the entry in LRU cache if found
pub fn get(&self, server_key: &str) -> Option<Arc<McpClient>> {
self.connections
.lock()
.get(server_key)
.map(|cached| Arc::clone(&cached.client))
}
}
impl Default for McpConnectionPool {
@@ -242,8 +214,7 @@ impl Default for McpConnectionPool {
#[derive(Debug, Clone)]
pub struct PoolStats {
pub total_connections: usize,
pub active_connections: usize,
pub idle_connections: usize,
pub capacity: usize,
}
#[cfg(test)]
@@ -271,114 +242,13 @@ mod tests {
assert!(pool.is_empty());
}
#[test]
#[allow(invalid_value)]
fn test_cached_connection_touch() {
let config = create_test_config("http://localhost:3000");
let client: Arc<McpClient> = Arc::new(unsafe {
// SAFETY: This is only for testing the CachedConnection struct
std::mem::MaybeUninit::zeroed().assume_init()
});
let mut cached = CachedConnection::new(client.clone(), config);
let first_time = cached.last_used;
std::thread::sleep(Duration::from_millis(10));
cached.touch();
assert!(cached.last_used > first_time);
// Prevent drop of invalid Arc (would segfault)
std::mem::forget(client);
}
#[test]
#[allow(invalid_value)]
fn test_cached_connection_is_idle() {
let config = create_test_config("http://localhost:3000");
let client: Arc<McpClient> = Arc::new(unsafe {
// SAFETY: This is only for testing the CachedConnection struct
std::mem::MaybeUninit::zeroed().assume_init()
});
let cached = CachedConnection::new(client.clone(), config);
// Fresh connection should not be idle
assert!(!cached.is_idle(Duration::from_secs(1)));
// Wait and check
std::thread::sleep(Duration::from_millis(100));
assert!(cached.is_idle(Duration::from_millis(50)));
// Prevent drop of invalid Arc (would segfault)
std::mem::forget(client);
}
#[test]
fn test_pool_stats() {
let pool = McpConnectionPool::with_config(Duration::from_millis(100), 10);
let pool = McpConnectionPool::with_capacity(10);
let stats = pool.stats();
assert_eq!(stats.total_connections, 0);
assert_eq!(stats.active_connections, 0);
assert_eq!(stats.idle_connections, 0);
}
#[test]
#[allow(invalid_value)]
fn test_cleanup_idle_connections() {
let pool = McpConnectionPool::with_config(Duration::from_millis(50), 10);
// Initially empty
assert_eq!(pool.len(), 0);
// Add a connection manually for testing
let config = create_test_config("http://localhost:3000");
let client: Arc<McpClient> =
Arc::new(unsafe { std::mem::MaybeUninit::zeroed().assume_init() });
let cached = CachedConnection::new(client.clone(), config);
pool.connections
.insert("http://localhost:3000".to_string(), cached);
assert_eq!(pool.len(), 1);
// Wait for TTL to expire
std::thread::sleep(Duration::from_millis(100));
// Cleanup should remove idle connection
pool.cleanup_idle_connections();
assert_eq!(pool.len(), 0);
// Prevent drop of invalid Arc (would segfault)
std::mem::forget(client);
}
#[test]
#[allow(invalid_value)]
fn test_find_oldest_connection() {
let pool = McpConnectionPool::new();
// Collect clients to forget at end
let mut clients = Vec::new();
// Add connections with different timestamps
for i in 0..3 {
let url = format!("http://localhost:{}", 3000 + i);
let config = create_test_config(&url);
let client: Arc<McpClient> =
Arc::new(unsafe { std::mem::MaybeUninit::zeroed().assume_init() });
let cached = CachedConnection::new(client.clone(), config);
pool.connections.insert(url, cached);
clients.push(client);
std::thread::sleep(Duration::from_millis(10));
}
// Oldest should be the first one
let oldest = pool.find_oldest_connection();
assert!(oldest.is_some());
assert_eq!(oldest.unwrap(), "http://localhost:3000");
// Prevent drop of invalid Arcs (would segfault)
for client in clients {
std::mem::forget(client);
}
assert_eq!(stats.capacity, 10);
}
#[test]
@@ -392,7 +262,8 @@ mod tests {
Arc::new(unsafe { std::mem::MaybeUninit::zeroed().assume_init() });
let cached = CachedConnection::new(client.clone(), config);
pool.connections
.insert("http://localhost:3000".to_string(), cached);
.lock()
.push("http://localhost:3000".to_string(), cached);
assert_eq!(pool.len(), 1);
@@ -418,8 +289,7 @@ mod tests {
};
// Create pool with proxy
let pool =
McpConnectionPool::with_full_config(Duration::from_secs(300), 100, Some(proxy.clone()));
let pool = McpConnectionPool::with_full_config(100, Some(proxy.clone()));
// Verify proxy is stored
assert!(pool.global_proxy.is_some());

5
sgl-router/src/mcp/error.rs
View File

@@ -1,3 +1,8 @@
//! MCP error types.
//!
//! Defines error variants for MCP operations including connection, tool execution,
//! and configuration errors.
use thiserror::Error;
pub type McpResult<T> = Result<T, McpError>;

316
sgl-router/src/mcp/inventory.rs
View File

@@ -1,10 +1,6 @@
// MCP Tool Inventory with TTL-based Caching
//
// This module provides TTL-based caching for MCP tools, prompts, and resources.
// Tools are cached with timestamps and automatically expire after the configured TTL.
// Background refresh tasks can proactively update the inventory.
use std::time::{Duration, Instant};
//! MCP tool, prompt, and resource inventory.
//!
//! Thread-safe cache for MCP capabilities across all connected servers.
use dashmap::DashMap;
@@ -15,7 +11,6 @@ use crate::mcp::config::{Prompt, RawResource, Tool};
pub struct CachedTool {
pub server_name: String,
pub tool: Tool,
pub cached_at: Instant,
}
/// Cached prompt with metadata
@@ -23,7 +18,6 @@ pub struct CachedTool {
pub struct CachedPrompt {
pub server_name: String,
pub prompt: Prompt,
pub cached_at: Instant,
}
/// Cached resource with metadata
@@ -31,13 +25,12 @@ pub struct CachedPrompt {
pub struct CachedResource {
pub server_name: String,
pub resource: RawResource,
pub cached_at: Instant,
}
/// Tool inventory with TTL-based caching
/// Tool inventory with periodic refresh
///
/// Provides thread-safe caching of MCP tools, prompts, and resources with automatic expiration.
/// Entries are timestamped and can be queried with TTL validation.
/// Provides thread-safe caching of MCP tools, prompts, and resources.
/// Entries are refreshed periodically by background tasks.
pub struct ToolInventory {
/// Map of tool_name -> cached tool
tools: DashMap<String, CachedTool>,
@@ -47,80 +40,59 @@ pub struct ToolInventory {
/// Map of resource_uri -> cached resource
resources: DashMap<String, CachedResource>,
/// Tool cache TTL
tool_ttl: Duration,
/// Last refresh time per server
server_refresh_times: DashMap<String, Instant>,
}
impl ToolInventory {
/// Create a new tool inventory with the specified TTL
pub fn new(tool_ttl: Duration) -> Self {
/// Create a new tool inventory
pub fn new() -> Self {
Self {
tools: DashMap::new(),
prompts: DashMap::new(),
resources: DashMap::new(),
tool_ttl,
server_refresh_times: DashMap::new(),
}
}
}
impl Default for ToolInventory {
fn default() -> Self {
Self::new()
}
}
impl ToolInventory {
// ============================================================================
// Tool Methods
// ============================================================================
/// Get a tool if it exists and is fresh (within TTL)
///
/// Returns None if the tool doesn't exist or has expired.
/// Get a tool if it exists
pub fn get_tool(&self, tool_name: &str) -> Option<(String, Tool)> {
self.tools.get(tool_name).and_then(|entry| {
let cached = entry.value();
// Check if still fresh
if cached.cached_at.elapsed() < self.tool_ttl {
Some((cached.server_name.clone(), cached.tool.clone()))
} else {
// Expired - will be removed by cleanup
None
}
})
self.tools
.get(tool_name)
.map(|entry| (entry.server_name.clone(), entry.tool.clone()))
}
/// Check if tool exists (regardless of TTL)
/// Check if tool exists
pub fn has_tool(&self, tool_name: &str) -> bool {
self.tools.contains_key(tool_name)
}
/// Insert or update a tool
pub fn insert_tool(&self, tool_name: String, server_name: String, tool: Tool) {
self.tools.insert(
tool_name,
CachedTool {
server_name,
tool,
cached_at: Instant::now(),
},
);
self.tools
.insert(tool_name, CachedTool { server_name, tool });
}
/// Get all tools (fresh only)
/// Get all tools
pub fn list_tools(&self) -> Vec<(String, String, Tool)> {
let now = Instant::now();
self.tools
.iter()
.filter_map(|entry| {
.map(|entry| {
let (name, cached) = entry.pair();
if now.duration_since(cached.cached_at) < self.tool_ttl {
Some((
name.clone(),
cached.server_name.clone(),
cached.tool.clone(),
))
} else {
None
}
(
name.clone(),
cached.server_name.clone(),
cached.tool.clone(),
)
})
.collect()
}
@@ -129,21 +101,14 @@ impl ToolInventory {
// Prompt Methods
// ============================================================================
/// Get a prompt if it exists and is fresh (within TTL)
/// Get a prompt if it exists
pub fn get_prompt(&self, prompt_name: &str) -> Option<(String, Prompt)> {
self.prompts.get(prompt_name).and_then(|entry| {
let cached = entry.value();
// Check if still fresh
if cached.cached_at.elapsed() < self.tool_ttl {
Some((cached.server_name.clone(), cached.prompt.clone()))
} else {
None
}
})
self.prompts
.get(prompt_name)
.map(|entry| (entry.server_name.clone(), entry.prompt.clone()))
}
/// Check if prompt exists (regardless of TTL)
/// Check if prompt exists
pub fn has_prompt(&self, prompt_name: &str) -> bool {
self.prompts.contains_key(prompt_name)
}
@@ -155,27 +120,21 @@ impl ToolInventory {
CachedPrompt {
server_name,
prompt,
cached_at: Instant::now(),
},
);
}
/// Get all prompts (fresh only)
/// Get all prompts
pub fn list_prompts(&self) -> Vec<(String, String, Prompt)> {
let now = Instant::now();
self.prompts
.iter()
.filter_map(|entry| {
.map(|entry| {
let (name, cached) = entry.pair();
if now.duration_since(cached.cached_at) < self.tool_ttl {
Some((
name.clone(),
cached.server_name.clone(),
cached.prompt.clone(),
))
} else {
None
}
(
name.clone(),
cached.server_name.clone(),
cached.prompt.clone(),
)
})
.collect()
}
@@ -184,21 +143,14 @@ impl ToolInventory {
// Resource Methods
// ============================================================================
/// Get a resource if it exists and is fresh (within TTL)
/// Get a resource if it exists
pub fn get_resource(&self, resource_uri: &str) -> Option<(String, RawResource)> {
self.resources.get(resource_uri).and_then(|entry| {
let cached = entry.value();
// Check if still fresh
if cached.cached_at.elapsed() < self.tool_ttl {
Some((cached.server_name.clone(), cached.resource.clone()))
} else {
None
}
})
self.resources
.get(resource_uri)
.map(|entry| (entry.server_name.clone(), entry.resource.clone()))
}
/// Check if resource exists (regardless of TTL)
/// Check if resource exists
pub fn has_resource(&self, resource_uri: &str) -> bool {
self.resources.contains_key(resource_uri)
}
@@ -215,27 +167,21 @@ impl ToolInventory {
CachedResource {
server_name,
resource,
cached_at: Instant::now(),
},
);
}
/// Get all resources (fresh only)
/// Get all resources
pub fn list_resources(&self) -> Vec<(String, String, RawResource)> {
let now = Instant::now();
self.resources
.iter()
.filter_map(|entry| {
.map(|entry| {
let (uri, cached) = entry.pair();
if now.duration_since(cached.cached_at) < self.tool_ttl {
Some((
uri.clone(),
cached.server_name.clone(),
cached.resource.clone(),
))
} else {
None
}
(
uri.clone(),
cached.server_name.clone(),
cached.resource.clone(),
)
})
.collect()
}
@@ -244,7 +190,7 @@ impl ToolInventory {
// Server Management Methods
// ============================================================================
/// Clear all cached items for a specific server (before refresh)
/// Clear all cached items for a specific server (called when LRU evicts client)
pub fn clear_server_tools(&self, server_name: &str) {
self.tools
.retain(|_, cached| cached.server_name != server_name);
@@ -254,51 +200,6 @@ impl ToolInventory {
.retain(|_, cached| cached.server_name != server_name);
}
/// Mark server as refreshed
pub fn mark_refreshed(&self, server_name: &str) {
self.server_refresh_times
.insert(server_name.to_string(), Instant::now());
}
/// Check if server needs refresh based on refresh interval
pub fn needs_refresh(&self, server_name: &str, refresh_interval: Duration) -> bool {
self.server_refresh_times
.get(server_name)
.map(|t| t.elapsed() > refresh_interval)
.unwrap_or(true) // Never refreshed = needs refresh
}
/// Get last refresh time for a server
pub fn last_refresh(&self, server_name: &str) -> Option<Instant> {
self.server_refresh_times
.get(server_name)
.map(|t| *t.value())
}
// ============================================================================
// Cleanup Methods
// ============================================================================
/// Cleanup expired entries
///
/// Removes all tools, prompts, and resources that have exceeded their TTL.
/// Should be called periodically by a background task.
pub fn cleanup_expired(&self) {
let now = Instant::now();
// Remove expired tools
self.tools
.retain(|_, cached| now.duration_since(cached.cached_at) < self.tool_ttl);
// Remove expired prompts
self.prompts
.retain(|_, cached| now.duration_since(cached.cached_at) < self.tool_ttl);
// Remove expired resources
self.resources
.retain(|_, cached| now.duration_since(cached.cached_at) < self.tool_ttl);
}
/// Get count of cached items
pub fn counts(&self) -> (usize, usize, usize) {
(self.tools.len(), self.prompts.len(), self.resources.len())
@@ -309,7 +210,6 @@ impl ToolInventory {
self.tools.clear();
self.prompts.clear();
self.resources.clear();
self.server_refresh_times.clear();
}
}
@@ -370,7 +270,7 @@ mod tests {
#[test]
fn test_tool_insert_and_get() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
let tool = create_test_tool("test_tool");
inventory.insert_tool("test_tool".to_string(), "server1".to_string(), tool.clone());
@@ -383,30 +283,9 @@ mod tests {
assert_eq!(retrieved_tool.name, "test_tool");
}
#[test]
fn test_tool_expiration() {
let inventory = ToolInventory::new(Duration::from_millis(100));
let tool = create_test_tool("expiring_tool");
inventory.insert_tool(
"expiring_tool".to_string(),
"server1".to_string(),
tool.clone(),
);
// Should be available immediately
assert!(inventory.get_tool("expiring_tool").is_some());
// Wait for expiration
std::thread::sleep(Duration::from_millis(150));
// Should be expired now
assert!(inventory.get_tool("expiring_tool").is_none());
}
#[test]
fn test_has_tool() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
let tool = create_test_tool("check_tool");
assert!(!inventory.has_tool("check_tool"));
@@ -418,7 +297,7 @@ mod tests {
#[test]
fn test_list_tools() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
inventory.insert_tool(
"tool1".to_string(),
@@ -440,29 +319,9 @@ mod tests {
assert_eq!(tools.len(), 3);
}
#[test]
fn test_list_tools_filters_expired() {
let inventory = ToolInventory::new(Duration::from_millis(100));
inventory.insert_tool(
"tool1".to_string(),
"server1".to_string(),
create_test_tool("tool1"),
);
// Should have 1 tool
assert_eq!(inventory.list_tools().len(), 1);
// Wait for expiration
std::thread::sleep(Duration::from_millis(150));
// Should have 0 tools (filtered out)
assert_eq!(inventory.list_tools().len(), 0);
}
#[test]
fn test_clear_server_tools() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
inventory.insert_tool(
"tool1".to_string(),
@@ -484,55 +343,9 @@ mod tests {
assert_eq!(tools[0].0, "tool2");
}
#[test]
fn test_server_refresh_tracking() {
let inventory = ToolInventory::new(Duration::from_secs(60));
// Never refreshed
assert!(inventory.needs_refresh("server1", Duration::from_secs(10)));
// Mark as refreshed
inventory.mark_refreshed("server1");
// Should not need refresh immediately
assert!(!inventory.needs_refresh("server1", Duration::from_secs(10)));
// Wait and check again
std::thread::sleep(Duration::from_millis(100));
assert!(inventory.needs_refresh("server1", Duration::from_millis(50)));
}
#[test]
fn test_cleanup_expired() {
let inventory = ToolInventory::new(Duration::from_millis(100));
inventory.insert_tool(
"tool1".to_string(),
"server1".to_string(),
create_test_tool("tool1"),
);
inventory.insert_tool(
"tool2".to_string(),
"server1".to_string(),
create_test_tool("tool2"),
);
let (tools, _, _) = inventory.counts();
assert_eq!(tools, 2);
// Wait for expiration
std::thread::sleep(Duration::from_millis(150));
// Cleanup expired entries
inventory.cleanup_expired();
let (tools, _, _) = inventory.counts();
assert_eq!(tools, 0);
}
#[test]
fn test_prompt_operations() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
let prompt = create_test_prompt("test_prompt");
inventory.insert_prompt(
@@ -553,7 +366,7 @@ mod tests {
#[test]
fn test_resource_operations() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
let resource = create_test_resource("file:///test.txt");
inventory.insert_resource(
@@ -576,7 +389,7 @@ mod tests {
async fn test_concurrent_access() {
use std::sync::Arc;
let inventory = Arc::new(ToolInventory::new(Duration::from_secs(60)));
let inventory = Arc::new(ToolInventory::new());
// Spawn multiple tasks that insert tools concurrently
let mut handles = vec![];
@@ -601,7 +414,7 @@ mod tests {
#[test]
fn test_clear_all() {
let inventory = ToolInventory::new(Duration::from_secs(60));
let inventory = ToolInventory::new();
inventory.insert_tool(
"tool1".to_string(),
@@ -619,8 +432,6 @@ mod tests {
create_test_resource("res1"),
);
inventory.mark_refreshed("server1");
let (tools, prompts, resources) = inventory.counts();
assert_eq!(tools, 1);
assert_eq!(prompts, 1);
@@ -632,6 +443,5 @@ mod tests {
assert_eq!(tools, 0);
assert_eq!(prompts, 0);
assert_eq!(resources, 0);
assert!(inventory.last_refresh("server1").is_none());
}
}

251
sgl-router/src/mcp/manager.rs
View File

@@ -1,13 +1,7 @@
//! Refactored MCP Manager - Single flat structure for all MCP operations
//! MCP client management and orchestration.
//!
//! This replaces the previous hierarchy:
//! - McpManager (wrapper for static/dynamic distinction)
//! - McpClientManager (manages multiple clients)
//! - McpClient (actual client)
//!
//! New flat structure:
//! - McpManager (single component handling all MCP concerns)
//! - McpClient (actual client to one server)
//! Manages both static MCP servers (from config) and dynamic MCP servers (from requests).
//! Static clients are never evicted; dynamic clients use LRU eviction via connection pool.
use std::{borrow::Cow, sync::Arc, time::Duration};
@@ -40,68 +34,47 @@ use crate::mcp::{
/// Type alias for MCP client
type McpClient = RunningService<RoleClient, ()>;
/// Unified MCP Manager - handles all MCP operations
///
/// This single component manages:
/// - Client connections (both static and dynamic)
/// - Tool inventory and caching
/// - Connection pooling
/// - Background refresh
/// - Tool/prompt/resource operations
pub struct McpManager {
/// All MCP clients (static + dynamic)
/// Key: server_name for static, server_url for dynamic
/// Using DashMap for concurrent access
clients: Arc<DashMap<String, Arc<McpClient>>>,
/// Track which servers are static (from config)
/// Using DashMap for thread-safe mutation during workflow registration
static_servers: Arc<DashMap<String, ()>>,
/// Shared tool inventory with TTL and caching
static_clients: Arc<DashMap<String, Arc<McpClient>>>,
inventory: Arc<ToolInventory>,
/// Connection pool for dynamic servers (TTL-based cleanup)
connection_pool: Arc<McpConnectionPool>,
/// Original config for static servers (kept for potential future use)
_config: McpConfig,
}
impl McpManager {
/// Create a new MCP manager with custom TTLs
pub async fn new(
config: McpConfig,
tool_ttl: Duration,
pool_idle_ttl: Duration,
pool_max_connections: usize,
) -> McpResult<Self> {
// Create shared inventory
let inventory = Arc::new(ToolInventory::new(tool_ttl));
const MAX_DYNAMIC_CLIENTS: usize = 200;
// Create connection pool
let connection_pool = Arc::new(McpConnectionPool::with_config(
pool_idle_ttl,
pool_max_connections,
));
pub async fn new(config: McpConfig, pool_max_connections: usize) -> McpResult<Self> {
let inventory = Arc::new(ToolInventory::new());
// Create manager structure
let clients = Arc::new(DashMap::new());
let static_servers = Arc::new(DashMap::new());
let mut connection_pool =
McpConnectionPool::with_full_config(pool_max_connections, config.proxy.clone());
let inventory_clone = Arc::clone(&inventory);
connection_pool.set_eviction_callback(move |server_key: &str| {
debug!(
"LRU evicted dynamic server '{}' - clearing tools from inventory",
server_key
);
inventory_clone.clear_server_tools(server_key);
});
let connection_pool = Arc::new(connection_pool);
// Create storage for static clients
let static_clients = Arc::new(DashMap::new());
// Get global proxy config for all servers
let global_proxy = config.proxy.as_ref();
// Connect to all static servers from config
for server_config in &config.servers {
static_servers.insert(server_config.name.clone(), ());
match Self::connect_server(server_config, global_proxy).await {
Ok(client) => {
let client_arc = Arc::new(client);
// Load inventory for this server
Self::load_server_inventory(&inventory, &server_config.name, &client_arc).await;
clients.insert(server_config.name.clone(), client_arc);
static_clients.insert(server_config.name.clone(), client_arc);
info!("Connected to static server '{}'", server_config.name);
}
Err(e) => {
@@ -113,52 +86,40 @@ impl McpManager {
}
}
if static_servers.is_empty() || clients.is_empty() {
if static_clients.is_empty() {
warn!("No static MCP servers connected");
}
Ok(Self {
clients,
static_servers,
static_clients,
inventory,
connection_pool,
_config: config,
})
}
/// Create with default settings (300s TTL, 300s idle, 100 max connections)
pub async fn with_defaults(config: McpConfig) -> McpResult<Self> {
Self::new(
config,
Duration::from_secs(300),
Duration::from_secs(300),
100,
)
.await
Self::new(config, Self::MAX_DYNAMIC_CLIENTS).await
}
// ========================================================================
// Client Management
// ========================================================================
/// Get a client by server name (static or dynamic)
pub async fn get_client(&self, server_name: &str) -> Option<Arc<McpClient>> {
self.clients.get(server_name).map(|e| Arc::clone(e.value()))
if let Some(client) = self.static_clients.get(server_name) {
return Some(Arc::clone(client.value()));
}
self.connection_pool.get(server_name)
}
/// Get or create a dynamic client from server config
pub async fn get_or_create_client(
&self,
server_config: McpServerConfig,
) -> McpResult<Arc<McpClient>> {
// Check if client already exists
let server_key = Self::server_key(&server_config);
let server_name = server_config.name.clone();
if let Some(client) = self.clients.get(&server_key) {
if let Some(client) = self.static_clients.get(&server_name) {
return Ok(Arc::clone(client.value()));
}
// Client doesn't exist, create new one via connection pool
let server_key = Self::server_key(&server_config);
let client = self
.connection_pool
.get_or_create(
@@ -170,47 +131,27 @@ impl McpManager {
)
.await?;
// Store in clients map
self.clients.insert(server_key, Arc::clone(&client));
self.inventory.clear_server_tools(&server_key);
Self::load_server_inventory(&self.inventory, &server_key, &client).await;
Ok(client)
}
/// List all static server names
pub fn list_static_servers(&self) -> Vec<String> {
self.static_servers
self.static_clients
.iter()
.map(|e| e.key().clone())
.collect()
}
/// Check if a server is static
pub fn is_static_server(&self, server_name: &str) -> bool {
self.static_servers.contains_key(server_name)
self.static_clients.contains_key(server_name)
}
/// Register a static server (called by workflow system)
///
/// This method registers a static MCP server that was configured and connected
/// via the workflow system. Static servers are never removed during runtime.
///
/// # Arguments
/// * `name` - Unique server name (from config)
/// * `client` - Connected MCP client
pub fn register_static_server(&self, name: String, client: Arc<McpClient>) {
// Insert into clients map
self.clients.insert(name.clone(), client);
// Mark as static server (for background refresh and stats)
self.static_servers.insert(name.clone(), ());
self.static_clients.insert(name.clone(), client);
info!("Registered static MCP server: {}", name);
}
// ========================================================================
// Tool Operations (delegate to clients via inventory)
// ========================================================================
/// List all available tools from all servers
pub fn list_tools(&self) -> Vec<Tool> {
self.inventory
@@ -267,10 +208,6 @@ impl McpManager {
.map(|(_server_name, tool_info)| tool_info)
}
// ========================================================================
// Prompt Operations
// ========================================================================
/// Get a prompt by name
pub async fn get_prompt(
&self,
@@ -310,10 +247,6 @@ impl McpManager {
.collect()
}
// ========================================================================
// Resource Operations
// ========================================================================
/// Read a resource by URI
pub async fn read_resource(&self, uri: &str) -> McpResult<ReadResourceResult> {
// Get server that owns this resource
@@ -348,10 +281,6 @@ impl McpManager {
.collect()
}
// ========================================================================
// Inventory Management
// ========================================================================
/// Refresh inventory for a specific server
pub async fn refresh_server_inventory(&self, server_name: &str) -> McpResult<()> {
let client = self
@@ -365,7 +294,8 @@ impl McpManager {
Ok(())
}
/// Start background refresh for all static servers
/// Start background refresh for ALL servers (static + dynamic)
/// Refreshes every 10-15 minutes to keep tool inventory up-to-date
pub fn spawn_background_refresh_all(
self: Arc<Self>,
refresh_interval: Duration,
@@ -377,17 +307,24 @@ impl McpManager {
loop {
interval.tick().await;
let server_names = self.list_static_servers();
// Get all static server keys
// Note: Dynamic clients in the connection pool are refreshed on-demand
// when they are accessed via get_or_create_client()
let server_keys: Vec<String> = self
.static_clients
.iter()
.map(|e| e.key().clone())
.collect();
if !server_names.is_empty() {
if !server_keys.is_empty() {
debug!(
"Background refresh: Refreshing {} static server(s)",
server_names.len()
server_keys.len()
);
for server_name in server_names {
if let Err(e) = self.refresh_server_inventory(&server_name).await {
warn!("Background refresh failed for '{}': {}", server_name, e);
for server_key in server_keys {
if let Err(e) = self.refresh_server_inventory(&server_key).await {
warn!("Background refresh failed for '{}': {}", server_key, e);
}
}
@@ -397,10 +334,6 @@ impl McpManager {
})
}
// ========================================================================
// Additional Tool/Prompt/Resource Methods
// ========================================================================
/// Check if a tool exists
pub fn has_tool(&self, name: &str) -> bool {
self.inventory.has_tool(name)
@@ -462,41 +395,56 @@ impl McpManager {
.map_err(|e| McpError::ToolExecution(format!("Failed to unsubscribe: {}", e)))
}
/// List all connected servers
/// List all connected servers (static + dynamic)
pub fn list_servers(&self) -> Vec<String> {
self.clients.iter().map(|e| e.key().clone()).collect()
let mut servers = Vec::new();
// Add static servers
servers.extend(self.static_clients.iter().map(|e| e.key().clone()));
// Add dynamic servers from connection pool
servers.extend(self.connection_pool.list_server_keys());
servers
}
/// Disconnect from all servers (for cleanup)
pub async fn shutdown(&self) {
let keys: Vec<String> = self.clients.iter().map(|e| e.key().clone()).collect();
for name in keys {
if let Some((_, client)) = self.clients.remove(&name) {
// Shutdown static servers
let static_keys: Vec<String> = self
.static_clients
.iter()
.map(|e| e.key().clone())
.collect();
for name in static_keys {
if let Some((_key, client)) = self.static_clients.remove(&name) {
// Try to unwrap Arc to call cancel
match Arc::try_unwrap(client) {
Ok(client) => {
if let Err(e) = client.cancel().await {
warn!("Error disconnecting from '{}': {}", name, e);
warn!("Error disconnecting from static server '{}': {}", name, e);
}
}
Err(_) => {
warn!("Could not shutdown '{}': client still in use", name);
warn!(
"Could not shutdown static server '{}': client still in use",
name
);
}
}
}
}
}
// ========================================================================
// Statistics and Accessors
// ========================================================================
// Clear dynamic clients from connection pool
// The pool will handle cleanup on drop
self.connection_pool.clear();
}
/// Get statistics about the manager
pub fn stats(&self) -> McpManagerStats {
let (tools, prompts, resources) = self.inventory.counts();
McpManagerStats {
static_server_count: self.static_servers.len(),
static_server_count: self.static_clients.len(),
pool_stats: self.connection_pool.stats(),
tool_count: tools,
prompt_count: prompts,
@@ -525,44 +473,41 @@ impl McpManager {
/// It discovers all tools, prompts, and resources from the client and caches them in the inventory.
pub async fn load_server_inventory(
inventory: &Arc<ToolInventory>,
server_name: &str,
server_key: &str,
client: &Arc<McpClient>,
) {
// Tools
match client.peer().list_all_tools().await {
Ok(ts) => {
info!("Discovered {} tools from '{}'", ts.len(), server_name);
info!("Discovered {} tools from '{}'", ts.len(), server_key);
for t in ts {
inventory.insert_tool(t.name.to_string(), server_name.to_string(), t);
inventory.insert_tool(t.name.to_string(), server_key.to_string(), t);
}
}
Err(e) => warn!("Failed to list tools from '{}': {}", server_name, e),
Err(e) => warn!("Failed to list tools from '{}': {}", server_key, e),
}
// Prompts
match client.peer().list_all_prompts().await {
Ok(ps) => {
info!("Discovered {} prompts from '{}'", ps.len(), server_name);
info!("Discovered {} prompts from '{}'", ps.len(), server_key);
for p in ps {
inventory.insert_prompt(p.name.clone(), server_name.to_string(), p);
inventory.insert_prompt(p.name.clone(), server_key.to_string(), p);
}
}
Err(e) => debug!("No prompts or failed to list on '{}': {}", server_name, e),
Err(e) => debug!("No prompts or failed to list on '{}': {}", server_key, e),
}
// Resources
match client.peer().list_all_resources().await {
Ok(rs) => {
info!("Discovered {} resources from '{}'", rs.len(), server_name);
info!("Discovered {} resources from '{}'", rs.len(), server_key);
for r in rs {
inventory.insert_resource(r.uri.clone(), server_name.to_string(), r.raw);
inventory.insert_resource(r.uri.clone(), server_key.to_string(), r.raw);
}
}
Err(e) => debug!("No resources or failed to list on '{}': {}", server_name, e),
Err(e) => debug!("No resources or failed to list on '{}': {}", server_key, e),
}
// Mark server as refreshed
inventory.mark_refreshed(server_name);
}
/// Discover and cache tools/prompts/resources for a connected server (internal wrapper)
@@ -602,9 +547,6 @@ impl McpManager {
}
Err(e) => debug!("No resources or failed to list on '{}': {}", server_name, e),
}
// Mark server as refreshed
self.inventory.mark_refreshed(server_name);
}
// ========================================================================
@@ -834,14 +776,7 @@ mod tests {
inventory: Default::default(),
};
let manager = McpManager::new(
config,
Duration::from_secs(300),
Duration::from_secs(300),
100,
)
.await
.unwrap();
let manager = McpManager::new(config, 100).await.unwrap();
assert_eq!(manager.list_static_servers().len(), 0);
}
}

12
sgl-router/src/mcp/mod.rs
View File

@@ -1,11 +1,7 @@
// MCP Client for SGLang Router
//
// This module provides a complete MCP (Model Context Protocol) client implementation
// supporting multiple transport types (stdio, SSE, HTTP) and all MCP features:
// - Tools: Discovery and execution
// - Prompts: Reusable templates for LLM interactions
// - Resources: File/data access with subscription support
// - OAuth: Secure authentication for remote servers
//! Model Context Protocol (MCP) client implementation.
//!
//! Provides MCP client functionality including tools, prompts, resources, and OAuth.
//! Supports stdio, SSE, and HTTP transports with connection pooling and caching.
pub mod config;
pub mod connection_pool;

4
sgl-router/src/mcp/oauth.rs
View File

@@ -1,4 +1,6 @@
// OAuth authentication support for MCP servers
//! OAuth authentication for MCP servers.
//!
//! Handles OAuth flow including callback server and token exchange.
use std::{net::SocketAddr, sync::Arc};

8
sgl-router/src/mcp/proxy.rs
View File

@@ -1,8 +1,6 @@
// MCP Proxy Configuration and Resolution
//
// This module provides proxy configuration resolution and HTTP client creation
// for MCP server connections. Proxy settings are MCP-specific and do NOT affect
// LLM API traffic.
//! HTTP proxy configuration for MCP connections.
//!
//! Resolves proxy settings and creates HTTP clients for MCP server connections.
use std::time::Duration;

5
sgl-router/src/mcp/tool_args.rs
View File

@@ -1,7 +1,6 @@
//! Tool arguments handling and type coercion
//! MCP tool argument handling.
//!
//! This module provides utilities for handling MCP tool arguments,
//! supporting both JSON strings and parsed Maps with automatic type coercion.
//! Supports both JSON strings and parsed Maps with automatic type coercion.
use serde_json::Map;

6
sgl-router/src/server.rs
View File

@@ -782,12 +782,12 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
info!("No MCP config provided, skipping MCP server initialization");
}
// Start background refresh for all registered static MCP servers
// Start background refresh for ALL MCP servers (static + dynamic in LRU cache)
if let Some(mcp_manager) = app_context.mcp_manager.get() {
let refresh_interval = Duration::from_secs(300); // 5 minutes, matches default TTL
let refresh_interval = Duration::from_secs(600); // 10 minutes
let _refresh_handle =
Arc::clone(mcp_manager).spawn_background_refresh_all(refresh_interval);
info!("Started background refresh for all static MCP servers");
info!("Started background refresh for all MCP servers (every 10 minutes)");
}
let worker_stats = app_context.worker_registry.stats();