diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index 7c1b81cfa..364a2b0c8 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -21,8 +21,8 @@ use crate::{ tokenizer::{ cache::{CacheConfig, CachedTokenizer}, factory as tokenizer_factory, + registry::TokenizerRegistry, traits::Tokenizer, - TokenizerRegistry, }, tool_parser::ParserFactory as ToolParserFactory, wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager}, diff --git a/sgl-model-gateway/src/auth/jwks.rs b/sgl-model-gateway/src/auth/jwks.rs index f7d34aedd..26611e53f 100644 --- a/sgl-model-gateway/src/auth/jwks.rs +++ b/sgl-model-gateway/src/auth/jwks.rs @@ -118,7 +118,7 @@ fn is_link_local_v6(ip: &Ipv6Addr) -> bool { /// /// For testing purposes, HTTP is allowed for localhost/127.0.0.1 only. /// In production, only HTTPS should be used. -pub fn validate_url(url_str: &str) -> Result { +pub(crate) fn validate_url(url_str: &str) -> Result { let url = Url::parse(url_str) .map_err(|e| JwksError::InvalidUrl(format!("Failed to parse URL: {}", e)))?; @@ -210,7 +210,7 @@ impl CachedJwks { } /// JWKS provider with caching and automatic refresh. -pub struct JwksProvider { +pub(crate) struct JwksProvider { /// HTTP client for fetching JWKS client: reqwest::Client, /// JWKS endpoint URL (validated) @@ -304,6 +304,7 @@ impl JwksProvider { } /// Get the JWKS URI. + #[allow(dead_code)] pub fn jwks_uri(&self) -> &str { &self.jwks_uri } @@ -411,6 +412,7 @@ impl JwksProvider { } /// Force refresh the JWKS cache. + #[allow(dead_code)] pub async fn refresh(&self) -> Result<(), JwksError> { let jwks = self.fetch_jwks().await?; let mut cache = self.cache.write(); diff --git a/sgl-model-gateway/src/auth/jwt.rs b/sgl-model-gateway/src/auth/jwt.rs index d43d0352a..edb1b51d6 100644 --- a/sgl-model-gateway/src/auth/jwt.rs +++ b/sgl-model-gateway/src/auth/jwt.rs @@ -75,7 +75,7 @@ pub enum JwtValidatorError { /// Standard JWT claims we extract. #[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StandardClaims { +pub(crate) struct StandardClaims { /// Subject (user ID) pub sub: Option, @@ -115,7 +115,7 @@ pub struct StandardClaims { /// Audience claim can be a single string or an array. #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(untagged)] -pub enum Audience { +pub(crate) enum Audience { Single(String), Multiple(Vec), #[default] @@ -123,6 +123,7 @@ pub enum Audience { } impl Audience { + #[allow(dead_code)] pub fn contains(&self, aud: &str) -> bool { match self { Audience::Single(s) => s == aud, @@ -151,7 +152,8 @@ pub struct ValidatedToken { pub name: Option, /// Full claims for additional processing - pub claims: StandardClaims, + #[allow(dead_code)] + pub(crate) claims: StandardClaims, } /// JTI (JWT ID) cache entry with expiration tracking. @@ -181,12 +183,13 @@ pub struct JwtValidator { impl JwtValidator { /// Create a new JWT validator with explicit JWKS URI. - pub fn new(config: JwtConfig, jwks_provider: Arc) -> Self { + #[allow(dead_code)] + pub(crate) fn new(config: JwtConfig, jwks_provider: Arc) -> Self { Self::new_with_options(config, jwks_provider, false) } /// Create a new JWT validator with optional JTI replay protection. - pub fn new_with_options( + pub(crate) fn new_with_options( config: JwtConfig, jwks_provider: Arc, enable_jti_check: bool, @@ -500,7 +503,8 @@ impl JwtValidator { } /// Get a reference to the JWKS provider. - pub fn jwks_provider(&self) -> &Arc { + #[allow(dead_code)] + pub(crate) fn jwks_provider(&self) -> &Arc { &self.jwks_provider } diff --git a/sgl-model-gateway/src/config/mod.rs b/sgl-model-gateway/src/config/mod.rs index 018683f66..4ed69c064 100644 --- a/sgl-model-gateway/src/config/mod.rs +++ b/sgl-model-gateway/src/config/mod.rs @@ -1,10 +1,9 @@ pub mod builder; pub mod types; -pub mod validation; +pub(crate) mod validation; pub use builder::*; pub use types::*; -pub use validation::*; #[derive(Debug, thiserror::Error)] pub enum ConfigError { diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 36eee6209..c85ebe13e 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -1,10 +1,10 @@ use super::*; /// Configuration validator -pub struct ConfigValidator; +pub(crate) struct ConfigValidator; impl ConfigValidator { - pub fn validate(config: &RouterConfig) -> ConfigResult<()> { + pub(crate) fn validate(config: &RouterConfig) -> ConfigResult<()> { Self::validate_mode(&config.mode)?; Self::validate_policy(&config.policy)?; Self::validate_server_settings(config)?; diff --git a/sgl-model-gateway/src/core/metrics_aggregator.rs b/sgl-model-gateway/src/core/metrics_aggregator.rs index f39c96a52..cf65b6cda 100644 --- a/sgl-model-gateway/src/core/metrics_aggregator.rs +++ b/sgl-model-gateway/src/core/metrics_aggregator.rs @@ -76,7 +76,7 @@ fn merge_family(a: PrometheusFamily, b: PrometheusFamily) -> anyhow::Result(iterable: I, f: F) -> Result, E> +fn try_reduce(iterable: I, f: F) -> Result, E> where I: IntoIterator, F: FnMut(T, T) -> Result, diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs index cbf6cda9f..587ec44cd 100644 --- a/sgl-model-gateway/src/core/mod.rs +++ b/sgl-model-gateway/src/core/mod.rs @@ -27,19 +27,17 @@ pub mod worker_manager; pub mod worker_registry; pub mod worker_service; -pub use circuit_breaker::{ - CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState, -}; +// Re-export commonly used types for convenience +pub use circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; pub use error::{WorkerError, WorkerResult}; pub use job_queue::{Job, JobQueue, JobQueueConfig}; pub use model_card::{ModelCard, ProviderType}; -pub use model_type::{Endpoint, ModelType}; -pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor}; +pub use retry::{is_retryable_status, RetryExecutor}; pub use worker::{ - attach_guards_to_response, worker_to_info, BasicWorker, ConnectionMode, DPAwareWorker, - HealthChecker, HealthConfig, RuntimeType, Worker, WorkerFactory, WorkerLoadGuard, WorkerType, + attach_guards_to_response, BasicWorker, ConnectionMode, HealthConfig, RuntimeType, Worker, + WorkerLoadGuard, WorkerType, }; pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder}; pub use worker_manager::{LoadMonitor, WorkerManager}; -pub use worker_registry::{HashRing, WorkerId, WorkerRegistry, WorkerRegistryStats}; -pub use worker_service::{WorkerService, WorkerServiceError}; +pub use worker_registry::{HashRing, WorkerRegistry}; +pub use worker_service::WorkerService; diff --git a/sgl-model-gateway/src/core/model_card.rs b/sgl-model-gateway/src/core/model_card.rs index 90fa6b9e4..e0c89351c 100644 --- a/sgl-model-gateway/src/core/model_card.rs +++ b/sgl-model-gateway/src/core/model_card.rs @@ -95,7 +95,7 @@ impl std::fmt::Display for ProviderType { /// # Example /// /// ``` -/// use smg::core::{ModelCard, ModelType, ProviderType}; +/// use smg::core::{model_type::ModelType, ModelCard, ProviderType}; /// /// let card = ModelCard::new("meta-llama/Llama-3.1-8B-Instruct") /// .with_display_name("Llama 3.1 8B Instruct") diff --git a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs index 8ab6285d8..015784198 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs @@ -8,8 +8,10 @@ use tracing::{debug, info}; use crate::{ app_context::AppContext, core::{ - model_card::ModelCard, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, - HealthConfig, RuntimeType, Worker, WorkerType, + circuit_breaker::CircuitBreakerConfig, + model_card::ModelCard, + worker::{HealthConfig, RuntimeType, WorkerType}, + BasicWorkerBuilder, ConnectionMode, Worker, }, protocols::worker_spec::WorkerConfigRequest, workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, diff --git a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs index 67614a808..ee1dcbc2c 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs @@ -9,8 +9,10 @@ use super::discover_dp::DpInfo; use crate::{ app_context::AppContext, core::{ - BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder, - HealthConfig, ModelCard, RuntimeType, Worker, WorkerType, UNKNOWN_MODEL_ID, + circuit_breaker::CircuitBreakerConfig, + model_card::ModelCard, + worker::{HealthConfig, RuntimeType, WorkerType}, + BasicWorkerBuilder, ConnectionMode, DPAwareWorkerBuilder, Worker, UNKNOWN_MODEL_ID, }, protocols::worker_spec::WorkerConfigRequest, workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, diff --git a/sgl-model-gateway/src/core/worker.rs b/sgl-model-gateway/src/core/worker.rs index 3a1544ab0..036e27597 100644 --- a/sgl-model-gateway/src/core/worker.rs +++ b/sgl-model-gateway/src/core/worker.rs @@ -4,7 +4,7 @@ use std::{ atomic::{AtomicBool, AtomicUsize, Ordering}, Arc, LazyLock, RwLock as StdRwLock, }, - time::{Duration, Instant}, + time::Duration, }; use async_trait::async_trait; @@ -13,11 +13,11 @@ use serde::{Deserialize, Serialize}; use tokio::{sync::OnceCell, time}; use super::{ - CircuitBreaker, Endpoint, ModelCard, ModelType, ProviderType, WorkerError, WorkerResult, - UNKNOWN_MODEL_ID, + model_card::{ModelCard, ProviderType}, + model_type::{Endpoint, ModelType}, + CircuitBreaker, WorkerError, WorkerResult, UNKNOWN_MODEL_ID, }; use crate::{ - core::{BasicWorkerBuilder, DPAwareWorkerBuilder}, observability::metrics::{metrics_labels, Metrics}, protocols::worker_spec::WorkerInfo, routers::grpc::client::GrpcClient, @@ -1000,94 +1000,6 @@ impl Worker for DPAwareWorker { } } -/// Worker factory for creating workers of different types -pub struct WorkerFactory; - -impl WorkerFactory { - /// Create a DP-aware worker of specified type - pub fn create_dp_aware( - base_url: String, - dp_rank: usize, - dp_size: usize, - worker_type: WorkerType, - api_key: Option, - ) -> Box { - let mut builder = - DPAwareWorkerBuilder::new(base_url, dp_rank, dp_size).worker_type(worker_type); - if let Some(api_key) = api_key { - builder = builder.api_key(api_key); - } - Box::new(builder.build()) - } - - /// Static health validation before creating a worker - /// This replaces wait_for_worker_health in handlers - pub async fn validate_health(url: &str, timeout_secs: u64) -> WorkerResult<()> { - let start_time = Instant::now(); - let timeout = Duration::from_secs(timeout_secs); - - loop { - if start_time.elapsed() > timeout { - return Err(WorkerError::HealthCheckFailed { - url: url.to_string(), - reason: format!( - "Timeout {}s waiting for worker to become healthy", - timeout_secs - ), - }); - } - - // Note: This static function doesn't have access to worker's API key - // API key authentication is handled in the worker instance's check_health_async method - match WORKER_CLIENT - .get(format!("{}/health", url)) - .timeout(Duration::from_secs(5)) - .send() - .await - { - Ok(res) if res.status().is_success() => { - tracing::info!("Worker {} is healthy", url); - return Ok(()); - } - Ok(res) => { - tracing::warn!( - "Worker {} health check failed with status: {}", - url, - res.status() - ); - } - Err(e) => { - tracing::warn!("Failed to contact worker {}: {}", url, e); - } - } - - time::sleep(Duration::from_secs(1)).await; - } - } -} - -/// Convert a list of worker URLs to worker trait objects -pub fn urls_to_workers(urls: Vec, api_key: Option) -> Vec> { - urls.into_iter() - .map(|url| { - let worker_builder = BasicWorkerBuilder::new(url).worker_type(WorkerType::Regular); - - let worker = if let Some(ref api_key) = api_key { - worker_builder.api_key(api_key.clone()).build() - } else { - worker_builder.build() - }; - - Box::new(worker) as Box - }) - .collect() -} - -/// Convert worker trait objects back to URLs -pub fn workers_to_urls(workers: &[Box]) -> Vec { - workers.iter().map(|w| w.url().to_string()).collect() -} - /// RAII guard for worker load management /// /// Automatically decrements worker load when dropped. Can be attached to @@ -1204,7 +1116,8 @@ impl http_body::Body for MultiGuardedBody { } /// Health checker handle with graceful shutdown -pub struct HealthChecker { +pub(crate) struct HealthChecker { + #[allow(dead_code)] handle: tokio::task::JoinHandle<()>, shutdown: Arc, } @@ -1224,6 +1137,7 @@ impl HealthChecker { } /// Shutdown the health checker gracefully + #[allow(dead_code)] pub async fn shutdown(self) { self.shutdown.store(true, Ordering::Release); let _ = self.handle.await; @@ -1280,7 +1194,10 @@ mod tests { use std::{thread, time::Duration}; use super::*; - use crate::core::{CircuitBreakerConfig, CircuitState}; + use crate::core::{ + circuit_breaker::{CircuitBreakerConfig, CircuitState}, + DPAwareWorkerBuilder, + }; #[test] fn test_worker_type_display() { @@ -1585,6 +1502,7 @@ mod tests { #[test] fn test_create_regular_worker() { + use crate::core::BasicWorkerBuilder; let worker: Box = Box::new( BasicWorkerBuilder::new("http://regular:8080") .worker_type(WorkerType::Regular) @@ -1596,6 +1514,7 @@ mod tests { #[test] fn test_create_prefill_worker() { + use crate::core::BasicWorkerBuilder; let worker1: Box = Box::new( BasicWorkerBuilder::new("http://prefill:8080") .worker_type(WorkerType::Prefill { @@ -1628,6 +1547,7 @@ mod tests { #[test] fn test_create_decode_worker() { + use crate::core::BasicWorkerBuilder; let worker: Box = Box::new( BasicWorkerBuilder::new("http://decode:8080") .worker_type(WorkerType::Decode) @@ -1637,36 +1557,6 @@ mod tests { assert_eq!(worker.worker_type(), &WorkerType::Decode); } - #[test] - fn test_urls_to_workers() { - let urls = vec!["http://w1:8080".to_string(), "http://w2:8080".to_string()]; - - let workers = urls_to_workers(urls, Some("test_api_key".to_string())); - assert_eq!(workers.len(), 2); - assert_eq!(workers[0].url(), "http://w1:8080"); - assert_eq!(workers[1].url(), "http://w2:8080"); - assert_eq!(workers[0].worker_type(), &WorkerType::Regular); - } - - #[test] - fn test_workers_to_urls() { - let workers: Vec> = vec![ - Box::new( - BasicWorkerBuilder::new("http://w1:8080") - .worker_type(WorkerType::Regular) - .build(), - ), - Box::new( - BasicWorkerBuilder::new("http://w2:8080") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - let urls = workers_to_urls(&workers); - assert_eq!(urls, vec!["http://w1:8080", "http://w2:8080"]); - } - #[tokio::test] async fn test_check_health_async() { use crate::core::BasicWorkerBuilder; @@ -1819,45 +1709,6 @@ mod tests { assert_eq!(dp_worker.processed_requests(), 1); } - #[tokio::test] - async fn test_factory_create_dp_aware() { - let worker = WorkerFactory::create_dp_aware( - "http://worker1:8080".to_string(), - 1, - 4, - WorkerType::Regular, - Some("test_api_key".to_string()), - ); - - assert_eq!(worker.url(), "http://worker1:8080@1"); - assert!(worker.is_dp_aware()); - assert_eq!(worker.dp_rank(), Some(1)); - assert_eq!(worker.dp_size(), Some(4)); - assert_eq!(worker.worker_type(), &WorkerType::Regular); - } - - #[tokio::test] - async fn test_factory_create_dp_aware_prefill() { - let worker = WorkerFactory::create_dp_aware( - "http://worker1:8080".to_string(), - 0, - 2, - WorkerType::Prefill { - bootstrap_port: Some(8090), - }, - Some("test_api_key".to_string()), - ); - - assert_eq!(worker.url(), "http://worker1:8080@0"); - assert!(worker.is_dp_aware()); - assert_eq!( - worker.worker_type(), - &WorkerType::Prefill { - bootstrap_port: Some(8090) - } - ); - } - #[test] fn test_worker_circuit_breaker() { use crate::core::BasicWorkerBuilder; @@ -1929,6 +1780,7 @@ mod tests { #[tokio::test] async fn test_mixed_worker_types() { + use crate::core::BasicWorkerBuilder; let regular: Box = Box::new( BasicWorkerBuilder::new("http://regular:8080") .worker_type(WorkerType::Regular) @@ -1946,28 +1798,25 @@ mod tests { .worker_type(WorkerType::Decode) .build(), ); - let dp_aware_regular = WorkerFactory::create_dp_aware( - "http://dp:8080".to_string(), - 0, - 2, - WorkerType::Regular, - Some("test_api_key".to_string()), + let dp_aware_regular: Box = Box::new( + DPAwareWorkerBuilder::new("http://dp:8080", 0, 2) + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), ); - let dp_aware_prefill = WorkerFactory::create_dp_aware( - "http://dp-prefill:8080".to_string(), - 1, - 2, - WorkerType::Prefill { - bootstrap_port: None, - }, - Some("test_api_key".to_string()), + let dp_aware_prefill: Box = Box::new( + DPAwareWorkerBuilder::new("http://dp-prefill:8080", 1, 2) + .worker_type(WorkerType::Prefill { + bootstrap_port: None, + }) + .api_key("test_api_key") + .build(), ); - let dp_aware_decode = WorkerFactory::create_dp_aware( - "http://dp-decode:8080".to_string(), - 0, - 4, - WorkerType::Decode, - Some("test_api_key".to_string()), + let dp_aware_decode: Box = Box::new( + DPAwareWorkerBuilder::new("http://dp-decode:8080", 0, 4) + .worker_type(WorkerType::Decode) + .api_key("test_api_key") + .build(), ); let workers: Vec> = vec![ diff --git a/sgl-model-gateway/src/core/worker_registry.rs b/sgl-model-gateway/src/core/worker_registry.rs index 14b86eb0d..71747ae28 100644 --- a/sgl-model-gateway/src/core/worker_registry.rs +++ b/sgl-model-gateway/src/core/worker_registry.rs @@ -17,7 +17,11 @@ use dashmap::DashMap; use uuid::Uuid; use crate::{ - core::{CircuitState, ConnectionMode, RuntimeType, Worker, WorkerType}, + core::{ + circuit_breaker::CircuitState, + worker::{HealthChecker, RuntimeType, WorkerType}, + ConnectionMode, Worker, + }, observability::metrics::Metrics, }; @@ -588,7 +592,7 @@ impl WorkerRegistry { /// Start a health checker for all workers in the registry /// This should be called once after the registry is populated with workers - pub fn start_health_checker(&self, check_interval_secs: u64) -> crate::core::HealthChecker { + pub(crate) fn start_health_checker(&self, check_interval_secs: u64) -> HealthChecker { use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, @@ -632,7 +636,7 @@ impl WorkerRegistry { } }); - crate::core::HealthChecker::new(handle, shutdown) + HealthChecker::new(handle, shutdown) } } @@ -676,7 +680,7 @@ mod tests { use std::collections::HashMap; use super::*; - use crate::core::{BasicWorkerBuilder, CircuitBreakerConfig}; + use crate::core::{circuit_breaker::CircuitBreakerConfig, BasicWorkerBuilder}; #[test] fn test_worker_registry() { @@ -697,7 +701,7 @@ mod tests { .build(), ); - // Register worker (WorkerFactory returns Box, convert to Arc) + // Register worker let worker_id = registry.register(Arc::from(worker)); assert!(registry.get(&worker_id).is_some()); diff --git a/sgl-model-gateway/src/core/worker_service.rs b/sgl-model-gateway/src/core/worker_service.rs index 40f7a83a6..6cac685d7 100644 --- a/sgl-model-gateway/src/core/worker_service.rs +++ b/sgl-model-gateway/src/core/worker_service.rs @@ -16,7 +16,7 @@ use tracing::warn; use crate::{ config::RouterConfig, - core::{worker_to_info, Job, JobQueue, WorkerId, WorkerRegistry}, + core::{worker::worker_to_info, worker_registry::WorkerId, Job, JobQueue, WorkerRegistry}, protocols::worker_spec::{ WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest, }, diff --git a/sgl-model-gateway/src/mcp/connection_pool.rs b/sgl-model-gateway/src/mcp/connection_pool.rs index cd2066d91..f0e4eef9f 100644 --- a/sgl-model-gateway/src/mcp/connection_pool.rs +++ b/sgl-model-gateway/src/mcp/connection_pool.rs @@ -20,10 +20,11 @@ type EvictionCallback = Arc; /// Cached MCP connection with metadata #[derive(Clone)] -pub struct CachedConnection { +pub(crate) struct CachedConnection { /// The MCP client instance pub client: Arc, /// Server configuration used to create this connection + #[allow(dead_code)] pub config: McpServerConfig, } @@ -220,7 +221,7 @@ pub struct PoolStats { #[cfg(test)] mod tests { use super::*; - use crate::mcp::McpTransport; + use crate::mcp::config::McpTransport; // Helper to create test server config fn create_test_config(url: &str) -> McpServerConfig { @@ -277,7 +278,7 @@ mod tests { #[test] fn test_pool_with_global_proxy() { - use crate::mcp::McpProxyConfig; + use crate::mcp::config::McpProxyConfig; // Create proxy config let proxy = McpProxyConfig { diff --git a/sgl-model-gateway/src/mcp/inventory.rs b/sgl-model-gateway/src/mcp/inventory.rs index e03e75515..f4da1b986 100644 --- a/sgl-model-gateway/src/mcp/inventory.rs +++ b/sgl-model-gateway/src/mcp/inventory.rs @@ -8,21 +8,21 @@ use crate::mcp::config::{Prompt, RawResource, Tool}; /// Cached tool with metadata #[derive(Clone)] -pub struct CachedTool { +pub(crate) struct CachedTool { pub server_name: String, pub tool: Tool, } /// Cached prompt with metadata #[derive(Clone)] -pub struct CachedPrompt { +pub(crate) struct CachedPrompt { pub server_name: String, pub prompt: Prompt, } /// Cached resource with metadata #[derive(Clone)] -pub struct CachedResource { +pub(crate) struct CachedResource { pub server_name: String, pub resource: RawResource, } diff --git a/sgl-model-gateway/src/mcp/mod.rs b/sgl-model-gateway/src/mcp/mod.rs index 3ba98e340..59b3be298 100644 --- a/sgl-model-gateway/src/mcp/mod.rs +++ b/sgl-model-gateway/src/mcp/mod.rs @@ -12,14 +12,6 @@ pub mod oauth; pub mod proxy; pub mod tool_args; -// Re-export the main types for convenience -pub use config::{ - InventoryConfig, McpConfig, McpPoolConfig, McpProxyConfig, McpServerConfig, McpTransport, - Prompt, RawResource, Tool, WarmupServer, -}; -pub use connection_pool::{CachedConnection, McpConnectionPool, PoolStats}; -pub use error::{McpError, McpResult}; -pub use inventory::ToolInventory; -pub use manager::{McpManager, McpManagerStats}; -pub use proxy::{create_http_client, resolve_proxy_config}; -pub use tool_args::ToolArgs; +// 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 index d8eaac4c7..bc564987b 100644 --- a/sgl-model-gateway/src/mcp/oauth.rs +++ b/sgl-model-gateway/src/mcp/oauth.rs @@ -69,7 +69,7 @@ const CALLBACK_HTML: &str = r#" "#; /// OAuth authentication helper for MCP servers -pub struct OAuthHelper { +pub(crate) struct OAuthHelper { server_url: String, redirect_uri: String, callback_port: u16, diff --git a/sgl-model-gateway/src/mcp/proxy.rs b/sgl-model-gateway/src/mcp/proxy.rs index daf9b8673..a6a2a1a8b 100644 --- a/sgl-model-gateway/src/mcp/proxy.rs +++ b/sgl-model-gateway/src/mcp/proxy.rs @@ -4,7 +4,10 @@ use std::time::Duration; -use crate::mcp::{McpError, McpProxyConfig, McpResult, McpServerConfig}; +use crate::mcp::{ + config::{McpProxyConfig, McpServerConfig}, + error::{McpError, McpResult}, +}; /// Resolve proxy configuration for a server /// Priority: server.proxy > global.proxy > None @@ -15,7 +18,7 @@ use crate::mcp::{McpError, McpProxyConfig, McpResult, McpServerConfig}; /// /// # Returns /// The resolved proxy configuration, or None for direct connection -pub fn resolve_proxy_config<'a>( +pub(crate) fn resolve_proxy_config<'a>( server_config: &'a McpServerConfig, global_proxy: Option<&'a McpProxyConfig>, ) -> Option<&'a McpProxyConfig> { @@ -42,7 +45,7 @@ pub fn resolve_proxy_config<'a>( /// /// # Returns /// The configured builder or error -pub fn apply_proxy_to_builder( +pub(super) fn apply_proxy_to_builder( mut builder: reqwest::ClientBuilder, proxy_cfg: &McpProxyConfig, ) -> McpResult { @@ -94,7 +97,9 @@ pub fn apply_proxy_to_builder( /// /// # Returns /// A configured reqwest::Client or error -pub fn create_http_client(proxy_config: Option<&McpProxyConfig>) -> McpResult { +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 @@ -110,7 +115,7 @@ pub fn create_http_client(proxy_config: Option<&McpProxyConfig>) -> McpResult>> = Lazy::new(DashMap::new /// This function is designed for high-throughput scenarios where the same /// strings (model IDs, worker URLs) appear repeatedly. The first call allocates, /// subsequent calls just clone the Arc (very cheap - just a ref count increment). -pub fn intern_string(s: &str) -> Arc { +pub(crate) fn intern_string(s: &str) -> Arc { // Fast path: check if already interned if let Some(entry) = STRING_INTERNER.get(s) { return Arc::clone(entry.value()); @@ -46,7 +46,8 @@ pub fn intern_string(s: &str) -> Arc { .clone() } -pub fn interner_size() -> usize { +#[allow(dead_code)] +pub(crate) fn interner_size() -> usize { STRING_INTERNER.len() } @@ -91,7 +92,7 @@ pub fn status_code_to_static_str(code: u16) -> Option<&'static str> { } /// Static HTTP method strings to avoid allocations on every request. -pub mod http_methods { +pub(crate) mod http_methods { pub const GET: &str = "GET"; pub const POST: &str = "POST"; pub const PUT: &str = "PUT"; @@ -144,7 +145,7 @@ impl Default for PrometheusConfig { } } -pub fn init_metrics() { +pub(crate) fn init_metrics() { // Layer 1: HTTP metrics describe_counter!( "smg_http_requests_total", diff --git a/sgl-model-gateway/src/observability/otel_trace.rs b/sgl-model-gateway/src/observability/otel_trace.rs index 76542c6c0..99ed6b004 100644 --- a/sgl-model-gateway/src/observability/otel_trace.rs +++ b/sgl-model-gateway/src/observability/otel_trace.rs @@ -52,7 +52,7 @@ fn get_allowed_targets() -> &'static [&'static str; 3] { /// Filter that only allows specific module targets to be exported to OTEL. #[derive(Clone, Copy, Default)] -pub struct CustomOtelFilter; +pub(crate) struct CustomOtelFilter; impl CustomOtelFilter { #[inline] diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 1b0a7e556..a5491628a 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -23,13 +23,16 @@ use crate::{ app_context::AppContext, config::{RouterConfig, RoutingMode}, core::{ + job_queue::{JobQueue, JobQueueConfig}, steps::{ create_external_worker_registration_workflow, create_mcp_registration_workflow, create_tokenizer_registration_workflow, create_wasm_module_registration_workflow, create_wasm_module_removal_workflow, create_worker_registration_workflow, create_worker_removal_workflow, create_worker_update_workflow, }, - Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType, + worker::WorkerType, + worker_manager::WorkerManager, + Job, }, middleware::{self, AuthConfig, QueuedRequest}, observability::{ diff --git a/sgl-model-gateway/src/tokenizer/mod.rs b/sgl-model-gateway/src/tokenizer/mod.rs index 6050b48bd..3e03562a0 100644 --- a/sgl-model-gateway/src/tokenizer/mod.rs +++ b/sgl-model-gateway/src/tokenizer/mod.rs @@ -23,20 +23,14 @@ pub mod tiktoken; #[cfg(test)] mod tests; -// Re-exports -pub use cache::{CacheConfig, CacheStats, CachedTokenizer, TokenizerFingerprint}; -pub use factory::{ - create_tokenizer, create_tokenizer_async, create_tokenizer_async_with_chat_template, - create_tokenizer_from_file, create_tokenizer_with_chat_template, - create_tokenizer_with_chat_template_blocking, TokenizerType, -}; +// Internal imports for Tokenizer struct +use factory::{create_tokenizer_from_file, create_tokenizer_with_chat_template}; +// Re-export types used outside this module pub use huggingface::HuggingFaceTokenizer; pub use registry::TokenizerRegistry; -pub use sequence::Sequence; -pub use stop::{SequenceDecoderOutput, StopSequenceConfig, StopSequenceDecoder}; +pub use stop::StopSequenceDecoder; pub use stream::DecodeStream; -pub use tiktoken::{TiktokenModel, TiktokenTokenizer}; -pub use traits::{Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait}; +pub use traits::{Decoder, Encoder, Encoding, SpecialTokens}; /// Main tokenizer wrapper that provides a unified interface for different tokenizer implementations #[derive(Clone)] diff --git a/sgl-model-gateway/src/tokenizer/tiktoken.rs b/sgl-model-gateway/src/tokenizer/tiktoken.rs index 615410c9f..00929c709 100644 --- a/sgl-model-gateway/src/tokenizer/tiktoken.rs +++ b/sgl-model-gateway/src/tokenizer/tiktoken.rs @@ -6,7 +6,7 @@ use super::traits::{ }; /// Tiktoken tokenizer wrapper for OpenAI GPT models -pub struct TiktokenTokenizer { +pub(crate) struct TiktokenTokenizer { tokenizer: CoreBPE, #[allow(dead_code)] model: TiktokenModel, diff --git a/sgl-model-gateway/src/tool_parser/mod.rs b/sgl-model-gateway/src/tool_parser/mod.rs index 7cfd1ffc4..227d7a6a5 100644 --- a/sgl-model-gateway/src/tool_parser/mod.rs +++ b/sgl-model-gateway/src/tool_parser/mod.rs @@ -5,7 +5,6 @@ pub mod errors; pub mod factory; pub mod partial_json; -pub mod state; pub mod traits; pub mod types; @@ -15,13 +14,11 @@ pub mod parsers; #[cfg(test)] mod tests; -// Re-export commonly used types -pub use errors::{ParserError, ParserResult}; -pub use factory::{ParserFactory, ParserRegistry, PooledParser}; -// Re-export parsers for convenience +// Re-export types used outside this module +pub use factory::{ParserFactory, PooledParser}; pub use parsers::{ DeepSeekParser, Glm4MoeParser, JsonParser, KimiK2Parser, LlamaParser, MinimaxM2Parser, MistralParser, PythonicParser, QwenParser, Step3Parser, }; -pub use traits::{PartialJsonParser, ToolParser}; +pub use traits::ToolParser; pub use types::{FunctionCall, PartialToolCall, StreamingParseResult, ToolCall}; diff --git a/sgl-model-gateway/src/tool_parser/parsers/helpers.rs b/sgl-model-gateway/src/tool_parser/parsers/helpers.rs index aa39fa8f5..6c073794f 100644 --- a/sgl-model-gateway/src/tool_parser/parsers/helpers.rs +++ b/sgl-model-gateway/src/tool_parser/parsers/helpers.rs @@ -186,7 +186,7 @@ pub fn normalize_arguments_field(mut obj: Value) -> Value { /// - `Ok(StreamingParseResult)` with any tool call items to stream /// - `Err(ParserError)` if JSON parsing or serialization fails #[allow(clippy::too_many_arguments)] -pub fn handle_json_tool_streaming( +pub(crate) fn handle_json_tool_streaming( current_text: &str, start_idx: usize, partial_json: &mut crate::tool_parser::partial_json::PartialJson, diff --git a/sgl-model-gateway/src/tool_parser/parsers/mod.rs b/sgl-model-gateway/src/tool_parser/parsers/mod.rs index 5b28d0b4b..4e8f16dad 100644 --- a/sgl-model-gateway/src/tool_parser/parsers/mod.rs +++ b/sgl-model-gateway/src/tool_parser/parsers/mod.rs @@ -26,7 +26,7 @@ pub use kimik2::KimiK2Parser; pub use llama::LlamaParser; pub use minimax_m2::MinimaxM2Parser; pub use mistral::MistralParser; -pub use passthrough::PassthroughParser; +pub(crate) use passthrough::PassthroughParser; pub use pythonic::PythonicParser; pub use qwen::QwenParser; pub use step3::Step3Parser; diff --git a/sgl-model-gateway/src/tool_parser/parsers/passthrough.rs b/sgl-model-gateway/src/tool_parser/parsers/passthrough.rs index 11170f9d3..c6b438e9a 100644 --- a/sgl-model-gateway/src/tool_parser/parsers/passthrough.rs +++ b/sgl-model-gateway/src/tool_parser/parsers/passthrough.rs @@ -17,7 +17,7 @@ use crate::{ /// Passthrough parser that returns text unchanged with no tool calls #[derive(Default)] -pub struct PassthroughParser; +pub(crate) struct PassthroughParser; impl PassthroughParser { pub fn new() -> Self { diff --git a/sgl-model-gateway/src/tool_parser/partial_json.rs b/sgl-model-gateway/src/tool_parser/partial_json.rs index 0764572e8..ee1f90144 100644 --- a/sgl-model-gateway/src/tool_parser/partial_json.rs +++ b/sgl-model-gateway/src/tool_parser/partial_json.rs @@ -531,23 +531,3 @@ impl<'a> Parser<'a> { } } } - -/// Utility function to check if a string contains complete JSON -pub fn is_complete_json(input: &str) -> bool { - serde_json::from_str::(input).is_ok() -} - -/// Utility function to find common prefix between two strings -pub fn find_common_prefix(s1: &str, s2: &str) -> usize { - s1.chars() - .zip(s2.chars()) - .take_while(|(a, b)| a == b) - .count() -} - -/// Utility function to compute diff between old and new strings -pub fn compute_diff(old: &str, new: &str) -> String { - let common_len = find_common_prefix(old, new); - // Convert character count to byte offset - new.chars().skip(common_len).collect() -} diff --git a/sgl-model-gateway/src/tool_parser/state.rs b/sgl-model-gateway/src/tool_parser/state.rs deleted file mode 100644 index 9345ccc04..000000000 --- a/sgl-model-gateway/src/tool_parser/state.rs +++ /dev/null @@ -1,16 +0,0 @@ -/// Placeholder for Harmony streaming metadata captured during token-aware parsing. -#[derive(Debug, Clone, Default)] -pub struct HarmonyStreamState { - /// All tokens observed so far for the current assistant response. - pub tokens: Vec, - /// Number of tokens that have already been processed by the Harmony parser. - pub processed_tokens: usize, - /// Number of tool calls emitted downstream. - pub emitted_calls: usize, - /// Pending analysis-channel content awaiting flush into normal text output. - pub analysis_buffer: String, - /// Whether the tool name has been surfaced for the current call. - pub emitted_name: bool, - /// Whether arguments have been surfaced for the current call. - pub emitted_args: bool, -} diff --git a/sgl-model-gateway/src/tool_parser/tests.rs b/sgl-model-gateway/src/tool_parser/tests.rs index bf43f9553..49d271723 100644 --- a/sgl-model-gateway/src/tool_parser/tests.rs +++ b/sgl-model-gateway/src/tool_parser/tests.rs @@ -1,9 +1,5 @@ use super::*; -use crate::tool_parser::{ - parsers::JsonParser, - partial_json::{compute_diff, find_common_prefix, is_complete_json, PartialJson}, - traits::ToolParser, -}; +use crate::tool_parser::{parsers::JsonParser, partial_json::PartialJson, traits::ToolParser}; #[tokio::test] async fn test_tool_parser_factory() { @@ -97,37 +93,6 @@ fn test_partial_json_depth_limit() { assert!(result.is_err()); } -#[test] -fn test_is_complete_json() { - assert!(is_complete_json(r#"{"name": "test"}"#)); - assert!(is_complete_json(r#"[1, 2, 3]"#)); - assert!(is_complete_json(r#""string""#)); - assert!(is_complete_json("42")); - assert!(is_complete_json("true")); - assert!(is_complete_json("null")); - - assert!(!is_complete_json(r#"{"name": "#)); - assert!(!is_complete_json(r#"[1, 2, "#)); - assert!(!is_complete_json(r#""unclosed"#)); -} - -#[test] -fn test_find_common_prefix() { - assert_eq!(find_common_prefix("hello", "hello"), 5); - assert_eq!(find_common_prefix("hello", "help"), 3); - assert_eq!(find_common_prefix("hello", "world"), 0); - assert_eq!(find_common_prefix("", "hello"), 0); - assert_eq!(find_common_prefix("hello", ""), 0); -} - -#[test] -fn test_compute_diff() { - assert_eq!(compute_diff("hello", "hello world"), " world"); - assert_eq!(compute_diff("", "hello"), "hello"); - assert_eq!(compute_diff("hello", "hello"), ""); - assert_eq!(compute_diff("test", "hello"), "hello"); -} - // NOTE: test_stream_result_variants removed - StreamResult enum replaced by StreamingParseResult #[test] diff --git a/sgl-model-gateway/src/tool_parser/types.rs b/sgl-model-gateway/src/tool_parser/types.rs index 8157a44e2..e716da88a 100644 --- a/sgl-model-gateway/src/tool_parser/types.rs +++ b/sgl-model-gateway/src/tool_parser/types.rs @@ -16,42 +16,6 @@ pub struct FunctionCall { pub arguments: String, } -/// Streaming parse result -#[derive(Debug, Clone)] -pub enum StreamResult { - /// Need more data to continue parsing - Incomplete, - /// Found a tool name (for streaming) - ToolName { index: usize, name: String }, - /// Found incremental arguments (for streaming) - ToolArguments { index: usize, arguments: String }, - /// Completed parsing a tool - ToolComplete(ToolCall), - /// Normal text (not part of tool call) - NormalText(String), -} - -/// Token configuration for parsing -#[derive(Debug, Clone)] -pub struct TokenConfig { - /// Start tokens for tool calls - pub start_tokens: Vec, - /// End tokens for tool calls - pub end_tokens: Vec, - /// Separator between multiple tool calls - pub separator: String, -} - -impl TokenConfig { - /// Iterate over start/end token pairs - pub fn iter_pairs(&self) -> impl Iterator { - self.start_tokens - .iter() - .zip(self.end_tokens.iter()) - .map(|(s, e)| (s.as_str(), e.as_str())) - } -} - /// Simple partial tool call for streaming #[derive(Debug, Clone)] pub struct PartialToolCall { diff --git a/sgl-model-gateway/tests/mcp_test.rs b/sgl-model-gateway/tests/mcp_test.rs index 4c7e50ca9..661639ef6 100644 --- a/sgl-model-gateway/tests/mcp_test.rs +++ b/sgl-model-gateway/tests/mcp_test.rs @@ -13,7 +13,7 @@ use std::collections::HashMap; use common::mock_mcp_server::MockMCPServer; use serde_json::json; -use smg::mcp::{McpConfig, McpError, McpManager, McpServerConfig, McpTransport}; +use smg::mcp::{error::McpError, McpConfig, McpManager, McpServerConfig, McpTransport}; /// Create a new mock server for testing (each test gets its own) async fn create_mock_server() -> MockMCPServer { diff --git a/sgl-model-gateway/tests/routing/test_openai_routing.rs b/sgl-model-gateway/tests/routing/test_openai_routing.rs index df2f05ed6..5ad6325bc 100644 --- a/sgl-model-gateway/tests/routing/test_openai_routing.rs +++ b/sgl-model-gateway/tests/routing/test_openai_routing.rs @@ -18,9 +18,7 @@ use axum::{ }; use serde_json::json; use smg::{ - config::{ - ConfigError, ConfigValidator, HistoryBackend, OracleConfig, RouterConfig, RoutingMode, - }, + config::{ConfigError, HistoryBackend, OracleConfig, RouterConfig, RoutingMode}, data_connector::{ResponseId, StoredResponse}, protocols::{ chat::{ChatCompletionRequest, ChatMessage, MessageContent}, @@ -857,8 +855,9 @@ fn oracle_config_validation_requires_config_when_enabled() { .history_backend(HistoryBackend::Oracle) .build_unchecked(); - let err = - ConfigValidator::validate(&config).expect_err("config should fail without oracle details"); + let err = config + .validate() + .expect_err("config should fail without oracle details"); match err { ConfigError::MissingRequired { field } => { @@ -883,7 +882,7 @@ fn oracle_config_validation_accepts_dsn_only() { }) .build_unchecked(); - ConfigValidator::validate(&config).expect("dsn-based config should validate"); + config.validate().expect("dsn-based config should validate"); } #[test] @@ -901,5 +900,7 @@ fn oracle_config_validation_accepts_wallet_alias() { }) .build_unchecked(); - ConfigValidator::validate(&config).expect("wallet-based config should validate"); + config + .validate() + .expect("wallet-based config should validate"); }