diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 83ba00a92..0aae7adcc 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -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" diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index 3ce0f5c8c..ee586a23a 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -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; diff --git a/sgl-model-gateway/src/mcp/config.rs b/sgl-model-gateway/src/mcp/config.rs deleted file mode 100644 index e5d913aa2..000000000 --- a/sgl-model-gateway/src/mcp/config.rs +++ /dev/null @@ -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, - - /// 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, - - /// Pre-warm these connections at startup - #[serde(default)] - pub warmup: Vec, - - /// 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, - - /// 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, - #[serde(default)] - envs: HashMap, - }, - Sse { - url: String, - #[serde(skip_serializing_if = "Option::is_none")] - token: Option, - }, - Streamable { - url: String, - #[serde(skip_serializing_if = "Option::is_none")] - token: Option, - }, -} - -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, - - /// HTTPS proxy URL - pub https: Option, - - /// Comma-separated hosts to exclude from proxying - /// Example: "localhost,127.0.0.1,*.internal,10.*" - pub no_proxy: Option, - - /// Custom proxy authentication (if needed) - #[serde(skip_serializing_if = "Option::is_none")] - pub username: Option, - - #[serde(skip_serializing_if = "Option::is_none")] - pub password: Option, -} - -/// 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, -} - -// 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 { - 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> { - 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 { - // 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"), - } - } -} diff --git a/sgl-model-gateway/src/mcp/connection_pool.rs b/sgl-model-gateway/src/mcp/connection_pool.rs deleted file mode 100644 index f0e4eef9f..000000000 --- a/sgl-model-gateway/src/mcp/connection_pool.rs +++ /dev/null @@ -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; - -/// Type alias for eviction callback -type EvictionCallback = Arc; - -/// Cached MCP connection with metadata -#[derive(Clone)] -pub(crate) struct CachedConnection { - /// The MCP client instance - pub client: Arc, - /// 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, 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>>, - - /// 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, - - /// Optional eviction callback (called when LRU evicts a connection) - /// Used to clean up tools from inventory - eviction_callback: Option, -} - -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) -> 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(&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( - &self, - server_url: &str, - server_config: McpServerConfig, - connect_fn: F, - ) -> McpResult> - where - F: FnOnce(McpServerConfig, Option) -> Fut, - Fut: std::future::Future>, - { - // 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 { - 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> { - 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 = - 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()); - } -} diff --git a/sgl-model-gateway/src/mcp/error.rs b/sgl-model-gateway/src/mcp/error.rs deleted file mode 100644 index d6af2d3ec..000000000 --- a/sgl-model-gateway/src/mcp/error.rs +++ /dev/null @@ -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 = Result; - -#[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), - - #[error(transparent)] - Io(#[from] std::io::Error), - - #[error(transparent)] - Http(#[from] reqwest::Error), -} diff --git a/sgl-model-gateway/src/mcp/inventory.rs b/sgl-model-gateway/src/mcp/inventory.rs deleted file mode 100644 index f4da1b986..000000000 --- a/sgl-model-gateway/src/mcp/inventory.rs +++ /dev/null @@ -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, - - /// Map of prompt_name -> cached prompt - prompts: DashMap, - - /// Map of resource_uri -> cached resource - resources: DashMap, -} - -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); - } -} diff --git a/sgl-model-gateway/src/mcp/manager.rs b/sgl-model-gateway/src/mcp/manager.rs deleted file mode 100644 index 33868e1a0..000000000 --- a/sgl-model-gateway/src/mcp/manager.rs +++ /dev/null @@ -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; - -pub struct McpManager { - static_clients: Arc>>, - inventory: Arc, - connection_pool: Arc, - _config: McpConfig, -} - -impl McpManager { - const MAX_DYNAMIC_CLIENTS: usize = 200; - - pub async fn new(config: McpConfig, pool_max_connections: usize) -> McpResult { - 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::new(config, Self::MAX_DYNAMIC_CLIENTS).await - } - - pub async fn get_client(&self, server_name: &str) -> Option> { - 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> { - 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 { - 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) { - 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 { - 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 { - 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, - ) -> McpResult { - // 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 { - 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>, - ) -> McpResult { - // 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 { - 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 { - // 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 { - 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, - 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 = 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 { - self.inventory.get_prompt(name).map(|(_server, info)| info) - } - - /// Get resource info by URI - pub fn get_resource_info(&self, uri: &str) -> Option { - 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 { - 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 = 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 { - Arc::clone(&self.inventory) - } - - /// Get the connection pool - pub fn connection_pool(&self) -> Arc { - 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, - server_key: &str, - client: &Arc, - ) { - // 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 { - 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 { - 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 { - 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); - } -} diff --git a/sgl-model-gateway/src/mcp/mod.rs b/sgl-model-gateway/src/mcp/mod.rs deleted file mode 100644 index 59b3be298..000000000 --- a/sgl-model-gateway/src/mcp/mod.rs +++ /dev/null @@ -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; diff --git a/sgl-model-gateway/src/mcp/oauth.rs b/sgl-model-gateway/src/mcp/oauth.rs deleted file mode 100644 index bc564987b..000000000 --- a/sgl-model-gateway/src/mcp/oauth.rs +++ /dev/null @@ -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, -} - -/// State for the callback server -#[derive(Clone)] -struct CallbackState { - code_receiver: Arc>>>, -} - -/// HTML page returned after successful OAuth callback -const CALLBACK_HTML: &str = r#" - - - - OAuth Success - - - -
-
✓
-

Authentication Successful!

-

You can now close this window and return to your application.

-
- - -"#; - -/// 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 { - // 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 { - let (code_sender, code_receiver) = oneshot::channel::(); - - 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, - State(state): State, - ) -> Html { - 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> { - 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) -} diff --git a/sgl-model-gateway/src/mcp/proxy.rs b/sgl-model-gateway/src/mcp/proxy.rs deleted file mode 100644 index a6a2a1a8b..000000000 --- a/sgl-model-gateway/src/mcp/proxy.rs +++ /dev/null @@ -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 { - // 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 { - 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"); - } -} diff --git a/sgl-model-gateway/src/mcp/tool_args.rs b/sgl-model-gateway/src/mcp/tool_args.rs deleted file mode 100644 index 3207b9acc..000000000 --- a/sgl-model-gateway/src/mcp/tool_args.rs +++ /dev/null @@ -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>), -} - -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>, 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::() { - *val = serde_json::json!(num); - } - } - } - } - Ok(()) - } -} - -// Implement From traits for convenient conversion -impl From 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>> for ToolArgs { - fn from(map: Option>) -> Self { - ToolArgs::Map(map) - } -}