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

remove self managed mcp as it has been replaced with official rmcp crate (#17740)

Browse Source
This commit is contained in:
Simo Lin
2026-01-26 01:17:06 -05:00
committed by GitHub
parent 52e0f65fce
commit 7890a96f97
11 changed files with 2 additions and 2726 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-model-gateway/Cargo.toml
View File

@@ -87,6 +87,7 @@ llm-tokenizer = "1.0.0"
smg-auth = "1.0.0"
wfaas = "1.0.0"
data-connector = "1.0.0"
smg-mcp = "1.0.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
rustls-pemfile = "2.2"
openssl = "0.10.73"

2
sgl-model-gateway/src/lib.rs
View File

@@ -4,7 +4,7 @@ pub mod config;
pub mod core;
pub use data_connector;
pub mod grpc_client;
pub mod mcp;
pub use smg_mcp as mcp;
pub mod mesh;
pub mod middleware;
pub mod multimodal;

528
sgl-model-gateway/src/mcp/config.rs
View File

@@ -1,528 +0,0 @@
//! MCP configuration types and utilities.
//!
//! Defines configuration structures for MCP servers, transports, proxies, and inventory.
use std::{collections::HashMap, fmt};
pub use rmcp::model::{Prompt, RawResource, Tool};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpConfig {
/// Static MCP servers (loaded at startup)
pub servers: Vec<McpServerConfig>,
/// Connection pool settings
#[serde(default)]
pub pool: McpPoolConfig,
/// Global MCP proxy configuration (default for all servers)
/// Can be overridden per-server
#[serde(default)]
pub proxy: Option<McpProxyConfig>,
/// Pre-warm these connections at startup
#[serde(default)]
pub warmup: Vec<WarmupServer>,
/// Tool inventory refresh settings
#[serde(default)]
pub inventory: InventoryConfig,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpServerConfig {
pub name: String,
#[serde(flatten)]
pub transport: McpTransport,
/// Per-server proxy override (overrides global proxy)
/// Set to `null` in YAML to force direct connection (no proxy)
#[serde(default)]
pub proxy: Option<McpProxyConfig>,
/// Whether this server is required for router startup
/// - true: Router startup fails if this server cannot be reached
/// - false: Log warning but continue (default)
#[serde(default)]
pub required: bool,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(tag = "protocol", rename_all = "lowercase")]
pub enum McpTransport {
Stdio {
command: String,
#[serde(default)]
args: Vec<String>,
#[serde(default)]
envs: HashMap<String, String>,
},
Sse {
url: String,
#[serde(skip_serializing_if = "Option::is_none")]
token: Option<String>,
},
Streamable {
url: String,
#[serde(skip_serializing_if = "Option::is_none")]
token: Option<String>,
},
}
impl fmt::Debug for McpTransport {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
McpTransport::Stdio {
command,
args,
envs,
} => f
.debug_struct("Stdio")
.field("command", command)
.field("args", args)
.field("envs", envs)
.finish(),
McpTransport::Sse { url, token } => f
.debug_struct("Sse")
.field("url", url)
.field("token", &token.as_ref().map(|_| "****"))
.finish(),
McpTransport::Streamable { url, token } => f
.debug_struct("Streamable")
.field("url", url)
.field("token", &token.as_ref().map(|_| "****"))
.finish(),
}
}
}
/// MCP-specific proxy configuration (does NOT affect LLM API traffic)
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpProxyConfig {
/// HTTP proxy URL (e.g., "http://proxy.internal:8080")
pub http: Option<String>,
/// HTTPS proxy URL
pub https: Option<String>,
/// Comma-separated hosts to exclude from proxying
/// Example: "localhost,127.0.0.1,*.internal,10.*"
pub no_proxy: Option<String>,
/// Custom proxy authentication (if needed)
#[serde(skip_serializing_if = "Option::is_none")]
pub username: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub password: Option<String>,
}
/// Connection pool configuration
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpPoolConfig {
/// Maximum cached connections per server URL
#[serde(default = "default_max_connections")]
pub max_connections: usize,
/// Idle timeout before closing connection (seconds)
#[serde(default = "default_idle_timeout")]
pub idle_timeout: u64,
}
/// Tool inventory refresh configuration
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct InventoryConfig {
/// Enable automatic tool inventory refresh
#[serde(default = "default_true")]
pub enable_refresh: bool,
/// Tool cache TTL (seconds) - how long tools are considered fresh
#[serde(default = "default_tool_ttl")]
pub tool_ttl: u64,
/// Background refresh interval (seconds) - proactive refresh
#[serde(default = "default_refresh_interval")]
pub refresh_interval: u64,
/// Refresh on tool call failure (try refreshing if tool not found)
#[serde(default = "default_true")]
pub refresh_on_error: bool,
}
/// Pre-warm server connections at startup
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct WarmupServer {
/// Server URL
pub url: String,
/// Server label/name
pub label: String,
/// Optional authentication token
#[serde(skip_serializing_if = "Option::is_none")]
pub token: Option<String>,
}
// Default value functions
fn default_max_connections() -> usize {
100
}
fn default_idle_timeout() -> u64 {
300 // 5 minutes
}
fn default_true() -> bool {
true
}
fn default_tool_ttl() -> u64 {
300 // 5 minutes
}
fn default_refresh_interval() -> u64 {
60 // 1 minute
}
// Default implementations
impl Default for McpPoolConfig {
fn default() -> Self {
Self {
max_connections: default_max_connections(),
idle_timeout: default_idle_timeout(),
}
}
}
impl Default for InventoryConfig {
fn default() -> Self {
Self {
enable_refresh: true,
tool_ttl: default_tool_ttl(),
refresh_interval: default_refresh_interval(),
refresh_on_error: true,
}
}
}
impl McpProxyConfig {
/// Load proxy config from standard environment variables
pub fn from_env() -> Option<Self> {
let http = std::env::var("MCP_HTTP_PROXY")
.ok()
.or_else(|| std::env::var("HTTP_PROXY").ok());
let https = std::env::var("MCP_HTTPS_PROXY")
.ok()
.or_else(|| std::env::var("HTTPS_PROXY").ok());
let no_proxy = std::env::var("MCP_NO_PROXY")
.ok()
.or_else(|| std::env::var("NO_PROXY").ok());
if http.is_some() || https.is_some() {
Some(Self {
http,
https,
no_proxy,
username: None,
password: None,
})
} else {
None
}
}
}
impl McpConfig {
/// Load configuration from a YAML file
pub async fn from_file(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
let content = tokio::fs::read_to_string(path).await?;
let config: Self = serde_yaml::from_str(&content)?;
Ok(config)
}
/// Load configuration from environment variables (optional)
pub fn from_env() -> Option<Self> {
// This could be expanded to read from env vars
// For now, return None to indicate env config not implemented
None
}
/// Merge with environment-based proxy config
pub fn with_env_proxy(mut self) -> Self {
if self.proxy.is_none() {
self.proxy = McpProxyConfig::from_env();
}
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_pool_config() {
let config = McpPoolConfig::default();
assert_eq!(config.max_connections, 100);
assert_eq!(config.idle_timeout, 300);
}
#[test]
fn test_default_inventory_config() {
let config = InventoryConfig::default();
assert!(config.enable_refresh);
assert_eq!(config.tool_ttl, 300);
assert_eq!(config.refresh_interval, 60);
assert!(config.refresh_on_error);
}
#[test]
fn test_proxy_from_env_empty() {
// Ensure no proxy env vars are set for this test
std::env::remove_var("MCP_HTTP_PROXY");
std::env::remove_var("MCP_HTTPS_PROXY");
std::env::remove_var("HTTP_PROXY");
std::env::remove_var("HTTPS_PROXY");
let proxy = McpProxyConfig::from_env();
assert!(proxy.is_none(), "Should return None when no env vars set");
}
#[test]
fn test_proxy_from_env_with_vars() {
std::env::set_var("MCP_HTTP_PROXY", "http://test-proxy:8080");
std::env::set_var("MCP_NO_PROXY", "localhost,127.0.0.1");
let proxy = McpProxyConfig::from_env();
assert!(proxy.is_some(), "Should return Some when env vars set");
let proxy = proxy.unwrap();
assert_eq!(proxy.http.as_ref().unwrap(), "http://test-proxy:8080");
assert_eq!(proxy.no_proxy.as_ref().unwrap(), "localhost,127.0.0.1");
// Cleanup
std::env::remove_var("MCP_HTTP_PROXY");
std::env::remove_var("MCP_NO_PROXY");
}
#[tokio::test]
async fn test_yaml_minimal_config() {
let yaml = r#"
servers:
- name: "test-server"
protocol: sse
url: "http://localhost:3000/sse"
"#;
let config: McpConfig = serde_yaml::from_str(yaml).expect("Failed to parse YAML");
assert_eq!(config.servers.len(), 1);
assert_eq!(config.servers[0].name, "test-server");
assert!(!config.servers[0].required); // Should default to false
assert!(config.servers[0].proxy.is_none()); // Should default to None
assert_eq!(config.pool.max_connections, 100); // Should use default
assert_eq!(config.inventory.tool_ttl, 300); // Should use default
}
#[tokio::test]
async fn test_yaml_full_config() {
let yaml = r#"
# Global proxy configuration
proxy:
http: "http://global-proxy:8080"
https: "http://global-proxy:8080"
no_proxy: "localhost,127.0.0.1,*.internal"
# Connection pool settings
pool:
max_connections: 50
idle_timeout: 600
# Tool inventory settings
inventory:
enable_refresh: true
tool_ttl: 600
refresh_interval: 120
refresh_on_error: true
# Static servers
servers:
- name: "required-server"
protocol: sse
url: "https://api.example.com/sse"
token: "secret-token"
required: true
- name: "optional-server"
protocol: stdio
command: "mcp-server"
args: ["--port", "3000"]
required: false
proxy:
http: "http://server-specific-proxy:9090"
# Pre-warm connections
warmup:
- url: "http://localhost:3000/sse"
label: "local-dev"
"#;
let config: McpConfig = serde_yaml::from_str(yaml).expect("Failed to parse YAML");
// Check global proxy
assert!(config.proxy.is_some());
let global_proxy = config.proxy.as_ref().unwrap();
assert_eq!(
global_proxy.http.as_ref().unwrap(),
"http://global-proxy:8080"
);
// Check pool config
assert_eq!(config.pool.max_connections, 50);
assert_eq!(config.pool.idle_timeout, 600);
// Check inventory config
assert_eq!(config.inventory.tool_ttl, 600);
assert_eq!(config.inventory.refresh_interval, 120);
// Check servers
assert_eq!(config.servers.len(), 2);
// Required server
assert_eq!(config.servers[0].name, "required-server");
assert!(config.servers[0].required);
assert!(config.servers[0].proxy.is_none()); // Inherits global proxy
// Optional server with custom proxy
assert_eq!(config.servers[1].name, "optional-server");
assert!(!config.servers[1].required);
assert!(config.servers[1].proxy.is_some());
assert_eq!(
config.servers[1]
.proxy
.as_ref()
.unwrap()
.http
.as_ref()
.unwrap(),
"http://server-specific-proxy:9090"
);
// Check warmup
assert_eq!(config.warmup.len(), 1);
assert_eq!(config.warmup[0].label, "local-dev");
}
#[tokio::test]
async fn test_yaml_backward_compatibility() {
// Old config format should still work
let yaml = r#"
servers:
- name: "legacy-server"
protocol: sse
url: "http://localhost:3000/sse"
"#;
let config: McpConfig = serde_yaml::from_str(yaml).expect("Failed to parse old format");
assert_eq!(config.servers.len(), 1);
assert_eq!(config.servers[0].name, "legacy-server");
assert!(!config.servers[0].required); // New field defaults to false
assert!(config.servers[0].proxy.is_none()); // New field defaults to None
assert!(config.proxy.is_none()); // New field defaults to None
assert!(config.warmup.is_empty()); // New field defaults to empty
}
#[tokio::test]
async fn test_yaml_null_proxy_override() {
// Test that explicit null in YAML sets proxy to None
let yaml = r#"
proxy:
http: "http://global-proxy:8080"
servers:
- name: "direct-connection"
protocol: sse
url: "http://localhost:3000/sse"
proxy: null
"#;
let config: McpConfig = serde_yaml::from_str(yaml).expect("Failed to parse YAML");
assert!(config.proxy.is_some()); // Global proxy set
assert_eq!(config.servers.len(), 1);
assert!(config.servers[0].proxy.is_none()); // Explicitly set to null
}
#[test]
fn test_transport_stdio() {
let yaml = r#"
name: "test"
protocol: stdio
command: "mcp-server"
args: ["--port", "3000"]
envs:
VAR1: "value1"
VAR2: "value2"
"#;
let config: McpServerConfig = serde_yaml::from_str(yaml).expect("Failed to parse stdio");
assert_eq!(config.name, "test");
match config.transport {
McpTransport::Stdio {
command,
args,
envs,
} => {
assert_eq!(command, "mcp-server");
assert_eq!(args.len(), 2);
assert_eq!(args[0], "--port");
assert_eq!(envs.get("VAR1").unwrap(), "value1");
}
_ => panic!("Expected Stdio transport"),
}
}
#[test]
fn test_transport_sse() {
let yaml = r#"
name: "test"
protocol: sse
url: "http://localhost:3000/sse"
token: "secret"
"#;
let config: McpServerConfig = serde_yaml::from_str(yaml).expect("Failed to parse sse");
assert_eq!(config.name, "test");
match config.transport {
McpTransport::Sse { url, token } => {
assert_eq!(url, "http://localhost:3000/sse");
assert_eq!(token.unwrap(), "secret");
}
_ => panic!("Expected Sse transport"),
}
}
#[test]
fn test_transport_streamable() {
let yaml = r#"
name: "test"
protocol: streamable
url: "http://localhost:3000"
"#;
let config: McpServerConfig =
serde_yaml::from_str(yaml).expect("Failed to parse streamable");
assert_eq!(config.name, "test");
match config.transport {
McpTransport::Streamable { url, token } => {
assert_eq!(url, "http://localhost:3000");
assert!(token.is_none());
}
_ => panic!("Expected Streamable transport"),
}
}
}

319
sgl-model-gateway/src/mcp/connection_pool.rs
View File

@@ -1,319 +0,0 @@
/// MCP Connection Pool
///
/// This module provides connection pooling for dynamic MCP servers (per-request).
use std::sync::Arc;
use lru::LruCache;
use parking_lot::Mutex;
use rmcp::{service::RunningService, RoleClient};
use crate::mcp::{
config::{McpProxyConfig, McpServerConfig},
error::McpResult,
};
/// 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(crate) struct CachedConnection {
/// The MCP client instance
pub client: Arc<McpClient>,
/// Server configuration used to create this connection
#[allow(dead_code)]
pub config: McpServerConfig,
}
impl CachedConnection {
/// Create a new cached connection
pub fn new(client: Arc<McpClient>, config: McpServerConfig) -> Self {
Self { client, config }
}
}
/// Connection pool for dynamic MCP servers
///
/// Provides thread-safe connection pooling with LRU eviction.
/// Connections are keyed by server URL and reused across requests.
pub struct McpConnectionPool {
/// LRU cache of server_url -> cached connection
connections: Arc<Mutex<LruCache<String, CachedConnection>>>,
/// 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:
/// - max_connections: 200
/// - global_proxy: Loaded from environment variables (MCP_HTTP_PROXY, etc.)
pub fn new() -> Self {
Self {
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 capacity
pub fn with_capacity(max_connections: usize) -> Self {
Self {
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(max_connections: usize, global_proxy: Option<McpProxyConfig>) -> Self {
Self {
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 (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)
/// * `server_config` - Server configuration (used to create new connection if needed)
/// * `connect_fn` - Async function to create a new client connection
///
/// # Returns
/// Arc to the MCP client, either from cache or newly created
pub async fn get_or_create<F, Fut>(
&self,
server_url: &str,
server_config: McpServerConfig,
connect_fn: F,
) -> McpResult<Arc<McpClient>>
where
F: FnOnce(McpServerConfig, Option<McpProxyConfig>) -> Fut,
Fut: std::future::Future<Output = McpResult<McpClient>>,
{
// 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));
}
}
// Slow path: Create new connection
let client = connect_fn(server_config.clone(), self.global_proxy.clone()).await?;
let client_arc = Arc::new(client);
// 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);
}
}
}
Ok(client_arc)
}
/// Get current number of cached connections
pub fn len(&self) -> usize {
self.connections.lock().len()
}
/// Check if pool is empty
pub fn is_empty(&self) -> bool {
self.connections.lock().is_empty()
}
/// Clear all connections
pub fn clear(&self) {
self.connections.lock().clear();
}
/// Get connection statistics
pub fn stats(&self) -> PoolStats {
let total = self.connections.lock().len();
PoolStats {
total_connections: total,
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 {
fn default() -> Self {
Self::new()
}
}
/// Connection pool statistics
#[derive(Debug, Clone)]
pub struct PoolStats {
pub total_connections: usize,
pub capacity: usize,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mcp::config::McpTransport;
// Helper to create test server config
fn create_test_config(url: &str) -> McpServerConfig {
McpServerConfig {
name: "test_server".to_string(),
transport: McpTransport::Streamable {
url: url.to_string(),
token: None,
},
proxy: None,
required: false,
}
}
#[tokio::test]
async fn test_pool_creation() {
let pool = McpConnectionPool::new();
assert_eq!(pool.len(), 0);
assert!(pool.is_empty());
}
#[test]
fn test_pool_stats() {
let pool = McpConnectionPool::with_capacity(10);
let stats = pool.stats();
assert_eq!(stats.total_connections, 0);
assert_eq!(stats.capacity, 10);
}
#[test]
#[allow(invalid_value)]
fn test_pool_clear() {
let pool = McpConnectionPool::new();
// Add a connection
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
.lock()
.push("http://localhost:3000".to_string(), cached);
assert_eq!(pool.len(), 1);
pool.clear();
assert_eq!(pool.len(), 0);
assert!(pool.is_empty());
// Prevent drop of invalid Arc (would segfault)
std::mem::forget(client);
}
#[test]
fn test_pool_with_global_proxy() {
use crate::mcp::config::McpProxyConfig;
// Create proxy config
let proxy = McpProxyConfig {
http: Some("http://proxy.example.com:8080".to_string()),
https: None,
no_proxy: Some("localhost,127.0.0.1".to_string()),
username: None,
password: None,
};
// Create pool with proxy
let pool = McpConnectionPool::with_full_config(100, Some(proxy.clone()));
// Verify proxy is stored
assert!(pool.global_proxy.is_some());
let stored_proxy = pool.global_proxy.as_ref().unwrap();
assert_eq!(
stored_proxy.http.as_ref().unwrap(),
"http://proxy.example.com:8080"
);
assert_eq!(
stored_proxy.no_proxy.as_ref().unwrap(),
"localhost,127.0.0.1"
);
}
#[test]
fn test_pool_proxy_from_env() {
// Note: This test depends on environment variables
// In production, proxy is loaded from MCP_HTTP_PROXY or HTTP_PROXY env vars
let pool = McpConnectionPool::new();
// Pool should either have proxy from env or None
// We can't assert specific value since it depends on test environment
// Just verify it doesn't panic
assert!(pool.global_proxy.is_some() || pool.global_proxy.is_none());
}
}

50
sgl-model-gateway/src/mcp/error.rs
View File

@@ -1,50 +0,0 @@
//! 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>;
#[derive(Debug, Error)]
pub enum McpError {
#[error("Server not found: {0}")]
ServerNotFound(String),
#[error("Tool not found: {0}")]
ToolNotFound(String),
#[error("Transport error: {0}")]
Transport(String),
#[error("Tool execution failed: {0}")]
ToolExecution(String),
#[error("Connection failed: {0}")]
ConnectionFailed(String),
#[error("Configuration error: {0}")]
Config(String),
#[error("Authentication error: {0}")]
Auth(String),
#[error("Resource not found: {0}")]
ResourceNotFound(String),
#[error("Prompt not found: {0}")]
PromptNotFound(String),
#[error("Invalid arguments: {0}")]
InvalidArguments(String),
#[error(transparent)]
Sdk(#[from] Box<rmcp::RmcpError>),
#[error(transparent)]
Io(#[from] std::io::Error),
#[error(transparent)]
Http(#[from] reqwest::Error),
}

447
sgl-model-gateway/src/mcp/inventory.rs
View File

@@ -1,447 +0,0 @@
//! MCP tool, prompt, and resource inventory.
//!
//! Thread-safe cache for MCP capabilities across all connected servers.
use dashmap::DashMap;
use crate::mcp::config::{Prompt, RawResource, Tool};
/// Cached tool with metadata
#[derive(Clone)]
pub(crate) struct CachedTool {
pub server_name: String,
pub tool: Tool,
}
/// Cached prompt with metadata
#[derive(Clone)]
pub(crate) struct CachedPrompt {
pub server_name: String,
pub prompt: Prompt,
}
/// Cached resource with metadata
#[derive(Clone)]
pub(crate) struct CachedResource {
pub server_name: String,
pub resource: RawResource,
}
/// Tool inventory with periodic refresh
///
/// 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>,
/// Map of prompt_name -> cached prompt
prompts: DashMap<String, CachedPrompt>,
/// Map of resource_uri -> cached resource
resources: DashMap<String, CachedResource>,
}
impl ToolInventory {
/// Create a new tool inventory
pub fn new() -> Self {
Self {
tools: DashMap::new(),
prompts: DashMap::new(),
resources: DashMap::new(),
}
}
}
impl Default for ToolInventory {
fn default() -> Self {
Self::new()
}
}
impl ToolInventory {
// ============================================================================
// Tool Methods
// ============================================================================
/// Get a tool if it exists
pub fn get_tool(&self, tool_name: &str) -> Option<(String, Tool)> {
self.tools
.get(tool_name)
.map(|entry| (entry.server_name.clone(), entry.tool.clone()))
}
/// 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 });
}
/// Get all tools
pub fn list_tools(&self) -> Vec<(String, String, Tool)> {
self.tools
.iter()
.map(|entry| {
let (name, cached) = entry.pair();
(
name.clone(),
cached.server_name.clone(),
cached.tool.clone(),
)
})
.collect()
}
// ============================================================================
// Prompt Methods
// ============================================================================
/// Get a prompt if it exists
pub fn get_prompt(&self, prompt_name: &str) -> Option<(String, Prompt)> {
self.prompts
.get(prompt_name)
.map(|entry| (entry.server_name.clone(), entry.prompt.clone()))
}
/// Check if prompt exists
pub fn has_prompt(&self, prompt_name: &str) -> bool {
self.prompts.contains_key(prompt_name)
}
/// Insert or update a prompt
pub fn insert_prompt(&self, prompt_name: String, server_name: String, prompt: Prompt) {
self.prompts.insert(
prompt_name,
CachedPrompt {
server_name,
prompt,
},
);
}
/// Get all prompts
pub fn list_prompts(&self) -> Vec<(String, String, Prompt)> {
self.prompts
.iter()
.map(|entry| {
let (name, cached) = entry.pair();
(
name.clone(),
cached.server_name.clone(),
cached.prompt.clone(),
)
})
.collect()
}
// ============================================================================
// Resource Methods
// ============================================================================
/// Get a resource if it exists
pub fn get_resource(&self, resource_uri: &str) -> Option<(String, RawResource)> {
self.resources
.get(resource_uri)
.map(|entry| (entry.server_name.clone(), entry.resource.clone()))
}
/// Check if resource exists
pub fn has_resource(&self, resource_uri: &str) -> bool {
self.resources.contains_key(resource_uri)
}
/// Insert or update a resource
pub fn insert_resource(
&self,
resource_uri: String,
server_name: String,
resource: RawResource,
) {
self.resources.insert(
resource_uri,
CachedResource {
server_name,
resource,
},
);
}
/// Get all resources
pub fn list_resources(&self) -> Vec<(String, String, RawResource)> {
self.resources
.iter()
.map(|entry| {
let (uri, cached) = entry.pair();
(
uri.clone(),
cached.server_name.clone(),
cached.resource.clone(),
)
})
.collect()
}
// ============================================================================
// Server Management Methods
// ============================================================================
/// 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);
self.prompts
.retain(|_, cached| cached.server_name != server_name);
self.resources
.retain(|_, cached| cached.server_name != server_name);
}
/// Get count of cached items
pub fn counts(&self) -> (usize, usize, usize) {
(self.tools.len(), self.prompts.len(), self.resources.len())
}
/// Clear all cached items
pub fn clear_all(&self) {
self.tools.clear();
self.prompts.clear();
self.resources.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mcp::config::{Prompt, RawResource, Tool};
// Helper to create a test tool
fn create_test_tool(name: &str) -> Tool {
use std::{borrow::Cow, sync::Arc};
let schema_obj = serde_json::json!({
"type": "object",
"properties": {}
});
let schema_map = if let serde_json::Value::Object(m) = schema_obj {
m
} else {
serde_json::Map::new()
};
Tool {
name: Cow::Owned(name.to_string()),
title: None,
description: Some(Cow::Owned(format!("Test tool: {}", name))),
input_schema: Arc::new(schema_map),
output_schema: None,
annotations: None,
icons: None,
}
}
// Helper to create a test prompt
fn create_test_prompt(name: &str) -> Prompt {
Prompt {
name: name.to_string(),
title: None,
description: Some(format!("Test prompt: {}", name)),
arguments: None,
icons: None,
}
}
// Helper to create a test resource
fn create_test_resource(uri: &str) -> RawResource {
RawResource {
uri: uri.to_string(),
name: uri.to_string(),
title: None,
description: Some(format!("Test resource: {}", uri)),
mime_type: Some("text/plain".to_string()),
size: None,
icons: None,
}
}
#[test]
fn test_tool_insert_and_get() {
let inventory = ToolInventory::new();
let tool = create_test_tool("test_tool");
inventory.insert_tool("test_tool".to_string(), "server1".to_string(), tool.clone());
let result = inventory.get_tool("test_tool");
assert!(result.is_some());
let (server_name, retrieved_tool) = result.unwrap();
assert_eq!(server_name, "server1");
assert_eq!(retrieved_tool.name, "test_tool");
}
#[test]
fn test_has_tool() {
let inventory = ToolInventory::new();
let tool = create_test_tool("check_tool");
assert!(!inventory.has_tool("check_tool"));
inventory.insert_tool("check_tool".to_string(), "server1".to_string(), tool);
assert!(inventory.has_tool("check_tool"));
}
#[test]
fn test_list_tools() {
let inventory = ToolInventory::new();
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"),
);
inventory.insert_tool(
"tool3".to_string(),
"server2".to_string(),
create_test_tool("tool3"),
);
let tools = inventory.list_tools();
assert_eq!(tools.len(), 3);
}
#[test]
fn test_clear_server_tools() {
let inventory = ToolInventory::new();
inventory.insert_tool(
"tool1".to_string(),
"server1".to_string(),
create_test_tool("tool1"),
);
inventory.insert_tool(
"tool2".to_string(),
"server2".to_string(),
create_test_tool("tool2"),
);
assert_eq!(inventory.list_tools().len(), 2);
inventory.clear_server_tools("server1");
let tools = inventory.list_tools();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].0, "tool2");
}
#[test]
fn test_prompt_operations() {
let inventory = ToolInventory::new();
let prompt = create_test_prompt("test_prompt");
inventory.insert_prompt(
"test_prompt".to_string(),
"server1".to_string(),
prompt.clone(),
);
assert!(inventory.has_prompt("test_prompt"));
let result = inventory.get_prompt("test_prompt");
assert!(result.is_some());
let (server_name, retrieved_prompt) = result.unwrap();
assert_eq!(server_name, "server1");
assert_eq!(retrieved_prompt.name, "test_prompt");
}
#[test]
fn test_resource_operations() {
let inventory = ToolInventory::new();
let resource = create_test_resource("file:///test.txt");
inventory.insert_resource(
"file:///test.txt".to_string(),
"server1".to_string(),
resource.clone(),
);
assert!(inventory.has_resource("file:///test.txt"));
let result = inventory.get_resource("file:///test.txt");
assert!(result.is_some());
let (server_name, retrieved_resource) = result.unwrap();
assert_eq!(server_name, "server1");
assert_eq!(retrieved_resource.uri, "file:///test.txt");
}
#[tokio::test]
async fn test_concurrent_access() {
use std::sync::Arc;
let inventory = Arc::new(ToolInventory::new());
// Spawn multiple tasks that insert tools concurrently
let mut handles = vec![];
for i in 0..10 {
let inv = Arc::clone(&inventory);
let handle = tokio::spawn(async move {
let tool = create_test_tool(&format!("tool_{}", i));
inv.insert_tool(format!("tool_{}", i), format!("server_{}", i % 3), tool);
});
handles.push(handle);
}
// Wait for all tasks to complete
for handle in handles {
handle.await.unwrap();
}
// Should have 10 tools
let (tools, _, _) = inventory.counts();
assert_eq!(tools, 10);
}
#[test]
fn test_clear_all() {
let inventory = ToolInventory::new();
inventory.insert_tool(
"tool1".to_string(),
"server1".to_string(),
create_test_tool("tool1"),
);
inventory.insert_prompt(
"prompt1".to_string(),
"server1".to_string(),
create_test_prompt("prompt1"),
);
inventory.insert_resource(
"res1".to_string(),
"server1".to_string(),
create_test_resource("res1"),
);
let (tools, prompts, resources) = inventory.counts();
assert_eq!(tools, 1);
assert_eq!(prompts, 1);
assert_eq!(resources, 1);
inventory.clear_all();
let (tools, prompts, resources) = inventory.counts();
assert_eq!(tools, 0);
assert_eq!(prompts, 0);
assert_eq!(resources, 0);
}
}

813
sgl-model-gateway/src/mcp/manager.rs
View File

@@ -1,813 +0,0 @@
//! MCP client management and orchestration.
//!
//! 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};
use backoff::ExponentialBackoffBuilder;
use dashmap::DashMap;
use rmcp::{
model::{
CallToolRequestParam, CallToolResult, GetPromptRequestParam, GetPromptResult,
ReadResourceRequestParam, ReadResourceResult, SubscribeRequestParam,
UnsubscribeRequestParam,
},
service::RunningService,
transport::{
sse_client::SseClientConfig, streamable_http_client::StreamableHttpClientTransportConfig,
ConfigureCommandExt, SseClientTransport, StreamableHttpClientTransport, TokioChildProcess,
},
RoleClient, ServiceExt,
};
use serde_json::Map;
use tracing::{debug, error, info, warn};
use crate::mcp::{
config::{McpConfig, McpProxyConfig, McpServerConfig, McpTransport, Prompt, RawResource, Tool},
connection_pool::McpConnectionPool,
error::{McpError, McpResult},
inventory::ToolInventory,
tool_args::ToolArgs,
};
/// Type alias for MCP client
type McpClient = RunningService<RoleClient, ()>;
pub struct McpManager {
static_clients: Arc<DashMap<String, Arc<McpClient>>>,
inventory: Arc<ToolInventory>,
connection_pool: Arc<McpConnectionPool>,
_config: McpConfig,
}
impl McpManager {
const MAX_DYNAMIC_CLIENTS: usize = 200;
pub async fn new(config: McpConfig, pool_max_connections: usize) -> McpResult<Self> {
let inventory = Arc::new(ToolInventory::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 {
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;
static_clients.insert(server_config.name.clone(), client_arc);
info!("Connected to static server '{}'", server_config.name);
}
Err(e) => {
error!(
"Failed to connect to static server '{}': {}",
server_config.name, e
);
}
}
}
if static_clients.is_empty() {
info!("No static MCP servers connected");
}
Ok(Self {
static_clients,
inventory,
connection_pool,
_config: config,
})
}
pub async fn with_defaults(config: McpConfig) -> McpResult<Self> {
Self::new(config, Self::MAX_DYNAMIC_CLIENTS).await
}
pub async fn get_client(&self, server_name: &str) -> Option<Arc<McpClient>> {
if let Some(client) = self.static_clients.get(server_name) {
return Some(Arc::clone(client.value()));
}
self.connection_pool.get(server_name)
}
pub async fn get_or_create_client(
&self,
server_config: McpServerConfig,
) -> McpResult<Arc<McpClient>> {
let server_name = server_config.name.clone();
if let Some(client) = self.static_clients.get(&server_name) {
return Ok(Arc::clone(client.value()));
}
let server_key = Self::server_key(&server_config);
let client = self
.connection_pool
.get_or_create(
&server_key,
server_config,
|config, global_proxy| async move {
Self::connect_server(&config, global_proxy.as_ref()).await
},
)
.await?;
self.inventory.clear_server_tools(&server_key);
Self::load_server_inventory(&self.inventory, &server_key, &client).await;
Ok(client)
}
pub fn list_static_servers(&self) -> Vec<String> {
self.static_clients
.iter()
.map(|e| e.key().clone())
.collect()
}
pub fn is_static_server(&self, server_name: &str) -> bool {
self.static_clients.contains_key(server_name)
}
pub fn register_static_server(&self, name: String, client: Arc<McpClient>) {
self.static_clients.insert(name.clone(), client);
info!("Registered static MCP server: {}", name);
}
/// List all available tools from all servers
pub fn list_tools(&self) -> Vec<Tool> {
self.inventory
.list_tools()
.into_iter()
.map(|(_tool_name, _server_name, tool_info)| tool_info)
.collect()
}
/// List tools only from specific servers plus all static servers
///
/// This method filters tools to only include:
/// 1. Tools from static servers (always visible)
/// 2. Tools from the specified dynamic servers
///
/// This provides request-scoped tool isolation while maintaining
/// global visibility for static servers.
pub fn list_tools_for_servers(&self, server_keys: &[String]) -> Vec<Tool> {
self.inventory
.list_tools()
.into_iter()
.filter(|(_tool_name, server_key, _tool_info)| {
// Include if:
// 1. It's a static server (check by name in static_clients)
// 2. It's in the requested servers list
self.is_static_server_by_key(server_key) || server_keys.contains(server_key)
})
.map(|(_tool_name, _server_key, tool_info)| tool_info)
.collect()
}
/// Check if a server key belongs to a static server
///
/// Static servers can be identified by checking if their name
/// exists in the static_clients map. We need to handle the fact
/// that static servers use name as key while dynamic use URL.
fn is_static_server_by_key(&self, server_key: &str) -> bool {
// For static servers, the server_key in inventory is the server name
// Check if this key exists in static_clients
self.static_clients.contains_key(server_key)
}
/// Call a tool by name with automatic type coercion
///
/// Accepts either JSON string or parsed Map as arguments.
/// Automatically converts string numbers to actual numbers based on tool schema.
pub async fn call_tool(
&self,
tool_name: &str,
args: impl Into<ToolArgs>,
) -> McpResult<CallToolResult> {
// Get tool info for schema and server
let (server_name, tool_info) = self
.inventory
.get_tool(tool_name)
.ok_or_else(|| McpError::ToolNotFound(tool_name.to_string()))?;
// Convert args with type coercion based on schema
let tool_schema = Some(serde_json::Value::Object((*tool_info.input_schema).clone()));
let args_map = args
.into()
.into_map(tool_schema.as_ref())
.map_err(McpError::InvalidArguments)?;
// Get client for that server
let client = self
.get_client(&server_name)
.await
.ok_or_else(|| McpError::ServerNotFound(server_name.clone()))?;
// Call the tool
let request = CallToolRequestParam {
name: Cow::Owned(tool_name.to_string()),
arguments: args_map,
};
client
.call_tool(request)
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to call tool: {}", e)))
}
/// Get a tool by name
pub fn get_tool(&self, tool_name: &str) -> Option<Tool> {
self.inventory
.get_tool(tool_name)
.map(|(_server_name, tool_info)| tool_info)
}
/// Get a prompt by name
pub async fn get_prompt(
&self,
prompt_name: &str,
args: Option<Map<String, serde_json::Value>>,
) -> McpResult<GetPromptResult> {
// Get server that owns this prompt
let (server_name, _prompt_info) = self
.inventory
.get_prompt(prompt_name)
.ok_or_else(|| McpError::PromptNotFound(prompt_name.to_string()))?;
// Get client for that server
let client = self
.get_client(&server_name)
.await
.ok_or_else(|| McpError::ServerNotFound(server_name.clone()))?;
// Get the prompt
let request = GetPromptRequestParam {
name: prompt_name.to_string(),
arguments: args,
};
client
.get_prompt(request)
.await
.map_err(|e| McpError::Transport(format!("Failed to get prompt: {}", e)))
}
/// List all available prompts
pub fn list_prompts(&self) -> Vec<Prompt> {
self.inventory
.list_prompts()
.into_iter()
.map(|(_prompt_name, _server_name, prompt_info)| prompt_info)
.collect()
}
/// Read a resource by URI
pub async fn read_resource(&self, uri: &str) -> McpResult<ReadResourceResult> {
// Get server that owns this resource
let (server_name, _resource_info) = self
.inventory
.get_resource(uri)
.ok_or_else(|| McpError::ResourceNotFound(uri.to_string()))?;
// Get client for that server
let client = self
.get_client(&server_name)
.await
.ok_or_else(|| McpError::ServerNotFound(server_name.clone()))?;
// Read the resource
let request = ReadResourceRequestParam {
uri: uri.to_string(),
};
client
.read_resource(request)
.await
.map_err(|e| McpError::Transport(format!("Failed to read resource: {}", e)))
}
/// List all available resources
pub fn list_resources(&self) -> Vec<RawResource> {
self.inventory
.list_resources()
.into_iter()
.map(|(_resource_uri, _server_name, resource_info)| resource_info)
.collect()
}
/// Refresh inventory for a specific server
pub async fn refresh_server_inventory(&self, server_name: &str) -> McpResult<()> {
let client = self
.get_client(server_name)
.await
.ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;
info!("Refreshing inventory for server: {}", server_name);
self.load_server_inventory_internal(server_name, &client)
.await;
Ok(())
}
/// 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,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut interval = tokio::time::interval(refresh_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
interval.tick().await;
// 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_keys.is_empty() {
debug!(
"Background refresh: Refreshing {} static server(s)",
server_keys.len()
);
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);
}
}
debug!("Background refresh: Completed refresh cycle");
}
}
})
}
/// Check if a tool exists
pub fn has_tool(&self, name: &str) -> bool {
self.inventory.has_tool(name)
}
/// Get prompt info by name
pub fn get_prompt_info(&self, name: &str) -> Option<Prompt> {
self.inventory.get_prompt(name).map(|(_server, info)| info)
}
/// Get resource info by URI
pub fn get_resource_info(&self, uri: &str) -> Option<RawResource> {
self.inventory.get_resource(uri).map(|(_server, info)| info)
}
/// Subscribe to resource changes
pub async fn subscribe_resource(&self, uri: &str) -> McpResult<()> {
let (server_name, _resource_info) = self
.inventory
.get_resource(uri)
.ok_or_else(|| McpError::ResourceNotFound(uri.to_string()))?;
let client = self
.get_client(&server_name)
.await
.ok_or_else(|| McpError::ServerNotFound(server_name.clone()))?;
debug!("Subscribing to '{}' on '{}'", uri, server_name);
client
.peer()
.subscribe(SubscribeRequestParam {
uri: uri.to_string(),
})
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to subscribe: {}", e)))
}
/// Unsubscribe from resource changes
pub async fn unsubscribe_resource(&self, uri: &str) -> McpResult<()> {
let (server_name, _resource_info) = self
.inventory
.get_resource(uri)
.ok_or_else(|| McpError::ResourceNotFound(uri.to_string()))?;
let client = self
.get_client(&server_name)
.await
.ok_or_else(|| McpError::ServerNotFound(server_name.clone()))?;
debug!("Unsubscribing from '{}' on '{}'", uri, server_name);
client
.peer()
.unsubscribe(UnsubscribeRequestParam {
uri: uri.to_string(),
})
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to unsubscribe: {}", e)))
}
/// List all connected servers (static + dynamic)
pub fn list_servers(&self) -> Vec<String> {
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) {
// 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 static server '{}': {}", name, e);
}
}
Err(_) => {
warn!(
"Could not shutdown static server '{}': client still in use",
name
);
}
}
}
}
// 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_clients.len(),
pool_stats: self.connection_pool.stats(),
tool_count: tools,
prompt_count: prompts,
resource_count: resources,
}
}
/// Get the shared tool inventory
pub fn inventory(&self) -> Arc<ToolInventory> {
Arc::clone(&self.inventory)
}
/// Get the connection pool
pub fn connection_pool(&self) -> Arc<McpConnectionPool> {
Arc::clone(&self.connection_pool)
}
// ========================================================================
// Internal Helper Methods
// ========================================================================
/// Static helper for loading inventory (for new())
/// Discover and cache tools/prompts/resources for a connected server
///
/// This method is public to allow workflow-based inventory loading.
/// 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_key: &str,
client: &Arc<McpClient>,
) {
// Tools
match client.peer().list_all_tools().await {
Ok(ts) => {
info!("Discovered {} tools from '{}'", ts.len(), server_key);
for t in ts {
inventory.insert_tool(t.name.to_string(), server_key.to_string(), t);
}
}
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_key);
for p in ps {
inventory.insert_prompt(p.name.clone(), server_key.to_string(), p);
}
}
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_key);
for r in rs {
inventory.insert_resource(r.uri.clone(), server_key.to_string(), r.raw);
}
}
Err(e) => debug!("No resources or failed to list on '{}': {}", server_key, e),
}
}
/// Discover and cache tools/prompts/resources for a connected server (internal wrapper)
async fn load_server_inventory_internal(&self, server_name: &str, client: &McpClient) {
// Tools
match client.peer().list_all_tools().await {
Ok(ts) => {
info!("Discovered {} tools from '{}'", ts.len(), server_name);
for t in ts {
self.inventory
.insert_tool(t.name.to_string(), server_name.to_string(), t);
}
}
Err(e) => warn!("Failed to list tools from '{}': {}", server_name, e),
}
// Prompts
match client.peer().list_all_prompts().await {
Ok(ps) => {
info!("Discovered {} prompts from '{}'", ps.len(), server_name);
for p in ps {
self.inventory
.insert_prompt(p.name.clone(), server_name.to_string(), p);
}
}
Err(e) => debug!("No prompts or failed to list on '{}': {}", server_name, e),
}
// Resources
match client.peer().list_all_resources().await {
Ok(rs) => {
info!("Discovered {} resources from '{}'", rs.len(), server_name);
for r in rs {
self.inventory
.insert_resource(r.uri.clone(), server_name.to_string(), r.raw);
}
}
Err(e) => debug!("No resources or failed to list on '{}': {}", server_name, e),
}
}
// ========================================================================
// Connection Logic (from client_manager.rs)
// ========================================================================
/// Connect to an MCP server
///
/// This method is public to allow workflow-based server registration at runtime.
/// It handles connection with automatic retry for network-based transports (SSE/Streamable).
pub async fn connect_server(
config: &McpServerConfig,
global_proxy: Option<&McpProxyConfig>,
) -> McpResult<McpClient> {
let needs_retry = matches!(
&config.transport,
McpTransport::Sse { .. } | McpTransport::Streamable { .. }
);
if needs_retry {
Self::connect_server_with_retry(config, global_proxy).await
} else {
Self::connect_server_impl(config, global_proxy).await
}
}
/// Connect with exponential backoff retry for remote servers
async fn connect_server_with_retry(
config: &McpServerConfig,
global_proxy: Option<&McpProxyConfig>,
) -> McpResult<McpClient> {
let backoff = ExponentialBackoffBuilder::new()
.with_initial_interval(Duration::from_secs(1))
.with_max_interval(Duration::from_secs(30))
.with_max_elapsed_time(Some(Duration::from_secs(30)))
.build();
backoff::future::retry(backoff, || async {
match Self::connect_server_impl(config, global_proxy).await {
Ok(client) => Ok(client),
Err(e) => {
if Self::is_permanent_error(&e) {
error!(
"Permanent error connecting to '{}': {} - not retrying",
config.name, e
);
Err(backoff::Error::permanent(e))
} else {
warn!("Failed to connect to '{}', retrying: {}", config.name, e);
Err(backoff::Error::transient(e))
}
}
}
})
.await
}
/// Determine if an error is permanent (should not retry) or transient
fn is_permanent_error(error: &McpError) -> bool {
match error {
McpError::Config(_) => true,
McpError::Auth(_) => true,
McpError::ServerNotFound(_) => true,
McpError::Transport(_) => true,
McpError::ConnectionFailed(msg) => {
msg.contains("initialize")
|| msg.contains("connection closed")
|| msg.contains("connection refused")
|| msg.contains("invalid URL")
|| msg.contains("not found")
}
_ => false,
}
}
/// Internal implementation of server connection (stdio/sse/streamable)
async fn connect_server_impl(
config: &McpServerConfig,
global_proxy: Option<&McpProxyConfig>,
) -> McpResult<McpClient> {
info!(
"Connecting to MCP server '{}' via {:?}",
config.name, config.transport
);
match &config.transport {
McpTransport::Stdio {
command,
args,
envs,
} => {
let transport = TokioChildProcess::new(
tokio::process::Command::new(command).configure(|cmd| {
cmd.args(args)
.envs(envs.iter())
.stderr(std::process::Stdio::inherit());
}),
)
.map_err(|e| McpError::Transport(format!("create stdio transport: {}", e)))?;
let client = ().serve(transport).await.map_err(|e| {
McpError::ConnectionFailed(format!("initialize stdio client: {}", e))
})?;
info!("Connected to stdio server '{}'", config.name);
Ok(client)
}
McpTransport::Sse { url, token } => {
// Resolve proxy configuration
let proxy_config = crate::mcp::proxy::resolve_proxy_config(config, global_proxy);
// Create HTTP client with proxy support
let client = if token.is_some() {
let mut builder =
reqwest::Client::builder().connect_timeout(Duration::from_secs(10));
// Apply proxy configuration using proxy.rs helper
if let Some(proxy_cfg) = proxy_config {
builder = crate::mcp::proxy::apply_proxy_to_builder(builder, proxy_cfg)?;
}
// Add Authorization header
builder = builder.default_headers({
let mut headers = reqwest::header::HeaderMap::new();
headers.insert(
reqwest::header::AUTHORIZATION,
format!("Bearer {}", token.as_ref().unwrap())
.parse()
.map_err(|e| McpError::Transport(format!("auth token: {}", e)))?,
);
headers
});
builder
.build()
.map_err(|e| McpError::Transport(format!("build HTTP client: {}", e)))?
} else {
crate::mcp::proxy::create_http_client(proxy_config)?
};
let cfg = SseClientConfig {
sse_endpoint: url.clone().into(),
..Default::default()
};
let transport = SseClientTransport::start_with_client(client, cfg)
.await
.map_err(|e| McpError::Transport(format!("create SSE transport: {}", e)))?;
let client = ().serve(transport).await.map_err(|e| {
McpError::ConnectionFailed(format!("initialize SSE client: {}", e))
})?;
info!("Connected to SSE server '{}' at {}", config.name, url);
Ok(client)
}
McpTransport::Streamable { url, token } => {
// Note: Streamable transport doesn't support proxy yet
let _proxy_config = crate::mcp::proxy::resolve_proxy_config(config, global_proxy);
if _proxy_config.is_some() {
warn!(
"Proxy configuration detected but not supported for Streamable transport on server '{}'",
config.name
);
}
let transport = if let Some(tok) = token {
let mut cfg = StreamableHttpClientTransportConfig::with_uri(url.as_str());
cfg.auth_header = Some(format!("Bearer {}", tok));
StreamableHttpClientTransport::from_config(cfg)
} else {
StreamableHttpClientTransport::from_uri(url.as_str())
};
let client = ().serve(transport).await.map_err(|e| {
McpError::ConnectionFailed(format!("initialize streamable client: {}", e))
})?;
info!(
"Connected to streamable HTTP server '{}' at {}",
config.name, url
);
Ok(client)
}
}
}
/// Generate a unique key for a server config based on its transport
pub fn server_key(config: &McpServerConfig) -> String {
match &config.transport {
McpTransport::Streamable { url, .. } => url.clone(),
McpTransport::Sse { url, .. } => url.clone(),
McpTransport::Stdio { command, .. } => command.clone(),
}
}
}
/// Statistics about the MCP manager
#[derive(Debug, Clone)]
pub struct McpManagerStats {
/// Number of static servers registered
pub static_server_count: usize,
/// Connection pool statistics
pub pool_stats: crate::mcp::connection_pool::PoolStats,
/// Number of cached tools
pub tool_count: usize,
/// Number of cached prompts
pub prompt_count: usize,
/// Number of cached resources
pub resource_count: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_manager_creation() {
let config = McpConfig {
servers: vec![],
pool: Default::default(),
proxy: None,
warmup: vec![],
inventory: Default::default(),
};
let manager = McpManager::new(config, 100).await.unwrap();
assert_eq!(manager.list_static_servers().len(), 0);
}
}

17
sgl-model-gateway/src/mcp/mod.rs
View File

@@ -1,17 +0,0 @@
//! 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;
pub mod error;
pub mod inventory;
pub mod manager;
pub mod oauth;
pub mod proxy;
pub mod tool_args;
// Re-export types used outside this module
pub use config::{McpConfig, McpServerConfig, McpTransport, Tool};
pub use manager::McpManager;

194
sgl-model-gateway/src/mcp/oauth.rs
View File

@@ -1,194 +0,0 @@
//! OAuth authentication for MCP servers.
//!
//! Handles OAuth flow including callback server and token exchange.
use std::{net::SocketAddr, sync::Arc};
use axum::{
extract::{Query, State},
response::Html,
routing::get,
Router,
};
use rmcp::transport::auth::OAuthState;
use serde::Deserialize;
use tokio::sync::{oneshot, Mutex};
use crate::mcp::error::{McpError, McpResult};
/// OAuth callback parameters
#[derive(Debug, Deserialize)]
struct CallbackParams {
code: String,
#[allow(dead_code)]
state: Option<String>,
}
/// State for the callback server
#[derive(Clone)]
struct CallbackState {
code_receiver: Arc<Mutex<Option<oneshot::Sender<String>>>>,
}
/// HTML page returned after successful OAuth callback
const CALLBACK_HTML: &str = r#"
<!DOCTYPE html>
<html>
<head>
<title>OAuth Success</title>
<style>
body {
font-family: Arial, sans-serif;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
margin: 0;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
}
.container {
background: white;
padding: 40px;
border-radius: 10px;
box-shadow: 0 10px 30px rgba(0,0,0,0.2);
text-align: center;
}
h1 { color: #333; }
p { color: #666; margin: 20px 0; }
.success { color: #4CAF50; font-size: 48px; }
</style>
</head>
<body>
<div class="container">
<div class="success">✓</div>
<h1>Authentication Successful!</h1>
<p>You can now close this window and return to your application.</p>
</div>
</body>
</html>
"#;
/// OAuth authentication helper for MCP servers
pub(crate) struct OAuthHelper {
server_url: String,
redirect_uri: String,
callback_port: u16,
}
impl OAuthHelper {
/// Create a new OAuth helper
pub fn new(server_url: String, redirect_uri: String, callback_port: u16) -> Self {
Self {
server_url,
redirect_uri,
callback_port,
}
}
/// Perform OAuth authentication flow
pub async fn authenticate(
&self,
scopes: &[&str],
) -> McpResult<rmcp::transport::auth::AuthorizationManager> {
// Initialize OAuth state machine
let mut oauth_state = OAuthState::new(&self.server_url, None)
.await
.map_err(|e| McpError::Auth(format!("Failed to initialize OAuth: {}", e)))?;
oauth_state
.start_authorization(scopes, &self.redirect_uri, None)
.await
.map_err(|e| McpError::Auth(format!("Failed to start authorization: {}", e)))?;
// Get authorization URL
let auth_url = oauth_state
.get_authorization_url()
.await
.map_err(|e| McpError::Auth(format!("Failed to get authorization URL: {}", e)))?;
tracing::info!("OAuth authorization URL: {}", auth_url);
// Start callback server and wait for code
let auth_code = self.start_callback_server().await?;
// Exchange code for token
oauth_state
.handle_callback(&auth_code, "")
.await
.map_err(|e| McpError::Auth(format!("Failed to handle OAuth callback: {}", e)))?;
// Get authorization manager
oauth_state
.into_authorization_manager()
.ok_or_else(|| McpError::Auth("Failed to get authorization manager".to_string()))
}
/// Start a local HTTP server to receive the OAuth callback
async fn start_callback_server(&self) -> McpResult<String> {
let (code_sender, code_receiver) = oneshot::channel::<String>();
let state = CallbackState {
code_receiver: Arc::new(Mutex::new(Some(code_sender))),
};
// Create router for callback
let app = Router::new()
.route("/callback", get(Self::callback_handler))
.with_state(state);
let addr = SocketAddr::from(([127, 0, 0, 1], self.callback_port));
// Start server in background
let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
McpError::Auth(format!(
"Failed to bind to callback port {}: {}",
self.callback_port, e
))
})?;
tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
tracing::info!(
"OAuth callback server started on port {}",
self.callback_port
);
// Wait for authorization code
code_receiver
.await
.map_err(|_| McpError::Auth("Failed to receive authorization code".to_string()))
}
/// Handle OAuth callback
async fn callback_handler(
Query(params): Query<CallbackParams>,
State(state): State<CallbackState>,
) -> Html<String> {
tracing::debug!("Received OAuth callback with code");
// Send code to waiting task
if let Some(sender) = state.code_receiver.lock().await.take() {
let _ = sender.send(params.code);
}
Html(CALLBACK_HTML.to_string())
}
}
/// Create an OAuth-authenticated client
pub async fn create_oauth_client(
server_url: String,
_sse_url: String,
redirect_uri: String,
callback_port: u16,
scopes: &[&str],
) -> McpResult<rmcp::transport::auth::AuthClient<reqwest::Client>> {
let helper = OAuthHelper::new(server_url, redirect_uri, callback_port);
let auth_manager = helper.authenticate(scopes).await?;
let client = rmcp::transport::auth::AuthClient::new(reqwest::Client::default(), auth_manager);
Ok(client)
}

254
sgl-model-gateway/src/mcp/proxy.rs
View File

@@ -1,254 +0,0 @@
//! HTTP proxy configuration for MCP connections.
//!
//! Resolves proxy settings and creates HTTP clients for MCP server connections.
use std::time::Duration;
use crate::mcp::{
config::{McpProxyConfig, McpServerConfig},
error::{McpError, McpResult},
};
/// Resolve proxy configuration for a server
/// Priority: server.proxy > global.proxy > None
///
/// # Arguments
/// * `server_config` - Server-specific configuration
/// * `global_proxy` - Global proxy configuration from McpConfig
///
/// # Returns
/// The resolved proxy configuration, or None for direct connection
pub(crate) fn resolve_proxy_config<'a>(
server_config: &'a McpServerConfig,
global_proxy: Option<&'a McpProxyConfig>,
) -> Option<&'a McpProxyConfig> {
// Priority 1: Check if server has explicit proxy config
// Note: server.proxy = Some(config) uses that config
// server.proxy = None (set explicitly in YAML as null) forces direct connection
// server.proxy not set (field missing) falls back to global
if server_config.proxy.is_some() {
server_config.proxy.as_ref()
} else {
// Priority 2: Fall back to global proxy
global_proxy
}
}
/// Apply proxy configuration to a ClientBuilder
///
/// This is a reusable helper that applies proxy settings without building the client,
/// allowing additional configuration (like auth headers) to be added afterward.
///
/// # Arguments
/// * `builder` - The reqwest::ClientBuilder to configure
/// * `proxy_config` - The proxy configuration to apply
///
/// # Returns
/// The configured builder or error
pub(super) fn apply_proxy_to_builder(
mut builder: reqwest::ClientBuilder,
proxy_cfg: &McpProxyConfig,
) -> McpResult<reqwest::ClientBuilder> {
// Configure HTTP proxy
if let Some(ref http_proxy) = proxy_cfg.http {
let mut proxy = reqwest::Proxy::http(http_proxy)
.map_err(|e| McpError::Config(format!("Invalid HTTP proxy: {}", e)))?;
// Apply no_proxy exclusions
if let Some(ref no_proxy) = proxy_cfg.no_proxy {
proxy = proxy.no_proxy(reqwest::NoProxy::from_string(no_proxy));
}
// Apply authentication if configured
if let (Some(ref username), Some(ref password)) = (&proxy_cfg.username, &proxy_cfg.password)
{
proxy = proxy.basic_auth(username, password);
}
builder = builder.proxy(proxy);
}
// Configure HTTPS proxy
if let Some(ref https_proxy) = proxy_cfg.https {
let mut proxy = reqwest::Proxy::https(https_proxy)
.map_err(|e| McpError::Config(format!("Invalid HTTPS proxy: {}", e)))?;
// Apply no_proxy exclusions
if let Some(ref no_proxy) = proxy_cfg.no_proxy {
proxy = proxy.no_proxy(reqwest::NoProxy::from_string(no_proxy));
}
// Apply authentication if configured
if let (Some(ref username), Some(ref password)) = (&proxy_cfg.username, &proxy_cfg.password)
{
proxy = proxy.basic_auth(username, password);
}
builder = builder.proxy(proxy);
}
Ok(builder)
}
/// Create HTTP client with MCP-specific proxy configuration
///
/// # Arguments
/// * `proxy_config` - Optional proxy configuration to apply
///
/// # Returns
/// A configured reqwest::Client or error
pub(crate) fn create_http_client(
proxy_config: Option<&McpProxyConfig>,
) -> McpResult<reqwest::Client> {
let mut builder = reqwest::Client::builder().connect_timeout(Duration::from_secs(10));
// Apply MCP-specific proxy if configured
if let Some(proxy_cfg) = proxy_config {
builder = apply_proxy_to_builder(builder, proxy_cfg)?;
}
builder
.build()
.map_err(|e| McpError::Transport(format!("Failed to build HTTP client: {}", e)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mcp::config::McpTransport;
#[test]
fn test_resolve_proxy_no_config() {
let server = McpServerConfig {
name: "test".to_string(),
transport: McpTransport::Sse {
url: "http://localhost:3000/sse".to_string(),
token: None,
},
proxy: None,
required: false,
};
let result = resolve_proxy_config(&server, None);
assert!(
result.is_none(),
"Should return None when no proxy configured"
);
}
#[test]
fn test_resolve_proxy_global_only() {
let server = McpServerConfig {
name: "test".to_string(),
transport: McpTransport::Sse {
url: "http://localhost:3000/sse".to_string(),
token: None,
},
proxy: None,
required: false,
};
let global = McpProxyConfig {
http: Some("http://global-proxy:8080".to_string()),
https: None,
no_proxy: None,
username: None,
password: None,
};
let result = resolve_proxy_config(&server, Some(&global));
assert!(result.is_some(), "Should use global proxy");
assert_eq!(
result.unwrap().http.as_ref().unwrap(),
"http://global-proxy:8080"
);
}
#[test]
fn test_resolve_proxy_server_override() {
let server_proxy = McpProxyConfig {
http: Some("http://server-proxy:9090".to_string()),
https: None,
no_proxy: None,
username: None,
password: None,
};
let server = McpServerConfig {
name: "test".to_string(),
transport: McpTransport::Sse {
url: "http://localhost:3000/sse".to_string(),
token: None,
},
proxy: Some(server_proxy),
required: false,
};
let global = McpProxyConfig {
http: Some("http://global-proxy:8080".to_string()),
https: None,
no_proxy: None,
username: None,
password: None,
};
let result = resolve_proxy_config(&server, Some(&global));
assert!(result.is_some(), "Should use server-specific proxy");
assert_eq!(
result.unwrap().http.as_ref().unwrap(),
"http://server-proxy:9090",
"Server proxy should override global"
);
}
#[test]
fn test_create_http_client_no_proxy() {
let client = create_http_client(None);
assert!(client.is_ok(), "Should create client without proxy");
}
#[test]
fn test_create_http_client_with_proxy() {
let proxy = McpProxyConfig {
http: Some("http://proxy.example.com:8080".to_string()),
https: None,
no_proxy: Some("localhost,127.0.0.1".to_string()),
username: None,
password: None,
};
let client = create_http_client(Some(&proxy));
assert!(client.is_ok(), "Should create client with proxy");
}
#[test]
fn test_create_http_client_with_auth() {
let proxy = McpProxyConfig {
http: Some("http://proxy.example.com:8080".to_string()),
https: None,
no_proxy: None,
username: Some("user".to_string()),
password: Some("pass".to_string()),
};
let client = create_http_client(Some(&proxy));
assert!(
client.is_ok(),
"Should create client with proxy authentication"
);
}
#[test]
fn test_create_http_client_invalid_proxy() {
let proxy = McpProxyConfig {
http: Some("://invalid".to_string()), // Invalid URL format
https: None,
no_proxy: None,
username: None,
password: None,
};
let client = create_http_client(Some(&proxy));
assert!(client.is_err(), "Should fail with invalid proxy URL");
}
}

103
sgl-model-gateway/src/mcp/tool_args.rs
View File

@@ -1,103 +0,0 @@
//! MCP tool argument handling.
//!
//! Supports both JSON strings and parsed Maps with automatic type coercion.
use serde_json::Map;
/// Tool arguments input - supports both JSON strings and parsed Maps
pub enum ToolArgs {
/// JSON string that needs parsing
JsonString(String),
/// Already parsed map
Map(Option<Map<String, serde_json::Value>>),
}
impl ToolArgs {
/// Convert to Map with type coercion based on tool schema
pub(crate) fn into_map(
self,
tool_schema: Option<&serde_json::Value>,
) -> Result<Option<Map<String, serde_json::Value>>, String> {
match self {
ToolArgs::JsonString(s) => {
if s.is_empty() || s == "{}" {
return Ok(None);
}
let mut value: serde_json::Value =
serde_json::from_str(&s).map_err(|e| format!("parse tool args: {}", e))?;
Self::coerce_types(&mut value, tool_schema)?;
let result = match value {
serde_json::Value::Object(m) => Some(m),
_ => None,
};
Ok(result)
}
ToolArgs::Map(map) => {
if let Some(m) = map {
let mut value = serde_json::Value::Object(m);
Self::coerce_types(&mut value, tool_schema)?;
let result = match value {
serde_json::Value::Object(m) => Some(m),
_ => None,
};
Ok(result)
} else {
Ok(None)
}
}
}
}
/// Coerce string numbers to actual numbers based on schema
/// LLMs often output numbers as strings, so we need to convert them
fn coerce_types(
value: &mut serde_json::Value,
tool_schema: Option<&serde_json::Value>,
) -> Result<(), String> {
let Some(params) = tool_schema else {
return Ok(());
};
let Some(props) = params.get("properties").and_then(|p| p.as_object()) else {
return Ok(());
};
let Some(args) = value.as_object_mut() else {
return Ok(());
};
for (key, val) in args.iter_mut() {
let should_be_number = props
.get(key)
.and_then(|s| s.get("type"))
.and_then(|t| t.as_str())
.is_some_and(|t| matches!(t, "number" | "integer"));
if should_be_number {
if let Some(s) = val.as_str() {
if let Ok(num) = s.parse::<f64>() {
*val = serde_json::json!(num);
}
}
}
}
Ok(())
}
}
// Implement From traits for convenient conversion
impl From<String> for ToolArgs {
fn from(s: String) -> Self {
ToolArgs::JsonString(s)
}
}
impl From<&str> for ToolArgs {
fn from(s: &str) -> Self {
ToolArgs::JsonString(s.to_string())
}
}
impl From<Option<Map<String, serde_json::Value>>> for ToolArgs {
fn from(map: Option<Map<String, serde_json::Value>>) -> Self {
ToolArgs::Map(map)
}
}