chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/src/app_context.rs

@@ -0,0 +1,532 @@
+use std::{
+    sync::{Arc, OnceLock},
+    time::Duration,
+};
+
+use data_connector::{
+    create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage,
+    StorageFactoryConfig,
+};
+use reqwest::Client;
+use smg_mcp::McpManager;
+use tracing::debug;
+
+use crate::{
+    config::RouterConfig,
+    core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService},
+    middleware::TokenBucket,
+    observability::inflight_tracker::InFlightRequestTracker,
+    policies::PolicyRegistry,
+    reasoning_parser::ParserFactory as ReasoningParserFactory,
+    routers::router_manager::RouterManager,
+    tokenizer::registry::TokenizerRegistry,
+    tool_parser::ParserFactory as ToolParserFactory,
+    wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
+};
+
+/// Error type for AppContext builder
+#[derive(Debug)]
+pub struct AppContextBuildError(&'static str);
+
+impl std::fmt::Display for AppContextBuildError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Missing required field: {}", self.0)
+    }
+}
+
+impl std::error::Error for AppContextBuildError {}
+
+#[derive(Clone)]
+pub struct AppContext {
+    pub client: Client,
+    pub router_config: RouterConfig,
+    pub rate_limiter: Option<Arc<TokenBucket>>,
+    pub tokenizer_registry: Arc<TokenizerRegistry>,
+    pub reasoning_parser_factory: Option<ReasoningParserFactory>,
+    pub tool_parser_factory: Option<ToolParserFactory>,
+    pub worker_registry: Arc<WorkerRegistry>,
+    pub policy_registry: Arc<PolicyRegistry>,
+    pub router_manager: Option<Arc<RouterManager>>,
+    pub response_storage: Arc<dyn ResponseStorage>,
+    pub conversation_storage: Arc<dyn ConversationStorage>,
+    pub conversation_item_storage: Arc<dyn ConversationItemStorage>,
+    pub load_monitor: Option<Arc<LoadMonitor>>,
+    pub configured_reasoning_parser: Option<String>,
+    pub configured_tool_parser: Option<String>,
+    pub worker_job_queue: Arc<OnceLock<Arc<JobQueue>>>,
+    pub workflow_engines: Arc<OnceLock<WorkflowEngines>>,
+    pub mcp_manager: Arc<OnceLock<Arc<McpManager>>>,
+    pub wasm_manager: Option<Arc<WasmModuleManager>>,
+    pub worker_service: Arc<WorkerService>,
+    pub inflight_tracker: Arc<InFlightRequestTracker>,
+}
+
+impl std::fmt::Debug for AppContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AppContext")
+            .field("router_config", &self.router_config)
+            .finish_non_exhaustive()
+    }
+}
+
+pub struct AppContextBuilder {
+    client: Option<Client>,
+    router_config: Option<RouterConfig>,
+    rate_limiter: Option<Arc<TokenBucket>>,
+    tokenizer_registry: Option<Arc<TokenizerRegistry>>,
+    reasoning_parser_factory: Option<ReasoningParserFactory>,
+    tool_parser_factory: Option<ToolParserFactory>,
+    worker_registry: Option<Arc<WorkerRegistry>>,
+    policy_registry: Option<Arc<PolicyRegistry>>,
+    router_manager: Option<Arc<RouterManager>>,
+    response_storage: Option<Arc<dyn ResponseStorage>>,
+    conversation_storage: Option<Arc<dyn ConversationStorage>>,
+    conversation_item_storage: Option<Arc<dyn ConversationItemStorage>>,
+    load_monitor: Option<Arc<LoadMonitor>>,
+    worker_job_queue: Option<Arc<OnceLock<Arc<JobQueue>>>>,
+    workflow_engines: Option<Arc<OnceLock<WorkflowEngines>>>,
+    mcp_manager: Option<Arc<OnceLock<Arc<McpManager>>>>,
+    wasm_manager: Option<Arc<WasmModuleManager>>,
+}
+
+impl AppContext {
+    pub fn builder() -> AppContextBuilder {
+        AppContextBuilder::new()
+    }
+
+    /// Create AppContext from config with all components initialized
+    /// This is the main entry point that replaces ~194 lines of initialization in server.rs
+    pub async fn from_config(
+        router_config: RouterConfig,
+        request_timeout_secs: u64,
+    ) -> Result<Self, String> {
+        AppContextBuilder::from_config(router_config, request_timeout_secs)
+            .await?
+            .build()
+            .map_err(|e| e.to_string())
+    }
+}
+
+impl AppContextBuilder {
+    pub fn new() -> Self {
+        Self {
+            client: None,
+            router_config: None,
+            rate_limiter: None,
+            tokenizer_registry: None,
+            reasoning_parser_factory: None,
+            tool_parser_factory: None,
+            worker_registry: None,
+            policy_registry: None,
+            router_manager: None,
+            response_storage: None,
+            conversation_storage: None,
+            conversation_item_storage: None,
+            load_monitor: None,
+            worker_job_queue: None,
+            workflow_engines: None,
+            mcp_manager: None,
+            wasm_manager: None,
+        }
+    }
+
+    pub fn client(mut self, client: Client) -> Self {
+        self.client = Some(client);
+        self
+    }
+
+    pub fn router_config(mut self, router_config: RouterConfig) -> Self {
+        self.router_config = Some(router_config);
+        self
+    }
+
+    pub fn rate_limiter(mut self, rate_limiter: Option<Arc<TokenBucket>>) -> Self {
+        self.rate_limiter = rate_limiter;
+        self
+    }
+
+    pub fn tokenizer_registry(mut self, tokenizer_registry: Arc<TokenizerRegistry>) -> Self {
+        self.tokenizer_registry = Some(tokenizer_registry);
+        self
+    }
+
+    pub fn reasoning_parser_factory(
+        mut self,
+        reasoning_parser_factory: Option<ReasoningParserFactory>,
+    ) -> Self {
+        self.reasoning_parser_factory = reasoning_parser_factory;
+        self
+    }
+
+    pub fn tool_parser_factory(mut self, tool_parser_factory: Option<ToolParserFactory>) -> Self {
+        self.tool_parser_factory = tool_parser_factory;
+        self
+    }
+
+    pub fn worker_registry(mut self, worker_registry: Arc<WorkerRegistry>) -> Self {
+        self.worker_registry = Some(worker_registry);
+        self
+    }
+
+    pub fn policy_registry(mut self, policy_registry: Arc<PolicyRegistry>) -> Self {
+        self.policy_registry = Some(policy_registry);
+        self
+    }
+
+    pub fn router_manager(mut self, router_manager: Option<Arc<RouterManager>>) -> Self {
+        self.router_manager = router_manager;
+        self
+    }
+
+    pub fn response_storage(mut self, response_storage: Arc<dyn ResponseStorage>) -> Self {
+        self.response_storage = Some(response_storage);
+        self
+    }
+
+    pub fn conversation_storage(
+        mut self,
+        conversation_storage: Arc<dyn ConversationStorage>,
+    ) -> Self {
+        self.conversation_storage = Some(conversation_storage);
+        self
+    }
+
+    pub fn conversation_item_storage(
+        mut self,
+        conversation_item_storage: Arc<dyn ConversationItemStorage>,
+    ) -> Self {
+        self.conversation_item_storage = Some(conversation_item_storage);
+        self
+    }
+
+    pub fn load_monitor(mut self, load_monitor: Option<Arc<LoadMonitor>>) -> Self {
+        self.load_monitor = load_monitor;
+        self
+    }
+
+    pub fn worker_job_queue(mut self, worker_job_queue: Arc<OnceLock<Arc<JobQueue>>>) -> Self {
+        self.worker_job_queue = Some(worker_job_queue);
+        self
+    }
+
+    pub fn workflow_engines(mut self, workflow_engines: Arc<OnceLock<WorkflowEngines>>) -> Self {
+        self.workflow_engines = Some(workflow_engines);
+        self
+    }
+
+    pub fn mcp_manager(mut self, mcp_manager: Arc<OnceLock<Arc<McpManager>>>) -> Self {
+        self.mcp_manager = Some(mcp_manager);
+        self
+    }
+
+    pub fn wasm_manager(mut self, wasm_manager: Option<Arc<WasmModuleManager>>) -> Self {
+        self.wasm_manager = wasm_manager;
+        self
+    }
+
+    pub fn build(self) -> Result<AppContext, AppContextBuildError> {
+        let router_config = self
+            .router_config
+            .ok_or(AppContextBuildError("router_config"))?;
+        let configured_reasoning_parser = router_config.reasoning_parser.clone();
+        let configured_tool_parser = router_config.tool_call_parser.clone();
+
+        let worker_registry = self
+            .worker_registry
+            .ok_or(AppContextBuildError("worker_registry"))?;
+        let worker_job_queue = self
+            .worker_job_queue
+            .ok_or(AppContextBuildError("worker_job_queue"))?;
+
+        // Create WorkerService from the already-built components
+        let worker_service = Arc::new(WorkerService::new(
+            worker_registry.clone(),
+            worker_job_queue.clone(),
+            router_config.clone(),
+        ));
+
+        Ok(AppContext {
+            client: self.client.ok_or(AppContextBuildError("client"))?,
+            router_config,
+            rate_limiter: self.rate_limiter,
+            tokenizer_registry: self
+                .tokenizer_registry
+                .ok_or(AppContextBuildError("tokenizer_registry"))?,
+            reasoning_parser_factory: self.reasoning_parser_factory,
+            tool_parser_factory: self.tool_parser_factory,
+            worker_registry,
+            policy_registry: self
+                .policy_registry
+                .ok_or(AppContextBuildError("policy_registry"))?,
+            router_manager: self.router_manager,
+            response_storage: self
+                .response_storage
+                .ok_or(AppContextBuildError("response_storage"))?,
+            conversation_storage: self
+                .conversation_storage
+                .ok_or(AppContextBuildError("conversation_storage"))?,
+            conversation_item_storage: self
+                .conversation_item_storage
+                .ok_or(AppContextBuildError("conversation_item_storage"))?,
+            load_monitor: self.load_monitor,
+            configured_reasoning_parser,
+            configured_tool_parser,
+            worker_job_queue,
+            workflow_engines: self
+                .workflow_engines
+                .ok_or(AppContextBuildError("workflow_engines"))?,
+            mcp_manager: self
+                .mcp_manager
+                .ok_or(AppContextBuildError("mcp_manager"))?,
+            wasm_manager: self.wasm_manager,
+            worker_service,
+            inflight_tracker: InFlightRequestTracker::new(),
+        })
+    }
+
+    /// Initialize AppContext from config - creates ALL components
+    /// This replaces ~194 lines of initialization logic from server.rs
+    pub async fn from_config(
+        router_config: RouterConfig,
+        request_timeout_secs: u64,
+    ) -> Result<Self, String> {
+        Ok(Self::new()
+            .with_client(&router_config, request_timeout_secs)?
+            .maybe_rate_limiter(&router_config)
+            .with_tokenizer_registry(&router_config)?
+            .with_reasoning_parser_factory()
+            .with_tool_parser_factory()
+            .with_worker_registry()
+            .with_policy_registry(&router_config)
+            .with_storage(&router_config)?
+            .with_load_monitor(&router_config)
+            .with_worker_job_queue()
+            .with_workflow_engines()
+            .with_mcp_manager(&router_config)
+            .await?
+            .with_wasm_manager(&router_config)?
+            .router_config(router_config))
+    }
+
+    /// Create HTTP client with TLS/mTLS configuration
+    fn with_client(mut self, config: &RouterConfig, timeout_secs: u64) -> Result<Self, String> {
+        // FIXME: Current implementation creates a single HTTP client for all workers.
+        // This works well for single security domain deployments where all workers share
+        // the same CA and can accept the same client certificate.
+        //
+        // For multi-domain deployments (e.g., different model families with different CAs),
+        // this architecture needs significant refactoring:
+        // 1. Move client creation into worker registration workflow (per-worker clients)
+        // 2. Store client per worker in WorkerRegistry
+        // 3. Update PDRouter and other routers to fetch client from worker
+        // 4. Add per-worker TLS spec in WorkerConfigRequest
+        //
+        // Current single-domain approach is sufficient for most deployments.
+        //
+        // Use rustls TLS backend when TLS/mTLS is configured (client cert or CA certs provided).
+        // This ensures proper PKCS#8 key format support. For plain HTTP workers, use default
+        // backend to avoid unnecessary TLS initialization overhead.
+        let has_tls_config = config.client_identity.is_some() || !config.ca_certificates.is_empty();
+
+        let mut client_builder = Client::builder()
+            .pool_idle_timeout(Some(Duration::from_secs(50)))
+            .pool_max_idle_per_host(500)
+            .timeout(Duration::from_secs(timeout_secs))
+            .connect_timeout(Duration::from_secs(10))
+            .tcp_nodelay(true)
+            .tcp_keepalive(Some(Duration::from_secs(30)));
+
+        // Force rustls backend when TLS is configured
+        if has_tls_config {
+            client_builder = client_builder.use_rustls_tls();
+            debug!("Using rustls TLS backend for TLS/mTLS connections");
+        }
+
+        // Configure mTLS client identity if provided (certificates already loaded during config creation)
+        if let Some(identity_pem) = &config.client_identity {
+            let identity = reqwest::Identity::from_pem(identity_pem)
+                .map_err(|e| format!("Failed to create client identity: {}", e))?;
+            client_builder = client_builder.identity(identity);
+            debug!("mTLS client authentication enabled");
+        }
+
+        // Add CA certificates for verifying worker TLS (certificates already loaded during config creation)
+        for ca_cert in &config.ca_certificates {
+            let cert = reqwest::Certificate::from_pem(ca_cert)
+                .map_err(|e| format!("Failed to add CA certificate: {}", e))?;
+            client_builder = client_builder.add_root_certificate(cert);
+        }
+        if !config.ca_certificates.is_empty() {
+            debug!(
+                "Added {} CA certificate(s) for worker verification",
+                config.ca_certificates.len()
+            );
+        }
+
+        let client = client_builder
+            .build()
+            .map_err(|e| format!("Failed to create HTTP client: {}", e))?;
+
+        self.client = Some(client);
+        Ok(self)
+    }
+
+    /// Create rate limiter based on config
+    fn maybe_rate_limiter(mut self, config: &RouterConfig) -> Self {
+        self.rate_limiter = match config.max_concurrent_requests {
+            n if n <= 0 => None,
+            n => {
+                let rate_limit_tokens = config
+                    .rate_limit_tokens_per_second
+                    .filter(|&t| t > 0)
+                    .unwrap_or(n);
+                Some(Arc::new(TokenBucket::new(
+                    n as usize,
+                    rate_limit_tokens as usize,
+                )))
+            }
+        };
+        self
+    }
+
+    /// Create reasoning parser factory for gRPC mode or IGW mode
+    fn with_reasoning_parser_factory(mut self) -> Self {
+        // Initialize reasoning parser factory
+        self.reasoning_parser_factory = Some(ReasoningParserFactory::new());
+        self
+    }
+
+    /// Create tool parser factory for gRPC mode or IGW mode
+    fn with_tool_parser_factory(mut self) -> Self {
+        // Initialize tool parser factory
+        self.tool_parser_factory = Some(ToolParserFactory::new());
+        self
+    }
+
+    /// Create empty tokenizer registry
+    ///
+    /// Tokenizers are loaded via the tokenizer_registration workflow, which is triggered:
+    /// - At startup (if --tokenizer-path or --model-path is provided)
+    /// - When workers connect (registers under model_id)
+    /// - Via POST /v1/tokenizers API (registers under user-specified name)
+    ///
+    /// This unified approach ensures consistent behavior (caching, validation) across all paths.
+    fn with_tokenizer_registry(mut self, _config: &RouterConfig) -> Result<Self, String> {
+        self.tokenizer_registry = Some(Arc::new(TokenizerRegistry::new()));
+        Ok(self)
+    }
+
+    /// Create worker registry
+    fn with_worker_registry(mut self) -> Self {
+        self.worker_registry = Some(Arc::new(WorkerRegistry::new()));
+        self
+    }
+
+    /// Create policy registry
+    fn with_policy_registry(mut self, config: &RouterConfig) -> Self {
+        self.policy_registry = Some(Arc::new(PolicyRegistry::new(config.policy.clone())));
+        self
+    }
+
+    /// Create all storage backends using the factory function
+    fn with_storage(mut self, config: &RouterConfig) -> Result<Self, String> {
+        let storage_config = StorageFactoryConfig {
+            backend: &config.history_backend,
+            oracle: config.oracle.as_ref(),
+            postgres: config.postgres.as_ref(),
+            redis: config.redis.as_ref(),
+        };
+        let (response_storage, conversation_storage, conversation_item_storage) =
+            create_storage(storage_config)?;
+
+        self.response_storage = Some(response_storage);
+        self.conversation_storage = Some(conversation_storage);
+        self.conversation_item_storage = Some(conversation_item_storage);
+
+        Ok(self)
+    }
+
+    /// Create load monitor
+    fn with_load_monitor(mut self, config: &RouterConfig) -> Self {
+        let client = self
+            .client
+            .as_ref()
+            .expect("client must be set before load monitor");
+        self.load_monitor = Some(Arc::new(LoadMonitor::new(
+            self.worker_registry
+                .as_ref()
+                .expect("worker_registry must be set")
+                .clone(),
+            self.policy_registry
+                .as_ref()
+                .expect("policy_registry must be set")
+                .clone(),
+            client.clone(),
+            config.worker_startup_check_interval_secs,
+        )));
+        self
+    }
+
+    /// Create worker job queue OnceLock container
+    fn with_worker_job_queue(mut self) -> Self {
+        self.worker_job_queue = Some(Arc::new(OnceLock::new()));
+        self
+    }
+
+    /// Create workflow engines OnceLock container
+    fn with_workflow_engines(mut self) -> Self {
+        self.workflow_engines = Some(Arc::new(OnceLock::new()));
+        self
+    }
+
+    /// Create and initialize MCP manager with empty config
+    ///
+    /// This initializes the MCP manager with an empty config and default settings.
+    /// MCP servers will be registered later via the InitializeMcpServers job.
+    async fn with_mcp_manager(mut self, _router_config: &RouterConfig) -> Result<Self, String> {
+        // Create OnceLock container
+        let mcp_manager_lock = Arc::new(OnceLock::new());
+
+        // Always create with empty config and defaults
+        debug!("Initializing MCP manager with empty config and default settings (5 min TTL, 100 max connections)");
+
+        let empty_config = smg_mcp::McpConfig {
+            servers: Vec::new(),
+            pool: Default::default(),
+            proxy: None,
+            warmup: Vec::new(),
+            inventory: Default::default(),
+        };
+
+        let manager = McpManager::with_defaults(empty_config)
+            .await
+            .map_err(|e| format!("Failed to initialize MCP manager with defaults: {}", e))?;
+
+        // Store the initialized manager in the OnceLock
+        mcp_manager_lock
+            .set(Arc::new(manager))
+            .map_err(|_| "Failed to set MCP manager in OnceLock".to_string())?;
+
+        self.mcp_manager = Some(mcp_manager_lock);
+        Ok(self)
+    }
+
+    /// Create wasm manager if enabled in config
+    fn with_wasm_manager(mut self, config: &RouterConfig) -> Result<Self, String> {
+        self.wasm_manager = if config.enable_wasm {
+            Some(Arc::new(
+                WasmModuleManager::new(WasmRuntimeConfig::default())
+                    .map_err(|e| format!("Failed to initialize WASM module manager: {}", e))?,
+            ))
+        } else {
+            None
+        };
+        Ok(self)
+    }
+}
+
+impl Default for AppContextBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}

third_party/sglang/sgl-model-gateway/src/config/builder.rs

@@ -0,0 +1,846 @@
+use smg_mcp::McpConfig;
+
+use super::{
+    CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig,
+    HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RedisConfig,
+    RetryConfig, RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig,
+};
+use crate::core::ConnectionMode;
+
+/// Builder for RouterConfig that wraps the config itself
+/// This eliminates field duplication and stays in sync automatically
+#[derive(Debug, Clone, Default)]
+pub struct RouterConfigBuilder {
+    config: RouterConfig,
+    // Temporary fields for certificate paths (read during build)
+    client_cert_path: Option<String>,
+    client_key_path: Option<String>,
+    ca_cert_paths: Vec<String>,
+    server_cert_path: Option<String>,
+    server_key_path: Option<String>,
+    mcp_config_path: Option<String>,
+}
+
+impl RouterConfigBuilder {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Takes ownership
+    pub fn from_config(config: RouterConfig) -> Self {
+        Self {
+            config,
+            client_cert_path: None,
+            client_key_path: None,
+            ca_cert_paths: Vec::new(),
+            server_cert_path: None,
+            server_key_path: None,
+            mcp_config_path: None,
+        }
+    }
+
+    pub fn from_config_ref(config: &RouterConfig) -> Self {
+        Self::from_config(config.clone())
+    }
+
+    // ==================== Routing Mode ====================
+
+    pub fn regular_mode(mut self, worker_urls: Vec<String>) -> Self {
+        self.config.mode = RoutingMode::Regular { worker_urls };
+        self
+    }
+
+    pub fn prefill_decode_mode(
+        mut self,
+        prefill_urls: Vec<(String, Option<u16>)>,
+        decode_urls: Vec<String>,
+    ) -> Self {
+        self.config.mode = RoutingMode::PrefillDecode {
+            prefill_urls,
+            decode_urls,
+            prefill_policy: None,
+            decode_policy: None,
+        };
+        self
+    }
+
+    /// With separate policies
+    pub fn prefill_decode_mode_with_policies(
+        mut self,
+        prefill_urls: Vec<(String, Option<u16>)>,
+        decode_urls: Vec<String>,
+        prefill_policy: Option<PolicyConfig>,
+        decode_policy: Option<PolicyConfig>,
+    ) -> Self {
+        self.config.mode = RoutingMode::PrefillDecode {
+            prefill_urls,
+            decode_urls,
+            prefill_policy,
+            decode_policy,
+        };
+        self
+    }
+
+    pub fn openai_mode(mut self, worker_urls: Vec<String>) -> Self {
+        self.config.mode = RoutingMode::OpenAI { worker_urls };
+        self
+    }
+
+    pub fn mode(mut self, mode: RoutingMode) -> Self {
+        self.config.mode = mode;
+        self
+    }
+
+    // ==================== Policy ====================
+
+    pub fn policy(mut self, policy: PolicyConfig) -> Self {
+        self.config.policy = policy;
+        self
+    }
+
+    pub fn random_policy(mut self) -> Self {
+        self.config.policy = PolicyConfig::Random;
+        self
+    }
+
+    pub fn round_robin_policy(mut self) -> Self {
+        self.config.policy = PolicyConfig::RoundRobin;
+        self
+    }
+
+    pub fn cache_aware_policy(
+        mut self,
+        cache_threshold: f32,
+        balance_abs_threshold: usize,
+        balance_rel_threshold: f32,
+        eviction_interval_secs: u64,
+        max_tree_size: usize,
+    ) -> Self {
+        self.config.policy = PolicyConfig::CacheAware {
+            cache_threshold,
+            balance_abs_threshold,
+            balance_rel_threshold,
+            eviction_interval_secs,
+            max_tree_size,
+        };
+        self
+    }
+
+    pub fn power_of_two_policy(mut self, load_check_interval_secs: u64) -> Self {
+        self.config.policy = PolicyConfig::PowerOfTwo {
+            load_check_interval_secs,
+        };
+        self
+    }
+
+    // ==================== Connection ====================
+
+    pub fn connection_mode(mut self, mode: ConnectionMode) -> Self {
+        self.config.connection_mode = mode;
+        self
+    }
+
+    pub fn http_connection(mut self) -> Self {
+        self.config.connection_mode = ConnectionMode::Http;
+        self
+    }
+
+    pub fn grpc_connection(mut self, port: Option<u16>) -> Self {
+        self.config.connection_mode = ConnectionMode::Grpc { port };
+        self
+    }
+
+    pub fn grpc_connection_default(mut self) -> Self {
+        self.config.connection_mode = ConnectionMode::Grpc { port: None };
+        self
+    }
+
+    pub fn host<S: Into<String>>(mut self, host: S) -> Self {
+        self.config.host = host.into();
+        self
+    }
+
+    pub fn port(mut self, port: u16) -> Self {
+        self.config.port = port;
+        self
+    }
+
+    // ==================== Request ====================
+
+    pub fn max_payload_size(mut self, size: usize) -> Self {
+        self.config.max_payload_size = size;
+        self
+    }
+
+    pub fn request_timeout_secs(mut self, timeout: u64) -> Self {
+        self.config.request_timeout_secs = timeout;
+        self
+    }
+
+    pub fn worker_startup_timeout_secs(mut self, timeout: u64) -> Self {
+        self.config.worker_startup_timeout_secs = timeout;
+        self
+    }
+
+    pub fn worker_startup_check_interval_secs(mut self, interval: u64) -> Self {
+        self.config.worker_startup_check_interval_secs = interval;
+        self
+    }
+
+    // ==================== Rate Limiting ====================
+
+    pub fn max_concurrent_requests(mut self, max: i32) -> Self {
+        self.config.max_concurrent_requests = max;
+        self
+    }
+
+    pub fn disable_rate_limiting(mut self) -> Self {
+        self.config.max_concurrent_requests = -1;
+        self
+    }
+
+    pub fn queue_size(mut self, size: usize) -> Self {
+        self.config.queue_size = size;
+        self
+    }
+
+    pub fn queue_timeout_secs(mut self, timeout: u64) -> Self {
+        self.config.queue_timeout_secs = timeout;
+        self
+    }
+
+    pub fn rate_limit_tokens_per_second(mut self, tokens: i32) -> Self {
+        self.config.rate_limit_tokens_per_second = Some(tokens);
+        self
+    }
+
+    // ==================== Security & CORS ====================
+
+    pub fn api_key<S: Into<String>>(mut self, key: S) -> Self {
+        self.config.api_key = Some(key.into());
+        self
+    }
+
+    pub fn cors_allowed_origins(mut self, origins: Vec<String>) -> Self {
+        self.config.cors_allowed_origins = origins;
+        self
+    }
+
+    pub fn add_cors_origin<S: Into<String>>(mut self, origin: S) -> Self {
+        self.config.cors_allowed_origins.push(origin.into());
+        self
+    }
+
+    // ==================== Retry ====================
+
+    pub fn retry_config(mut self, retry: RetryConfig) -> Self {
+        self.config.retry = retry;
+        self
+    }
+
+    pub fn disable_retries(mut self) -> Self {
+        self.config.disable_retries = true;
+        self
+    }
+
+    pub fn enable_retries(mut self) -> Self {
+        self.config.disable_retries = false;
+        self
+    }
+
+    // ==================== Circuit Breaker ====================
+
+    pub fn circuit_breaker_config(mut self, circuit_breaker: CircuitBreakerConfig) -> Self {
+        self.config.circuit_breaker = circuit_breaker;
+        self
+    }
+
+    pub fn disable_circuit_breaker(mut self) -> Self {
+        self.config.disable_circuit_breaker = true;
+        self
+    }
+
+    pub fn enable_circuit_breaker(mut self) -> Self {
+        self.config.disable_circuit_breaker = false;
+        self
+    }
+
+    // ==================== Health Check ====================
+
+    pub fn health_check_config(mut self, health_check: HealthCheckConfig) -> Self {
+        self.config.health_check = health_check;
+        self
+    }
+
+    // ==================== Discovery ====================
+
+    pub fn discovery_config(mut self, discovery: DiscoveryConfig) -> Self {
+        self.config.discovery = Some(discovery);
+        self
+    }
+
+    /// With default settings
+    pub fn enable_discovery(mut self) -> Self {
+        self.config.discovery = Some(DiscoveryConfig {
+            enabled: true,
+            ..Default::default()
+        });
+        self
+    }
+
+    // ==================== Metrics ====================
+
+    pub fn metrics_config(mut self, metrics: MetricsConfig) -> Self {
+        self.config.metrics = Some(metrics);
+        self
+    }
+
+    pub fn enable_metrics<S: Into<String>>(mut self, host: S, port: u16) -> Self {
+        self.config.metrics = Some(MetricsConfig {
+            host: host.into(),
+            port,
+        });
+        self
+    }
+
+    // ===================== Otel Trace ====================
+
+    pub fn enable_trace<S: Into<String>>(mut self, endpoint: S) -> Self {
+        self.config.trace_config = Some(TraceConfig {
+            enable_trace: true,
+            otlp_traces_endpoint: endpoint.into(),
+        });
+        self
+    }
+
+    pub fn disable_trace(mut self) -> Self {
+        self.config.trace_config = Some(TraceConfig {
+            enable_trace: false,
+            otlp_traces_endpoint: "".to_string(),
+        });
+        self
+    }
+
+    // ==================== Logging ====================
+
+    pub fn log_dir<S: Into<String>>(mut self, dir: S) -> Self {
+        self.config.log_dir = Some(dir.into());
+        self
+    }
+
+    pub fn log_level<S: Into<String>>(mut self, level: S) -> Self {
+        self.config.log_level = Some(level.into());
+        self
+    }
+
+    pub fn request_id_headers(mut self, headers: Vec<String>) -> Self {
+        self.config.request_id_headers = Some(headers);
+        self
+    }
+
+    // ==================== IGW Mode ====================
+
+    pub fn enable_igw(mut self) -> Self {
+        self.config.enable_igw = true;
+        self
+    }
+
+    /// Use proxy mode
+    pub fn disable_igw(mut self) -> Self {
+        self.config.enable_igw = false;
+        self
+    }
+
+    // ==================== WASM ====================
+
+    pub fn enable_wasm(mut self, enable: bool) -> Self {
+        self.config.enable_wasm = enable;
+        self
+    }
+
+    pub fn model_path<S: Into<String>>(mut self, path: S) -> Self {
+        self.config.model_path = Some(path.into());
+        self
+    }
+
+    /// Overrides model_path tokenizer
+    pub fn tokenizer_path<S: Into<String>>(mut self, path: S) -> Self {
+        self.config.tokenizer_path = Some(path.into());
+        self
+    }
+
+    pub fn chat_template<S: Into<String>>(mut self, path: S) -> Self {
+        self.config.chat_template = Some(path.into());
+        self
+    }
+
+    // ==================== History Backend ====================
+
+    pub fn history_backend(mut self, backend: HistoryBackend) -> Self {
+        self.config.history_backend = backend;
+        self
+    }
+
+    pub fn memory_history(mut self) -> Self {
+        self.config.history_backend = HistoryBackend::Memory;
+        self
+    }
+
+    pub fn no_history(mut self) -> Self {
+        self.config.history_backend = HistoryBackend::None;
+        self
+    }
+
+    pub fn oracle_history(mut self, oracle_config: OracleConfig) -> Self {
+        self.config.history_backend = HistoryBackend::Oracle;
+        self.config.oracle = Some(oracle_config);
+        self
+    }
+
+    pub fn redis_history(mut self, redis_config: RedisConfig) -> Self {
+        self.config.history_backend = HistoryBackend::Redis;
+        self.config.redis = Some(redis_config);
+        self
+    }
+
+    // ==================== Parsers ====================
+
+    pub fn reasoning_parser<S: Into<String>>(mut self, parser: S) -> Self {
+        self.config.reasoning_parser = Some(parser.into());
+        self
+    }
+
+    pub fn tool_call_parser<S: Into<String>>(mut self, parser: S) -> Self {
+        self.config.tool_call_parser = Some(parser.into());
+        self
+    }
+
+    // ==================== Tokenizer Cache ====================
+
+    pub fn tokenizer_cache(mut self, cache: TokenizerCacheConfig) -> Self {
+        self.config.tokenizer_cache = cache;
+        self
+    }
+
+    pub fn enable_l0_cache(mut self, max_entries: usize) -> Self {
+        self.config.tokenizer_cache.enable_l0 = true;
+        self.config.tokenizer_cache.l0_max_entries = max_entries;
+        self
+    }
+
+    pub fn enable_l1_cache(mut self, max_memory: usize) -> Self {
+        self.config.tokenizer_cache.enable_l1 = true;
+        self.config.tokenizer_cache.l1_max_memory = max_memory;
+        self
+    }
+
+    // ==================== Data Parallelism ====================
+
+    pub fn enable_dp_aware(mut self) -> Self {
+        self.config.dp_aware = true;
+        self
+    }
+
+    pub fn disable_dp_aware(mut self) -> Self {
+        self.config.dp_aware = false;
+        self
+    }
+
+    // ==================== Boolean Setters ====================
+    // Accept bool parameters to conditionally set flags without if statements
+
+    pub fn dp_aware(mut self, enable: bool) -> Self {
+        self.config.dp_aware = enable;
+        self
+    }
+
+    /// Inverse of disable_retries field
+    pub fn retries(mut self, enable: bool) -> Self {
+        self.config.disable_retries = !enable;
+        self
+    }
+
+    /// Inverse of disable_circuit_breaker field
+    pub fn circuit_breaker(mut self, enable: bool) -> Self {
+        self.config.disable_circuit_breaker = !enable;
+        self
+    }
+
+    pub fn igw(mut self, enable: bool) -> Self {
+        self.config.enable_igw = enable;
+        self
+    }
+
+    // ==================== Option Setters ====================
+    // Accept Option<T> and only set if Some
+
+    pub fn maybe_api_key(mut self, key: Option<impl Into<String>>) -> Self {
+        if let Some(k) = key {
+            self.config.api_key = Some(k.into());
+        }
+        self
+    }
+
+    pub fn maybe_discovery(mut self, discovery: Option<DiscoveryConfig>) -> Self {
+        self.config.discovery = discovery;
+        self
+    }
+
+    pub fn maybe_metrics(mut self, metrics: Option<MetricsConfig>) -> Self {
+        self.config.metrics = metrics;
+        self
+    }
+
+    pub fn maybe_trace(mut self, trace_config: Option<TraceConfig>) -> Self {
+        self.config.trace_config = trace_config;
+        self
+    }
+
+    pub fn maybe_log_dir(mut self, dir: Option<impl Into<String>>) -> Self {
+        self.config.log_dir = dir.map(|d| d.into());
+        self
+    }
+
+    pub fn maybe_log_level(mut self, level: Option<impl Into<String>>) -> Self {
+        self.config.log_level = level.map(|l| l.into());
+        self
+    }
+
+    pub fn maybe_request_id_headers(mut self, headers: Option<Vec<String>>) -> Self {
+        self.config.request_id_headers = headers;
+        self
+    }
+
+    pub fn maybe_rate_limit_tokens_per_second(mut self, tokens: Option<i32>) -> Self {
+        self.config.rate_limit_tokens_per_second = tokens;
+        self
+    }
+
+    pub fn maybe_model_path(mut self, path: Option<impl Into<String>>) -> Self {
+        self.config.model_path = path.map(|p| p.into());
+        self
+    }
+
+    pub fn maybe_tokenizer_path(mut self, path: Option<impl Into<String>>) -> Self {
+        self.config.tokenizer_path = path.map(|p| p.into());
+        self
+    }
+
+    pub fn maybe_chat_template(mut self, template: Option<impl Into<String>>) -> Self {
+        self.config.chat_template = template.map(|t| t.into());
+        self
+    }
+
+    pub fn maybe_oracle(mut self, oracle: Option<OracleConfig>) -> Self {
+        if let Some(cfg) = oracle {
+            self.config.history_backend = HistoryBackend::Oracle;
+            self.config.oracle = Some(cfg);
+        }
+        self
+    }
+
+    pub fn maybe_postgres(mut self, postgres: Option<PostgresConfig>) -> Self {
+        if let Some(cfg) = postgres {
+            self.config.history_backend = HistoryBackend::Postgres;
+            self.config.postgres = Some(cfg);
+        }
+        self
+    }
+
+    pub fn maybe_redis(mut self, redis: Option<RedisConfig>) -> Self {
+        if let Some(cfg) = redis {
+            self.config.history_backend = HistoryBackend::Redis;
+            self.config.redis = Some(cfg);
+        }
+        self
+    }
+
+    pub fn maybe_reasoning_parser(mut self, parser: Option<impl Into<String>>) -> Self {
+        self.config.reasoning_parser = parser.map(|p| p.into());
+        self
+    }
+
+    pub fn maybe_tool_call_parser(mut self, parser: Option<impl Into<String>>) -> Self {
+        self.config.tool_call_parser = parser.map(|p| p.into());
+        self
+    }
+
+    // ==================== mTLS ====================
+
+    /// Both paths must be provided together. Files read during build()
+    pub fn client_cert_and_key<S1: Into<String>, S2: Into<String>>(
+        mut self,
+        cert_path: S1,
+        key_path: S2,
+    ) -> Self {
+        self.client_cert_path = Some(cert_path.into());
+        self.client_key_path = Some(key_path.into());
+        self
+    }
+
+    /// Files read during build()
+    pub fn maybe_client_cert_and_key(
+        mut self,
+        cert_path: Option<impl Into<String>>,
+        key_path: Option<impl Into<String>>,
+    ) -> Self {
+        self.client_cert_path = cert_path.map(|p| p.into());
+        self.client_key_path = key_path.map(|p| p.into());
+        self
+    }
+
+    /// File read during build()
+    pub fn add_ca_certificate<S: Into<String>>(mut self, ca_cert_path: S) -> Self {
+        self.ca_cert_paths.push(ca_cert_path.into());
+        self
+    }
+
+    /// Files read during build()
+    pub fn add_ca_certificates<S: Into<String>>(mut self, ca_cert_paths: Vec<S>) -> Self {
+        self.ca_cert_paths
+            .extend(ca_cert_paths.into_iter().map(|p| p.into()));
+        self
+    }
+
+    // ==================== Server TLS ====================
+
+    /// Both paths must be provided together. Files read during build()
+    pub fn server_cert_and_key<S1: Into<String>, S2: Into<String>>(
+        mut self,
+        cert_path: S1,
+        key_path: S2,
+    ) -> Self {
+        self.server_cert_path = Some(cert_path.into());
+        self.server_key_path = Some(key_path.into());
+        self
+    }
+
+    /// Files read during build()
+    pub fn maybe_server_cert_and_key(
+        mut self,
+        cert_path: Option<impl Into<String>>,
+        key_path: Option<impl Into<String>>,
+    ) -> Self {
+        self.server_cert_path = cert_path.map(|p| p.into());
+        self.server_key_path = key_path.map(|p| p.into());
+        self
+    }
+
+    // ==================== MCP ====================
+
+    /// Config file loaded during build()
+    pub fn mcp_config_path<S: Into<String>>(mut self, path: S) -> Self {
+        self.mcp_config_path = Some(path.into());
+        self
+    }
+
+    /// Config file loaded during build()
+    pub fn maybe_mcp_config_path(mut self, path: Option<impl Into<String>>) -> Self {
+        self.mcp_config_path = path.map(|p| p.into());
+        self
+    }
+
+    // ==================== Build ====================
+
+    pub fn build(self) -> ConfigResult<RouterConfig> {
+        self.build_with_validation(true)
+    }
+
+    pub fn build_unchecked(self) -> RouterConfig {
+        self.into()
+    }
+
+    pub fn build_with_validation(mut self, validate: bool) -> ConfigResult<RouterConfig> {
+        // Read mTLS certificates from paths if provided
+        self = self.read_mtls_certificates()?;
+
+        // Read Server TLS certificates from paths if provided
+        self = self.read_server_certificates()?;
+
+        // Read MCP config from path if provided
+        self = self.read_mcp_config()?;
+
+        let config: RouterConfig = self.into();
+        if validate {
+            config.validate()?;
+        }
+        Ok(config)
+    }
+
+    /// Internal method to read mTLS certificates from paths
+    fn read_mtls_certificates(mut self) -> ConfigResult<Self> {
+        // Read client certificate and key
+        match (&self.client_cert_path, &self.client_key_path) {
+            (Some(cert_path), Some(key_path)) => {
+                let cert = std::fs::read(cert_path).map_err(|e| ConfigError::ValidationFailed {
+                    reason: format!(
+                        "Failed to read client certificate from {}: {}",
+                        cert_path, e
+                    ),
+                })?;
+                let key = std::fs::read(key_path).map_err(|e| ConfigError::ValidationFailed {
+                    reason: format!("Failed to read client key from {}: {}", key_path, e),
+                })?;
+
+                // Combine cert and key into single PEM for reqwest::Identity
+                // When using rustls, certificate must come first, then key
+                // Ensure proper PEM formatting with newlines
+                let mut combined = cert;
+                if !combined.ends_with(b"\n") {
+                    combined.push(b'\n');
+                }
+                combined.extend_from_slice(&key);
+                if !combined.ends_with(b"\n") {
+                    combined.push(b'\n');
+                }
+
+                self.config.client_identity = Some(combined);
+            }
+            (None, None) => {
+                // No client cert configured, that's fine
+            }
+            _ => {
+                return Err(ConfigError::ValidationFailed {
+                    reason:
+                        "Both --client-cert-path and --client-key-path must be specified together"
+                            .to_string(),
+                });
+            }
+        }
+
+        // Read CA certificates
+        for path in &self.ca_cert_paths {
+            let cert = std::fs::read(path).map_err(|e| ConfigError::ValidationFailed {
+                reason: format!("Failed to read CA certificate from {}: {}", path, e),
+            })?;
+            self.config.ca_certificates.push(cert);
+        }
+
+        Ok(self)
+    }
+
+    /// Internal method to read Server TLS certificates from paths
+    fn read_server_certificates(mut self) -> ConfigResult<Self> {
+        match (&self.server_cert_path, &self.server_key_path) {
+            (Some(cert_path), Some(key_path)) => {
+                let cert = std::fs::read(cert_path).map_err(|e| ConfigError::ValidationFailed {
+                    reason: format!(
+                        "Failed to read server certificate from {}: {}",
+                        cert_path, e
+                    ),
+                })?;
+                let key = std::fs::read(key_path).map_err(|e| ConfigError::ValidationFailed {
+                    reason: format!("Failed to read server key from {}: {}", key_path, e),
+                })?;
+                self.config.server_cert = Some(cert);
+                self.config.server_key = Some(key);
+            }
+            (None, None) => {}
+            _ => {
+                return Err(ConfigError::ValidationFailed {
+                    reason: "Both --tls-cert-path and --tls-key-path must be specified together"
+                        .to_string(),
+                });
+            }
+        }
+        Ok(self)
+    }
+
+    /// Internal method to read MCP config from path
+    fn read_mcp_config(mut self) -> ConfigResult<Self> {
+        if let Some(mcp_config_path) = &self.mcp_config_path {
+            let contents = std::fs::read_to_string(mcp_config_path).map_err(|e| {
+                ConfigError::ValidationFailed {
+                    reason: format!("Failed to read MCP config from {}: {}", mcp_config_path, e),
+                }
+            })?;
+            let mcp_config: McpConfig =
+                serde_yaml::from_str(&contents).map_err(|e| ConfigError::ValidationFailed {
+                    reason: format!("Failed to parse MCP config from {}: {}", mcp_config_path, e),
+                })?;
+            self.config.mcp_config = Some(mcp_config);
+        }
+
+        Ok(self)
+    }
+}
+
+impl From<RouterConfigBuilder> for RouterConfig {
+    fn from(builder: RouterConfigBuilder) -> Self {
+        builder.config
+    }
+}
+
+impl RouterConfig {
+    /// Create a builder for RouterConfig
+    pub fn builder() -> RouterConfigBuilder {
+        RouterConfigBuilder::new()
+    }
+
+    /// Create a builder from this configuration
+    pub fn to_builder(&self) -> RouterConfigBuilder {
+        RouterConfigBuilder::from_config_ref(self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Test that .to_builder() round-trip conversion works correctly
+    #[test]
+    fn test_builder_from_existing_config() {
+        let original = RouterConfigBuilder::new()
+            .regular_mode(vec!["http://worker1:8000".to_string()])
+            .port(3000)
+            .build()
+            .unwrap();
+
+        let modified = original
+            .to_builder()
+            .port(4000)
+            .enable_metrics("0.0.0.0", 29000)
+            .enable_trace("localhost:4317")
+            .build()
+            .unwrap();
+
+        assert_eq!(modified.port, 4000);
+        assert!(modified.metrics.is_some());
+        assert!(modified.trace_config.is_some());
+    }
+
+    /// Test complex routing mode helper method
+    #[test]
+    fn test_builder_prefill_decode_mode() {
+        let config = RouterConfigBuilder::new()
+            .prefill_decode_mode(
+                vec![("http://prefill:8000".to_string(), Some(8001))],
+                vec!["http://decode:8000".to_string()],
+            )
+            .power_of_two_policy(60)
+            .build()
+            .unwrap();
+
+        assert!(config.mode.is_pd_mode());
+        assert_eq!(config.mode.worker_count(), 2);
+    }
+
+    /// Test complex policy helper method with multiple parameters
+    #[test]
+    fn test_builder_cache_aware_policy() {
+        let config = RouterConfigBuilder::new()
+            .regular_mode(vec!["http://worker1:8000".to_string()])
+            .cache_aware_policy(0.8, 10, 1.5, 300, 1000)
+            .build()
+            .unwrap();
+
+        match config.policy {
+            PolicyConfig::CacheAware {
+                cache_threshold, ..
+            } => {
+                assert!((cache_threshold - 0.8).abs() < 0.0001);
+            }
+            _ => panic!("Expected CacheAware policy"),
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/config/mod.rs

@@ -0,0 +1,27 @@
+pub mod builder;
+pub mod types;
+pub(crate) mod validation;
+
+pub use builder::*;
+pub use types::*;
+
+#[derive(Debug, thiserror::Error)]
+pub enum ConfigError {
+    #[error("Validation failed: {reason}")]
+    ValidationFailed { reason: String },
+
+    #[error("Invalid value for field '{field}': {value} - {reason}")]
+    InvalidValue {
+        field: String,
+        value: String,
+        reason: String,
+    },
+
+    #[error("Incompatible configuration: {reason}")]
+    IncompatibleConfig { reason: String },
+
+    #[error("Missing required field: {field}")]
+    MissingRequired { field: String },
+}
+
+pub type ConfigResult<T> = Result<T, ConfigError>;

third_party/sglang/sgl-model-gateway/src/core/circuit_breaker.rs

@@ -0,0 +1,626 @@
+use std::{
+    sync::atomic::{AtomicU32, AtomicU64, AtomicU8, Ordering},
+    time::{Duration, Instant},
+};
+
+use tracing::info;
+
+use crate::observability::metrics::Metrics;
+
+/// Circuit breaker configuration
+#[derive(Debug, Clone)]
+pub struct CircuitBreakerConfig {
+    /// Number of consecutive failures to open the circuit
+    pub failure_threshold: u32,
+    /// Success threshold to close circuit from half-open
+    pub success_threshold: u32,
+    /// Duration to wait before attempting half-open
+    pub timeout_duration: Duration,
+    /// Time window for failure counting
+    pub window_duration: Duration,
+}
+
+impl Default for CircuitBreakerConfig {
+    fn default() -> Self {
+        Self {
+            failure_threshold: 5,
+            success_threshold: 2,
+            timeout_duration: Duration::from_secs(30),
+            window_duration: Duration::from_secs(60),
+        }
+    }
+}
+
+/// Circuit breaker state constants for atomic storage
+const STATE_CLOSED: u8 = 0;
+const STATE_OPEN: u8 = 1;
+const STATE_HALF_OPEN: u8 = 2;
+
+/// Circuit breaker state
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CircuitState {
+    /// Normal operation - requests are allowed
+    Closed,
+    /// Circuit is open - requests are rejected
+    Open,
+    /// Testing if service has recovered - limited requests allowed
+    HalfOpen,
+}
+
+impl std::fmt::Display for CircuitState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CircuitState::Closed => write!(f, "Closed"),
+            CircuitState::Open => write!(f, "Open"),
+            CircuitState::HalfOpen => write!(f, "HalfOpen"),
+        }
+    }
+}
+
+impl CircuitState {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            CircuitState::Closed => "closed",
+            CircuitState::Open => "open",
+            CircuitState::HalfOpen => "half_open",
+        }
+    }
+
+    pub fn to_int(&self) -> u8 {
+        match self {
+            CircuitState::Closed => STATE_CLOSED,
+            CircuitState::Open => STATE_OPEN,
+            CircuitState::HalfOpen => STATE_HALF_OPEN,
+        }
+    }
+
+    fn from_int(v: u8) -> Self {
+        match v {
+            STATE_CLOSED => CircuitState::Closed,
+            STATE_OPEN => CircuitState::Open,
+            STATE_HALF_OPEN => CircuitState::HalfOpen,
+            _ => CircuitState::Closed, // Default to closed for safety
+        }
+    }
+}
+
+/// Get current time as milliseconds since an arbitrary epoch.
+/// Uses Instant for monotonic time, converting to ms for atomic storage.
+#[inline]
+fn now_ms() -> u64 {
+    // Use a static reference point for consistent timing
+    static START: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
+    let start = START.get_or_init(Instant::now);
+    start.elapsed().as_millis() as u64
+}
+
+/// Circuit breaker implementation using lock-free atomics for hot paths.
+///
+/// This implementation avoids RwLock contention by using atomic operations
+/// for state checks (the most common operation). Only state transitions
+/// use compare-and-swap which is still lock-free.
+#[derive(Debug)]
+pub struct CircuitBreaker {
+    /// Circuit state stored as atomic u8 (0=Closed, 1=Open, 2=HalfOpen)
+    state: AtomicU8,
+    consecutive_failures: AtomicU32,
+    consecutive_successes: AtomicU32,
+    total_failures: AtomicU64,
+    total_successes: AtomicU64,
+    /// Last failure time in milliseconds (from now_ms())
+    last_failure_time_ms: AtomicU64,
+    /// Last state change time in milliseconds (from now_ms())
+    last_state_change_ms: AtomicU64,
+    config: CircuitBreakerConfig,
+    metric_label: String,
+}
+
+impl CircuitBreaker {
+    /// Create a new circuit breaker with default configuration
+    pub fn new() -> Self {
+        Self::with_config_and_label(CircuitBreakerConfig::default(), String::new())
+    }
+
+    /// Create a new circuit breaker with custom configuration and metric label
+    pub fn with_config_and_label(config: CircuitBreakerConfig, metric_label: String) -> Self {
+        let init_state = CircuitState::Closed;
+        Metrics::set_worker_cb_state(&metric_label, init_state.to_int());
+        Self {
+            state: AtomicU8::new(STATE_CLOSED),
+            consecutive_failures: AtomicU32::new(0),
+            consecutive_successes: AtomicU32::new(0),
+            total_failures: AtomicU64::new(0),
+            total_successes: AtomicU64::new(0),
+            last_failure_time_ms: AtomicU64::new(0),
+            last_state_change_ms: AtomicU64::new(now_ms()),
+            config,
+            metric_label,
+        }
+    }
+
+    /// Get the metric label
+    pub fn metric_label(&self) -> &str {
+        &self.metric_label
+    }
+
+    /// Check if a request can be executed (lock-free hot path)
+    #[inline]
+    pub fn can_execute(&self) -> bool {
+        let state = self.state();
+        match state {
+            CircuitState::Closed => true,
+            CircuitState::Open => false,
+            CircuitState::HalfOpen => true,
+        }
+    }
+
+    /// Get the current state (lock-free)
+    #[inline]
+    pub fn state(&self) -> CircuitState {
+        self.check_and_update_state_returning()
+    }
+
+    /// Check and update state, returning the current state (lock-free)
+    #[inline]
+    fn check_and_update_state_returning(&self) -> CircuitState {
+        let current_state_int = self.state.load(Ordering::Acquire);
+        let current_state = CircuitState::from_int(current_state_int);
+
+        if current_state == CircuitState::Open {
+            let last_change_ms = self.last_state_change_ms.load(Ordering::Acquire);
+            let elapsed_ms = now_ms().saturating_sub(last_change_ms);
+            let timeout_ms = self.config.timeout_duration.as_millis() as u64;
+
+            if elapsed_ms >= timeout_ms {
+                // Try to transition to HalfOpen using CAS
+                if self
+                    .state
+                    .compare_exchange(
+                        STATE_OPEN,
+                        STATE_HALF_OPEN,
+                        Ordering::AcqRel,
+                        Ordering::Acquire,
+                    )
+                    .is_ok()
+                {
+                    self.last_state_change_ms.store(now_ms(), Ordering::Release);
+                    self.consecutive_failures.store(0, Ordering::Release);
+                    self.consecutive_successes.store(0, Ordering::Release);
+
+                    info!("Circuit breaker state transition: open -> half_open");
+                    Metrics::record_worker_cb_transition(&self.metric_label, "open", "half_open");
+                    Metrics::set_worker_cb_state(&self.metric_label, STATE_HALF_OPEN);
+                    self.publish_gauge_metrics();
+                    return CircuitState::HalfOpen;
+                }
+                // Another thread already transitioned, re-read the state
+                return CircuitState::from_int(self.state.load(Ordering::Acquire));
+            }
+        }
+        current_state
+    }
+
+    /// Record the outcome of a request
+    pub fn record_outcome(&self, success: bool) {
+        if success {
+            self.record_success();
+        } else {
+            self.record_failure();
+        }
+
+        let outcome_str = if success { "success" } else { "failure" };
+        Metrics::record_worker_cb_outcome(&self.metric_label, outcome_str);
+        self.publish_gauge_metrics();
+    }
+
+    /// Record a successful request
+    pub fn record_success(&self) {
+        self.total_successes.fetch_add(1, Ordering::Relaxed);
+        self.consecutive_failures.store(0, Ordering::Release);
+        let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1;
+
+        let current_state = CircuitState::from_int(self.state.load(Ordering::Acquire));
+
+        match current_state {
+            CircuitState::HalfOpen => {
+                if successes >= self.config.success_threshold {
+                    self.transition_to(CircuitState::Closed);
+                }
+            }
+            CircuitState::Closed => {}
+            CircuitState::Open => {
+                tracing::warn!("Success recorded while circuit is open");
+            }
+        }
+    }
+
+    /// Record a failed request
+    pub fn record_failure(&self) {
+        self.total_failures.fetch_add(1, Ordering::Relaxed);
+        self.consecutive_successes.store(0, Ordering::Release);
+        let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1;
+
+        // Update last failure time atomically
+        self.last_failure_time_ms.store(now_ms(), Ordering::Release);
+
+        let current_state = CircuitState::from_int(self.state.load(Ordering::Acquire));
+
+        match current_state {
+            CircuitState::Closed => {
+                if failures >= self.config.failure_threshold {
+                    self.transition_to(CircuitState::Open);
+                }
+            }
+            CircuitState::HalfOpen => {
+                self.transition_to(CircuitState::Open);
+            }
+            CircuitState::Open => {}
+        }
+    }
+
+    /// Transition to a new state (uses CAS for lock-free operation)
+    fn transition_to(&self, new_state: CircuitState) {
+        let new_state_int = new_state.to_int();
+        let old_state_int = self.state.swap(new_state_int, Ordering::AcqRel);
+        let old_state = CircuitState::from_int(old_state_int);
+
+        if old_state != new_state {
+            self.last_state_change_ms.store(now_ms(), Ordering::Release);
+
+            match new_state {
+                CircuitState::Closed => {
+                    self.consecutive_failures.store(0, Ordering::Release);
+                    self.consecutive_successes.store(0, Ordering::Release);
+                }
+                CircuitState::Open => {
+                    self.consecutive_successes.store(0, Ordering::Release);
+                }
+                CircuitState::HalfOpen => {
+                    self.consecutive_failures.store(0, Ordering::Release);
+                    self.consecutive_successes.store(0, Ordering::Release);
+                }
+            }
+
+            let from = old_state.as_str();
+            let to = new_state.as_str();
+            info!("Circuit breaker state transition: {} -> {}", from, to);
+            Metrics::record_worker_cb_transition(&self.metric_label, from, to);
+            Metrics::set_worker_cb_state(&self.metric_label, new_state.to_int());
+            self.publish_gauge_metrics();
+        }
+    }
+
+    /// Get the number of consecutive failures
+    pub fn consecutive_failures(&self) -> u32 {
+        self.consecutive_failures.load(Ordering::Acquire)
+    }
+
+    /// Get the number of consecutive successes
+    pub fn consecutive_successes(&self) -> u32 {
+        self.consecutive_successes.load(Ordering::Acquire)
+    }
+
+    /// Get total failures
+    pub fn total_failures(&self) -> u64 {
+        self.total_failures.load(Ordering::Relaxed)
+    }
+
+    /// Get total successes
+    pub fn total_successes(&self) -> u64 {
+        self.total_successes.load(Ordering::Relaxed)
+    }
+
+    /// Get time since last failure
+    pub fn time_since_last_failure(&self) -> Option<Duration> {
+        let last_ms = self.last_failure_time_ms.load(Ordering::Acquire);
+        if last_ms == 0 {
+            None
+        } else {
+            let elapsed_ms = now_ms().saturating_sub(last_ms);
+            Some(Duration::from_millis(elapsed_ms))
+        }
+    }
+
+    /// Get time since last state change
+    pub fn time_since_last_state_change(&self) -> Duration {
+        let last_ms = self.last_state_change_ms.load(Ordering::Acquire);
+        let elapsed_ms = now_ms().saturating_sub(last_ms);
+        Duration::from_millis(elapsed_ms)
+    }
+
+    /// Check if the circuit is in a half-open state
+    pub fn is_half_open(&self) -> bool {
+        self.state() == CircuitState::HalfOpen
+    }
+
+    /// Record a test success (for health check probing)
+    pub fn record_test_success(&self) {
+        if self.is_half_open() {
+            self.record_success();
+        }
+    }
+
+    /// Record a test failure (for health check probing)
+    pub fn record_test_failure(&self) {
+        if self.is_half_open() {
+            self.record_failure();
+        }
+    }
+
+    /// Reset the circuit breaker to closed state
+    pub fn reset(&self) {
+        self.transition_to(CircuitState::Closed);
+        self.consecutive_failures.store(0, Ordering::Release);
+        self.consecutive_successes.store(0, Ordering::Release);
+        self.publish_gauge_metrics();
+    }
+
+    /// Force the circuit to open (for manual intervention)
+    pub fn force_open(&self) {
+        self.transition_to(CircuitState::Open);
+    }
+
+    /// Get circuit breaker statistics
+    pub fn stats(&self) -> CircuitBreakerStats {
+        CircuitBreakerStats {
+            state: self.state(),
+            consecutive_failures: self.consecutive_failures(),
+            consecutive_successes: self.consecutive_successes(),
+            total_failures: self.total_failures(),
+            total_successes: self.total_successes(),
+            time_since_last_failure: self.time_since_last_failure(),
+            time_since_last_state_change: self.time_since_last_state_change(),
+        }
+    }
+
+    fn publish_gauge_metrics(&self) {
+        Metrics::set_worker_cb_consecutive_failures(
+            &self.metric_label,
+            self.consecutive_failures(),
+        );
+        Metrics::set_worker_cb_consecutive_successes(
+            &self.metric_label,
+            self.consecutive_successes(),
+        );
+    }
+}
+
+impl Clone for CircuitBreaker {
+    fn clone(&self) -> Self {
+        Self {
+            state: AtomicU8::new(self.state.load(Ordering::Acquire)),
+            consecutive_failures: AtomicU32::new(self.consecutive_failures.load(Ordering::Acquire)),
+            consecutive_successes: AtomicU32::new(
+                self.consecutive_successes.load(Ordering::Acquire),
+            ),
+            total_failures: AtomicU64::new(self.total_failures.load(Ordering::Relaxed)),
+            total_successes: AtomicU64::new(self.total_successes.load(Ordering::Relaxed)),
+            last_failure_time_ms: AtomicU64::new(self.last_failure_time_ms.load(Ordering::Acquire)),
+            last_state_change_ms: AtomicU64::new(self.last_state_change_ms.load(Ordering::Acquire)),
+            config: self.config.clone(),
+            metric_label: self.metric_label.clone(),
+        }
+    }
+}
+
+impl Default for CircuitBreaker {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Circuit breaker statistics
+#[derive(Debug, Clone)]
+pub struct CircuitBreakerStats {
+    pub state: CircuitState,
+    pub consecutive_failures: u32,
+    pub consecutive_successes: u32,
+    pub total_failures: u64,
+    pub total_successes: u64,
+    pub time_since_last_failure: Option<Duration>,
+    pub time_since_last_state_change: Duration,
+}
+
+#[cfg(test)]
+mod tests {
+    use std::thread;
+
+    use super::*;
+
+    #[test]
+    fn test_circuit_breaker_initial_state() {
+        let cb = CircuitBreaker::new();
+        assert_eq!(cb.state(), CircuitState::Closed);
+        assert!(cb.can_execute());
+        assert_eq!(cb.consecutive_failures(), 0);
+        assert_eq!(cb.consecutive_successes(), 0);
+    }
+
+    #[test]
+    fn test_circuit_opens_on_threshold() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 3,
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        assert_eq!(cb.state(), CircuitState::Closed);
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Closed);
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Closed);
+        cb.record_failure();
+
+        assert_eq!(cb.state(), CircuitState::Open);
+        assert!(!cb.can_execute());
+        assert_eq!(cb.consecutive_failures(), 3);
+    }
+
+    #[test]
+    fn test_circuit_half_open_after_timeout() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 1,
+            timeout_duration: Duration::from_millis(100),
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Open);
+
+        thread::sleep(Duration::from_millis(150));
+
+        assert_eq!(cb.state(), CircuitState::HalfOpen);
+        assert!(cb.can_execute());
+    }
+
+    #[test]
+    fn test_circuit_closes_on_success_threshold() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 1,
+            success_threshold: 2,
+            timeout_duration: Duration::from_millis(50),
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Open);
+
+        thread::sleep(Duration::from_millis(100));
+        assert_eq!(cb.state(), CircuitState::HalfOpen);
+
+        cb.record_success();
+        assert_eq!(cb.state(), CircuitState::HalfOpen);
+        cb.record_success();
+
+        assert_eq!(cb.state(), CircuitState::Closed);
+        assert!(cb.can_execute());
+    }
+
+    #[test]
+    fn test_circuit_reopens_on_half_open_failure() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 1,
+            timeout_duration: Duration::from_millis(50),
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Open);
+
+        thread::sleep(Duration::from_millis(100));
+        assert_eq!(cb.state(), CircuitState::HalfOpen);
+
+        cb.record_failure();
+
+        assert_eq!(cb.state(), CircuitState::Open);
+        assert!(!cb.can_execute());
+    }
+
+    #[test]
+    fn test_success_resets_failure_count() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 3,
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        cb.record_failure();
+        cb.record_failure();
+        assert_eq!(cb.consecutive_failures(), 2);
+
+        cb.record_success();
+        assert_eq!(cb.consecutive_failures(), 0);
+        assert_eq!(cb.consecutive_successes(), 1);
+
+        cb.record_failure();
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Closed);
+    }
+
+    #[test]
+    fn test_manual_reset() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 1,
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        cb.record_failure();
+        assert_eq!(cb.state(), CircuitState::Open);
+
+        cb.reset();
+        assert_eq!(cb.state(), CircuitState::Closed);
+        assert_eq!(cb.consecutive_failures(), 0);
+        assert_eq!(cb.consecutive_successes(), 0);
+    }
+
+    #[test]
+    fn test_force_open() {
+        let cb = CircuitBreaker::new();
+        assert_eq!(cb.state(), CircuitState::Closed);
+
+        cb.force_open();
+        assert_eq!(cb.state(), CircuitState::Open);
+        assert!(!cb.can_execute());
+    }
+
+    #[test]
+    fn test_stats() {
+        let config = CircuitBreakerConfig {
+            failure_threshold: 2,
+            ..Default::default()
+        };
+        let cb = CircuitBreaker::with_config_and_label(config, String::new());
+
+        cb.record_success();
+        cb.record_failure();
+        cb.record_failure();
+
+        let stats = cb.stats();
+        assert_eq!(stats.state, CircuitState::Open);
+        assert_eq!(stats.consecutive_failures, 2);
+        assert_eq!(stats.consecutive_successes, 0);
+        assert_eq!(stats.total_failures, 2);
+        assert_eq!(stats.total_successes, 1);
+    }
+
+    #[test]
+    fn test_clone() {
+        let cb1 = CircuitBreaker::new();
+        cb1.record_failure();
+
+        let cb2 = cb1.clone();
+        assert_eq!(cb2.consecutive_failures(), 1);
+
+        cb1.record_failure();
+        assert_eq!(cb1.consecutive_failures(), 2);
+        assert_eq!(cb2.consecutive_failures(), 1); // cb2 is unchanged
+    }
+
+    #[test]
+    fn test_thread_safety() {
+        use std::sync::Arc;
+
+        let cb = Arc::new(CircuitBreaker::new());
+        let mut handles = vec![];
+
+        for _ in 0..10 {
+            let cb_clone = Arc::clone(&cb);
+            let handle = thread::spawn(move || {
+                for _ in 0..100 {
+                    cb_clone.record_failure();
+                }
+            });
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.join().unwrap();
+        }
+
+        assert_eq!(cb.total_failures(), 1000);
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/error.rs

@@ -0,0 +1,210 @@
+//! Error types for the SGLang router core
+//!
+//! This module defines error types used throughout the router for worker operations.
+
+/// Worker-related errors
+#[derive(Debug, thiserror::Error)]
+pub enum WorkerError {
+    #[error("Health check failed for worker {url}: {reason}")]
+    HealthCheckFailed { url: String, reason: String },
+
+    #[error("Worker not found: {url}")]
+    WorkerNotFound { url: String },
+
+    #[error("Invalid worker configuration: {message}")]
+    InvalidConfiguration { message: String },
+
+    #[error("Network error for worker {url}: {error}")]
+    NetworkError { url: String, error: String },
+
+    #[error("Worker at capacity: {url}")]
+    WorkerAtCapacity { url: String },
+
+    #[error("Invalid URL format: {url}")]
+    InvalidUrl { url: String },
+
+    #[error("Connection failed for worker {url}: {reason}")]
+    ConnectionFailed { url: String, reason: String },
+}
+
+/// Result type for worker operations
+pub type WorkerResult<T> = Result<T, WorkerError>;
+
+/// Convert from reqwest errors to worker errors
+impl From<reqwest::Error> for WorkerError {
+    fn from(err: reqwest::Error) -> Self {
+        WorkerError::NetworkError {
+            url: err.url().map(|u| u.to_string()).unwrap_or_default(),
+            error: err.to_string(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::error::Error;
+
+    use super::*;
+
+    #[test]
+    fn test_health_check_failed_display() {
+        let error = WorkerError::HealthCheckFailed {
+            url: "http://worker1:8080".to_string(),
+            reason: "Connection refused".to_string(),
+        };
+        assert_eq!(
+            error.to_string(),
+            "Health check failed for worker http://worker1:8080: Connection refused"
+        );
+    }
+
+    #[test]
+    fn test_worker_not_found_display() {
+        let error = WorkerError::WorkerNotFound {
+            url: "http://worker2:8080".to_string(),
+        };
+        assert_eq!(error.to_string(), "Worker not found: http://worker2:8080");
+    }
+
+    #[test]
+    fn test_invalid_configuration_display() {
+        let error = WorkerError::InvalidConfiguration {
+            message: "Missing port number".to_string(),
+        };
+        assert_eq!(
+            error.to_string(),
+            "Invalid worker configuration: Missing port number"
+        );
+    }
+
+    #[test]
+    fn test_network_error_display() {
+        let error = WorkerError::NetworkError {
+            url: "http://worker3:8080".to_string(),
+            error: "Timeout after 30s".to_string(),
+        };
+        assert_eq!(
+            error.to_string(),
+            "Network error for worker http://worker3:8080: Timeout after 30s"
+        );
+    }
+
+    #[test]
+    fn test_worker_at_capacity_display() {
+        let error = WorkerError::WorkerAtCapacity {
+            url: "http://worker4:8080".to_string(),
+        };
+        assert_eq!(error.to_string(), "Worker at capacity: http://worker4:8080");
+    }
+
+    #[test]
+    fn test_worker_error_implements_std_error() {
+        let error = WorkerError::WorkerNotFound {
+            url: "http://test".to_string(),
+        };
+        let _: &dyn Error = &error;
+        assert!(error.source().is_none());
+    }
+
+    #[test]
+    fn test_error_send_sync() {
+        fn assert_send_sync<T: Send + Sync>() {}
+        assert_send_sync::<WorkerError>();
+    }
+
+    #[test]
+    fn test_worker_result_type_alias() {
+        let result: WorkerResult<i32> = Ok(42);
+        assert!(matches!(result, Ok(42)));
+
+        let error = WorkerError::WorkerNotFound {
+            url: "test".to_string(),
+        };
+        let result: WorkerResult<i32> = Err(error);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_empty_url_handling() {
+        let error1 = WorkerError::HealthCheckFailed {
+            url: "".to_string(),
+            reason: "No connection".to_string(),
+        };
+        assert_eq!(
+            error1.to_string(),
+            "Health check failed for worker : No connection"
+        );
+
+        let error2 = WorkerError::NetworkError {
+            url: "".to_string(),
+            error: "DNS failure".to_string(),
+        };
+        assert_eq!(error2.to_string(), "Network error for worker : DNS failure");
+
+        let error3 = WorkerError::WorkerNotFound {
+            url: "".to_string(),
+        };
+        assert_eq!(error3.to_string(), "Worker not found: ");
+    }
+
+    #[test]
+    fn test_special_characters_in_messages() {
+        let error = WorkerError::InvalidConfiguration {
+            message: "Invalid JSON: {\"error\": \"test\"}".to_string(),
+        };
+        assert_eq!(
+            error.to_string(),
+            "Invalid worker configuration: Invalid JSON: {\"error\": \"test\"}"
+        );
+
+        let error2 = WorkerError::HealthCheckFailed {
+            url: "http://测试:8080".to_string(),
+            reason: "连接被拒绝".to_string(),
+        };
+        assert_eq!(
+            error2.to_string(),
+            "Health check failed for worker http://测试:8080: 连接被拒绝"
+        );
+    }
+
+    #[test]
+    fn test_very_long_error_messages() {
+        let long_message = "A".repeat(10000);
+        let error = WorkerError::InvalidConfiguration {
+            message: long_message.clone(),
+        };
+        let display = error.to_string();
+        assert!(display.contains(&long_message));
+        assert_eq!(
+            display.len(),
+            "Invalid worker configuration: ".len() + long_message.len()
+        );
+    }
+
+    #[test]
+    fn test_reqwest_error_conversion() {
+        let network_error = WorkerError::NetworkError {
+            url: "http://example.com".to_string(),
+            error: "connection timeout".to_string(),
+        };
+
+        match network_error {
+            WorkerError::NetworkError { url, error } => {
+                assert_eq!(url, "http://example.com");
+                assert_eq!(error, "connection timeout");
+            }
+            _ => panic!("Expected NetworkError variant"),
+        }
+    }
+
+    #[test]
+    fn test_error_equality() {
+        let error1 = WorkerError::WorkerNotFound {
+            url: "http://test".to_string(),
+        };
+        let error2 = WorkerError::WorkerNotFound {
+            url: "http://test".to_string(),
+        };
+        assert_eq!(error1.to_string(), error2.to_string());
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/job_queue.rs

@@ -0,0 +1,798 @@
+//! Async job queue for control plane operations
+//!
+//! Provides non-blocking worker management by queuing operations and processing
+//! them asynchronously in background worker tasks.
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, Weak},
+    time::{Duration, SystemTime},
+};
+
+use dashmap::DashMap;
+use smg_mcp::McpConfig;
+use tokio::sync::{mpsc, Semaphore};
+use tracing::{debug, error, info, warn};
+use wfaas::WorkflowId;
+
+use crate::{
+    app_context::AppContext,
+    config::{RouterConfig, RoutingMode},
+    core::steps::{
+        create_external_worker_workflow_data, create_local_worker_workflow_data,
+        create_mcp_workflow_data, create_tokenizer_workflow_data,
+        create_wasm_registration_workflow_data, create_wasm_removal_workflow_data,
+        create_worker_removal_workflow_data, create_worker_update_workflow_data,
+        McpServerConfigRequest, TokenizerConfigRequest, TokenizerRemovalRequest,
+        WasmModuleConfigRequest, WasmModuleRemovalRequest,
+    },
+    protocols::worker_spec::{JobStatus, WorkerConfigRequest, WorkerUpdateRequest},
+};
+
+/// Job types for control plane operations
+#[derive(Debug, Clone)]
+pub enum Job {
+    AddWorker {
+        config: Box<WorkerConfigRequest>,
+    },
+    UpdateWorker {
+        url: String,
+        update: Box<WorkerUpdateRequest>,
+    },
+    RemoveWorker {
+        url: String,
+    },
+    InitializeWorkersFromConfig {
+        router_config: Box<RouterConfig>,
+    },
+    InitializeMcpServers {
+        mcp_config: Box<McpConfig>,
+    },
+    RegisterMcpServer {
+        config: Box<McpServerConfigRequest>,
+    },
+    AddWasmModule {
+        config: Box<WasmModuleConfigRequest>,
+    },
+    RemoveWasmModule {
+        request: Box<WasmModuleRemovalRequest>,
+    },
+    AddTokenizer {
+        config: Box<TokenizerConfigRequest>,
+    },
+    RemoveTokenizer {
+        request: Box<TokenizerRemovalRequest>,
+    },
+}
+
+impl Job {
+    /// Get job type as string for logging
+    pub fn job_type(&self) -> &'static str {
+        match self {
+            Job::AddWorker { .. } => "AddWorker",
+            Job::UpdateWorker { .. } => "UpdateWorker",
+            Job::RemoveWorker { .. } => "RemoveWorker",
+            Job::InitializeWorkersFromConfig { .. } => "InitializeWorkersFromConfig",
+            Job::InitializeMcpServers { .. } => "InitializeMcpServers",
+            Job::RegisterMcpServer { .. } => "RegisterMcpServer",
+            Job::AddWasmModule { .. } => "AddWasmModule",
+            Job::RemoveWasmModule { .. } => "RemoveWasmModule",
+            Job::AddTokenizer { .. } => "AddTokenizer",
+            Job::RemoveTokenizer { .. } => "RemoveTokenizer",
+        }
+    }
+
+    /// Get worker URL, MCP server name, WASM module, or tokenizer identifier for logging and status tracking
+    pub fn worker_url(&self) -> &str {
+        match self {
+            Job::AddWorker { config } => &config.url,
+            Job::UpdateWorker { url, .. } => url,
+            Job::RemoveWorker { url } => url,
+            Job::InitializeWorkersFromConfig { .. } => "startup",
+            Job::InitializeMcpServers { .. } => "startup",
+            Job::RegisterMcpServer { config } => &config.name,
+            Job::AddWasmModule { config } => &config.descriptor.name,
+            Job::RemoveWasmModule { request } => &request.uuid_string,
+            Job::AddTokenizer { config } => &config.id,
+            Job::RemoveTokenizer { request } => &request.id,
+        }
+    }
+}
+
+/// Job queue configuration
+#[derive(Clone, Debug)]
+pub struct JobQueueConfig {
+    /// Maximum pending jobs in queue
+    pub queue_capacity: usize,
+    /// Maximum number of jobs executing concurrently
+    pub max_concurrent_jobs: usize,
+}
+
+impl Default for JobQueueConfig {
+    fn default() -> Self {
+        Self {
+            queue_capacity: 1000,
+            max_concurrent_jobs: 200,
+        }
+    }
+}
+
+/// Job queue manager for worker validation and removal operations
+pub struct JobQueue {
+    /// Channel for submitting jobs
+    tx: mpsc::Sender<Job>,
+    /// Weak reference to AppContext to avoid circular dependencies
+    context: Weak<AppContext>,
+    /// Job status tracking by worker URL
+    status_map: Arc<DashMap<String, JobStatus>>,
+    /// Semaphore to limit concurrent job execution
+    concurrency_limit: Arc<Semaphore>,
+}
+
+impl std::fmt::Debug for JobQueue {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("JobQueue")
+            .field("status_count", &self.status_map.len())
+            .finish()
+    }
+}
+
+impl JobQueue {
+    /// Create a new job queue with semaphore-based concurrency control
+    ///
+    /// Takes a Weak reference to AppContext to avoid circular strong references.
+    /// Spawns a single dispatcher task that spawns individual job tasks with semaphore control.
+    pub fn new(config: JobQueueConfig, context: Weak<AppContext>) -> Arc<Self> {
+        let (tx, mut rx) = mpsc::channel(config.queue_capacity);
+
+        debug!(
+            "Initializing job queue: capacity={}, max_concurrent={}",
+            config.queue_capacity, config.max_concurrent_jobs
+        );
+
+        let status_map = Arc::new(DashMap::new());
+        let concurrency_limit = Arc::new(Semaphore::new(config.max_concurrent_jobs));
+
+        let queue = Arc::new(Self {
+            tx,
+            context: context.clone(),
+            status_map: status_map.clone(),
+            concurrency_limit: concurrency_limit.clone(),
+        });
+
+        // Single dispatcher task
+        let ctx = context.clone();
+        let status = status_map.clone();
+        let sem = concurrency_limit.clone();
+
+        tokio::spawn(async move {
+            while let Some(job) = rx.recv().await {
+                // Acquire permit (blocks if at concurrency limit)
+                let Ok(permit) = sem.clone().acquire_owned().await else {
+                    error!("Semaphore closed, stopping dispatcher");
+                    break;
+                };
+
+                let ctx_clone = ctx.clone();
+                let status_clone = status.clone();
+
+                tokio::spawn(async move {
+                    Self::process_job(job, ctx_clone, status_clone, permit).await;
+                });
+            }
+
+            debug!("Job dispatcher stopped");
+        });
+
+        // Spawn cleanup task for old job statuses (TTL 5 minutes)
+        let cleanup_status_map = status_map.clone();
+        tokio::spawn(async move {
+            Self::cleanup_old_statuses(cleanup_status_map).await;
+        });
+
+        queue
+    }
+
+    /// Get current queue and concurrency status
+    pub fn get_load_info(&self) -> (usize, usize) {
+        let queue_depth = self.tx.max_capacity() - self.tx.capacity();
+        let available_permits = self.concurrency_limit.available_permits();
+        (queue_depth, available_permits)
+    }
+
+    /// Submit a job with detailed queue status
+    pub async fn submit(&self, job: Job) -> Result<(), String> {
+        // Check if context is still alive before accepting jobs
+        if self.context.upgrade().is_none() {
+            return Err("Job queue shutting down: AppContext dropped".to_string());
+        }
+
+        // Extract values before moving job
+        let job_type = job.job_type();
+        let worker_url = job.worker_url().to_string();
+
+        // Record pending status
+        self.status_map.insert(
+            worker_url.clone(),
+            JobStatus::pending(job_type, &worker_url),
+        );
+
+        match self.tx.send(job).await {
+            Ok(_) => {
+                let (queue_depth, available_permits) = self.get_load_info();
+                debug!(
+                    "Job submitted: type={}, worker={}, queue_depth={}, available_slots={}",
+                    job_type, worker_url, queue_depth, available_permits
+                );
+                Ok(())
+            }
+            Err(_) => {
+                self.status_map.remove(&worker_url);
+                let (queue_depth, _) = self.get_load_info();
+                Err(format!(
+                    "Job queue full: {} jobs pending (capacity: {})",
+                    queue_depth,
+                    self.tx.max_capacity()
+                ))
+            }
+        }
+    }
+
+    /// Get job status by worker URL
+    pub fn get_status(&self, worker_url: &str) -> Option<JobStatus> {
+        self.status_map.get(worker_url).map(|entry| entry.clone())
+    }
+
+    /// Remove job status (called when worker is deleted)
+    pub fn remove_status(&self, worker_url: &str) {
+        self.status_map.remove(worker_url);
+    }
+
+    /// Process a single job with status tracking and error handling
+    async fn process_job(
+        job: Job,
+        context: Weak<AppContext>,
+        status_map: Arc<DashMap<String, JobStatus>>,
+        _permit: tokio::sync::OwnedSemaphorePermit,
+    ) {
+        let job_type = job.job_type();
+        let worker_url = job.worker_url().to_string();
+        let start = std::time::Instant::now();
+
+        // Update to processing
+        status_map.insert(
+            worker_url.clone(),
+            JobStatus::processing(job_type, &worker_url),
+        );
+
+        debug!("Processing job: type={}, worker={}", job_type, worker_url);
+
+        // Execute job
+        match context.upgrade() {
+            Some(ctx) => {
+                let result = Self::execute_job(&job, &ctx).await;
+                let duration = start.elapsed();
+                Self::record_job_completion(job_type, &worker_url, duration, &result, &status_map);
+            }
+            None => {
+                let error_msg = "AppContext dropped".to_string();
+                status_map.insert(
+                    worker_url.clone(),
+                    JobStatus::failed(job_type, &worker_url, error_msg),
+                );
+                error!(
+                    "AppContext dropped, cannot process job: type={}, worker={}",
+                    job_type, worker_url
+                );
+            }
+        }
+
+        // Permit automatically released when dropped
+    }
+
+    /// Execute a specific job
+    async fn execute_job(job: &Job, context: &Arc<AppContext>) -> Result<String, String> {
+        match job {
+            Job::AddWorker { config } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let timeout_duration =
+                    Duration::from_secs(context.router_config.worker_startup_timeout_secs + 30);
+
+                // Select workflow based on runtime field
+                match config.runtime.as_deref() {
+                    Some("external") => {
+                        let workflow_data = create_external_worker_workflow_data(
+                            (**config).clone(),
+                            Arc::clone(context),
+                        );
+                        let instance_id = engines
+                            .external_worker
+                            .start_workflow(
+                                WorkflowId::new("external_worker_registration"),
+                                workflow_data,
+                            )
+                            .await
+                            .map_err(|e| {
+                                format!(
+                                    "Failed to start external worker registration workflow: {:?}",
+                                    e
+                                )
+                            })?;
+
+                        debug!(
+                            "Started external worker registration workflow for {} (instance: {})",
+                            config.url, instance_id
+                        );
+
+                        engines
+                            .external_worker
+                            .wait_for_completion(instance_id, &config.url, timeout_duration)
+                            .await
+                    }
+                    _ => {
+                        let workflow_data = create_local_worker_workflow_data(
+                            (**config).clone(),
+                            Arc::clone(context),
+                        );
+                        let instance_id = engines
+                            .local_worker
+                            .start_workflow(
+                                WorkflowId::new("local_worker_registration"),
+                                workflow_data,
+                            )
+                            .await
+                            .map_err(|e| {
+                                format!(
+                                    "Failed to start local worker registration workflow: {:?}",
+                                    e
+                                )
+                            })?;
+
+                        debug!(
+                            "Started local worker registration workflow for {} (instance: {})",
+                            config.url, instance_id
+                        );
+
+                        engines
+                            .local_worker
+                            .wait_for_completion(instance_id, &config.url, timeout_duration)
+                            .await
+                    }
+                }
+            }
+            Job::UpdateWorker { url, update } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let workflow_data = create_worker_update_workflow_data(
+                    url.to_string(),
+                    (**update).clone(),
+                    Arc::clone(context),
+                );
+
+                let instance_id = engines
+                    .worker_update
+                    .start_workflow(WorkflowId::new("worker_update"), workflow_data)
+                    .await
+                    .map_err(|e| format!("Failed to start worker update workflow: {:?}", e))?;
+
+                debug!(
+                    "Started worker update workflow for {} (instance: {})",
+                    url, instance_id
+                );
+
+                let timeout_duration = Duration::from_secs(30);
+
+                engines
+                    .worker_update
+                    .wait_for_completion(instance_id, url, timeout_duration)
+                    .await
+            }
+            Job::RemoveWorker { url } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let workflow_data = create_worker_removal_workflow_data(
+                    url.to_string(),
+                    context.router_config.dp_aware,
+                    Arc::clone(context),
+                );
+
+                let instance_id = engines
+                    .worker_removal
+                    .start_workflow(WorkflowId::new("worker_removal"), workflow_data)
+                    .await
+                    .map_err(|e| format!("Failed to start worker removal workflow: {:?}", e))?;
+
+                debug!(
+                    "Started worker removal workflow for {} (instance: {})",
+                    url, instance_id
+                );
+
+                let timeout_duration = Duration::from_secs(30);
+
+                let result = engines
+                    .worker_removal
+                    .wait_for_completion(instance_id, url, timeout_duration)
+                    .await;
+
+                // Clean up job status when removing worker
+                if let Some(queue) = context.worker_job_queue.get() {
+                    queue.remove_status(url);
+                }
+
+                result
+            }
+            Job::AddWasmModule { config } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let workflow_data =
+                    create_wasm_registration_workflow_data(*config.clone(), Arc::clone(context));
+
+                let instance_id = engines
+                    .wasm_registration
+                    .start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
+                    .await
+                    .map_err(|e| {
+                        format!("Failed to start WASM module registration workflow: {:?}", e)
+                    })?;
+
+                debug!(
+                    "Started WASM module registration workflow for {} (instance: {})",
+                    config.descriptor.name, instance_id
+                );
+
+                let timeout_duration = Duration::from_secs(300); // 5 minutes
+
+                engines
+                    .wasm_registration
+                    .wait_for_completion(instance_id, &config.descriptor.name, timeout_duration)
+                    .await
+            }
+            Job::RemoveWasmModule { request } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let workflow_data =
+                    create_wasm_removal_workflow_data(*request.clone(), Arc::clone(context));
+
+                let instance_id = engines
+                    .wasm_removal
+                    .start_workflow(WorkflowId::new("wasm_module_removal"), workflow_data)
+                    .await
+                    .map_err(|e| {
+                        format!("Failed to start WASM module removal workflow: {:?}", e)
+                    })?;
+
+                debug!(
+                    "Started WASM module removal workflow for {} (instance: {})",
+                    request.module_uuid, instance_id
+                );
+
+                let timeout_duration = Duration::from_secs(60); // 1 minute
+
+                engines
+                    .wasm_removal
+                    .wait_for_completion(
+                        instance_id,
+                        &request.module_uuid.to_string(),
+                        timeout_duration,
+                    )
+                    .await
+            }
+            Job::InitializeWorkersFromConfig { router_config } => {
+                let api_key = router_config.api_key.clone();
+                let mut worker_count = 0;
+
+                // Create iterator of (url, worker_type, bootstrap_port) tuples based on mode
+                let workers: Vec<(String, &str, Option<u16>)> = match &router_config.mode {
+                    RoutingMode::Regular { worker_urls } => worker_urls
+                        .iter()
+                        .map(|url| (url.clone(), "regular", None))
+                        .collect(),
+                    RoutingMode::PrefillDecode {
+                        prefill_urls,
+                        decode_urls,
+                        ..
+                    } => {
+                        let prefill_workers = prefill_urls
+                            .iter()
+                            .map(|(url, port)| (url.clone(), "prefill", *port));
+
+                        let decode_workers =
+                            decode_urls.iter().map(|url| (url.clone(), "decode", None));
+
+                        prefill_workers.chain(decode_workers).collect()
+                    }
+                    RoutingMode::OpenAI { worker_urls } => {
+                        // OpenAI mode: submit AddWorker jobs with runtime: "external"
+                        // The external_worker_registration workflow handles model discovery
+                        let api_key = router_config.api_key.clone();
+                        let mut submitted_count = 0;
+
+                        for url in worker_urls {
+                            let url_for_error = url.clone();
+                            let config = WorkerConfigRequest {
+                                url: url.clone(),
+                                api_key: api_key.clone(),
+                                worker_type: Some("regular".to_string()),
+                                labels: HashMap::new(),
+                                model_id: None,
+                                priority: None,
+                                cost: None,
+                                runtime: Some("external".to_string()),
+                                tokenizer_path: None,
+                                reasoning_parser: None,
+                                tool_parser: None,
+                                chat_template: router_config.chat_template.clone(),
+                                bootstrap_port: None,
+                                health_check_timeout_secs: router_config.health_check.timeout_secs,
+                                health_check_interval_secs: router_config
+                                    .health_check
+                                    .check_interval_secs,
+                                health_success_threshold: router_config
+                                    .health_check
+                                    .success_threshold,
+                                health_failure_threshold: router_config
+                                    .health_check
+                                    .failure_threshold,
+                                disable_health_check: router_config
+                                    .health_check
+                                    .disable_health_check,
+                                max_connection_attempts: router_config
+                                    .health_check
+                                    .success_threshold
+                                    * 10,
+                                dp_aware: false,
+                            };
+
+                            let job = Job::AddWorker {
+                                config: Box::new(config),
+                            };
+
+                            if let Some(queue) = context.worker_job_queue.get() {
+                                queue.submit(job).await.map_err(|e| {
+                                    format!(
+                                        "Failed to submit AddWorker job for external endpoint {}: {}",
+                                        url_for_error, e
+                                    )
+                                })?;
+                                submitted_count += 1;
+                            } else {
+                                return Err("JobQueue not available".to_string());
+                            }
+                        }
+
+                        if submitted_count == 0 {
+                            info!("OpenAI mode: no worker URLs provided");
+                            return Ok("OpenAI mode: no worker URLs to initialize".to_string());
+                        }
+
+                        return Ok(format!(
+                            "Submitted {} AddWorker jobs for external endpoints",
+                            submitted_count
+                        ));
+                    }
+                };
+
+                debug!(
+                    "Creating AddWorker jobs for {} workers from config",
+                    workers.len()
+                );
+
+                // Process all workers with unified loop
+                for (url, worker_type, bootstrap_port) in workers {
+                    let url_for_error = url.clone(); // Clone for error message
+                    let config = WorkerConfigRequest {
+                        url,
+                        api_key: api_key.clone(),
+                        worker_type: Some(worker_type.to_string()),
+                        labels: HashMap::new(),
+                        model_id: None,
+                        priority: None,
+                        cost: None,
+                        runtime: None,
+                        tokenizer_path: None,
+                        reasoning_parser: None,
+                        tool_parser: None,
+                        chat_template: router_config.chat_template.clone(),
+                        bootstrap_port,
+                        health_check_timeout_secs: router_config.health_check.timeout_secs,
+                        health_check_interval_secs: router_config.health_check.check_interval_secs,
+                        health_success_threshold: router_config.health_check.success_threshold,
+                        health_failure_threshold: router_config.health_check.failure_threshold,
+                        disable_health_check: router_config.health_check.disable_health_check,
+                        max_connection_attempts: router_config.health_check.success_threshold * 10,
+                        dp_aware: router_config.dp_aware,
+                    };
+
+                    let job = Job::AddWorker {
+                        config: Box::new(config),
+                    };
+
+                    if let Some(queue) = context.worker_job_queue.get() {
+                        queue.submit(job).await.map_err(|e| {
+                            format!(
+                                "Failed to submit AddWorker job for {} worker {}: {}",
+                                worker_type, url_for_error, e
+                            )
+                        })?;
+                        worker_count += 1;
+                    } else {
+                        return Err("JobQueue not available".to_string());
+                    }
+                }
+
+                Ok(format!("Submitted {} AddWorker jobs", worker_count))
+            }
+            Job::InitializeMcpServers { mcp_config } => {
+                let mut server_count = 0;
+
+                debug!(
+                    "Creating RegisterMcpServer jobs for {} MCP servers from config",
+                    mcp_config.servers.len()
+                );
+
+                // Submit RegisterMcpServer jobs for each server in the config
+                for server_config in &mcp_config.servers {
+                    let mcp_server_request = McpServerConfigRequest {
+                        name: server_config.name.clone(),
+                        config: server_config.clone(),
+                    };
+
+                    let job = Job::RegisterMcpServer {
+                        config: Box::new(mcp_server_request),
+                    };
+
+                    if let Some(queue) = context.worker_job_queue.get() {
+                        queue.submit(job).await.map_err(|e| {
+                            format!(
+                                "Failed to submit RegisterMcpServer job for '{}': {}",
+                                server_config.name, e
+                            )
+                        })?;
+                        server_count += 1;
+                    } else {
+                        return Err("JobQueue not available".to_string());
+                    }
+                }
+
+                Ok(format!("Submitted {} RegisterMcpServer jobs", server_count))
+            }
+            Job::RegisterMcpServer { config } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let workflow_data =
+                    create_mcp_workflow_data((**config).clone(), Arc::clone(context));
+
+                let instance_id = engines
+                    .mcp
+                    .start_workflow(WorkflowId::new("mcp_registration"), workflow_data)
+                    .await
+                    .map_err(|e| format!("Failed to start MCP registration workflow: {:?}", e))?;
+
+                debug!(
+                    "Started MCP registration workflow for {} (instance: {})",
+                    config.name, instance_id
+                );
+
+                let timeout_duration = Duration::from_secs(7200 + 30); // 2hr + margin
+
+                engines
+                    .mcp
+                    .wait_for_completion(instance_id, &config.name, timeout_duration)
+                    .await
+            }
+            Job::AddTokenizer { config } => {
+                let engines = context
+                    .workflow_engines
+                    .get()
+                    .ok_or_else(|| "Workflow engines not initialized".to_string())?;
+
+                let workflow_data =
+                    create_tokenizer_workflow_data(*config.clone(), Arc::clone(context));
+
+                let instance_id = engines
+                    .tokenizer
+                    .start_workflow(WorkflowId::new("tokenizer_registration"), workflow_data)
+                    .await
+                    .map_err(|e| {
+                        format!("Failed to start tokenizer registration workflow: {:?}", e)
+                    })?;
+
+                debug!(
+                    "Started tokenizer registration workflow for '{}' id={} (instance: {})",
+                    config.name, config.id, instance_id
+                );
+
+                // Allow up to 10 minutes for HuggingFace downloads
+                let timeout_duration = Duration::from_secs(600);
+
+                engines
+                    .tokenizer
+                    .wait_for_completion(instance_id, &config.id, timeout_duration)
+                    .await
+            }
+            Job::RemoveTokenizer { request } => {
+                // Tokenizer removal is synchronous and fast
+                if let Some(entry) = context.tokenizer_registry.remove_by_id(&request.id) {
+                    info!(
+                        "Successfully removed tokenizer '{}' (id: {})",
+                        entry.name, entry.id
+                    );
+                    Ok(format!("Tokenizer '{}' removed successfully", entry.name))
+                } else {
+                    Err(format!("Tokenizer with id '{}' not found", request.id))
+                }
+            }
+        }
+    }
+
+    /// Update job status on completion
+    fn record_job_completion(
+        job_type: &'static str,
+        worker_url: &str,
+        _duration: Duration,
+        result: &Result<String, String>,
+        status_map: &Arc<DashMap<String, JobStatus>>,
+    ) {
+        match result {
+            Ok(message) => {
+                status_map.remove(worker_url);
+                debug!(
+                    "Completed job: type={}, worker={}, result={}",
+                    job_type, worker_url, message
+                );
+            }
+            Err(error) => {
+                status_map.insert(
+                    worker_url.to_string(),
+                    JobStatus::failed(job_type, worker_url, error.clone()),
+                );
+                warn!(
+                    "Failed job: type={}, worker={}, error={}",
+                    job_type, worker_url, error
+                );
+            }
+        }
+    }
+
+    /// Cleanup old job statuses (TTL 5 minutes)
+    async fn cleanup_old_statuses(status_map: Arc<DashMap<String, JobStatus>>) {
+        const CLEANUP_INTERVAL: Duration = Duration::from_secs(60); // Run every minute
+        const STATUS_TTL: u64 = 300; // 5 minutes in seconds
+
+        loop {
+            tokio::time::sleep(CLEANUP_INTERVAL).await;
+
+            let now = SystemTime::now()
+                .duration_since(SystemTime::UNIX_EPOCH)
+                .unwrap()
+                .as_secs();
+
+            // Remove statuses older than TTL
+            status_map.retain(|_key, value| now - value.timestamp < STATUS_TTL);
+
+            debug!(
+                "Cleaned up old job statuses, remaining: {}",
+                status_map.len()
+            );
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/metrics_aggregator.rs

@@ -0,0 +1,91 @@
+use anyhow::ensure;
+use openmetrics_parser::{MetricFamily, MetricsExposition, PrometheusType, PrometheusValue};
+use tracing::warn;
+
+#[derive(Debug)]
+pub struct MetricPack {
+    pub labels: Vec<(String, String)>,
+    pub metrics_text: String,
+}
+
+type PrometheusExposition = MetricsExposition<PrometheusType, PrometheusValue>;
+type PrometheusFamily = MetricFamily<PrometheusType, PrometheusValue>;
+
+/// Aggregate Prometheus metrics scraped from multiple sources into a unified one
+pub fn aggregate_metrics(metric_packs: Vec<MetricPack>) -> anyhow::Result<String> {
+    let mut expositions = vec![];
+    for metric_pack in metric_packs {
+        let metrics_text = &metric_pack.metrics_text;
+        // openmetrics_parser doesn't handle colons in metric names; replace with underscores
+        let metrics_text = metrics_text.replace(":", "_");
+
+        let exposition = match openmetrics_parser::prometheus::parse_prometheus(&metrics_text) {
+            Ok(x) => x,
+            Err(err) => {
+                warn!(
+                    "aggregate_metrics error when parsing text: pack={:?} err={:?}",
+                    metric_pack, err
+                );
+                continue;
+            }
+        };
+        let exposition = transform_metrics(exposition, &metric_pack.labels);
+        expositions.push(exposition);
+    }
+
+    let text = try_reduce(expositions.into_iter(), merge_exposition)?
+        .map(|x| format!("{x}"))
+        .unwrap_or_default();
+    Ok(text)
+}
+
+fn transform_metrics(
+    mut exposition: PrometheusExposition,
+    extra_labels: &[(String, String)],
+) -> PrometheusExposition {
+    for family in exposition.families.values_mut() {
+        *family = family.with_labels(extra_labels.iter().map(|(k, v)| (k.as_str(), v.as_str())));
+    }
+    exposition
+}
+
+fn merge_exposition(
+    a: PrometheusExposition,
+    b: PrometheusExposition,
+) -> anyhow::Result<PrometheusExposition> {
+    let mut ans = a;
+    for (name, family_b) in b.families.into_iter() {
+        let family_merged = if let Some(family_a) = ans.families.remove(&name) {
+            merge_family(family_a, family_b)?
+        } else {
+            family_b
+        };
+        ans.families.insert(name, family_merged);
+    }
+    Ok(ans)
+}
+
+fn merge_family(a: PrometheusFamily, b: PrometheusFamily) -> anyhow::Result<PrometheusFamily> {
+    ensure!(
+        a.get_label_names() == b.get_label_names(),
+        "Label names should agree a={:?} b={:?}",
+        a.get_label_names(),
+        b.get_label_names()
+    );
+    a.with_samples(b.into_iter_samples())
+        .map_err(|e| anyhow::anyhow!("failed to merge samples: {e:?}"))
+}
+
+fn try_reduce<I, T, E, F>(iterable: I, f: F) -> Result<Option<T>, E>
+where
+    I: IntoIterator<Item = T>,
+    F: FnMut(T, T) -> Result<T, E>,
+{
+    let mut it = iterable.into_iter();
+    let first = match it.next() {
+        None => return Ok(None),
+        Some(x) => x,
+    };
+
+    Ok(Some(it.try_fold(first, f)?))
+}

third_party/sglang/sgl-model-gateway/src/core/mod.rs

@@ -0,0 +1,43 @@
+//! Core abstractions for the SGLang router
+//!
+//! This module contains the fundamental types and traits used throughout the router:
+//! - Worker trait and implementations
+//! - Model types and endpoint definitions
+//! - Error types
+//! - Circuit breaker for reliability
+//! - Token buckets for rate limiting
+//! - Workflow steps for multi-step operations
+//! - Common utilities
+
+// Re-export UNKNOWN_MODEL_ID from protocols for use throughout core
+pub use crate::protocols::UNKNOWN_MODEL_ID;
+
+pub mod circuit_breaker;
+pub mod error;
+pub mod job_queue;
+pub mod metrics_aggregator;
+pub mod model_card;
+pub mod model_type;
+pub mod retry;
+pub mod steps;
+pub mod token_bucket;
+pub mod worker;
+pub mod worker_builder;
+pub mod worker_manager;
+pub mod worker_registry;
+pub mod worker_service;
+
+// 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 retry::{is_retryable_status, RetryExecutor};
+pub use worker::{
+    AttachedBody, BasicWorker, ConnectionMode, HealthConfig, RuntimeType, Worker, WorkerLoadGuard,
+    WorkerType,
+};
+pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder};
+pub use worker_manager::{LoadMonitor, WorkerManager};
+pub use worker_registry::{HashRing, WorkerRegistry};
+pub use worker_service::WorkerService;

third_party/sglang/sgl-model-gateway/src/core/model_card.rs

@@ -0,0 +1,395 @@
+//! Model card definitions for worker model configuration.
+//!
+//! This module defines [`ModelCard`] which consolidates model-related configuration
+//! that was previously scattered in `WorkerMetadata.labels` HashMap.
+//!
+//! Also defines [`ProviderType`] for vendor-specific API transformations.
+//!
+//! Inspired by Dynamo's ModelDeploymentCard but simplified for router needs.
+
+use std::collections::HashMap;
+
+use serde::{Deserialize, Serialize};
+
+use super::{
+    model_type::{Endpoint, ModelType},
+    UNKNOWN_MODEL_ID,
+};
+
+/// Provider type for external API transformations.
+///
+/// Different providers have different API formats and requirements.
+/// This enum identifies which vendor's API format to use for transformations.
+///
+/// Note: `None` (when used as `Option<ProviderType>`) means native/passthrough -
+/// no transformation needed. This is the case for local SGLang backends.
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum ProviderType {
+    /// OpenAI API - strip SGLang-specific fields
+    #[serde(alias = "openai")]
+    OpenAI,
+    /// xAI/Grok - special handling for input items
+    #[serde(alias = "xai", alias = "grok")]
+    XAI,
+    /// Anthropic Claude - different API format
+    #[serde(alias = "anthropic", alias = "claude")]
+    Anthropic,
+    /// Google Gemini - special logprobs handling
+    #[serde(alias = "gemini", alias = "google")]
+    Gemini,
+    /// Custom provider with string identifier
+    #[serde(untagged)]
+    Custom(String),
+}
+
+impl ProviderType {
+    /// Get provider name as string
+    pub fn as_str(&self) -> &str {
+        match self {
+            Self::OpenAI => "openai",
+            Self::XAI => "xai",
+            Self::Anthropic => "anthropic",
+            Self::Gemini => "gemini",
+            Self::Custom(s) => s.as_str(),
+        }
+    }
+
+    /// Detect provider from model name (heuristic fallback).
+    /// Returns `None` for models that don't match known external providers
+    /// (i.e., models served by local/native backends).
+    pub fn from_model_name(model: &str) -> Option<Self> {
+        let model_lower = model.to_lowercase();
+        if model_lower.starts_with("grok") {
+            Some(Self::XAI)
+        } else if model_lower.starts_with("gemini") {
+            Some(Self::Gemini)
+        } else if model_lower.starts_with("claude") {
+            Some(Self::Anthropic)
+        } else if model_lower.starts_with("gpt")
+            || model_lower.starts_with("o1")
+            || model_lower.starts_with("o3")
+        {
+            Some(Self::OpenAI)
+        } else {
+            None // Native/local model, no provider transformation needed
+        }
+    }
+}
+
+impl std::fmt::Display for ProviderType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.as_str())
+    }
+}
+
+/// Model card containing model configuration and capabilities.
+///
+/// Consolidates fields previously scattered in `WorkerMetadata.labels`:
+/// - `model_id` -> `id`
+/// - `tokenizer_path` -> `tokenizer_path`
+/// - `chat_template` -> `chat_template`
+/// - `reasoning_parser` -> `reasoning_parser`
+/// - `tool_parser` -> `tool_parser`
+///
+/// # Example
+///
+/// ```
+/// 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")
+///     .with_alias("llama-3.1-8b")
+///     .with_model_type(ModelType::VISION_LLM)
+///     .with_context_length(128_000)
+///     .with_tokenizer_path("meta-llama/Llama-3.1-8B-Instruct");
+///
+/// assert!(card.matches("llama-3.1-8b"));
+/// assert!(card.model_type.supports_vision());
+/// assert!(card.provider.is_none()); // Local model, no external provider
+/// ```
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ModelCard {
+    // === Identity ===
+    /// Primary model ID (e.g., "meta-llama/Llama-3.1-8B-Instruct")
+    /// Previously: labels.get("model_id")
+    pub id: String,
+
+    /// Optional display name (e.g., "Llama 3.1 8B Instruct")
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub display_name: Option<String>,
+
+    /// Alternative names/aliases for this model
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub aliases: Vec<String>,
+
+    // === Capabilities ===
+    /// Supported endpoint types (bitflags)
+    #[serde(default = "default_model_type")]
+    pub model_type: ModelType,
+
+    /// HuggingFace model type string (e.g., "llama", "qwen2", "gpt-oss")
+    /// This is different from `model_type` which is capability bitflags.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub hf_model_type: Option<String>,
+
+    /// Model architectures from HuggingFace config (e.g., ["LlamaForCausalLM"])
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub architectures: Vec<String>,
+
+    /// Provider hint for API transformations.
+    /// `None` means native/passthrough (no transformation needed).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub provider: Option<ProviderType>,
+
+    /// Maximum context length in tokens
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context_length: Option<u32>,
+
+    // === Tokenization & Parsing (previously in labels) ===
+    /// Path to tokenizer (e.g., HuggingFace model ID or local path)
+    /// Previously: labels.get("tokenizer_path")
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub tokenizer_path: Option<String>,
+
+    /// Chat template (Jinja2 template string or path)
+    /// Previously: labels.get("chat_template")
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub chat_template: Option<String>,
+
+    /// Reasoning parser type (e.g., "deepseek", "qwen")
+    /// Previously: labels.get("reasoning_parser")
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub reasoning_parser: Option<String>,
+
+    /// Tool/function calling parser type (e.g., "llama", "mistral")
+    /// Previously: labels.get("tool_parser")
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub tool_parser: Option<String>,
+
+    /// User-defined metadata (for fields not covered above)
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<serde_json::Value>,
+
+    // === Classification Support ===
+    /// Classification label mapping (class index -> label name).
+    /// Empty if not a classification model.
+    /// Example: {0: "negative", 1: "positive"}
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub id2label: HashMap<u32, String>,
+
+    /// Number of classification labels (0 if not a classifier).
+    #[serde(default, skip_serializing_if = "is_zero")]
+    pub num_labels: u32,
+}
+
+fn is_zero(n: &u32) -> bool {
+    *n == 0
+}
+
+fn default_model_type() -> ModelType {
+    ModelType::LLM
+}
+
+impl ModelCard {
+    /// Create a new model card with minimal configuration.
+    ///
+    /// Defaults to `ModelType::LLM` and no provider (native/passthrough).
+    pub fn new(id: impl Into<String>) -> Self {
+        Self {
+            id: id.into(),
+            display_name: None,
+            aliases: Vec::new(),
+            model_type: ModelType::LLM,
+            hf_model_type: None,
+            architectures: Vec::new(),
+            provider: None,
+            context_length: None,
+            tokenizer_path: None,
+            chat_template: None,
+            reasoning_parser: None,
+            tool_parser: None,
+            metadata: None,
+            id2label: HashMap::new(),
+            num_labels: 0,
+        }
+    }
+
+    // === Builder-style methods ===
+
+    /// Set the display name
+    pub fn with_display_name(mut self, name: impl Into<String>) -> Self {
+        self.display_name = Some(name.into());
+        self
+    }
+
+    /// Add a single alias
+    pub fn with_alias(mut self, alias: impl Into<String>) -> Self {
+        self.aliases.push(alias.into());
+        self
+    }
+
+    /// Add multiple aliases
+    pub fn with_aliases(mut self, aliases: impl IntoIterator<Item = impl Into<String>>) -> Self {
+        self.aliases.extend(aliases.into_iter().map(|a| a.into()));
+        self
+    }
+
+    /// Set the model type (capabilities)
+    pub fn with_model_type(mut self, model_type: ModelType) -> Self {
+        self.model_type = model_type;
+        self
+    }
+
+    /// Set the HuggingFace model type string
+    pub fn with_hf_model_type(mut self, hf_model_type: impl Into<String>) -> Self {
+        self.hf_model_type = Some(hf_model_type.into());
+        self
+    }
+
+    /// Set the model architectures
+    pub fn with_architectures(mut self, architectures: Vec<String>) -> Self {
+        self.architectures = architectures;
+        self
+    }
+
+    /// Set the provider type (for external API transformations)
+    pub fn with_provider(mut self, provider: ProviderType) -> Self {
+        self.provider = Some(provider);
+        self
+    }
+
+    /// Set the context length
+    pub fn with_context_length(mut self, length: u32) -> Self {
+        self.context_length = Some(length);
+        self
+    }
+
+    /// Set the tokenizer path
+    pub fn with_tokenizer_path(mut self, path: impl Into<String>) -> Self {
+        self.tokenizer_path = Some(path.into());
+        self
+    }
+
+    /// Set the chat template
+    pub fn with_chat_template(mut self, template: impl Into<String>) -> Self {
+        self.chat_template = Some(template.into());
+        self
+    }
+
+    /// Set the reasoning parser type
+    pub fn with_reasoning_parser(mut self, parser: impl Into<String>) -> Self {
+        self.reasoning_parser = Some(parser.into());
+        self
+    }
+
+    /// Set the tool parser type
+    pub fn with_tool_parser(mut self, parser: impl Into<String>) -> Self {
+        self.tool_parser = Some(parser.into());
+        self
+    }
+
+    /// Set custom metadata
+    pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
+        self.metadata = Some(metadata);
+        self
+    }
+
+    /// Set the id2label mapping for classification models
+    pub fn with_id2label(mut self, id2label: HashMap<u32, String>) -> Self {
+        self.num_labels = id2label.len() as u32;
+        self.id2label = id2label;
+        self
+    }
+
+    /// Set num_labels directly (alternative to with_id2label)
+    pub fn with_num_labels(mut self, num_labels: u32) -> Self {
+        self.num_labels = num_labels;
+        self
+    }
+
+    // === Query methods ===
+
+    /// Check if this model matches the given ID (including aliases)
+    pub fn matches(&self, model_id: &str) -> bool {
+        self.id == model_id || self.aliases.iter().any(|a| a == model_id)
+    }
+
+    /// Check if this model supports a given endpoint
+    pub fn supports_endpoint(&self, endpoint: Endpoint) -> bool {
+        self.model_type.supports_endpoint(endpoint)
+    }
+
+    /// Get the display name or fall back to ID
+    pub fn name(&self) -> &str {
+        self.display_name.as_deref().unwrap_or(&self.id)
+    }
+
+    /// Check if this is a native/local model (no external provider)
+    #[inline]
+    pub fn is_native(&self) -> bool {
+        self.provider.is_none()
+    }
+
+    /// Check if this model uses an external provider
+    #[inline]
+    pub fn has_external_provider(&self) -> bool {
+        self.provider.is_some()
+    }
+
+    /// Check if this is an LLM (supports chat)
+    #[inline]
+    pub fn is_llm(&self) -> bool {
+        self.model_type.is_llm()
+    }
+
+    /// Check if this is an embedding model
+    #[inline]
+    pub fn is_embedding_model(&self) -> bool {
+        self.model_type.is_embedding_model()
+    }
+
+    /// Check if this model supports vision/multimodal
+    #[inline]
+    pub fn supports_vision(&self) -> bool {
+        self.model_type.supports_vision()
+    }
+
+    /// Check if this model supports tools/function calling
+    #[inline]
+    pub fn supports_tools(&self) -> bool {
+        self.model_type.supports_tools()
+    }
+
+    /// Check if this model supports reasoning
+    #[inline]
+    pub fn supports_reasoning(&self) -> bool {
+        self.model_type.supports_reasoning()
+    }
+
+    /// Check if this is a classification model
+    #[inline]
+    pub fn is_classifier(&self) -> bool {
+        self.num_labels > 0
+    }
+
+    /// Get label for a class index, with fallback to generic label (LABEL_N)
+    pub fn get_label(&self, class_idx: u32) -> String {
+        self.id2label
+            .get(&class_idx)
+            .cloned()
+            .unwrap_or_else(|| format!("LABEL_{}", class_idx))
+    }
+}
+
+impl Default for ModelCard {
+    fn default() -> Self {
+        Self::new(UNKNOWN_MODEL_ID)
+    }
+}
+
+impl std::fmt::Display for ModelCard {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/model_type.rs

@@ -0,0 +1,332 @@
+//! Model type definitions using bitflags for endpoint support.
+//!
+//! This module defines [`ModelType`] using bitflags to represent which endpoints
+//! a model can support. This allows combining capabilities like
+//! `ModelType::CHAT | ModelType::COMPLETIONS`.
+//!
+//! Inspired by Dynamo's model_type.rs implementation.
+
+use bitflags::bitflags;
+use serde::{Deserialize, Serialize};
+
+bitflags! {
+    #[derive(Copy, Debug, Default, Clone, Eq, PartialEq, Hash)]
+    pub struct ModelType: u16 {
+        /// OpenAI Chat Completions API (/v1/chat/completions)
+        const CHAT        = 1 << 0;
+        /// OpenAI Completions API - legacy (/v1/completions)
+        const COMPLETIONS = 1 << 1;
+        /// OpenAI Responses API (/v1/responses)
+        const RESPONSES   = 1 << 2;
+        /// Embeddings API (/v1/embeddings)
+        const EMBEDDINGS  = 1 << 3;
+        /// Rerank API (/v1/rerank)
+        const RERANK      = 1 << 4;
+        /// SGLang Generate API (/generate)
+        const GENERATE    = 1 << 5;
+        /// Vision/multimodal support (images in input)
+        const VISION      = 1 << 6;
+        /// Tool/function calling support
+        const TOOLS       = 1 << 7;
+        /// Reasoning/thinking support (e.g., o1, DeepSeek-R1)
+        const REASONING   = 1 << 8;
+        /// Image generation (DALL-E, Sora, gpt-image)
+        const IMAGE_GEN   = 1 << 9;
+        /// Audio models (TTS, Whisper, realtime, transcribe)
+        const AUDIO       = 1 << 10;
+        /// Content moderation models
+        const MODERATION  = 1 << 11;
+
+        /// Standard LLM: chat + completions + responses + tools
+        const LLM = Self::CHAT.bits() | Self::COMPLETIONS.bits()
+                  | Self::RESPONSES.bits() | Self::TOOLS.bits();
+
+        /// Vision-capable LLM: LLM + vision
+        const VISION_LLM = Self::LLM.bits() | Self::VISION.bits();
+
+        /// Reasoning LLM: LLM + reasoning (e.g., o1, o3, DeepSeek-R1)
+        const REASONING_LLM = Self::LLM.bits() | Self::REASONING.bits();
+
+        /// Full-featured LLM: all text generation capabilities
+        const FULL_LLM = Self::VISION_LLM.bits() | Self::REASONING.bits();
+
+        /// Embedding model only
+        const EMBED_MODEL = Self::EMBEDDINGS.bits();
+
+        /// Reranker model only
+        const RERANK_MODEL = Self::RERANK.bits();
+
+        /// Image generation model only (DALL-E, Sora, gpt-image)
+        const IMAGE_MODEL = Self::IMAGE_GEN.bits();
+
+        /// Audio model only (TTS, Whisper, realtime)
+        const AUDIO_MODEL = Self::AUDIO.bits();
+
+        /// Content moderation model only
+        const MODERATION_MODEL = Self::MODERATION.bits();
+    }
+}
+
+/// Mapping of individual capability flags to their names.
+/// Used by `as_capability_names()` for a data-driven approach.
+const CAPABILITY_NAMES: &[(ModelType, &str)] = &[
+    (ModelType::CHAT, "chat"),
+    (ModelType::COMPLETIONS, "completions"),
+    (ModelType::RESPONSES, "responses"),
+    (ModelType::EMBEDDINGS, "embeddings"),
+    (ModelType::RERANK, "rerank"),
+    (ModelType::GENERATE, "generate"),
+    (ModelType::VISION, "vision"),
+    (ModelType::TOOLS, "tools"),
+    (ModelType::REASONING, "reasoning"),
+    (ModelType::IMAGE_GEN, "image_gen"),
+    (ModelType::AUDIO, "audio"),
+    (ModelType::MODERATION, "moderation"),
+];
+
+impl ModelType {
+    /// Check if this model type supports the chat completions endpoint
+    #[inline]
+    pub fn supports_chat(&self) -> bool {
+        self.contains(Self::CHAT)
+    }
+
+    /// Check if this model type supports the legacy completions endpoint
+    #[inline]
+    pub fn supports_completions(&self) -> bool {
+        self.contains(Self::COMPLETIONS)
+    }
+
+    /// Check if this model type supports the responses endpoint
+    #[inline]
+    pub fn supports_responses(&self) -> bool {
+        self.contains(Self::RESPONSES)
+    }
+
+    /// Check if this model type supports the embeddings endpoint
+    #[inline]
+    pub fn supports_embeddings(&self) -> bool {
+        self.contains(Self::EMBEDDINGS)
+    }
+
+    /// Check if this model type supports the rerank endpoint
+    #[inline]
+    pub fn supports_rerank(&self) -> bool {
+        self.contains(Self::RERANK)
+    }
+
+    /// Check if this model type supports the generate endpoint
+    #[inline]
+    pub fn supports_generate(&self) -> bool {
+        self.contains(Self::GENERATE)
+    }
+
+    /// Check if this model type supports vision/multimodal input
+    #[inline]
+    pub fn supports_vision(&self) -> bool {
+        self.contains(Self::VISION)
+    }
+
+    /// Check if this model type supports tool/function calling
+    #[inline]
+    pub fn supports_tools(&self) -> bool {
+        self.contains(Self::TOOLS)
+    }
+
+    /// Check if this model type supports reasoning/thinking
+    #[inline]
+    pub fn supports_reasoning(&self) -> bool {
+        self.contains(Self::REASONING)
+    }
+
+    /// Check if this model type supports image generation
+    #[inline]
+    pub fn supports_image_gen(&self) -> bool {
+        self.contains(Self::IMAGE_GEN)
+    }
+
+    /// Check if this model type supports audio (TTS, Whisper, etc.)
+    #[inline]
+    pub fn supports_audio(&self) -> bool {
+        self.contains(Self::AUDIO)
+    }
+
+    /// Check if this model type supports content moderation
+    #[inline]
+    pub fn supports_moderation(&self) -> bool {
+        self.contains(Self::MODERATION)
+    }
+
+    /// Check if this model type supports a given endpoint
+    pub fn supports_endpoint(&self, endpoint: Endpoint) -> bool {
+        match endpoint {
+            Endpoint::Chat => self.supports_chat(),
+            Endpoint::Completions => self.supports_completions(),
+            Endpoint::Responses => self.supports_responses(),
+            Endpoint::Embeddings => self.supports_embeddings(),
+            Endpoint::Rerank => self.supports_rerank(),
+            Endpoint::Generate => self.supports_generate(),
+            Endpoint::Models => true, // Models endpoint is always supported
+        }
+    }
+
+    /// Convert to a list of supported capability names
+    pub fn as_capability_names(&self) -> Vec<&'static str> {
+        let mut result = Vec::with_capacity(CAPABILITY_NAMES.len());
+        for &(flag, name) in CAPABILITY_NAMES {
+            if self.contains(flag) {
+                result.push(name);
+            }
+        }
+        result
+    }
+
+    /// Check if this is an LLM (supports at least chat)
+    #[inline]
+    pub fn is_llm(&self) -> bool {
+        self.supports_chat()
+    }
+
+    /// Check if this is an embedding model
+    #[inline]
+    pub fn is_embedding_model(&self) -> bool {
+        self.supports_embeddings() && !self.supports_chat()
+    }
+
+    /// Check if this is a reranker model
+    #[inline]
+    pub fn is_reranker(&self) -> bool {
+        self.supports_rerank() && !self.supports_chat()
+    }
+
+    /// Check if this is an image generation model
+    #[inline]
+    pub fn is_image_model(&self) -> bool {
+        self.supports_image_gen() && !self.supports_chat()
+    }
+
+    /// Check if this is an audio model
+    #[inline]
+    pub fn is_audio_model(&self) -> bool {
+        self.supports_audio() && !self.supports_chat()
+    }
+
+    /// Check if this is a moderation model
+    #[inline]
+    pub fn is_moderation_model(&self) -> bool {
+        self.supports_moderation() && !self.supports_chat()
+    }
+}
+
+impl std::fmt::Display for ModelType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let names = self.as_capability_names();
+        if names.is_empty() {
+            write!(f, "none")
+        } else {
+            write!(f, "{}", names.join(","))
+        }
+    }
+}
+
+// Custom Serialize/Deserialize for ModelType to handle bitflags properly
+impl Serialize for ModelType {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        // Serialize as the underlying u16 bits
+        serializer.serialize_u16(self.bits())
+    }
+}
+
+impl<'de> Deserialize<'de> for ModelType {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let bits = u16::deserialize(deserializer)?;
+        ModelType::from_bits(bits)
+            .ok_or_else(|| serde::de::Error::custom(format!("invalid ModelType bits: {}", bits)))
+    }
+}
+
+/// Endpoint types for routing decisions.
+///
+/// This enum represents the different API endpoints that can be routed to workers.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum Endpoint {
+    /// Chat completions endpoint (/v1/chat/completions)
+    Chat,
+    /// Legacy completions endpoint (/v1/completions)
+    Completions,
+    /// Responses endpoint (/v1/responses)
+    Responses,
+    /// Embeddings endpoint (/v1/embeddings)
+    Embeddings,
+    /// Rerank endpoint (/v1/rerank)
+    Rerank,
+    /// SGLang generate endpoint (/generate)
+    Generate,
+    /// Models listing endpoint (/v1/models)
+    Models,
+}
+
+impl Endpoint {
+    /// Get the URL path for this endpoint
+    pub fn path(&self) -> &'static str {
+        match self {
+            Endpoint::Chat => "/v1/chat/completions",
+            Endpoint::Completions => "/v1/completions",
+            Endpoint::Responses => "/v1/responses",
+            Endpoint::Embeddings => "/v1/embeddings",
+            Endpoint::Rerank => "/v1/rerank",
+            Endpoint::Generate => "/generate",
+            Endpoint::Models => "/v1/models",
+        }
+    }
+
+    /// Parse an endpoint from a URL path
+    pub fn from_path(path: &str) -> Option<Self> {
+        // Normalize: strip trailing slash and match
+        let path = path.trim_end_matches('/');
+        match path {
+            "/v1/chat/completions" => Some(Endpoint::Chat),
+            "/v1/completions" => Some(Endpoint::Completions),
+            "/v1/responses" => Some(Endpoint::Responses),
+            "/v1/embeddings" => Some(Endpoint::Embeddings),
+            "/v1/rerank" => Some(Endpoint::Rerank),
+            "/generate" => Some(Endpoint::Generate),
+            "/v1/models" => Some(Endpoint::Models),
+            _ => None,
+        }
+    }
+
+    /// Get the required ModelType flag for this endpoint
+    pub fn required_capability(&self) -> Option<ModelType> {
+        match self {
+            Endpoint::Chat => Some(ModelType::CHAT),
+            Endpoint::Completions => Some(ModelType::COMPLETIONS),
+            Endpoint::Responses => Some(ModelType::RESPONSES),
+            Endpoint::Embeddings => Some(ModelType::EMBEDDINGS),
+            Endpoint::Rerank => Some(ModelType::RERANK),
+            Endpoint::Generate => Some(ModelType::GENERATE),
+            Endpoint::Models => None, // No specific capability required
+        }
+    }
+}
+
+impl std::fmt::Display for Endpoint {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Endpoint::Chat => write!(f, "chat"),
+            Endpoint::Completions => write!(f, "completions"),
+            Endpoint::Responses => write!(f, "responses"),
+            Endpoint::Embeddings => write!(f, "embeddings"),
+            Endpoint::Rerank => write!(f, "rerank"),
+            Endpoint::Generate => write!(f, "generate"),
+            Endpoint::Models => write!(f, "models"),
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/retry.rs

@@ -0,0 +1,320 @@
+use std::time::Duration;
+
+use axum::{http::StatusCode, response::Response};
+use rand::Rng;
+use tracing::debug;
+
+use crate::config::types::RetryConfig;
+
+/// Check if an HTTP status code indicates a retryable error
+pub fn is_retryable_status(status: StatusCode) -> bool {
+    matches!(
+        status,
+        StatusCode::REQUEST_TIMEOUT
+            | StatusCode::TOO_MANY_REQUESTS
+            | StatusCode::INTERNAL_SERVER_ERROR
+            | StatusCode::BAD_GATEWAY
+            | StatusCode::SERVICE_UNAVAILABLE
+            | StatusCode::GATEWAY_TIMEOUT
+    )
+}
+
+/// Computes exponential backoff with optional jitter.
+#[derive(Debug, Clone)]
+pub struct BackoffCalculator;
+
+impl BackoffCalculator {
+    /// Calculate backoff delay for a given attempt index (0-based).
+    pub fn calculate_delay(config: &RetryConfig, attempt: u32) -> Duration {
+        let pow = config.backoff_multiplier.powi(attempt as i32);
+        let delay_ms = ((config.initial_backoff_ms as f32 * pow) as u64).min(config.max_backoff_ms);
+
+        let jitter = config.jitter_factor.clamp(0.0, 1.0);
+        if jitter > 0.0 {
+            let mut rng = rand::rng();
+            let jitter_scale: f32 = rng.random_range(-jitter..=jitter);
+            let jitter_ms = (delay_ms as f32 * jitter_scale)
+                .round()
+                .max(-(delay_ms as f32));
+            let adjusted = (delay_ms as i64 + jitter_ms as i64).max(0) as u64;
+            return Duration::from_millis(adjusted);
+        }
+
+        Duration::from_millis(delay_ms)
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum RetryError {
+    #[error("maximum retry attempts exceeded")]
+    MaxRetriesExceeded,
+}
+
+/// A thin async retry executor for generic operations.
+#[derive(Debug, Clone, Default)]
+pub struct RetryExecutor;
+
+impl RetryExecutor {
+    /// Execute an operation that returns an HTTP Response with retries and backoff.
+    ///
+    /// Usage pattern:
+    /// - `operation(attempt)`: perform one attempt (0-based). Construct and send the request,
+    ///   then return the `Response`. Do any per-attempt bookkeeping (e.g., load tracking,
+    ///   circuit-breaker outcome recording) inside this closure.
+    /// - `should_retry(&response, attempt)`: decide if the given response should be retried
+    ///   (e.g., based on HTTP status). Returning false short-circuits and returns the response.
+    /// - `on_backoff(delay, next_attempt)`: called before sleeping between attempts.
+    ///   Use this to record metrics.
+    /// - `on_exhausted()`: called when the executor has exhausted all retry attempts.
+    ///
+    /// Example:
+    /// ```ignore
+    /// let resp = RetryExecutor::execute_response_with_retry(
+    ///     &retry_cfg,
+    ///     |attempt| async move {
+    ///         let worker = select_cb_aware_worker()?;
+    ///         let resp = send_request(worker).await;
+    ///         worker.record_outcome(resp.status().is_success());
+    ///         resp
+    ///     },
+    ///     |res, _| matches!(res.status(), StatusCode::REQUEST_TIMEOUT | StatusCode::TOO_MANY_REQUESTS | StatusCode::INTERNAL_SERVER_ERROR | StatusCode::BAD_GATEWAY | StatusCode::SERVICE_UNAVAILABLE | StatusCode::GATEWAY_TIMEOUT),
+    ///     |delay, _attempt| { /* record backoff metrics */ },
+    ///     || { /* record retries exhausted */ },
+    /// ).await;
+    /// ```
+    pub async fn execute_response_with_retry<Op, Fut, ShouldRetry, OnBackoff, OnExhausted>(
+        config: &RetryConfig,
+        mut operation: Op,
+        should_retry: ShouldRetry,
+        on_backoff: OnBackoff,
+        mut on_exhausted: OnExhausted,
+    ) -> Response
+    where
+        Op: FnMut(u32) -> Fut,
+        Fut: std::future::Future<Output = Response>,
+        ShouldRetry: Fn(&Response, u32) -> bool,
+        OnBackoff: Fn(Duration, u32),
+        OnExhausted: FnMut(),
+    {
+        let max = config.max_retries.max(1);
+
+        let mut attempt: u32 = 0;
+        loop {
+            let response = operation(attempt).await;
+            let is_last = attempt + 1 >= max;
+
+            if !should_retry(&response, attempt) {
+                return response;
+            }
+
+            if is_last {
+                on_exhausted();
+                return response;
+            }
+
+            let next_attempt = attempt + 1;
+            let delay = BackoffCalculator::calculate_delay(config, attempt);
+            debug!(
+                attempt = attempt,
+                next_attempt = next_attempt,
+                delay_ms = delay.as_millis() as u64,
+                "Retry backoff"
+            );
+            on_backoff(delay, next_attempt);
+            tokio::time::sleep(delay).await;
+
+            attempt = next_attempt;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    };
+
+    use axum::{http::StatusCode, response::IntoResponse};
+
+    use super::*;
+
+    fn base_retry_config() -> RetryConfig {
+        RetryConfig {
+            max_retries: 3,
+            initial_backoff_ms: 1,
+            max_backoff_ms: 4,
+            backoff_multiplier: 2.0,
+            jitter_factor: 0.0,
+        }
+    }
+
+    #[test]
+    fn test_backoff_no_jitter_progression_and_cap() {
+        let cfg = RetryConfig {
+            max_retries: 10,
+            initial_backoff_ms: 100,
+            max_backoff_ms: 250,
+            backoff_multiplier: 2.0,
+            jitter_factor: 0.0,
+        };
+        assert_eq!(
+            BackoffCalculator::calculate_delay(&cfg, 0),
+            Duration::from_millis(100)
+        );
+        assert_eq!(
+            BackoffCalculator::calculate_delay(&cfg, 1),
+            Duration::from_millis(200)
+        );
+        assert_eq!(
+            BackoffCalculator::calculate_delay(&cfg, 2),
+            Duration::from_millis(250)
+        );
+        assert_eq!(
+            BackoffCalculator::calculate_delay(&cfg, 10),
+            Duration::from_millis(250)
+        );
+    }
+
+    #[test]
+    fn test_backoff_with_jitter_within_bounds() {
+        let cfg = RetryConfig {
+            max_retries: 5,
+            initial_backoff_ms: 100,
+            max_backoff_ms: 10_000,
+            backoff_multiplier: 2.0,
+            jitter_factor: 0.5,
+        };
+        let base = 400.0;
+        for _ in 0..50 {
+            let d = BackoffCalculator::calculate_delay(&cfg, 2).as_millis() as f32;
+            assert!(d >= base * 0.5 - 1.0 && d <= base * 1.5 + 1.0);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_execute_response_with_retry_success_path_and_hooks() {
+        let cfg = base_retry_config();
+        let remaining = Arc::new(AtomicU32::new(2));
+        let calls = Arc::new(AtomicU32::new(0));
+        let backoffs = Arc::new(AtomicU32::new(0));
+        let exhausted = Arc::new(AtomicU32::new(0));
+
+        let response = RetryExecutor::execute_response_with_retry(
+            &cfg,
+            {
+                let remaining = remaining.clone();
+                let calls = calls.clone();
+                move |_attempt| {
+                    calls.fetch_add(1, Ordering::Relaxed);
+                    let remaining = remaining.clone();
+                    async move {
+                        if remaining
+                            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |v| v.checked_sub(1))
+                            .is_ok()
+                        {
+                            (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response()
+                        } else {
+                            (StatusCode::OK, "ok").into_response()
+                        }
+                    }
+                }
+            },
+            |res, _attempt| !res.status().is_success(),
+            {
+                let backoffs = backoffs.clone();
+                move |_delay, _next_attempt| {
+                    backoffs.fetch_add(1, Ordering::Relaxed);
+                }
+            },
+            {
+                let exhausted = exhausted.clone();
+                move || {
+                    exhausted.fetch_add(1, Ordering::Relaxed);
+                }
+            },
+        )
+        .await;
+
+        assert_eq!(response.status(), StatusCode::OK);
+        assert_eq!(calls.load(Ordering::Relaxed), 3);
+        assert_eq!(backoffs.load(Ordering::Relaxed), 2);
+        assert_eq!(exhausted.load(Ordering::Relaxed), 0);
+    }
+
+    #[tokio::test]
+    async fn test_execute_response_with_retry_non_retryable_short_circuit() {
+        let cfg = base_retry_config();
+        let calls = Arc::new(AtomicU32::new(0));
+        let backoffs = Arc::new(AtomicU32::new(0));
+        let exhausted = Arc::new(AtomicU32::new(0));
+
+        let response = RetryExecutor::execute_response_with_retry(
+            &cfg,
+            {
+                let calls = calls.clone();
+                move |_attempt| {
+                    calls.fetch_add(1, Ordering::Relaxed);
+                    async move { (StatusCode::BAD_REQUEST, "bad").into_response() }
+                }
+            },
+            |_res, _attempt| false,
+            {
+                let backoffs = backoffs.clone();
+                move |_delay, _next_attempt| {
+                    backoffs.fetch_add(1, Ordering::Relaxed);
+                }
+            },
+            {
+                let exhausted = exhausted.clone();
+                move || {
+                    exhausted.fetch_add(1, Ordering::Relaxed);
+                }
+            },
+        )
+        .await;
+
+        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+        assert_eq!(calls.load(Ordering::Relaxed), 1);
+        assert_eq!(backoffs.load(Ordering::Relaxed), 0);
+        assert_eq!(exhausted.load(Ordering::Relaxed), 0);
+    }
+
+    #[tokio::test]
+    async fn test_execute_response_with_retry_exhausted_hooks() {
+        let cfg = base_retry_config();
+        let calls = Arc::new(AtomicU32::new(0));
+        let backoffs = Arc::new(AtomicU32::new(0));
+        let exhausted = Arc::new(AtomicU32::new(0));
+
+        let response = RetryExecutor::execute_response_with_retry(
+            &cfg,
+            {
+                let calls = calls.clone();
+                move |_attempt| {
+                    calls.fetch_add(1, Ordering::Relaxed);
+                    async move { (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response() }
+                }
+            },
+            |_res, _attempt| true,
+            {
+                let backoffs = backoffs.clone();
+                move |_delay, _next_attempt| {
+                    backoffs.fetch_add(1, Ordering::Relaxed);
+                }
+            },
+            {
+                let exhausted = exhausted.clone();
+                move || {
+                    exhausted.fetch_add(1, Ordering::Relaxed);
+                }
+            },
+        )
+        .await;
+
+        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
+        assert_eq!(calls.load(Ordering::Relaxed), cfg.max_retries);
+        assert_eq!(backoffs.load(Ordering::Relaxed), cfg.max_retries - 1);
+        assert_eq!(exhausted.load(Ordering::Relaxed), 1);
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/mcp_registration.rs

@@ -0,0 +1,326 @@
+use std::{sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use smg_mcp::{config::McpServerConfig, manager::McpManager};
+use tracing::{debug, error, info, warn};
+use wfaas::{
+    BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult,
+    WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
+};
+
+use super::workflow_data::McpWorkflowData;
+use crate::{app_context::AppContext, observability::metrics::Metrics};
+
+/// MCP server connection configuration
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct McpServerConfigRequest {
+    /// Server name (unique identifier)
+    pub name: String,
+    /// Server configuration (transport, proxy, etc.)
+    pub config: McpServerConfig,
+}
+
+impl McpServerConfigRequest {
+    /// Check if this server is required for router startup
+    pub fn is_required(&self) -> bool {
+        self.config.required
+    }
+}
+
+/// Step 1: Connect to MCP server
+///
+/// This step establishes a connection to the MCP server using the flat manager architecture.
+/// The connection is retried aggressively (100 attempts) with a long timeout (2 hours)
+/// to handle slow-starting servers or network issues.
+pub struct ConnectMcpServerStep;
+
+#[async_trait]
+impl StepExecutor<McpWorkflowData> for ConnectMcpServerStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<McpWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config_request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+
+        debug!("Connecting to MCP server: {}", config_request.name);
+
+        // Get proxy config from router_config if available, otherwise fall back to env
+        let proxy_config = app_context
+            .router_config
+            .mcp_config
+            .as_ref()
+            .and_then(|cfg| cfg.proxy.as_ref());
+
+        // Connect to MCP server
+        let client = McpManager::connect_server(&config_request.config, proxy_config)
+            .await
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("connect_mcp_server"),
+                message: format!(
+                    "Failed to connect to MCP server {}: {}",
+                    config_request.name, e
+                ),
+            })?;
+
+        info!(
+            "Successfully connected to MCP server: {}",
+            config_request.name
+        );
+
+        // Store client in typed data
+        context.data.mcp_client = Some(Arc::new(client));
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true // Connection failures are retryable
+    }
+}
+
+/// Step 2: Discover MCP inventory (tools, prompts, resources)
+///
+/// This step queries the MCP server for its capabilities using McpManager::load_server_inventory().
+/// - Tools: Available function calls
+/// - Prompts: Reusable prompt templates
+/// - Resources: Accessible files/data
+pub struct DiscoverMcpInventoryStep;
+
+#[async_trait]
+impl StepExecutor<McpWorkflowData> for DiscoverMcpInventoryStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<McpWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config_request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let mcp_client = context
+            .data
+            .mcp_client
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
+
+        debug!(
+            "Discovering inventory for MCP server: {}",
+            config_request.name
+        );
+
+        // Get shared ToolInventory from McpManager
+        let mcp_manager =
+            app_context
+                .mcp_manager
+                .get()
+                .ok_or_else(|| WorkflowError::StepFailed {
+                    step_id: StepId::new("discover_mcp_inventory"),
+                    message: "MCP manager not initialized".to_string(),
+                })?;
+
+        let inventory = mcp_manager.inventory();
+
+        // Use the public load_server_inventory method
+        McpManager::load_server_inventory(&inventory, &config_request.name, mcp_client).await;
+
+        info!("Completed inventory discovery for {}", config_request.name);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true // Discovery failures are retryable
+    }
+}
+
+/// Step 3: Register MCP server in manager
+///
+/// This step adds the MCP client to the McpManager's client map so it can be
+/// used for tool calls and inventory management.
+pub struct RegisterMcpServerStep;
+
+#[async_trait]
+impl StepExecutor<McpWorkflowData> for RegisterMcpServerStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<McpWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config_request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let mcp_client = context
+            .data
+            .mcp_client
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?
+            .clone();
+
+        debug!("Registering MCP server: {}", config_request.name);
+
+        // Get MCP manager from app context
+        let mcp_manager =
+            app_context
+                .mcp_manager
+                .get()
+                .ok_or_else(|| WorkflowError::StepFailed {
+                    step_id: StepId::new("register_mcp_server"),
+                    message: "MCP manager not initialized".to_string(),
+                })?;
+
+        // Register the client in the manager's client map
+        mcp_manager.register_static_server(config_request.name.clone(), mcp_client);
+
+        // Update active MCP servers metric
+        Metrics::set_mcp_servers_active(mcp_manager.list_servers().len());
+
+        info!("Registered MCP server: {}", config_request.name);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Registration is a simple operation, not retryable
+    }
+}
+
+/// Step 4: Validate registration based on required flag
+///
+/// This step checks if the server is marked as required. If the server is required
+/// but wasn't successfully registered (client not in context), this step fails the workflow.
+/// For optional servers, this step always succeeds, allowing the workflow to complete
+/// even if earlier steps failed.
+pub struct ValidateRegistrationStep;
+
+#[async_trait]
+impl StepExecutor<McpWorkflowData> for ValidateRegistrationStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<McpWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config_request = &context.data.config;
+        let client_registered = context.data.mcp_client.is_some();
+
+        if client_registered {
+            info!(
+                "MCP server '{}' registered successfully",
+                config_request.name
+            );
+
+            // Mark as validated
+            context.data.validated = true;
+
+            return Ok(StepResult::Success);
+        }
+
+        if config_request.is_required() {
+            error!(
+                "Required MCP server '{}' failed to register",
+                config_request.name
+            );
+            Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_registration"),
+                message: format!(
+                    "Required MCP server '{}' failed to register",
+                    config_request.name
+                ),
+            })
+        } else {
+            warn!(
+                "Optional MCP server '{}' failed to register, continuing workflow",
+                config_request.name
+            );
+            Ok(StepResult::Success)
+        }
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}
+
+/// Create MCP server registration workflow
+///
+/// This workflow adapts its failure behavior based on the `required` field in the server config:
+/// - If `required == true`: Uses FailWorkflow - router startup fails if server cannot be reached
+/// - If `required == false` (default): Uses ContinueNextStep - logs warning but continues
+///
+/// Workflow configuration:
+/// - ConnectMcpServer: 100 retries, 2hr timeout (aggressive retry for slow servers)
+/// - DiscoverMcpInventory: 3 retries, 10s timeout (discovery + caching)
+/// - RegisterMcpServer: No retry, 5s timeout (fast registration)
+/// - ValidateRegistration: Final validation step
+pub fn create_mcp_registration_workflow() -> WorkflowDefinition<McpWorkflowData> {
+    WorkflowDefinition::new("mcp_registration", "MCP Server Registration")
+        .add_step(
+            StepDefinition::new(
+                "connect_mcp_server",
+                "Connect to MCP Server",
+                Arc::new(ConnectMcpServerStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 100,
+                backoff: BackoffStrategy::Linear {
+                    increment: Duration::from_secs(1),
+                    max: Duration::from_secs(5),
+                },
+            })
+            .with_timeout(Duration::from_secs(7200)) // 2 hours
+            .with_failure_action(FailureAction::ContinueNextStep),
+        )
+        .add_step(
+            StepDefinition::new(
+                "discover_mcp_inventory",
+                "Discover and Cache MCP Inventory",
+                Arc::new(DiscoverMcpInventoryStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 3,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(1)),
+            })
+            .with_timeout(Duration::from_secs(10))
+            .with_failure_action(FailureAction::ContinueNextStep)
+            .depends_on(&["connect_mcp_server"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "register_mcp_server",
+                "Register MCP Server",
+                Arc::new(RegisterMcpServerStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::ContinueNextStep)
+            .depends_on(&["discover_mcp_inventory"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "validate_registration",
+                "Validate MCP Registration",
+                Arc::new(ValidateRegistrationStep),
+            )
+            .with_timeout(Duration::from_secs(1))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["register_mcp_server"]),
+        )
+}
+
+/// Helper to create initial workflow data for MCP registration
+pub fn create_mcp_workflow_data(
+    config: McpServerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> McpWorkflowData {
+    McpWorkflowData {
+        config,
+        validated: false,
+        app_context: Some(app_context),
+        mcp_client: None,
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/mod.rs

@@ -0,0 +1,87 @@
+//! Workflow step implementations
+//!
+//! This module contains concrete step implementations for various workflows:
+//! - Worker management (registration, removal, updates)
+//! - MCP server registration
+//! - WASM module registration and removal
+//! - Tokenizer registration
+
+pub mod mcp_registration;
+pub mod tokenizer_registration;
+pub mod wasm_module_registration;
+pub mod wasm_module_removal;
+pub mod worker;
+pub mod workflow_data;
+pub mod workflow_engines;
+
+// Worker management (registration, removal)
+pub use mcp_registration::{
+    create_mcp_registration_workflow, create_mcp_workflow_data, ConnectMcpServerStep,
+    DiscoverMcpInventoryStep, McpServerConfigRequest, RegisterMcpServerStep,
+    ValidateRegistrationStep,
+};
+pub use tokenizer_registration::{
+    create_tokenizer_registration_workflow, create_tokenizer_workflow_data, LoadTokenizerStep,
+    TokenizerConfigRequest, TokenizerRemovalRequest,
+};
+pub use wasm_module_registration::{
+    create_wasm_module_registration_workflow, create_wasm_registration_workflow_data,
+    CalculateHashStep, CheckDuplicateStep, LoadWasmBytesStep, RegisterModuleStep,
+    ValidateDescriptorStep, ValidateWasmComponentStep, WasmModuleConfigRequest,
+};
+pub use wasm_module_removal::{
+    create_wasm_module_removal_workflow, create_wasm_removal_workflow_data, FindModuleToRemoveStep,
+    RemoveModuleStep, WasmModuleRemovalRequest,
+};
+pub use worker::{
+    // Workflow builders
+    create_external_worker_workflow,
+    // Workflow data helpers
+    create_external_worker_workflow_data,
+    create_local_worker_workflow,
+    create_local_worker_workflow_data,
+    create_worker_removal_workflow,
+    create_worker_removal_workflow_data,
+    create_worker_update_workflow,
+    create_worker_update_workflow_data,
+    // Utility functions
+    group_models_into_cards,
+    infer_model_type_from_id,
+    // Shared steps
+    ActivateWorkersStep,
+    // External registration steps
+    CreateExternalWorkersStep,
+    // Local registration steps
+    CreateLocalWorkerStep,
+    DetectConnectionModeStep,
+    DiscoverDPInfoStep,
+    DiscoverMetadataStep,
+    DiscoverModelsStep,
+    DpInfo,
+    // Update steps
+    FindWorkerToUpdateStep,
+    // Removal steps
+    FindWorkersToRemoveStep,
+    ModelInfo,
+    ModelsResponse,
+    RegisterWorkersStep,
+    RemoveFromPolicyRegistryStep,
+    RemoveFromWorkerRegistryStep,
+    UpdatePoliciesForWorkerStep,
+    UpdatePoliciesStep,
+    UpdateRemainingPoliciesStep,
+    UpdateWorkerPropertiesStep,
+    WorkerList,
+    WorkerRemovalRequest,
+};
+// Typed workflow data structures
+pub use workflow_data::{
+    ExternalWorkerWorkflowData, LocalWorkerWorkflowData, McpWorkflowData, ProtocolUpdateRequest,
+    TokenizerWorkflowData, WasmRegistrationWorkflowData, WasmRemovalWorkflowData,
+    WorkerConfigRequest, WorkerList as WorkflowWorkerList, WorkerRegistrationData,
+    WorkerRemovalWorkflowData, WorkerUpdateWorkflowData,
+};
+// Typed workflow engines
+pub use workflow_engines::WorkflowEngines;
+
+pub use crate::config::TokenizerCacheConfig;

third_party/sglang/sgl-model-gateway/src/core/steps/tokenizer_registration.rs

@@ -0,0 +1,330 @@
+//! Tokenizer registration workflow
+//!
+//! This module provides a workflow for registering tokenizers asynchronously.
+//! Tokenizers can be loaded from local paths or downloaded from HuggingFace.
+//!
+//! This is the **single source of truth** for tokenizer registration. All paths
+//! (startup, worker connection, API) should use this workflow to ensure consistent
+//! behavior (validation, caching, deduplication).
+
+use std::{sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
+use tracing::{error, info};
+use wfaas::{
+    BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult,
+    WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
+};
+
+use super::workflow_data::TokenizerWorkflowData;
+use crate::{
+    app_context::AppContext,
+    config::TokenizerCacheConfig,
+    tokenizer::{
+        cache::{CacheConfig, CachedTokenizer},
+        factory,
+        registry::LoadOutcome,
+        traits::Tokenizer,
+    },
+};
+
+/// Configuration for adding a tokenizer
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TokenizerConfigRequest {
+    /// Pre-generated UUID for this tokenizer
+    pub id: String,
+    /// User-provided name (what to register under in the registry)
+    pub name: String,
+    /// Source: either a local path or HuggingFace model ID
+    pub source: String,
+    /// Optional path to chat template file
+    pub chat_template_path: Option<String>,
+    /// Optional cache configuration. If provided, wraps tokenizer with CachedTokenizer.
+    #[serde(default)]
+    pub cache_config: Option<TokenizerCacheConfig>,
+    /// If true, the workflow fails when a tokenizer with the same name already exists.
+    /// If false (default), the workflow succeeds and returns the existing tokenizer's ID.
+    /// API callers should set this to true.
+    #[serde(default)]
+    pub fail_on_duplicate: bool,
+}
+
+/// Configuration for removing a tokenizer
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TokenizerRemovalRequest {
+    /// UUID of the tokenizer to remove
+    pub id: String,
+}
+
+// ============================================================================
+// Workflow Steps
+// ============================================================================
+
+/// Load the tokenizer from source (local path or HuggingFace)
+///
+/// This step handles:
+/// - Input validation (via registry.load())
+/// - Deduplication (returns success if already exists)
+/// - Loading from local path or HuggingFace
+/// - Optional caching layer wrapping
+pub struct LoadTokenizerStep;
+
+#[async_trait]
+impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<TokenizerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+
+        info!(
+            "Loading tokenizer '{}' (id: {}) from source: {}{}",
+            config.name,
+            config.id,
+            config.source,
+            if config.cache_config.is_some() {
+                " with caching"
+            } else {
+                ""
+            }
+        );
+
+        // Clone needed values before async move
+        let id = config.id.clone();
+        let name = config.name.clone();
+        let source = config.source.clone();
+        let chat_template = config.chat_template_path.clone();
+        let cache_config = config.cache_config.clone();
+
+        // Load the tokenizer using the registry's load method
+        // This handles: validation, deduplication, and loading
+        let result = app_context
+            .tokenizer_registry
+            .load(&id, &name, &source, || {
+                let source = source.clone();
+                let chat_template = chat_template.clone();
+                let cache_cfg = cache_config.clone();
+                async move {
+                    // Load base tokenizer
+                    let base_tokenizer = factory::create_tokenizer_async_with_chat_template(
+                        &source,
+                        chat_template.as_deref(),
+                    )
+                    .await
+                    .map_err(|e| format!("Failed to load tokenizer: {}", e))?;
+
+                    // Wrap with caching layer if configured
+                    let tokenizer: Arc<dyn Tokenizer> = match cache_cfg {
+                        Some(cfg) if cfg.enable_l0 || cfg.enable_l1 => {
+                            let cache_config = CacheConfig {
+                                enable_l0: cfg.enable_l0,
+                                l0_max_entries: cfg.l0_max_entries,
+                                enable_l1: cfg.enable_l1,
+                                l1_max_memory: cfg.l1_max_memory,
+                            };
+                            Arc::new(CachedTokenizer::new(base_tokenizer, cache_config))
+                        }
+                        _ => base_tokenizer,
+                    };
+
+                    Ok(tokenizer)
+                }
+            })
+            .await;
+
+        match result {
+            Ok(outcome) => {
+                let loaded_id = outcome.id();
+
+                // Get vocab size for logging
+                let vocab_size = app_context
+                    .tokenizer_registry
+                    .get_by_id(loaded_id)
+                    .map(|e| e.tokenizer.vocab_size());
+
+                match &outcome {
+                    LoadOutcome::Loaded { id } => {
+                        info!(
+                            "Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
+                            name, id, vocab_size
+                        );
+                    }
+                    LoadOutcome::AlreadyExists { id } => {
+                        if config.fail_on_duplicate {
+                            return Err(WorkflowError::StepFailed {
+                                step_id: StepId::new("load_tokenizer"),
+                                message: format!(
+                                    "Tokenizer '{}' already exists (id: {})",
+                                    name, id
+                                ),
+                            });
+                        }
+                        info!(
+                            "Tokenizer '{}' already exists (id: {}), skipping load",
+                            name, id
+                        );
+                    }
+                }
+
+                // Store vocab size in typed data
+                if let Some(size) = vocab_size {
+                    context.data.vocab_size = Some(size);
+                }
+
+                Ok(StepResult::Success)
+            }
+            Err(e) => {
+                error!("Failed to load tokenizer '{}': {}", name, e);
+                Err(WorkflowError::StepFailed {
+                    step_id: StepId::new("load_tokenizer"),
+                    message: e.to_string(),
+                })
+            }
+        }
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true // Network/IO errors are retryable
+    }
+}
+
+// ============================================================================
+// Workflow Builder
+// ============================================================================
+
+/// Create the tokenizer registration workflow
+///
+/// This workflow loads and registers a tokenizer. The single LoadTokenizerStep handles:
+/// - Input validation (empty name/source)
+/// - Deduplication (returns success if already exists)
+/// - Loading from local path or HuggingFace
+/// - Optional caching layer wrapping
+///
+/// Configuration:
+/// - 3 retries with 2s backoff (for network issues)
+/// - 5 minute timeout (HuggingFace downloads can be slow)
+pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition<TokenizerWorkflowData> {
+    WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration").add_step(
+        StepDefinition::new(
+            "load_tokenizer",
+            "Load Tokenizer",
+            Arc::new(LoadTokenizerStep),
+        )
+        .with_retry(RetryPolicy {
+            max_attempts: 3,
+            backoff: BackoffStrategy::Fixed(Duration::from_secs(2)),
+        })
+        .with_timeout(Duration::from_secs(300)) // 5 min for HuggingFace downloads
+        .with_failure_action(FailureAction::FailWorkflow),
+    )
+}
+
+/// Helper to create initial workflow data for tokenizer registration
+pub fn create_tokenizer_workflow_data(
+    config: TokenizerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> TokenizerWorkflowData {
+    TokenizerWorkflowData {
+        config,
+        vocab_size: None,
+        app_context: Some(app_context),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_tokenizer_config_request_serialization() {
+        let config = TokenizerConfigRequest {
+            id: "test-uuid-1234".to_string(),
+            name: "test-model".to_string(),
+            source: "meta-llama/Llama-2-7b-hf".to_string(),
+            chat_template_path: None,
+            cache_config: None,
+            fail_on_duplicate: false,
+        };
+
+        let json = serde_json::to_string(&config).unwrap();
+        let parsed: TokenizerConfigRequest = serde_json::from_str(&json).unwrap();
+
+        assert_eq!(parsed.id, "test-uuid-1234");
+        assert_eq!(parsed.name, "test-model");
+        assert_eq!(parsed.source, "meta-llama/Llama-2-7b-hf");
+        assert!(parsed.chat_template_path.is_none());
+        assert!(parsed.cache_config.is_none());
+        assert!(!parsed.fail_on_duplicate);
+    }
+
+    #[test]
+    fn test_tokenizer_config_request_fail_on_duplicate_defaults_to_false() {
+        // Test that fail_on_duplicate defaults to false when not specified in JSON
+        let json = r#"{
+            "id": "test-uuid",
+            "name": "test-model",
+            "source": "/path/to/tokenizer"
+        }"#;
+        let parsed: TokenizerConfigRequest = serde_json::from_str(json).unwrap();
+        assert!(!parsed.fail_on_duplicate);
+    }
+
+    #[test]
+    fn test_tokenizer_config_request_fail_on_duplicate_true() {
+        let config = TokenizerConfigRequest {
+            id: "test-uuid-1234".to_string(),
+            name: "test-model".to_string(),
+            source: "meta-llama/Llama-2-7b-hf".to_string(),
+            chat_template_path: None,
+            cache_config: None,
+            fail_on_duplicate: true,
+        };
+
+        let json = serde_json::to_string(&config).unwrap();
+        let parsed: TokenizerConfigRequest = serde_json::from_str(&json).unwrap();
+        assert!(parsed.fail_on_duplicate);
+    }
+
+    #[test]
+    fn test_tokenizer_config_request_with_cache() {
+        let config = TokenizerConfigRequest {
+            id: "test-uuid-1234".to_string(),
+            name: "test-model".to_string(),
+            source: "meta-llama/Llama-2-7b-hf".to_string(),
+            chat_template_path: None,
+            cache_config: Some(TokenizerCacheConfig {
+                enable_l0: true,
+                l0_max_entries: 1000,
+                enable_l1: false,
+                l1_max_memory: 0,
+            }),
+            fail_on_duplicate: false,
+        };
+
+        let json = serde_json::to_string(&config).unwrap();
+        let parsed: TokenizerConfigRequest = serde_json::from_str(&json).unwrap();
+
+        assert!(parsed.cache_config.is_some());
+        let cache = parsed.cache_config.unwrap();
+        assert!(cache.enable_l0);
+        assert_eq!(cache.l0_max_entries, 1000);
+        assert!(!cache.enable_l1);
+    }
+
+    #[test]
+    fn test_workflow_creation() {
+        let mut workflow = create_tokenizer_registration_workflow();
+        assert_eq!(workflow.id.to_string(), "tokenizer_registration");
+        // Validate the workflow DAG
+        workflow
+            .validate()
+            .expect("Workflow validation should pass");
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/wasm_module_registration.rs

@@ -0,0 +1,637 @@
+use std::{
+    path::{Component as PathComponent, Path},
+    sync::Arc,
+    time::Duration,
+};
+
+use async_trait::async_trait;
+use sha2::{Digest, Sha256};
+use tracing::{debug, info, warn};
+use uuid::Uuid;
+use wasmtime::{component::Component, Config, Engine};
+use wfaas::{
+    BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult,
+    WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
+};
+
+use super::workflow_data::WasmRegistrationWorkflowData;
+use crate::{
+    app_context::AppContext,
+    wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta},
+};
+
+/// WASM module registration request
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct WasmModuleConfigRequest {
+    /// Module descriptor containing name, file_path, attach_points, etc.
+    pub descriptor: WasmModuleDescriptor,
+}
+
+/// Sensitive system directories that WASM modules cannot be loaded from.
+/// These are blocked to prevent information disclosure attacks.
+const BLOCKED_PATH_PREFIXES: &[&str] = &[
+    "/etc/",
+    "/proc/",
+    "/sys/",
+    "/dev/",
+    "/boot/",
+    "/root/",
+    "/var/log/",
+    "/var/run/",
+];
+
+/// Check if a path starts with any blocked prefix.
+/// Returns the matched prefix if found, None otherwise.
+fn find_blocked_prefix(path: &str) -> Option<&'static str> {
+    BLOCKED_PATH_PREFIXES
+        .iter()
+        .find(|&&prefix| path.starts_with(prefix))
+        .copied()
+}
+
+/// Check if a path has a .wasm extension (case-insensitive).
+fn has_wasm_extension(path: &Path) -> bool {
+    path.extension()
+        .is_some_and(|ext| ext.eq_ignore_ascii_case("wasm"))
+}
+
+/// Step 1: Validate module descriptor
+///
+/// Validates that the module descriptor has all required fields:
+/// - Module name is not empty
+/// - File path is not empty
+/// - File exists and is readable
+/// - File size is not zero
+pub struct ValidateDescriptorStep;
+
+#[async_trait]
+impl StepExecutor<WasmRegistrationWorkflowData> for ValidateDescriptorStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRegistrationWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let descriptor = &context.data.config.descriptor;
+
+        debug!("Validating WASM module descriptor: {}", descriptor.name);
+
+        // Validate name
+        if descriptor.name.is_empty() {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: "Module name cannot be empty".to_string(),
+            });
+        }
+
+        // Validate file path
+        if descriptor.file_path.is_empty() {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: "Module file path cannot be empty".to_string(),
+            });
+        }
+
+        // Security: Validate path to prevent path traversal attacks
+        let path = Path::new(&descriptor.file_path);
+
+        // Must be an absolute path
+        if !path.is_absolute() {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: format!(
+                    "Module file path must be absolute, got: {}",
+                    descriptor.file_path
+                ),
+            });
+        }
+
+        // Check for path traversal components (.. or symbolic links that could escape)
+        for component in path.components() {
+            match component {
+                PathComponent::ParentDir => {
+                    warn!(
+                        "Path traversal attempt detected in WASM module path: {}",
+                        descriptor.file_path
+                    );
+                    return Err(WorkflowError::StepFailed {
+                        step_id: StepId::new("validate_descriptor"),
+                        message: "Path traversal (..) not allowed in module file path".to_string(),
+                    });
+                }
+                PathComponent::CurDir => {
+                    return Err(WorkflowError::StepFailed {
+                        step_id: StepId::new("validate_descriptor"),
+                        message: "Current directory (.) not allowed in module file path"
+                            .to_string(),
+                    });
+                }
+                _ => {}
+            }
+        }
+
+        // Require .wasm extension to prevent loading arbitrary files
+        if !has_wasm_extension(path) {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: "Module file must have .wasm extension".to_string(),
+            });
+        }
+
+        // Block access to sensitive system directories
+        if let Some(prefix) = find_blocked_prefix(&descriptor.file_path) {
+            warn!(
+                "Attempt to access blocked directory in WASM module path: {}",
+                descriptor.file_path
+            );
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: format!("Access to {} directory is not allowed", prefix),
+            });
+        }
+
+        // Check if file exists and get size
+        let metadata = tokio::fs::metadata(&descriptor.file_path)
+            .await
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: format!("Failed to access file {}: {}", descriptor.file_path, e),
+            })?;
+
+        if metadata.len() == 0 {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: "Module file size cannot be 0".to_string(),
+            });
+        }
+
+        // Canonicalize the path to resolve symlinks and verify final location is safe
+        let canonical_path = tokio::fs::canonicalize(&descriptor.file_path)
+            .await
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: format!(
+                    "Failed to canonicalize path {}: {}",
+                    descriptor.file_path, e
+                ),
+            })?;
+
+        // Re-check blocked directories after symlink resolution
+        let canonical_str = canonical_path.to_string_lossy();
+        if let Some(prefix) = find_blocked_prefix(&canonical_str) {
+            warn!(
+                "Symlink resolved to blocked directory: {} -> {}",
+                descriptor.file_path, canonical_str
+            );
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: format!(
+                    "Path resolves to blocked directory {} (via symlink)",
+                    prefix
+                ),
+            });
+        }
+
+        // Ensure canonicalized path still has .wasm extension (symlink target check)
+        if !has_wasm_extension(&canonical_path) {
+            warn!(
+                "Symlink target is not a .wasm file: {} -> {}",
+                descriptor.file_path, canonical_str
+            );
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("validate_descriptor"),
+                message: "Symlink target must be a .wasm file".to_string(),
+            });
+        }
+
+        // Clone name for logging before mutable borrow
+        let module_name = descriptor.name.clone();
+
+        // Store file size in typed data
+        context.data.file_size_bytes = Some(metadata.len());
+
+        info!(
+            "Descriptor validated successfully for module: {}",
+            module_name
+        );
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Validation errors are not retryable (invalid input)
+    }
+}
+
+/// Step 2: Calculate SHA256 hash of the module file
+///
+/// Reads the file and calculates its SHA256 hash for deduplication.
+/// This step is I/O intensive and may take time for large files.
+pub struct CalculateHashStep;
+
+#[async_trait]
+impl StepExecutor<WasmRegistrationWorkflowData> for CalculateHashStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRegistrationWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let file_path = &context.data.config.descriptor.file_path;
+
+        debug!("Calculating SHA256 hash for: {}", file_path);
+
+        // Read file in chunks to handle large files efficiently
+        let mut file =
+            tokio::fs::File::open(file_path)
+                .await
+                .map_err(|e| WorkflowError::StepFailed {
+                    step_id: StepId::new("calculate_hash"),
+                    message: format!("Failed to open file {}: {}", file_path, e),
+                })?;
+
+        let mut hasher = Sha256::new();
+        let mut buffer = vec![0u8; 1024 * 1024]; // 1MB buffer
+
+        loop {
+            use tokio::io::AsyncReadExt;
+            let bytes_read =
+                file.read(&mut buffer)
+                    .await
+                    .map_err(|e| WorkflowError::StepFailed {
+                        step_id: StepId::new("calculate_hash"),
+                        message: format!("Failed to read file {}: {}", file_path, e),
+                    })?;
+
+            if bytes_read == 0 {
+                break;
+            }
+
+            hasher.update(&buffer[..bytes_read]);
+        }
+
+        let hash: [u8; 32] = hasher.finalize().into();
+
+        // Clone path for logging before mutable borrow
+        let path_for_log = file_path.clone();
+
+        // Store hash in typed data
+        context.data.sha256_hash = Some(hash);
+
+        info!("SHA256 hash calculated for: {}", path_for_log);
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true // File I/O errors are retryable (network filesystem, etc.)
+    }
+}
+
+/// Step 3: Check for duplicate SHA256 hash
+///
+/// Checks if a module with the same SHA256 hash already exists in the manager.
+/// This prevents duplicate modules from being registered.
+pub struct CheckDuplicateStep;
+
+#[async_trait]
+impl StepExecutor<WasmRegistrationWorkflowData> for CheckDuplicateStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRegistrationWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let sha256_hash = context
+            .data
+            .sha256_hash
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
+
+        debug!(
+            "Checking for duplicate SHA256 hash for module: {}",
+            context.data.config.descriptor.name
+        );
+
+        // Get WASM module manager from app context
+        let wasm_manager =
+            app_context
+                .wasm_manager
+                .as_ref()
+                .ok_or_else(|| WorkflowError::StepFailed {
+                    step_id: StepId::new("check_duplicate"),
+                    message: "WASM module manager not initialized".to_string(),
+                })?;
+
+        // Check for duplicate hash using manager's internal method
+        wasm_manager
+            .check_duplicate_sha256_hash(sha256_hash)
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("check_duplicate"),
+                message: format!("Duplicate SHA256 hash detected: {}", e),
+            })?;
+
+        info!(
+            "No duplicate found for module: {}",
+            context.data.config.descriptor.name
+        );
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Duplicate check failures are not retryable
+    }
+}
+
+/// Step 4: Load WASM bytes into memory
+///
+/// Reads the entire WASM file into memory for faster execution.
+/// This is an I/O operation that may take time for large files.
+pub struct LoadWasmBytesStep;
+
+#[async_trait]
+impl StepExecutor<WasmRegistrationWorkflowData> for LoadWasmBytesStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRegistrationWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let file_path = &context.data.config.descriptor.file_path;
+
+        debug!("Loading WASM bytes from: {}", file_path);
+
+        // Clone path for logging before mutable borrow
+        let path_for_log = file_path.clone();
+
+        let wasm_bytes =
+            tokio::fs::read(file_path)
+                .await
+                .map_err(|e| WorkflowError::StepFailed {
+                    step_id: StepId::new("load_wasm_bytes"),
+                    message: format!("Failed to read WASM file {}: {}", file_path, e),
+                })?;
+
+        // Store WASM bytes in typed data
+        context.data.wasm_bytes = Some(wasm_bytes);
+
+        info!("WASM bytes loaded from: {}", path_for_log);
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true // File read errors are retryable
+    }
+}
+
+/// Step 5: Validate WASM component format
+///
+/// Validates that the loaded WASM bytes represent a valid component.
+/// This catches format errors early during registration rather than during execution.
+pub struct ValidateWasmComponentStep;
+
+#[async_trait]
+impl StepExecutor<WasmRegistrationWorkflowData> for ValidateWasmComponentStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRegistrationWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let wasm_bytes = context
+            .data
+            .wasm_bytes
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?;
+
+        debug!(
+            "Validating WASM component format for module: {}",
+            context.data.config.descriptor.name
+        );
+
+        // Create a temporary engine to validate the component
+        let mut config = Config::new();
+        config.async_support(true);
+        config.wasm_component_model(true);
+
+        let engine = Engine::new(&config).map_err(|e| WorkflowError::StepFailed {
+            step_id: StepId::new("validate_wasm_component"),
+            message: format!("Failed to create WASM engine: {}", e),
+        })?;
+
+        // Attempt to compile the component to validate it
+        Component::new(&engine, wasm_bytes)
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("validate_wasm_component"),
+                message: format!(
+                    "Invalid WASM component: {}. \
+                     Hint: The WASM file must be in component format. \
+                     If you're using wit-bindgen, use 'wasm-tools component new' to wrap the WASM module into a component.",
+                    e
+                ),
+            })?;
+
+        info!(
+            "WASM component validated successfully for module: {}",
+            context.data.config.descriptor.name
+        );
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Validation errors are not retryable (invalid format)
+    }
+}
+
+/// Step 6: Register module in WasmModuleManager
+///
+/// Creates the WasmModule object and registers it in the manager's module map.
+/// This is the final step that makes the module available for execution.
+pub struct RegisterModuleStep;
+
+#[async_trait]
+impl StepExecutor<WasmRegistrationWorkflowData> for RegisterModuleStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRegistrationWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let sha256_hash = context
+            .data
+            .sha256_hash
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
+        let file_size_bytes = context
+            .data
+            .file_size_bytes
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("file_size_bytes".to_string()))?;
+        let wasm_bytes = context
+            .data
+            .wasm_bytes
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?
+            .clone();
+
+        let descriptor = &context.data.config.descriptor;
+
+        debug!("Registering WASM module in manager: {}", descriptor.name);
+
+        // Get WASM module manager from app context
+        let wasm_manager =
+            app_context
+                .wasm_manager
+                .as_ref()
+                .ok_or_else(|| WorkflowError::StepFailed {
+                    step_id: StepId::new("register_module"),
+                    message: "WASM module manager not initialized".to_string(),
+                })?;
+
+        // Create module metadata
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_else(|_| Duration::from_nanos(0))
+            .as_nanos() as u64;
+
+        let module_uuid = Uuid::new_v4();
+
+        let module = WasmModule {
+            module_uuid,
+            module_meta: WasmModuleMeta {
+                name: descriptor.name.clone(),
+                file_path: descriptor.file_path.clone(),
+                sha256_hash,
+                size_bytes: file_size_bytes,
+                created_at: now,
+                last_accessed_at: now,
+                access_count: 0,
+                attach_points: descriptor.attach_points.clone(),
+                wasm_bytes,
+            },
+        };
+
+        // Clone name for logging before mutable borrow
+        let module_name = descriptor.name.clone();
+
+        // Register module in manager
+        wasm_manager
+            .register_module_internal(module)
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("register_module"),
+                message: format!("Failed to register module: {}", e),
+            })?;
+
+        // Store module UUID in typed data
+        context.data.module_uuid = Some(module_uuid);
+
+        info!(
+            "WASM module registered successfully: {} (UUID: {})",
+            module_name, module_uuid
+        );
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Registration is a simple operation, not retryable
+    }
+}
+
+/// Create WASM module registration workflow
+///
+/// This workflow handles the complete process of registering a WASM module:
+/// - Validates the module descriptor
+/// - Calculates SHA256 hash for deduplication
+/// - Checks for duplicates
+/// - Loads WASM bytes into memory
+/// - Validates WASM component format
+/// - Registers the module in the manager
+///
+/// Workflow configuration:
+/// - ValidateDescriptor: No retry, 5s timeout (fast validation)
+/// - CalculateHash: 3 retries, 60s timeout (I/O intensive, may need retry)
+/// - CheckDuplicate: No retry, 5s timeout (fast check)
+/// - LoadWasmBytes: 3 retries, 60s timeout (I/O intensive)
+/// - ValidateWasmComponent: No retry, 30s timeout (CPU intensive validation)
+/// - RegisterModule: No retry, 5s timeout (fast registration)
+pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition<WasmRegistrationWorkflowData>
+{
+    WorkflowDefinition::new("wasm_module_registration", "WASM Module Registration")
+        .add_step(
+            StepDefinition::new(
+                "validate_descriptor",
+                "Validate Descriptor",
+                Arc::new(ValidateDescriptorStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow),
+        )
+        .add_step(
+            StepDefinition::new(
+                "calculate_hash",
+                "Calculate SHA256 Hash",
+                Arc::new(CalculateHashStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 3,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(1)),
+            })
+            .with_timeout(Duration::from_secs(60))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["validate_descriptor"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "check_duplicate",
+                "Check Duplicate Hash",
+                Arc::new(CheckDuplicateStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["calculate_hash"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "load_wasm_bytes",
+                "Load WASM Bytes",
+                Arc::new(LoadWasmBytesStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 3,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(1)),
+            })
+            .with_timeout(Duration::from_secs(60))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["check_duplicate"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "validate_wasm_component",
+                "Validate WASM Component",
+                Arc::new(ValidateWasmComponentStep),
+            )
+            .with_timeout(Duration::from_secs(30))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["load_wasm_bytes"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "register_module",
+                "Register Module",
+                Arc::new(RegisterModuleStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["validate_wasm_component"]),
+        )
+}
+
+/// Helper to create initial workflow data for WASM module registration
+pub fn create_wasm_registration_workflow_data(
+    config: WasmModuleConfigRequest,
+    app_context: Arc<AppContext>,
+) -> WasmRegistrationWorkflowData {
+    WasmRegistrationWorkflowData {
+        config,
+        wasm_bytes: None,
+        sha256_hash: None,
+        file_size_bytes: None,
+        module_uuid: None,
+        app_context: Some(app_context),
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/wasm_module_removal.rs

@@ -0,0 +1,180 @@
+use std::{sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use tracing::{debug, info};
+use uuid::Uuid;
+use wfaas::{
+    FailureAction, StepDefinition, StepExecutor, StepId, StepResult, WorkflowContext,
+    WorkflowDefinition, WorkflowError, WorkflowResult,
+};
+
+use super::workflow_data::WasmRemovalWorkflowData;
+use crate::app_context::AppContext;
+
+/// WASM module removal request
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct WasmModuleRemovalRequest {
+    /// Module UUID to remove
+    pub module_uuid: Uuid,
+    /// Cached UUID string for worker_url() method
+    pub(crate) uuid_string: String,
+}
+
+impl WasmModuleRemovalRequest {
+    pub fn new(module_uuid: Uuid) -> Self {
+        Self {
+            module_uuid,
+            uuid_string: module_uuid.to_string(),
+        }
+    }
+}
+
+/// Step 1: Find module to remove
+///
+/// Verifies that the module exists before attempting removal.
+pub struct FindModuleToRemoveStep;
+
+#[async_trait]
+impl StepExecutor<WasmRemovalWorkflowData> for FindModuleToRemoveStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRemovalWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let removal_request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+
+        debug!("Finding module to remove: {}", removal_request.module_uuid);
+
+        // Get WASM module manager from app context
+        let wasm_manager =
+            app_context
+                .wasm_manager
+                .as_ref()
+                .ok_or_else(|| WorkflowError::StepFailed {
+                    step_id: StepId::new("find_module_to_remove"),
+                    message: "WASM module manager not initialized".to_string(),
+                })?;
+
+        // Check if module exists
+        let module = wasm_manager
+            .get_module(removal_request.module_uuid)
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("find_module_to_remove"),
+                message: format!("Failed to get module: {}", e),
+            })?;
+
+        if module.is_none() {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("find_module_to_remove"),
+                message: format!("Module with UUID {} not found", removal_request.module_uuid),
+            });
+        }
+
+        // Clone uuid for logging before mutable borrow
+        let module_uuid = removal_request.module_uuid;
+
+        // Store the module ID in typed data
+        context.data.module_id = Some(module_uuid.to_string());
+
+        info!("Module found for removal: {}", module_uuid);
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Module not found is not retryable
+    }
+}
+
+/// Step 2: Remove module from WasmModuleManager
+///
+/// Removes the module from the manager's module map.
+pub struct RemoveModuleStep;
+
+#[async_trait]
+impl StepExecutor<WasmRemovalWorkflowData> for RemoveModuleStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WasmRemovalWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let removal_request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+
+        debug!("Removing WASM module: {}", removal_request.module_uuid);
+
+        // Get WASM module manager from app context
+        let wasm_manager =
+            app_context
+                .wasm_manager
+                .as_ref()
+                .ok_or_else(|| WorkflowError::StepFailed {
+                    step_id: StepId::new("remove_module"),
+                    message: "WASM module manager not initialized".to_string(),
+                })?;
+
+        // Remove module from manager
+        wasm_manager
+            .remove_module_internal(removal_request.module_uuid)
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("remove_module"),
+                message: format!("Failed to remove module: {}", e),
+            })?;
+
+        info!(
+            "WASM module removed successfully: {}",
+            removal_request.module_uuid
+        );
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Removal is not retryable
+    }
+}
+
+/// Create WASM module removal workflow
+///
+/// This workflow handles the process of removing a WASM module:
+/// - Finds the module to remove
+/// - Removes it from the manager
+///
+/// Workflow configuration:
+/// - FindModuleToRemove: No retry, 5s timeout (fast lookup)
+/// - RemoveModule: No retry, 5s timeout (fast removal)
+pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition<WasmRemovalWorkflowData> {
+    WorkflowDefinition::new("wasm_module_removal", "WASM Module Removal")
+        .add_step(
+            StepDefinition::new(
+                "find_module_to_remove",
+                "Find Module to Remove",
+                Arc::new(FindModuleToRemoveStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow),
+        )
+        .add_step(
+            StepDefinition::new("remove_module", "Remove Module", Arc::new(RemoveModuleStep))
+                .with_timeout(Duration::from_secs(5))
+                .with_failure_action(FailureAction::FailWorkflow)
+                .depends_on(&["find_module_to_remove"]),
+        )
+}
+
+/// Helper to create initial workflow data for WASM module removal
+pub fn create_wasm_removal_workflow_data(
+    config: WasmModuleRemovalRequest,
+    app_context: Arc<AppContext>,
+) -> WasmRemovalWorkflowData {
+    WasmRemovalWorkflowData {
+        config,
+        module_id: None,
+        app_context: Some(app_context),
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs

@@ -0,0 +1,169 @@
+//! External worker creation step.
+
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use tracing::{debug, info};
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::core::{
+    circuit_breaker::CircuitBreakerConfig,
+    steps::workflow_data::{ExternalWorkerWorkflowData, WorkerList},
+    worker::{HealthConfig, RuntimeType, WorkerType},
+    BasicWorkerBuilder, ConnectionMode, Worker,
+};
+
+/// Normalize URL for external APIs (ensure https://).
+fn normalize_external_url(url: &str) -> String {
+    if url.starts_with("http://") || url.starts_with("https://") {
+        url.to_string()
+    } else {
+        format!("https://{}", url)
+    }
+}
+
+/// Step 2: Create worker objects for each discovered model.
+pub struct CreateExternalWorkersStep;
+
+#[async_trait]
+impl StepExecutor<ExternalWorkerWorkflowData> for CreateExternalWorkersStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<ExternalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let model_cards = &context.data.model_cards;
+
+        // Build configs from router settings
+        let circuit_breaker_config = {
+            let cfg = app_context.router_config.effective_circuit_breaker_config();
+            CircuitBreakerConfig {
+                failure_threshold: cfg.failure_threshold,
+                success_threshold: cfg.success_threshold,
+                timeout_duration: Duration::from_secs(cfg.timeout_duration_secs),
+                window_duration: Duration::from_secs(cfg.window_duration_secs),
+            }
+        };
+
+        let health_config = {
+            let cfg = &app_context.router_config.health_check;
+            HealthConfig {
+                timeout_secs: cfg.timeout_secs,
+                check_interval_secs: cfg.check_interval_secs,
+                endpoint: cfg.endpoint.clone(),
+                failure_threshold: cfg.failure_threshold,
+                success_threshold: cfg.success_threshold,
+                disable_health_check: cfg.disable_health_check || config.disable_health_check,
+            }
+        };
+
+        // Build labels from config
+        let mut labels: HashMap<String, String> = config.labels.clone();
+        if let Some(priority) = config.priority {
+            labels.insert("priority".to_string(), priority.to_string());
+        }
+        if let Some(cost) = config.cost {
+            labels.insert("cost".to_string(), cost.to_string());
+        }
+
+        // Normalize URL (ensure https:// for external APIs)
+        let normalized_url = normalize_external_url(&config.url);
+
+        let mut workers = Vec::new();
+
+        // Handle wildcard mode: create a single worker with empty models list
+        if model_cards.is_empty() {
+            debug!("Creating wildcard worker (no models) for {}", config.url);
+
+            let mut builder = BasicWorkerBuilder::new(normalized_url.clone())
+                .models(vec![]) // Empty models = accepts any model
+                .worker_type(WorkerType::Regular)
+                .connection_mode(ConnectionMode::Http)
+                .runtime_type(RuntimeType::External)
+                .circuit_breaker_config(circuit_breaker_config.clone())
+                .health_config(health_config.clone());
+
+            if let Some(ref api_key) = config.api_key {
+                builder = builder.api_key(api_key.clone());
+            }
+
+            if !labels.is_empty() {
+                builder = builder.labels(labels.clone());
+            }
+
+            let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
+            if health_config.disable_health_check {
+                worker.set_healthy(true);
+            } else {
+                worker.set_healthy(false);
+            }
+
+            info!(
+                "Created wildcard worker at {} (accepts any model, user auth forwarded)",
+                normalized_url
+            );
+
+            workers.push(worker);
+        } else {
+            debug!(
+                "Creating {} external workers for {}",
+                model_cards.len(),
+                config.url
+            );
+
+            // Create a worker for each model
+            for model_card in model_cards.iter() {
+                let mut builder = BasicWorkerBuilder::new(normalized_url.clone())
+                    .model(model_card.clone())
+                    .worker_type(WorkerType::Regular)
+                    .connection_mode(ConnectionMode::Http)
+                    .runtime_type(RuntimeType::External)
+                    .circuit_breaker_config(circuit_breaker_config.clone())
+                    .health_config(health_config.clone());
+
+                if let Some(ref api_key) = config.api_key {
+                    builder = builder.api_key(api_key.clone());
+                }
+
+                if !labels.is_empty() {
+                    builder = builder.labels(labels.clone());
+                }
+
+                let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
+                if health_config.disable_health_check {
+                    worker.set_healthy(true);
+                } else {
+                    worker.set_healthy(false);
+                }
+
+                debug!(
+                    "Created external worker for model {} at {}",
+                    model_card.id, normalized_url
+                );
+
+                workers.push(worker);
+            }
+
+            info!(
+                "Created {} external workers from {}",
+                workers.len(),
+                config.url
+            );
+        }
+
+        // Store results in workflow data
+        context.data.workers = Some(WorkerList::from_workers(&workers));
+        context.data.actual_workers = Some(workers);
+        context.data.labels = labels;
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs

@@ -0,0 +1,274 @@
+//! Model discovery step for external API endpoints.
+
+use std::{collections::HashMap, time::Duration};
+
+use async_trait::async_trait;
+use once_cell::sync::Lazy;
+use regex::Regex;
+use reqwest::Client;
+use serde::Deserialize;
+use tracing::{debug, info};
+use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::core::{
+    model_card::{ModelCard, ProviderType},
+    model_type::ModelType,
+    steps::workflow_data::ExternalWorkerWorkflowData,
+};
+
+// HTTP client for API calls
+static HTTP_CLIENT: Lazy<Client> = Lazy::new(|| {
+    Client::builder()
+        .timeout(Duration::from_secs(30))
+        .build()
+        .expect("Failed to create HTTP client")
+});
+
+// Regex to strip date suffix: -YYYY-MM-DD or -YYYY-MM
+static DATE_SUFFIX_PATTERN: Lazy<Regex> =
+    Lazy::new(|| Regex::new(r"-\d{4}-\d{2}(-\d{2})?$").expect("Invalid date regex"));
+
+/// OpenAI /v1/models response format.
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelsResponse {
+    pub data: Vec<ModelInfo>,
+    #[serde(default)]
+    pub object: String,
+}
+
+/// Individual model information from /v1/models.
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelInfo {
+    pub id: String,
+    #[serde(default)]
+    pub object: String,
+    #[serde(default)]
+    pub created: Option<u64>,
+    #[serde(default)]
+    pub owned_by: Option<String>,
+}
+
+/// Group models by base name (stripping date suffixes) and create ModelCards with aliases.
+///
+/// # Example
+/// Input:  `["gpt-4o", "gpt-4o-2024-05-13", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"]`
+/// Output: `ModelCard { id: "gpt-4o", aliases: ["gpt-4o-2024-05-13", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"] }`
+pub fn group_models_into_cards(models: Vec<ModelInfo>) -> Vec<ModelCard> {
+    // Group model IDs by base name (with date stripped)
+    let mut groups: HashMap<String, Vec<String>> = HashMap::new();
+    for model in &models {
+        let base = DATE_SUFFIX_PATTERN.replace(&model.id, "").to_string();
+        groups.entry(base).or_default().push(model.id.clone());
+    }
+
+    // Create ModelCard for each group
+    groups
+        .into_values()
+        .map(|mut variants| {
+            // Sort: shortest first (base name), then alphabetically
+            variants.sort_by(|a, b| a.len().cmp(&b.len()).then_with(|| a.cmp(b)));
+
+            let primary_id = variants.remove(0); // shortest = primary ID
+            let aliases = variants; // rest = aliases
+
+            let model_type = infer_model_type_from_id(&primary_id);
+            let provider = infer_provider_from_id(&primary_id);
+
+            let mut card = ModelCard::new(&primary_id)
+                .with_aliases(aliases)
+                .with_model_type(model_type);
+
+            if let Some(p) = provider {
+                card = card.with_provider(p);
+            }
+
+            card
+        })
+        .collect()
+}
+
+/// Infer ModelType from model ID string.
+pub fn infer_model_type_from_id(id: &str) -> ModelType {
+    let id_lower = id.to_lowercase();
+
+    // Embedding models
+    if id_lower.contains("embed") || id_lower.contains("ada-002") {
+        return ModelType::EMBED_MODEL;
+    }
+
+    // Rerank models
+    if id_lower.contains("rerank") {
+        return ModelType::RERANK_MODEL;
+    }
+
+    // Image generation models
+    if id_lower.starts_with("dall-e")
+        || id_lower.starts_with("sora")
+        || (id_lower.contains("image") && !id_lower.contains("vision"))
+    {
+        return ModelType::IMAGE_MODEL;
+    }
+
+    // Audio models
+    if id_lower.starts_with("tts")
+        || id_lower.starts_with("whisper")
+        || id_lower.contains("audio")
+        || id_lower.contains("realtime")
+        || id_lower.contains("transcribe")
+    {
+        return ModelType::AUDIO_MODEL;
+    }
+
+    // Moderation models
+    if id_lower.contains("moderation") {
+        return ModelType::MODERATION_MODEL;
+    }
+
+    // Vision LLM
+    if id_lower.contains("vision") || id_lower.contains("4o") {
+        return ModelType::VISION_LLM;
+    }
+
+    // Reasoning models
+    if id_lower.starts_with("o1") || id_lower.starts_with("o3") {
+        return ModelType::REASONING_LLM;
+    }
+
+    // Default to standard LLM
+    ModelType::LLM
+}
+
+/// Infer provider type from model ID string.
+fn infer_provider_from_id(id: &str) -> Option<ProviderType> {
+    let id_lower = id.to_lowercase();
+
+    // OpenAI models
+    if id_lower.starts_with("gpt")
+        || id_lower.starts_with("o1")
+        || id_lower.starts_with("o3")
+        || id_lower.starts_with("dall-e")
+        || id_lower.starts_with("whisper")
+        || id_lower.starts_with("tts")
+        || id_lower.starts_with("text-embedding")
+        || id_lower.starts_with("babbage")
+        || id_lower.starts_with("davinci")
+        || id_lower.contains("omni")
+    {
+        return Some(ProviderType::OpenAI);
+    }
+
+    // xAI/Grok models
+    if id_lower.starts_with("grok") {
+        return Some(ProviderType::XAI);
+    }
+
+    // Anthropic Claude models
+    if id_lower.starts_with("claude") {
+        return Some(ProviderType::Anthropic);
+    }
+
+    // Google Gemini models
+    if id_lower.starts_with("gemini") {
+        return Some(ProviderType::Gemini);
+    }
+
+    None
+}
+
+/// Fetch models from /v1/models endpoint.
+async fn fetch_models(url: &str, api_key: Option<&str>) -> Result<Vec<ModelCard>, String> {
+    let base_url = url.trim_end_matches('/');
+    let models_url = format!("{}/v1/models", base_url);
+
+    let mut req = HTTP_CLIENT.get(&models_url);
+    if let Some(key) = api_key {
+        req = req.bearer_auth(key);
+    }
+
+    let response = req
+        .send()
+        .await
+        .map_err(|e| format!("Failed to connect to {}: {}", models_url, e))?;
+
+    if !response.status().is_success() {
+        return Err(format!(
+            "Server returned status {} from {}",
+            response.status(),
+            models_url
+        ));
+    }
+
+    let models_response: ModelsResponse = response
+        .json()
+        .await
+        .map_err(|e| format!("Failed to parse models response: {}", e))?;
+
+    debug!(
+        "Fetched {} raw models from {}",
+        models_response.data.len(),
+        url
+    );
+
+    let model_cards = group_models_into_cards(models_response.data);
+
+    debug!(
+        "Grouped into {} model cards with aliases",
+        model_cards.len()
+    );
+
+    Ok(model_cards)
+}
+
+/// Step 1: Discover models from external /v1/models endpoint.
+pub struct DiscoverModelsStep;
+
+#[async_trait]
+impl StepExecutor<ExternalWorkerWorkflowData> for DiscoverModelsStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<ExternalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+
+        // If no API key is provided, skip model discovery and use wildcard mode.
+        if config.api_key.as_ref().is_none_or(|k| k.is_empty()) {
+            info!(
+                "No API key provided for {} - using wildcard mode (accepts any model). \
+                 User's Authorization header will be forwarded to backend.",
+                config.url
+            );
+            // Leave model_cards empty for wildcard mode
+            return Ok(StepResult::Success);
+        }
+
+        debug!("Discovering models from external endpoint {}", config.url);
+
+        let model_cards = fetch_models(&config.url, config.api_key.as_deref())
+            .await
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("discover_models"),
+                message: format!("Failed to discover models from {}: {}", config.url, e),
+            })?;
+
+        if model_cards.is_empty() {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("discover_models"),
+                message: format!("No models discovered from {}", config.url),
+            });
+        }
+
+        info!(
+            "Discovered {} models from {}: {:?}",
+            model_cards.len(),
+            config.url,
+            model_cards.iter().map(|c| &c.id).collect::<Vec<_>>()
+        );
+
+        context.data.model_cards = model_cards;
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/external/mod.rs

@@ -0,0 +1,121 @@
+//! External worker registration steps for OpenAI-compatible API endpoints.
+//!
+//! These steps handle the discovery and creation of workers that connect to
+//! external API providers (OpenAI, Anthropic, etc.) via HTTP.
+
+mod create_workers;
+mod discover_models;
+
+use std::{sync::Arc, time::Duration};
+
+pub use create_workers::CreateExternalWorkersStep;
+pub use discover_models::{
+    group_models_into_cards, infer_model_type_from_id, DiscoverModelsStep, ModelInfo,
+    ModelsResponse,
+};
+use wfaas::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition};
+
+use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
+use crate::{
+    app_context::AppContext, core::steps::workflow_data::ExternalWorkerWorkflowData,
+    protocols::worker_spec::WorkerConfigRequest,
+};
+
+/// Create external worker registration workflow definition.
+///
+/// DAG structure with parallel execution opportunities:
+/// ```text
+///              discover_models
+///                    │
+///              create_workers
+///                    │
+///             register_workers
+///                    │
+///       ┌────────────┴────────────┐
+///       │                         │
+///  update_policies         activate_workers
+///       │                         │
+///       └────────────┴────────────┘
+/// ```
+pub fn create_external_worker_workflow() -> WorkflowDefinition<ExternalWorkerWorkflowData> {
+    WorkflowDefinition::new(
+        "external_worker_registration",
+        "External Worker Registration",
+    )
+    // Step 1: Discover models from /v1/models endpoint
+    .add_step(
+        StepDefinition::new(
+            "discover_models",
+            "Discover Models",
+            Arc::new(DiscoverModelsStep),
+        )
+        .with_retry(RetryPolicy {
+            max_attempts: 3,
+            backoff: BackoffStrategy::Exponential {
+                base: Duration::from_secs(1),
+                max: Duration::from_secs(10),
+            },
+        })
+        .with_timeout(Duration::from_secs(30))
+        .with_failure_action(FailureAction::FailWorkflow),
+    )
+    // Step 2: Create workers for each model
+    .add_step(
+        StepDefinition::new(
+            "create_workers",
+            "Create Workers",
+            Arc::new(CreateExternalWorkersStep),
+        )
+        .with_timeout(Duration::from_secs(5))
+        .with_failure_action(FailureAction::FailWorkflow)
+        .depends_on(&["discover_models"]),
+    )
+    // Step 3: Register workers (shared step)
+    .add_step(
+        StepDefinition::new(
+            "register_workers",
+            "Register Workers",
+            Arc::new(RegisterWorkersStep),
+        )
+        .with_timeout(Duration::from_secs(5))
+        .with_failure_action(FailureAction::FailWorkflow)
+        .depends_on(&["create_workers"]),
+    )
+    // Step 4a: Update policies (parallel with activation)
+    .add_step(
+        StepDefinition::new(
+            "update_policies",
+            "Update Policies",
+            Arc::new(UpdatePoliciesStep),
+        )
+        .with_timeout(Duration::from_secs(5))
+        .with_failure_action(FailureAction::ContinueNextStep)
+        .depends_on(&["register_workers"]),
+    )
+    // Step 4b: Activate workers (parallel with policy update)
+    .add_step(
+        StepDefinition::new(
+            "activate_workers",
+            "Activate Workers",
+            Arc::new(ActivateWorkersStep),
+        )
+        .with_timeout(Duration::from_secs(5))
+        .with_failure_action(FailureAction::FailWorkflow)
+        .depends_on(&["register_workers"]),
+    )
+}
+
+/// Helper to create initial workflow data for external worker registration
+pub fn create_external_worker_workflow_data(
+    config: WorkerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> ExternalWorkerWorkflowData {
+    ExternalWorkerWorkflowData {
+        config,
+        model_cards: Vec::new(),
+        workers: None,
+        labels: std::collections::HashMap::new(),
+        app_context: Some(app_context),
+        actual_workers: None,
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs

@@ -0,0 +1,398 @@
+//! Local worker creation step.
+
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::{
+    app_context::AppContext,
+    core::{
+        circuit_breaker::CircuitBreakerConfig,
+        model_card::ModelCard,
+        steps::workflow_data::LocalWorkerWorkflowData,
+        worker::{HealthConfig, RuntimeType, WorkerType},
+        BasicWorkerBuilder, ConnectionMode, DPAwareWorkerBuilder, Worker, UNKNOWN_MODEL_ID,
+    },
+    protocols::worker_spec::WorkerConfigRequest,
+};
+
+/// Step 3: Create worker object(s) with merged configuration + metadata.
+///
+/// This step:
+/// 1. Merges discovered labels with config labels
+/// 2. Determines the model ID from various sources
+/// 3. Creates ModelCard with metadata
+/// 4. Builds worker(s) - either single worker or multiple DP-aware workers
+/// 5. Outputs unified `workers: Vec<Arc<dyn Worker>>` for downstream steps
+pub struct CreateLocalWorkerStep;
+
+#[async_trait]
+impl StepExecutor<LocalWorkerWorkflowData> for CreateLocalWorkerStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<LocalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let connection_mode =
+            context.data.connection_mode.as_ref().ok_or_else(|| {
+                WorkflowError::ContextValueNotFound("connection_mode".to_string())
+            })?;
+        let discovered_labels = &context.data.discovered_labels;
+
+        // Check if worker already exists
+        if app_context
+            .worker_registry
+            .get_by_url(&config.url)
+            .is_some()
+        {
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("create_worker"),
+                message: format!("Worker {} already exists", config.url),
+            });
+        }
+
+        // Build labels from config
+        let mut config_labels = config.labels.clone();
+        if let Some(priority) = config.priority {
+            config_labels.insert("priority".to_string(), priority.to_string());
+        }
+        if let Some(cost) = config.cost {
+            config_labels.insert("cost".to_string(), cost.to_string());
+        }
+
+        // Merge: discovered labels first, then config labels (config takes precedence)
+        let mut final_labels = discovered_labels.clone();
+        for (key, value) in &config_labels {
+            final_labels.insert(key.clone(), value.clone());
+        }
+
+        // Determine model_id: config > served_model_name > model_path > UNKNOWN_MODEL_ID
+        let model_id = config
+            .model_id
+            .clone()
+            .or_else(|| final_labels.get("served_model_name").cloned())
+            .or_else(|| final_labels.get("model_path").cloned())
+            .unwrap_or_else(|| UNKNOWN_MODEL_ID.to_string());
+
+        if model_id != UNKNOWN_MODEL_ID {
+            debug!("Using model_id: {}", model_id);
+        }
+
+        // Create ModelCard
+        let model_card = build_model_card(&model_id, config, &final_labels);
+
+        debug!(
+            "Creating worker {} with {} discovered + {} config = {} final labels",
+            config.url,
+            discovered_labels.len(),
+            config_labels.len(),
+            final_labels.len()
+        );
+
+        // Parse worker type
+        let worker_type = parse_worker_type(config);
+
+        // Get runtime type (for gRPC workers)
+        let runtime_type = determine_runtime_type(connection_mode, &context.data, config);
+
+        // Build circuit breaker config
+        let circuit_breaker_config = build_circuit_breaker_config(app_context);
+
+        // Build health config
+        let health_config = build_health_config(app_context, config);
+
+        // Normalize URL
+        let normalized_url = normalize_url(&config.url, connection_mode);
+
+        if normalized_url != config.url {
+            debug!(
+                "Normalized worker URL: {} -> {} ({:?})",
+                config.url, normalized_url, connection_mode
+            );
+        }
+
+        // Create workers - always output as Vec for unified downstream handling
+        let workers = if config.dp_aware {
+            create_dp_aware_workers(
+                &context.data,
+                &normalized_url,
+                model_card,
+                worker_type,
+                connection_mode,
+                runtime_type,
+                circuit_breaker_config,
+                health_config,
+                config,
+                &final_labels,
+            )?
+        } else {
+            create_single_worker(
+                &normalized_url,
+                model_card,
+                worker_type,
+                connection_mode,
+                runtime_type,
+                circuit_breaker_config,
+                health_config,
+                config,
+                &final_labels,
+            )
+        };
+
+        // Update workflow data
+        context.data.actual_workers = Some(workers);
+        context.data.final_labels = final_labels;
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}
+
+fn build_model_card(
+    model_id: &str,
+    config: &WorkerConfigRequest,
+    labels: &HashMap<String, String>,
+) -> ModelCard {
+    let mut card = ModelCard::new(model_id);
+
+    if let Some(ref tokenizer_path) = config.tokenizer_path {
+        card = card.with_tokenizer_path(tokenizer_path.clone());
+    }
+    if let Some(ref reasoning_parser) = config.reasoning_parser {
+        card = card.with_reasoning_parser(reasoning_parser.clone());
+    }
+    if let Some(ref tool_parser) = config.tool_parser {
+        card = card.with_tool_parser(tool_parser.clone());
+    }
+    if let Some(ref chat_template) = config.chat_template {
+        card = card.with_chat_template(chat_template.clone());
+    }
+    if let Some(model_type_str) = labels.get("model_type") {
+        card = card.with_hf_model_type(model_type_str.clone());
+    }
+    if let Some(architectures_json) = labels.get("architectures") {
+        if let Ok(architectures) = serde_json::from_str::<Vec<String>>(architectures_json) {
+            card = card.with_architectures(architectures);
+        }
+    }
+
+    // Parse classification model id2label mapping
+    // The proto field is id2label_json: JSON string like {"0": "negative", "1": "positive"}
+    if let Some(id2label_json) = labels.get("id2label_json") {
+        if !id2label_json.is_empty() {
+            // Parse JSON: keys are string indices, values are label names
+            if let Ok(string_map) = serde_json::from_str::<HashMap<String, String>>(id2label_json) {
+                // Convert string keys ("0", "1") to u32 keys (0, 1)
+                let id2label: HashMap<u32, String> = string_map
+                    .into_iter()
+                    .filter_map(|(k, v)| k.parse::<u32>().ok().map(|idx| (idx, v)))
+                    .collect();
+
+                if !id2label.is_empty() {
+                    card = card.with_id2label(id2label);
+                    debug!("Parsed id2label with {} classes", card.num_labels);
+                }
+            }
+        }
+    }
+    // Fallback: if num_labels is set but id2label wasn't parsed, create default labels
+    // Match logic in serving_classify.py::_get_id2label_mapping
+    else if let Some(num_labels_str) = labels.get("num_labels") {
+        if let Ok(num_labels) = num_labels_str.parse::<u32>() {
+            if num_labels > 0 {
+                // Create default mapping: {0: "LABEL_0", 1: "LABEL_1", ...}
+                let id2label: HashMap<u32, String> = (0..num_labels)
+                    .map(|i| (i, format!("LABEL_{}", i)))
+                    .collect();
+                card = card.with_id2label(id2label);
+                debug!("Created default id2label with {} classes", num_labels);
+            }
+        }
+    }
+
+    card
+}
+
+fn parse_worker_type(config: &WorkerConfigRequest) -> WorkerType {
+    config
+        .worker_type
+        .as_ref()
+        .map(|t| match t.as_str() {
+            "prefill" => WorkerType::Prefill {
+                bootstrap_port: config.bootstrap_port,
+            },
+            "decode" => WorkerType::Decode,
+            _ => WorkerType::Regular,
+        })
+        .unwrap_or(WorkerType::Regular)
+}
+
+fn determine_runtime_type(
+    connection_mode: &ConnectionMode,
+    data: &LocalWorkerWorkflowData,
+    config: &WorkerConfigRequest,
+) -> RuntimeType {
+    if !matches!(connection_mode, ConnectionMode::Grpc { .. }) {
+        return RuntimeType::Sglang;
+    }
+
+    if let Some(ref detected_runtime) = data.detected_runtime_type {
+        match detected_runtime.as_str() {
+            "vllm" => RuntimeType::Vllm,
+            _ => RuntimeType::Sglang,
+        }
+    } else if let Some(ref runtime) = config.runtime {
+        match runtime.as_str() {
+            "vllm" => RuntimeType::Vllm,
+            _ => RuntimeType::Sglang,
+        }
+    } else {
+        RuntimeType::Sglang
+    }
+}
+
+fn build_circuit_breaker_config(app_context: &AppContext) -> CircuitBreakerConfig {
+    let cfg = app_context.router_config.effective_circuit_breaker_config();
+    CircuitBreakerConfig {
+        failure_threshold: cfg.failure_threshold,
+        success_threshold: cfg.success_threshold,
+        timeout_duration: Duration::from_secs(cfg.timeout_duration_secs),
+        window_duration: Duration::from_secs(cfg.window_duration_secs),
+    }
+}
+
+fn build_health_config(app_context: &AppContext, config: &WorkerConfigRequest) -> HealthConfig {
+    let cfg = &app_context.router_config.health_check;
+    HealthConfig {
+        timeout_secs: cfg.timeout_secs,
+        check_interval_secs: cfg.check_interval_secs,
+        endpoint: cfg.endpoint.clone(),
+        failure_threshold: cfg.failure_threshold,
+        success_threshold: cfg.success_threshold,
+        disable_health_check: cfg.disable_health_check || config.disable_health_check,
+    }
+}
+
+fn normalize_url(url: &str, connection_mode: &ConnectionMode) -> String {
+    if url.starts_with("http://") || url.starts_with("https://") || url.starts_with("grpc://") {
+        url.to_string()
+    } else {
+        match connection_mode {
+            ConnectionMode::Http => format!("http://{}", url),
+            ConnectionMode::Grpc { .. } => format!("grpc://{}", url),
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn create_dp_aware_workers(
+    data: &LocalWorkerWorkflowData,
+    normalized_url: &str,
+    model_card: ModelCard,
+    worker_type: WorkerType,
+    connection_mode: &ConnectionMode,
+    runtime_type: RuntimeType,
+    circuit_breaker_config: CircuitBreakerConfig,
+    health_config: HealthConfig,
+    config: &WorkerConfigRequest,
+    final_labels: &HashMap<String, String>,
+) -> Result<Vec<Arc<dyn Worker>>, WorkflowError> {
+    let dp_info = data
+        .dp_info
+        .as_ref()
+        .ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?;
+
+    debug!(
+        "Creating {} DP-aware workers for {} (dp_size: {})",
+        dp_info.dp_size, normalized_url, dp_info.dp_size
+    );
+
+    let mut workers = Vec::with_capacity(dp_info.dp_size);
+    for rank in 0..dp_info.dp_size {
+        let mut builder =
+            DPAwareWorkerBuilder::new(normalized_url.to_string(), rank, dp_info.dp_size)
+                .model(model_card.clone())
+                .worker_type(worker_type.clone())
+                .connection_mode(connection_mode.clone())
+                .runtime_type(runtime_type.clone())
+                .circuit_breaker_config(circuit_breaker_config.clone())
+                .health_config(health_config.clone());
+
+        if let Some(ref api_key) = config.api_key {
+            builder = builder.api_key(api_key.clone());
+        }
+        if !final_labels.is_empty() {
+            builder = builder.labels(final_labels.clone());
+        }
+
+        let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
+        if health_config.disable_health_check {
+            worker.set_healthy(true);
+        } else {
+            worker.set_healthy(false);
+        }
+        workers.push(worker);
+
+        debug!(
+            "Created DP-aware worker {}@{}/{} ({:?})",
+            normalized_url, rank, dp_info.dp_size, connection_mode
+        );
+    }
+
+    Ok(workers)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn create_single_worker(
+    normalized_url: &str,
+    model_card: ModelCard,
+    worker_type: WorkerType,
+    connection_mode: &ConnectionMode,
+    runtime_type: RuntimeType,
+    circuit_breaker_config: CircuitBreakerConfig,
+    health_config: HealthConfig,
+    config: &WorkerConfigRequest,
+    final_labels: &HashMap<String, String>,
+) -> Vec<Arc<dyn Worker>> {
+    let health_check_disabled = health_config.disable_health_check;
+
+    let mut builder = BasicWorkerBuilder::new(normalized_url.to_string())
+        .model(model_card)
+        .worker_type(worker_type)
+        .connection_mode(connection_mode.clone())
+        .runtime_type(runtime_type)
+        .circuit_breaker_config(circuit_breaker_config)
+        .health_config(health_config);
+
+    if let Some(ref api_key) = config.api_key {
+        builder = builder.api_key(api_key.clone());
+    }
+    if !final_labels.is_empty() {
+        builder = builder.labels(final_labels.clone());
+    }
+
+    let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
+    if health_check_disabled {
+        worker.set_healthy(true);
+    } else {
+        worker.set_healthy(false);
+    }
+
+    debug!(
+        "Created worker object for {} ({:?}) with {} labels",
+        normalized_url,
+        connection_mode,
+        final_labels.len()
+    );
+
+    vec![worker]
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs

@@ -0,0 +1,144 @@
+//! Connection mode detection step.
+
+use std::time::Duration;
+
+use async_trait::async_trait;
+use reqwest::Client;
+use tracing::debug;
+use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use super::strip_protocol;
+use crate::{
+    core::{steps::workflow_data::LocalWorkerWorkflowData, ConnectionMode},
+    routers::grpc::client::GrpcClient,
+};
+
+/// Try HTTP health check.
+async fn try_http_health_check(
+    url: &str,
+    timeout_secs: u64,
+    client: &Client,
+) -> Result<(), String> {
+    let is_https = url.starts_with("https://");
+    let protocol = if is_https { "https" } else { "http" };
+    let clean_url = strip_protocol(url);
+    let health_url = format!("{}://{}/health", protocol, clean_url);
+
+    client
+        .get(&health_url)
+        .timeout(Duration::from_secs(timeout_secs))
+        .send()
+        .await
+        .and_then(reqwest::Response::error_for_status)
+        .map_err(|e| format!("Health check failed: {}", e))?;
+
+    Ok(())
+}
+
+/// Perform gRPC health check with runtime type.
+async fn do_grpc_health_check(
+    grpc_url: &str,
+    timeout_secs: u64,
+    runtime_type: &str,
+) -> Result<(), String> {
+    let connect_future = GrpcClient::connect(grpc_url, runtime_type);
+    let client = tokio::time::timeout(Duration::from_secs(timeout_secs), connect_future)
+        .await
+        .map_err(|_| "gRPC connection timeout".to_string())?
+        .map_err(|e| format!("gRPC connection failed: {}", e))?;
+
+    let health_future = client.health_check();
+    tokio::time::timeout(Duration::from_secs(timeout_secs), health_future)
+        .await
+        .map_err(|_| "gRPC health check timeout".to_string())?
+        .map_err(|e| format!("gRPC health check failed: {}", e))?;
+
+    Ok(())
+}
+
+/// Try gRPC health check (tries SGLang first, then vLLM if not specified).
+async fn try_grpc_health_check(
+    url: &str,
+    timeout_secs: u64,
+    runtime_type: Option<&str>,
+) -> Result<(), String> {
+    let grpc_url = if url.starts_with("grpc://") {
+        url.to_string()
+    } else {
+        format!("grpc://{}", strip_protocol(url))
+    };
+
+    match runtime_type {
+        Some(runtime) => do_grpc_health_check(&grpc_url, timeout_secs, runtime).await,
+        None => {
+            // Try SGLang first, then vLLM as fallback
+            if let Ok(()) = do_grpc_health_check(&grpc_url, timeout_secs, "sglang").await {
+                return Ok(());
+            }
+            do_grpc_health_check(&grpc_url, timeout_secs, "vllm")
+                .await
+                .map_err(|e| format!("gRPC failed (tried SGLang and vLLM): {}", e))
+        }
+    }
+}
+
+/// Step 1: Detect connection mode by probing HTTP and gRPC.
+pub struct DetectConnectionModeStep;
+
+#[async_trait]
+impl StepExecutor<LocalWorkerWorkflowData> for DetectConnectionModeStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<LocalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+
+        debug!(
+            "Detecting connection mode for {} (timeout: {}s, max_attempts: {})",
+            config.url, config.health_check_timeout_secs, config.max_connection_attempts
+        );
+
+        // Try both protocols in parallel
+        let url = config.url.clone();
+        let timeout = config.health_check_timeout_secs;
+        let client = &app_context.client;
+        let runtime_type = config.runtime.as_deref();
+
+        let (http_result, grpc_result) = tokio::join!(
+            try_http_health_check(&url, timeout, client),
+            try_grpc_health_check(&url, timeout, runtime_type)
+        );
+
+        let connection_mode = match (http_result, grpc_result) {
+            (Ok(_), _) => {
+                debug!("{} detected as HTTP", config.url);
+                ConnectionMode::Http
+            }
+            (_, Ok(_)) => {
+                debug!("{} detected as gRPC", config.url);
+                ConnectionMode::Grpc { port: None }
+            }
+            (Err(http_err), Err(grpc_err)) => {
+                return Err(WorkflowError::StepFailed {
+                    step_id: StepId::new("detect_connection_mode"),
+                    message: format!(
+                        "Both HTTP and gRPC health checks failed for {}: HTTP: {}, gRPC: {}",
+                        config.url, http_err, grpc_err
+                    ),
+                });
+            }
+        };
+
+        context.data.connection_mode = Some(connection_mode);
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs

@@ -0,0 +1,78 @@
+//! Data Parallel (DP) information discovery step.
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use super::discover_metadata::get_server_info;
+use crate::core::{steps::workflow_data::LocalWorkerWorkflowData, UNKNOWN_MODEL_ID};
+
+/// DP (Data Parallel) information for a worker.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct DpInfo {
+    pub dp_size: usize,
+    pub model_id: String,
+}
+
+/// Get DP info for a worker URL.
+pub async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result<DpInfo, String> {
+    let info = get_server_info(url, api_key).await?;
+
+    let dp_size = info
+        .dp_size
+        .ok_or_else(|| format!("No dp_size in response from {}", url))?;
+
+    let model_id = info
+        .model_id
+        .filter(|s| !s.is_empty())
+        .or(info.served_model_name.filter(|s| !s.is_empty()))
+        .or_else(|| {
+            info.model_path
+                .and_then(|path| path.split('/').next_back().map(|s| s.to_string()))
+        })
+        .unwrap_or_else(|| UNKNOWN_MODEL_ID.to_string());
+
+    Ok(DpInfo { dp_size, model_id })
+}
+
+/// Step 2b: Discover DP (Data Parallel) information (only for DP-aware workers).
+pub struct DiscoverDPInfoStep;
+
+#[async_trait]
+impl StepExecutor<LocalWorkerWorkflowData> for DiscoverDPInfoStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<LocalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+
+        if !config.dp_aware {
+            debug!(
+                "Worker {} is not DP-aware, skipping DP discovery",
+                config.url
+            );
+            return Ok(StepResult::Success);
+        }
+
+        debug!("Discovering DP info for {} (DP-aware)", config.url);
+
+        let dp_info = get_dp_info(&config.url, config.api_key.as_deref())
+            .await
+            .map_err(|e| WorkflowError::StepFailed {
+                step_id: StepId::new("discover_dp_info"),
+                message: format!("Failed to get DP info: {}", e),
+            })?;
+
+        debug!(
+            "Discovered DP size {} for {} (model: {})",
+            dp_info.dp_size, config.url, dp_info.model_id
+        );
+
+        context.data.dp_info = Some(dp_info);
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs

@@ -0,0 +1,320 @@
+//! Metadata discovery step for local workers.
+
+use std::{collections::HashMap, time::Duration};
+
+use async_trait::async_trait;
+use once_cell::sync::Lazy;
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use tracing::{debug, warn};
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use super::strip_protocol;
+use crate::{
+    core::{steps::workflow_data::LocalWorkerWorkflowData, ConnectionMode},
+    routers::grpc::client::GrpcClient,
+};
+
+// HTTP client for metadata fetching
+static HTTP_CLIENT: Lazy<Client> = Lazy::new(|| {
+    Client::builder()
+        .timeout(Duration::from_secs(10))
+        .build()
+        .expect("Failed to create HTTP client")
+});
+
+/// Server information returned from /server_info endpoint.
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct ServerInfo {
+    #[serde(alias = "model")]
+    pub model_id: Option<String>,
+    pub model_path: Option<String>,
+    pub served_model_name: Option<String>,
+    pub tp_size: Option<usize>,
+    pub dp_size: Option<usize>,
+    pub load_balance_method: Option<String>,
+    pub disaggregation_mode: Option<String>,
+    pub version: Option<String>,
+    pub max_batch_size: Option<usize>,
+    pub max_total_tokens: Option<usize>,
+    pub max_prefill_tokens: Option<usize>,
+    pub max_running_requests: Option<usize>,
+    pub max_num_reqs: Option<usize>,
+}
+
+/// Model information returned from /model_info endpoint.
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct ModelInfo {
+    pub model_path: Option<String>,
+    pub tokenizer_path: Option<String>,
+    pub is_generation: Option<bool>,
+    pub model_type: Option<String>,
+    pub architectures: Option<Vec<String>>,
+}
+
+/// Fallback function to GET JSON from old endpoint (with "get_" prefix) for backward compatibility.
+async fn get_json_fallback(
+    base_url: &str,
+    endpoint: &str,
+    api_key: Option<&str>,
+) -> Result<Value, String> {
+    // FIXME: This fallback logic should be removed together with /get_server_info
+    // and /get_model_info endpoints in http_server.py
+    warn!(
+        concat!(
+            "Endpoint '/{}' returned 404, falling back to '/get_{}' for backward compatibility. ",
+            "The '/get_{}' endpoint is deprecated and will be removed in a future version. ",
+            "Please use '/{}' instead."
+        ),
+        endpoint, endpoint, endpoint, endpoint
+    );
+
+    let old_url = format!("{}/get_{}", base_url, endpoint);
+    let mut req = HTTP_CLIENT.get(&old_url);
+    if let Some(key) = api_key {
+        req = req.bearer_auth(key);
+    }
+
+    let response = req
+        .send()
+        .await
+        .map_err(|e| format!("Failed to connect to {}: {}", old_url, e))?;
+
+    if !response.status().is_success() {
+        return Err(format!(
+            "Server returned status {} from {}",
+            response.status(),
+            old_url
+        ));
+    }
+
+    response
+        .json::<Value>()
+        .await
+        .map_err(|e| format!("Failed to parse response from {}: {}", old_url, e))
+}
+
+/// Get server info from /server_info endpoint.
+pub async fn get_server_info(url: &str, api_key: Option<&str>) -> Result<ServerInfo, String> {
+    let base_url = url.trim_end_matches('/');
+    let server_info_url = format!("{}/server_info", base_url);
+
+    let mut req = HTTP_CLIENT.get(&server_info_url);
+    if let Some(key) = api_key {
+        req = req.bearer_auth(key);
+    }
+
+    let response = req
+        .send()
+        .await
+        .map_err(|e| format!("Failed to connect to {}: {}", server_info_url, e))?;
+
+    // If /server_info returns 404, fallback to /get_server_info for backward compatibility
+    if response.status() == reqwest::StatusCode::NOT_FOUND {
+        let json = get_json_fallback(base_url, "server_info", api_key).await?;
+        return serde_json::from_value(json)
+            .map_err(|e| format!("Failed to parse server info: {}", e));
+    }
+
+    if !response.status().is_success() {
+        return Err(format!(
+            "Server returned status {} from {}",
+            response.status(),
+            server_info_url
+        ));
+    }
+
+    let json = response
+        .json::<Value>()
+        .await
+        .map_err(|e| format!("Failed to parse response from {}: {}", server_info_url, e))?;
+
+    serde_json::from_value(json).map_err(|e| format!("Failed to parse server info: {}", e))
+}
+
+/// Get model info from /model_info endpoint.
+pub async fn get_model_info(url: &str, api_key: Option<&str>) -> Result<ModelInfo, String> {
+    let base_url = url.trim_end_matches('/');
+    let model_info_url = format!("{}/model_info", base_url);
+
+    let mut req = HTTP_CLIENT.get(&model_info_url);
+    if let Some(key) = api_key {
+        req = req.bearer_auth(key);
+    }
+
+    let response = req
+        .send()
+        .await
+        .map_err(|e| format!("Failed to connect to {}: {}", model_info_url, e))?;
+
+    // If /model_info returns 404, fallback to /get_model_info for backward compatibility
+    if response.status() == reqwest::StatusCode::NOT_FOUND {
+        let json = get_json_fallback(base_url, "model_info", api_key).await?;
+        return serde_json::from_value(json)
+            .map_err(|e| format!("Failed to parse model info: {}", e));
+    }
+
+    if !response.status().is_success() {
+        return Err(format!(
+            "Server returned status {} from {}",
+            response.status(),
+            model_info_url
+        ));
+    }
+
+    response
+        .json::<ModelInfo>()
+        .await
+        .map_err(|e| format!("Failed to parse response from {}: {}", model_info_url, e))
+}
+
+/// Fetch gRPC metadata (returns labels and detected runtime type).
+async fn fetch_grpc_metadata(
+    url: &str,
+    runtime_type: Option<&str>,
+) -> Result<(HashMap<String, String>, String), String> {
+    let grpc_url = if url.starts_with("grpc://") {
+        url.to_string()
+    } else {
+        format!("grpc://{}", strip_protocol(url))
+    };
+
+    async fn do_fetch(
+        grpc_url: &str,
+        runtime_type: &str,
+    ) -> Result<HashMap<String, String>, String> {
+        let client = GrpcClient::connect(grpc_url, runtime_type)
+            .await
+            .map_err(|e| format!("Failed to connect to gRPC: {}", e))?;
+
+        let model_info = client
+            .get_model_info()
+            .await
+            .map_err(|e| format!("Failed to fetch gRPC metadata: {}", e))?;
+
+        Ok(model_info.to_labels())
+    }
+
+    match runtime_type {
+        Some(runtime) => {
+            let labels = do_fetch(&grpc_url, runtime).await?;
+            Ok((labels, runtime.to_string()))
+        }
+        None => {
+            // Try SGLang first, then vLLM as fallback
+            if let Ok(labels) = do_fetch(&grpc_url, "sglang").await {
+                return Ok((labels, "sglang".to_string()));
+            }
+            let labels = do_fetch(&grpc_url, "vllm")
+                .await
+                .map_err(|e| format!("gRPC metadata failed (tried SGLang and vLLM): {}", e))?;
+            Ok((labels, "vllm".to_string()))
+        }
+    }
+}
+
+/// Step 2a: Discover metadata from worker.
+pub struct DiscoverMetadataStep;
+
+#[async_trait]
+impl StepExecutor<LocalWorkerWorkflowData> for DiscoverMetadataStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<LocalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let config = &context.data.config;
+        let connection_mode =
+            context.data.connection_mode.as_ref().ok_or_else(|| {
+                WorkflowError::ContextValueNotFound("connection_mode".to_string())
+            })?;
+
+        debug!(
+            "Discovering metadata for {} ({:?})",
+            config.url, connection_mode
+        );
+
+        let (discovered_labels, detected_runtime) = match connection_mode {
+            ConnectionMode::Http => {
+                let mut labels = HashMap::new();
+
+                // Fetch from /server_info for server-related metadata
+                if let Ok(server_info) =
+                    get_server_info(&config.url, config.api_key.as_deref()).await
+                {
+                    if let Some(model_path) = server_info.model_path.filter(|s| !s.is_empty()) {
+                        labels.insert("model_path".to_string(), model_path);
+                    }
+                    if let Some(served_model_name) =
+                        server_info.served_model_name.filter(|s| !s.is_empty())
+                    {
+                        labels.insert("served_model_name".to_string(), served_model_name);
+                    }
+                    if let Some(tp_size) = server_info.tp_size {
+                        labels.insert("tp_size".to_string(), tp_size.to_string());
+                    }
+                    if let Some(dp_size) = server_info.dp_size {
+                        labels.insert("dp_size".to_string(), dp_size.to_string());
+                    }
+                    if let Some(load_balance_method) = server_info.load_balance_method {
+                        labels.insert("load_balance_method".to_string(), load_balance_method);
+                    }
+                    if let Some(disaggregation_mode) = server_info.disaggregation_mode {
+                        labels.insert("disaggregation_mode".to_string(), disaggregation_mode);
+                    }
+                }
+
+                // Fetch from /model_info for model-related metadata
+                if let Ok(model_info) = get_model_info(&config.url, config.api_key.as_deref()).await
+                {
+                    if let Some(tokenizer_path) =
+                        model_info.tokenizer_path.filter(|s| !s.is_empty())
+                    {
+                        labels.insert("tokenizer_path".to_string(), tokenizer_path);
+                    }
+                    if let Some(model_type) = model_info.model_type.filter(|s| !s.is_empty()) {
+                        labels.insert("model_type".to_string(), model_type);
+                    }
+                    if let Some(architectures) = model_info.architectures.filter(|a| !a.is_empty())
+                    {
+                        if let Ok(json_str) = serde_json::to_string(&architectures) {
+                            labels.insert("architectures".to_string(), json_str);
+                        }
+                    }
+                }
+
+                Ok((labels, None))
+            }
+            ConnectionMode::Grpc { .. } => {
+                let runtime_type = config.runtime.as_deref();
+                fetch_grpc_metadata(&config.url, runtime_type)
+                    .await
+                    .map(|(labels, runtime)| (labels, Some(runtime)))
+            }
+        }
+        .unwrap_or_else(|e| {
+            warn!("Failed to fetch metadata for {}: {}", config.url, e);
+            (HashMap::new(), None)
+        });
+
+        let url = config.url.clone();
+        debug!(
+            "Discovered {} metadata labels for {}",
+            discovered_labels.len(),
+            url
+        );
+
+        // Update workflow data
+        context.data.discovered_labels = discovered_labels;
+        if let Some(runtime) = detected_runtime {
+            debug!("Detected runtime type: {}", runtime);
+            context.data.detected_runtime_type = Some(runtime);
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        true
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs

@@ -0,0 +1,59 @@
+//! Step to find a worker to update based on URL.
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use super::find_workers_by_url;
+use crate::core::steps::workflow_data::WorkerUpdateWorkflowData;
+
+/// Step to find workers to update based on URL.
+///
+/// For DP-aware workers, finds all workers with matching URL prefix.
+/// For regular workers, finds the single worker with exact URL match.
+pub struct FindWorkerToUpdateStep;
+
+#[async_trait]
+impl StepExecutor<WorkerUpdateWorkflowData> for FindWorkerToUpdateStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerUpdateWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let worker_url = &context.data.worker_url;
+        let dp_aware = context.data.dp_aware;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+
+        let workers_to_update =
+            find_workers_by_url(&app_context.worker_registry, worker_url, dp_aware);
+
+        if workers_to_update.is_empty() {
+            let error_msg = if dp_aware {
+                format!("No workers found with prefix {}@", worker_url)
+            } else {
+                format!("Worker {} not found", worker_url)
+            };
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("find_worker_to_update"),
+                message: error_msg,
+            });
+        }
+
+        debug!(
+            "Found {} worker(s) to update for {}",
+            workers_to_update.len(),
+            worker_url
+        );
+
+        context.data.workers_to_update = Some(workers_to_update);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs

@@ -0,0 +1,82 @@
+//! Step to find workers to remove based on URL.
+
+use std::collections::HashSet;
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use super::find_workers_by_url;
+use crate::core::steps::workflow_data::{WorkerList, WorkerRemovalWorkflowData};
+
+/// Request structure for worker removal.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct WorkerRemovalRequest {
+    pub url: String,
+    pub dp_aware: bool,
+}
+
+/// Step to find workers to remove based on URL.
+///
+/// For DP-aware workers, finds all workers with matching URL prefix.
+/// For regular workers, finds the single worker with exact URL match.
+pub struct FindWorkersToRemoveStep;
+
+#[async_trait]
+impl StepExecutor<WorkerRemovalWorkflowData> for FindWorkersToRemoveStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerRemovalWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+
+        let workers_to_remove =
+            find_workers_by_url(&app_context.worker_registry, &request.url, request.dp_aware);
+
+        if workers_to_remove.is_empty() {
+            let error_msg = if request.dp_aware {
+                format!("No workers found with prefix {}@", request.url)
+            } else {
+                format!("Worker {} not found", request.url)
+            };
+            return Err(WorkflowError::StepFailed {
+                step_id: StepId::new("find_workers_to_remove"),
+                message: error_msg,
+            });
+        }
+
+        debug!(
+            "Found {} worker(s) to remove for {}",
+            workers_to_remove.len(),
+            request.url
+        );
+
+        // Store workers and their model IDs for subsequent steps
+        let worker_urls: Vec<String> = workers_to_remove
+            .iter()
+            .map(|w| w.url().to_string())
+            .collect();
+
+        let affected_models: HashSet<String> = workers_to_remove
+            .iter()
+            .map(|w| w.model_id().to_string())
+            .collect();
+
+        // Update workflow data
+        context.data.workers_to_remove = Some(WorkerList::from_workers(&workers_to_remove));
+        context.data.actual_workers_to_remove = Some(workers_to_remove);
+        context.data.worker_urls = worker_urls;
+        context.data.affected_models = affected_models;
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/mod.rs

@@ -0,0 +1,363 @@
+mod create_worker;
+mod detect_connection;
+mod discover_dp;
+mod discover_metadata;
+mod find_worker_to_update;
+mod find_workers_to_remove;
+mod remove_from_policy_registry;
+mod remove_from_worker_registry;
+mod submit_tokenizer_job;
+mod update_policies_for_worker;
+mod update_remaining_policies;
+mod update_worker_properties;
+
+use std::{sync::Arc, time::Duration};
+
+/// Strip protocol prefix (http://, https://, grpc://) from URL.
+pub(crate) fn strip_protocol(url: &str) -> String {
+    url.trim_start_matches("http://")
+        .trim_start_matches("https://")
+        .trim_start_matches("grpc://")
+        .to_string()
+}
+
+pub use create_worker::CreateLocalWorkerStep;
+pub use detect_connection::DetectConnectionModeStep;
+pub use discover_dp::{get_dp_info, DiscoverDPInfoStep, DpInfo};
+pub use discover_metadata::DiscoverMetadataStep;
+pub use find_worker_to_update::FindWorkerToUpdateStep;
+pub use find_workers_to_remove::{FindWorkersToRemoveStep, WorkerRemovalRequest};
+pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep;
+pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep;
+pub use submit_tokenizer_job::SubmitTokenizerJobStep;
+pub use update_policies_for_worker::UpdatePoliciesForWorkerStep;
+pub use update_remaining_policies::UpdateRemainingPoliciesStep;
+pub use update_worker_properties::UpdateWorkerPropertiesStep;
+use wfaas::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition};
+
+use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
+use crate::{
+    app_context::AppContext,
+    config::RouterConfig,
+    core::{
+        steps::workflow_data::{
+            LocalWorkerWorkflowData, WorkerRemovalWorkflowData, WorkerUpdateWorkflowData,
+        },
+        Worker, WorkerRegistry,
+    },
+    protocols::worker_spec::{WorkerConfigRequest, WorkerUpdateRequest},
+};
+
+/// Find workers by URL, supporting both DP-aware (prefix match) and regular (exact match) modes.
+///
+/// For DP-aware workers, finds all workers with URL prefix `{url}@`.
+/// For regular workers, finds the single worker with exact URL match.
+pub(crate) fn find_workers_by_url(
+    registry: &WorkerRegistry,
+    url: &str,
+    dp_aware: bool,
+) -> Vec<Arc<dyn Worker>> {
+    if dp_aware {
+        let worker_url_prefix = format!("{}@", url);
+        registry
+            .get_all()
+            .iter()
+            .filter(|worker| worker.url().starts_with(&worker_url_prefix))
+            .cloned()
+            .collect()
+    } else {
+        match registry.get_by_url(url) {
+            Some(worker) => vec![worker],
+            None => Vec::new(),
+        }
+    }
+}
+
+pub fn create_local_worker_workflow(
+    router_config: &RouterConfig,
+) -> WorkflowDefinition<LocalWorkerWorkflowData> {
+    let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs);
+
+    // Calculate max_attempts based on timeout
+    let timeout_secs = detect_timeout.as_secs() as f64;
+    let effective_timeout = timeout_secs * 0.9;
+    let max_attempts = if effective_timeout > 10.0 {
+        (5 + ((effective_timeout - 10.0) / 5.0).ceil() as u32).max(3)
+    } else {
+        3
+    };
+
+    WorkflowDefinition::new("local_worker_registration", "Local Worker Registration")
+        // Step 1: Detect connection mode (HTTP vs gRPC)
+        .add_step(
+            StepDefinition::new(
+                "detect_connection_mode",
+                "Detect Connection Mode",
+                Arc::new(DetectConnectionModeStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts,
+                backoff: BackoffStrategy::Linear {
+                    increment: Duration::from_secs(1),
+                    max: Duration::from_secs(5),
+                },
+            })
+            .with_timeout(detect_timeout)
+            .with_failure_action(FailureAction::FailWorkflow),
+        )
+        // Step 2a: Discover metadata (parallel with DP discovery)
+        .add_step(
+            StepDefinition::new(
+                "discover_metadata",
+                "Discover Metadata",
+                Arc::new(DiscoverMetadataStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 3,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(1)),
+            })
+            .with_timeout(Duration::from_secs(10))
+            .with_failure_action(FailureAction::ContinueNextStep)
+            .depends_on(&["detect_connection_mode"]),
+        )
+        // Step 2b: Discover DP info (after metadata to avoid concurrent /server_info calls)
+        .add_step(
+            StepDefinition::new(
+                "discover_dp_info",
+                "Discover DP Info",
+                Arc::new(DiscoverDPInfoStep),
+            )
+            .with_retry(RetryPolicy {
+                max_attempts: 3,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(1)),
+            })
+            .with_timeout(Duration::from_secs(10))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["discover_metadata"]),
+        )
+        // Step 3: Create worker(s)
+        .add_step(
+            StepDefinition::new(
+                "create_worker",
+                "Create Worker",
+                Arc::new(CreateLocalWorkerStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["discover_dp_info"]),
+        )
+        // Step 4: Register workers (shared step)
+        .add_step(
+            StepDefinition::new(
+                "register_workers",
+                "Register Workers",
+                Arc::new(RegisterWorkersStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["create_worker"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "submit_tokenizer_job",
+                "Submit Tokenizer Job",
+                Arc::new(SubmitTokenizerJobStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::ContinueNextStep)
+            .depends_on(&["register_workers"]),
+        )
+        // Step 5a: Update policies (parallel with activation)
+        .add_step(
+            StepDefinition::new(
+                "update_policies",
+                "Update Policies",
+                Arc::new(UpdatePoliciesStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::ContinueNextStep)
+            .depends_on(&["register_workers"]),
+        )
+        // Step 5b: Activate workers (parallel with policy update)
+        .add_step(
+            StepDefinition::new(
+                "activate_workers",
+                "Activate Workers",
+                Arc::new(ActivateWorkersStep),
+            )
+            .with_timeout(Duration::from_secs(5))
+            .with_failure_action(FailureAction::FailWorkflow)
+            .depends_on(&["register_workers"]),
+        )
+}
+
+/// Create a worker removal workflow definition.
+///
+/// DAG structure:
+/// ```text
+///     find_workers_to_remove
+///              │
+///     remove_from_policy_registry
+///              │
+///     remove_from_worker_registry
+///              │
+///     update_remaining_policies
+/// ```
+pub fn create_worker_removal_workflow() -> WorkflowDefinition<WorkerRemovalWorkflowData> {
+    WorkflowDefinition::new("worker_removal", "Remove worker from router")
+        .add_step(
+            StepDefinition::new(
+                "find_workers_to_remove",
+                "Find workers to remove",
+                Arc::new(FindWorkersToRemoveStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            }),
+        )
+        .add_step(
+            StepDefinition::new(
+                "remove_from_policy_registry",
+                "Remove workers from policy registry",
+                Arc::new(RemoveFromPolicyRegistryStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["find_workers_to_remove"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "remove_from_worker_registry",
+                "Remove workers from worker registry",
+                Arc::new(RemoveFromWorkerRegistryStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["remove_from_policy_registry"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "update_remaining_policies",
+                "Update cache-aware policies for remaining workers",
+                Arc::new(UpdateRemainingPoliciesStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["remove_from_worker_registry"]),
+        )
+}
+
+/// Create a worker update workflow definition.
+///
+/// DAG structure:
+/// ```text
+///     find_worker_to_update
+///              │
+///     update_worker_properties
+///              │
+///     update_policies_for_worker
+/// ```
+pub fn create_worker_update_workflow() -> WorkflowDefinition<WorkerUpdateWorkflowData> {
+    WorkflowDefinition::new("worker_update", "Update worker properties")
+        .add_step(
+            StepDefinition::new(
+                "find_worker_to_update",
+                "Find worker to update",
+                Arc::new(FindWorkerToUpdateStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            }),
+        )
+        .add_step(
+            StepDefinition::new(
+                "update_worker_properties",
+                "Update worker properties",
+                Arc::new(UpdateWorkerPropertiesStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["find_worker_to_update"]),
+        )
+        .add_step(
+            StepDefinition::new(
+                "update_policies_for_worker",
+                "Update policies for updated worker",
+                Arc::new(UpdatePoliciesForWorkerStep),
+            )
+            .with_timeout(Duration::from_secs(10))
+            .with_retry(RetryPolicy {
+                max_attempts: 1,
+                backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
+            })
+            .depends_on(&["update_worker_properties"]),
+        )
+}
+
+/// Helper to create initial workflow data for local worker registration
+pub fn create_local_worker_workflow_data(
+    config: WorkerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> LocalWorkerWorkflowData {
+    LocalWorkerWorkflowData {
+        config,
+        connection_mode: None,
+        discovered_labels: std::collections::HashMap::new(),
+        dp_info: None,
+        workers: None,
+        final_labels: std::collections::HashMap::new(),
+        detected_runtime_type: None,
+        app_context: Some(app_context),
+        actual_workers: None,
+    }
+}
+
+/// Helper to create initial workflow data for worker removal
+pub fn create_worker_removal_workflow_data(
+    url: String,
+    dp_aware: bool,
+    app_context: Arc<AppContext>,
+) -> WorkerRemovalWorkflowData {
+    WorkerRemovalWorkflowData {
+        config: WorkerRemovalRequest { url, dp_aware },
+        workers_to_remove: None,
+        worker_urls: Vec::new(),
+        affected_models: std::collections::HashSet::new(),
+        app_context: Some(app_context),
+        actual_workers_to_remove: None,
+    }
+}
+
+/// Helper to create initial workflow data for worker update
+pub fn create_worker_update_workflow_data(
+    worker_url: String,
+    update_config: WorkerUpdateRequest,
+    app_context: Arc<AppContext>,
+) -> WorkerUpdateWorkflowData {
+    // Determine if this is a DP-aware update based on URL pattern
+    let dp_aware = worker_url.contains('@');
+    WorkerUpdateWorkflowData {
+        config: update_config,
+        worker_url,
+        dp_aware,
+        app_context: Some(app_context),
+        workers_to_update: None,
+        updated_workers: None,
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs

@@ -0,0 +1,61 @@
+//! Step to remove workers from policy registry.
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::core::steps::workflow_data::WorkerRemovalWorkflowData;
+
+/// Step to remove workers from the policy registry.
+///
+/// Removes each worker from cache-aware policies and notifies
+/// the policy registry of worker removal.
+pub struct RemoveFromPolicyRegistryStep;
+
+#[async_trait]
+impl StepExecutor<WorkerRemovalWorkflowData> for RemoveFromPolicyRegistryStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerRemovalWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let workers_to_remove = context
+            .data
+            .actual_workers_to_remove
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_remove".to_string()))?;
+
+        debug!(
+            "Removing {} worker(s) from policy registry",
+            workers_to_remove.len()
+        );
+
+        for worker in workers_to_remove.iter() {
+            let model_id = worker.model_id().to_string();
+            let worker_url = worker.url();
+
+            // Remove from cache-aware policy
+            app_context
+                .policy_registry
+                .remove_worker_from_cache_aware(&model_id, worker_url);
+
+            // Notify policy registry
+            app_context.policy_registry.on_worker_removed(&model_id);
+        }
+
+        debug!(
+            "Removed {} worker(s) from policy registry",
+            workers_to_remove.len()
+        );
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs

@@ -0,0 +1,103 @@
+//! Step to remove workers from worker registry.
+
+use std::collections::HashSet;
+
+use async_trait::async_trait;
+use tracing::{debug, warn};
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::{
+    core::steps::workflow_data::WorkerRemovalWorkflowData, observability::metrics::Metrics,
+};
+
+/// Step to remove workers from the worker registry.
+///
+/// Removes each worker by URL from the central worker registry.
+pub struct RemoveFromWorkerRegistryStep;
+
+#[async_trait]
+impl StepExecutor<WorkerRemovalWorkflowData> for RemoveFromWorkerRegistryStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerRemovalWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let worker_urls = &context.data.worker_urls;
+
+        debug!(
+            "Removing {} worker(s) from worker registry",
+            worker_urls.len()
+        );
+
+        // Collect unique worker configurations before removal for pool size updates
+        let unique_configs: HashSet<_> = worker_urls
+            .iter()
+            .filter_map(|url| app_context.worker_registry.get_by_url(url))
+            .map(|w| {
+                let meta = w.metadata();
+                (
+                    meta.worker_type.clone(),
+                    meta.connection_mode.clone(),
+                    w.model_id().to_string(),
+                )
+            })
+            .collect();
+
+        let mut removed_count = 0;
+        for worker_url in worker_urls.iter() {
+            if app_context
+                .worker_registry
+                .remove_by_url(worker_url)
+                .is_some()
+            {
+                removed_count += 1;
+            }
+        }
+
+        // Log if some workers were already removed (e.g., by another process)
+        if removed_count != worker_urls.len() {
+            warn!(
+                "Removed {} of {} workers (some may have been removed by another process)",
+                removed_count,
+                worker_urls.len()
+            );
+        } else {
+            debug!("Removed {} worker(s) from registry", removed_count);
+        }
+
+        // Update Layer 3 worker pool size metrics for unique configurations
+        for (worker_type, connection_mode, model_id) in unique_configs {
+            // Get labels before moving values into get_workers_filtered
+            let worker_type_label = worker_type.as_metric_label();
+            let connection_mode_label = connection_mode.as_metric_label();
+
+            let pool_size = app_context
+                .worker_registry
+                .get_workers_filtered(
+                    Some(&model_id),
+                    Some(worker_type),
+                    Some(connection_mode),
+                    None,
+                    false,
+                )
+                .len();
+
+            Metrics::set_worker_pool_size(
+                worker_type_label,
+                connection_mode_label,
+                &model_id,
+                pool_size,
+            );
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs

@@ -0,0 +1,135 @@
+//! Tokenizer registration step for local workers.
+//!
+//! This step submits a Job::AddTokenizer to the job queue, which triggers the
+//! tokenizer_registration workflow. The workflow handles validation, deduplication,
+//! and caching - this step just submits the job.
+
+use async_trait::async_trait;
+use tracing::{debug, info, warn};
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::{
+    core::{
+        steps::{workflow_data::LocalWorkerWorkflowData, TokenizerConfigRequest},
+        Job,
+    },
+    tokenizer::TokenizerRegistry,
+};
+
+/// Step: Submit tokenizer registration job for the worker's model
+///
+/// This step submits a Job::AddTokenizer to the job queue rather than loading
+/// the tokenizer directly. This ensures tokenizer registration goes through
+/// the unified tokenizer_registration workflow.
+pub struct SubmitTokenizerJobStep;
+
+#[async_trait]
+impl StepExecutor<LocalWorkerWorkflowData> for SubmitTokenizerJobStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<LocalWorkerWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let labels = &context.data.final_labels;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let workers = context
+            .data
+            .actual_workers
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
+
+        // Get job queue
+        let job_queue = match app_context.worker_job_queue.get() {
+            Some(queue) => queue,
+            None => {
+                warn!("Job queue not available, skipping tokenizer registration");
+                return Ok(StepResult::Success);
+            }
+        };
+
+        // Get chat_template: worker config > global router config
+        let chat_template = context
+            .data
+            .config
+            .chat_template
+            .clone()
+            .or_else(|| app_context.router_config.chat_template.clone());
+
+        // Get cache config from router config
+        let cache_config = app_context.router_config.tokenizer_cache.to_option();
+
+        for worker in workers.iter() {
+            let model_id = worker.model_id().to_string();
+
+            // Get tokenizer path with fallback chain:
+            // 1. Worker labels: tokenizer_path
+            // 2. Worker labels: model_path
+            // 3. Router config (CLI args): --tokenizer-path
+            // 4. Router config (CLI args): --model-path
+            let tokenizer_path: String = if let Some(path) = labels
+                .get("tokenizer_path")
+                .or_else(|| labels.get("model_path"))
+            {
+                path.clone()
+            } else if let Some(path) = app_context
+                .router_config
+                .tokenizer_path
+                .as_ref()
+                .or(app_context.router_config.model_path.as_ref())
+            {
+                debug!(
+                    "Using router config tokenizer path '{}' for model {}",
+                    path, model_id
+                );
+                path.clone()
+            } else {
+                warn!(
+                    "No tokenizer_path or model_path found for model {} (checked worker labels and router config)",
+                    model_id
+                );
+                continue;
+            };
+
+            // Note: We don't check if tokenizer already exists here.
+            // The registry.load() handles deduplication gracefully (returns AlreadyExists).
+            // This simplifies the code and ensures consistent behavior.
+
+            info!(
+                "Submitting tokenizer registration job for model {} from {}",
+                model_id, tokenizer_path
+            );
+
+            // Create tokenizer config request
+            let config = TokenizerConfigRequest {
+                id: TokenizerRegistry::generate_id(),
+                name: model_id.clone(),
+                source: tokenizer_path,
+                chat_template_path: chat_template.clone(),
+                cache_config: cache_config.clone(),
+                fail_on_duplicate: false,
+            };
+
+            // Submit job (fire-and-forget, don't wait for completion)
+            if let Err(e) = job_queue
+                .submit(Job::AddTokenizer {
+                    config: Box::new(config),
+                })
+                .await
+            {
+                warn!(
+                    "Failed to submit tokenizer job for model {}: {}",
+                    model_id, e
+                );
+            }
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false // Job submission failures are not retryable at this level
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs

@@ -0,0 +1,72 @@
+//! Step to update policies for updated workers.
+
+use std::collections::HashSet;
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::core::steps::workflow_data::WorkerUpdateWorkflowData;
+
+/// Step to update policies for updated workers.
+///
+/// After workers are updated, this step re-initializes cache-aware policies
+/// for the affected models to reflect any priority or label changes.
+pub struct UpdatePoliciesForWorkerStep;
+
+#[async_trait]
+impl StepExecutor<WorkerUpdateWorkflowData> for UpdatePoliciesForWorkerStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerUpdateWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let updated_workers =
+            context.data.updated_workers.as_ref().ok_or_else(|| {
+                WorkflowError::ContextValueNotFound("updated_workers".to_string())
+            })?;
+
+        // Collect affected models
+        let affected_models: HashSet<String> = updated_workers
+            .iter()
+            .map(|w| w.model_id().to_string())
+            .collect();
+
+        debug!(
+            "Updating policies for {} affected model(s) after worker update",
+            affected_models.len()
+        );
+
+        for model_id in &affected_models {
+            let workers = app_context.worker_registry.get_by_model(model_id);
+
+            if let Some(policy) = app_context.policy_registry.get_policy(model_id) {
+                if policy.name() == "cache_aware" && !workers.is_empty() {
+                    // Re-initialize cache-aware policy with updated workers
+                    app_context
+                        .policy_registry
+                        .init_cache_aware_policy(model_id, &workers);
+
+                    debug!(
+                        "Updated cache-aware policy for model {} ({} workers)",
+                        model_id,
+                        workers.len()
+                    );
+                }
+            }
+
+            // Notify policy registry of the update
+            app_context.policy_registry.on_worker_added(model_id, None);
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs

@@ -0,0 +1,69 @@
+//! Step to update cache-aware policies for remaining workers after removal.
+
+use async_trait::async_trait;
+use tracing::{debug, info};
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::core::steps::workflow_data::WorkerRemovalWorkflowData;
+
+/// Step to update cache-aware policies for remaining workers.
+///
+/// After workers are removed, this step re-initializes cache-aware policies
+/// for the affected models using the remaining workers.
+pub struct UpdateRemainingPoliciesStep;
+
+#[async_trait]
+impl StepExecutor<WorkerRemovalWorkflowData> for UpdateRemainingPoliciesStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerRemovalWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let affected_models = &context.data.affected_models;
+        let worker_urls = &context.data.worker_urls;
+
+        debug!(
+            "Updating cache-aware policies for {} affected model(s)",
+            affected_models.len()
+        );
+
+        for model_id in affected_models.iter() {
+            let remaining_workers = app_context.worker_registry.get_by_model(model_id);
+
+            if let Some(policy) = app_context.policy_registry.get_policy(model_id) {
+                if policy.name() == "cache_aware" && !remaining_workers.is_empty() {
+                    app_context
+                        .policy_registry
+                        .init_cache_aware_policy(model_id, &remaining_workers);
+
+                    debug!(
+                        "Updated cache-aware policy for model {} ({} remaining workers)",
+                        model_id,
+                        remaining_workers.len()
+                    );
+                }
+            }
+        }
+
+        // Log final result at info level
+        if worker_urls.len() == 1 {
+            info!("Removed worker {}", worker_urls[0]);
+        } else {
+            info!(
+                "Removed {} DP-aware workers: {:?}",
+                worker_urls.len(),
+                worker_urls
+            );
+        }
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs

@@ -0,0 +1,152 @@
+//! Step to update worker properties.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tracing::{debug, info};
+use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+
+use crate::core::{
+    steps::workflow_data::WorkerUpdateWorkflowData, BasicWorkerBuilder, HealthConfig, Worker,
+};
+
+/// Step to update worker properties.
+///
+/// This step creates new worker instances with updated properties and
+/// re-registers them to replace the old workers in the registry.
+pub struct UpdateWorkerPropertiesStep;
+
+#[async_trait]
+impl StepExecutor<WorkerUpdateWorkflowData> for UpdateWorkerPropertiesStep {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<WorkerUpdateWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let request = &context.data.config;
+        let app_context = context
+            .data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+        let workers_to_update = context
+            .data
+            .workers_to_update
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_update".to_string()))?
+            .clone();
+
+        debug!(
+            "Updating properties for {} worker(s)",
+            workers_to_update.len()
+        );
+
+        let mut updated_workers: Vec<Arc<dyn Worker>> = Vec::with_capacity(workers_to_update.len());
+
+        for worker in workers_to_update.iter() {
+            // Build updated labels - merge new labels into existing ones
+            let mut updated_labels = worker.metadata().labels.clone();
+            if let Some(ref new_labels) = request.labels {
+                for (key, value) in new_labels {
+                    updated_labels.insert(key.clone(), value.clone());
+                }
+            }
+
+            // Update priority if specified (stored in labels)
+            if let Some(priority) = request.priority {
+                updated_labels.insert("priority".to_string(), priority.to_string());
+            }
+
+            // Update cost if specified (stored in labels)
+            if let Some(cost) = request.cost {
+                updated_labels.insert("cost".to_string(), cost.to_string());
+            }
+
+            // Build updated health config
+            let existing_health = &worker.metadata().health_config;
+            let updated_health_config = HealthConfig {
+                timeout_secs: request
+                    .health_check_timeout_secs
+                    .unwrap_or(existing_health.timeout_secs),
+                check_interval_secs: request
+                    .health_check_interval_secs
+                    .unwrap_or(existing_health.check_interval_secs),
+                endpoint: existing_health.endpoint.clone(),
+                failure_threshold: request
+                    .health_failure_threshold
+                    .unwrap_or(existing_health.failure_threshold),
+                success_threshold: request
+                    .health_success_threshold
+                    .unwrap_or(existing_health.success_threshold),
+                disable_health_check: request
+                    .disable_health_check
+                    .unwrap_or(existing_health.disable_health_check),
+            };
+
+            // Determine API key: use new one if provided, otherwise keep existing
+            let updated_api_key = request
+                .api_key
+                .clone()
+                .or_else(|| worker.metadata().api_key.clone());
+
+            // Create a new worker with updated properties
+            let new_worker: Arc<dyn Worker> = if worker.is_dp_aware() {
+                // For DP-aware workers, extract DP info and rebuild
+                let dp_rank = worker.dp_rank().unwrap_or(0);
+                let dp_size = worker.dp_size().unwrap_or(1);
+                let base_url = worker.base_url().to_string();
+
+                let mut builder =
+                    crate::core::DPAwareWorkerBuilder::new(base_url, dp_rank, dp_size)
+                        .worker_type(worker.worker_type().clone())
+                        .connection_mode(worker.connection_mode().clone())
+                        .runtime_type(worker.metadata().runtime_type.clone())
+                        .labels(updated_labels)
+                        .health_config(updated_health_config.clone())
+                        .models(worker.metadata().models.clone());
+
+                if let Some(ref api_key) = updated_api_key {
+                    builder = builder.api_key(api_key.clone());
+                }
+
+                Arc::new(builder.build())
+            } else {
+                // For basic workers, rebuild with updated properties
+                let mut builder = BasicWorkerBuilder::new(worker.url())
+                    .worker_type(worker.worker_type().clone())
+                    .connection_mode(worker.connection_mode().clone())
+                    .runtime_type(worker.metadata().runtime_type.clone())
+                    .labels(updated_labels)
+                    .health_config(updated_health_config.clone())
+                    .models(worker.metadata().models.clone());
+
+                if let Some(ref api_key) = updated_api_key {
+                    builder = builder.api_key(api_key.clone());
+                }
+
+                Arc::new(builder.build())
+            };
+
+            // Re-register the worker (this replaces the old one)
+            app_context.worker_registry.register(new_worker.clone());
+
+            updated_workers.push(new_worker);
+        }
+
+        // Log result
+        if updated_workers.len() == 1 {
+            info!("Updated worker {}", updated_workers[0].url());
+        } else {
+            info!("Updated {} workers", updated_workers.len());
+        }
+
+        // Store updated workers for subsequent steps
+        context.data.updated_workers = Some(updated_workers);
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/mod.rs

@@ -0,0 +1,19 @@
+pub mod external;
+pub mod local;
+pub mod shared;
+
+pub use external::{
+    create_external_worker_workflow, create_external_worker_workflow_data, group_models_into_cards,
+    infer_model_type_from_id, CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo,
+    ModelsResponse,
+};
+pub use local::{
+    create_local_worker_workflow, create_local_worker_workflow_data,
+    create_worker_removal_workflow, create_worker_removal_workflow_data,
+    create_worker_update_workflow, create_worker_update_workflow_data, CreateLocalWorkerStep,
+    DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, DpInfo,
+    FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
+    RemoveFromWorkerRegistryStep, UpdatePoliciesForWorkerStep, UpdateRemainingPoliciesStep,
+    UpdateWorkerPropertiesStep, WorkerRemovalRequest,
+};
+pub use shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep, WorkerList};

third_party/sglang/sgl-model-gateway/src/core/steps/worker/shared/activate.rs

@@ -0,0 +1,37 @@
+//! Unified worker activation step.
+
+use async_trait::async_trait;
+use tracing::info;
+use wfaas::{
+    StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult,
+};
+
+use crate::core::steps::workflow_data::WorkerRegistrationData;
+
+/// Unified step to activate workers by marking them as healthy.
+///
+/// This is the final step in any worker registration workflow.
+/// Works with any workflow data type that implements `WorkerRegistrationData`.
+pub struct ActivateWorkersStep;
+
+#[async_trait]
+impl<D: WorkerRegistrationData + WorkflowData> StepExecutor<D> for ActivateWorkersStep {
+    async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult> {
+        let workers = context
+            .data
+            .get_actual_workers()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
+
+        for worker in workers.iter() {
+            worker.set_healthy(true);
+        }
+
+        info!("Activated {} worker(s) (marked as healthy)", workers.len());
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/shared/mod.rs

@@ -0,0 +1,21 @@
+//! Shared worker registration steps used by both local and external workflows.
+//!
+//! These steps are designed to work with any worker type and can be composed
+//! into different workflows using the DAG-based workflow engine.
+
+mod activate;
+mod register;
+mod update_policies;
+
+use std::sync::Arc;
+
+pub use activate::ActivateWorkersStep;
+pub use register::RegisterWorkersStep;
+pub use update_policies::UpdatePoliciesStep;
+
+use crate::core::Worker;
+
+/// Type alias for a collection of workers in workflow context.
+/// Both local (single/DP-aware) and external (multi-model) workflows
+/// use this unified type for consistency.
+pub type WorkerList = Vec<Arc<dyn Worker>>;

third_party/sglang/sgl-model-gateway/src/core/steps/worker/shared/register.rs

@@ -0,0 +1,99 @@
+//! Unified worker registration step.
+
+use std::{collections::HashSet, sync::Arc};
+
+use async_trait::async_trait;
+use tracing::debug;
+use wfaas::{
+    StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult,
+};
+
+use crate::{core::steps::workflow_data::WorkerRegistrationData, observability::metrics::Metrics};
+
+/// Unified step to register workers in the registry.
+///
+/// Works with both single workers and batches. Always expects `workers` key
+/// in context containing `Vec<Arc<dyn Worker>>`.
+/// Works with any workflow data type that implements `WorkerRegistrationData`.
+pub struct RegisterWorkersStep;
+
+#[async_trait]
+impl<D: WorkerRegistrationData + WorkflowData> StepExecutor<D> for RegisterWorkersStep {
+    async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .get_app_context()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+
+        let workers = context
+            .data
+            .get_actual_workers()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
+
+        let mut worker_ids = Vec::with_capacity(workers.len());
+
+        for worker in workers.iter() {
+            let worker_id = app_context.worker_registry.register(Arc::clone(worker));
+            debug!(
+                "Registered worker {} (model: {}) with ID {:?}",
+                worker.url(),
+                worker.model_id(),
+                worker_id
+            );
+            worker_ids.push(worker_id);
+        }
+
+        // Collect unique worker configurations to avoid redundant metric updates
+        let unique_configs: HashSet<_> = workers
+            .iter()
+            .map(|w| {
+                let meta = w.metadata();
+                (
+                    meta.worker_type.clone(),
+                    meta.connection_mode.clone(),
+                    w.model_id().to_string(),
+                )
+            })
+            .collect();
+
+        // Update Layer 3 worker pool size metrics per unique type/connection/model
+        for (worker_type, connection_mode, model_id) in unique_configs {
+            // Get labels before moving values into get_workers_filtered
+            let worker_type_label = worker_type.as_metric_label();
+            let connection_mode_label = connection_mode.as_metric_label();
+
+            let pool_size = app_context
+                .worker_registry
+                .get_workers_filtered(
+                    Some(&model_id),
+                    Some(worker_type),
+                    Some(connection_mode),
+                    None,
+                    false,
+                )
+                .len();
+
+            Metrics::set_worker_pool_size(
+                worker_type_label,
+                connection_mode_label,
+                &model_id,
+                pool_size,
+            );
+        }
+
+        // Note: worker_ids are stored for potential future use but not persisted
+        // as they are internal registry identifiers
+        debug!(
+            "Registered {} workers with IDs: {:?}",
+            worker_ids.len(),
+            worker_ids
+        );
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs

@@ -0,0 +1,156 @@
+//! Unified policy update step.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tracing::{debug, warn};
+use wfaas::{
+    StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult,
+};
+
+use crate::core::{steps::workflow_data::WorkerRegistrationData, Worker};
+
+/// Unified step to update policy registry for registered workers.
+///
+/// Handles both local workers (same model, possibly DP-aware) and
+/// external workers (different models per worker).
+pub struct UpdatePoliciesStep;
+
+impl UpdatePoliciesStep {
+    /// Check for conflicts between prefill and decode worker configurations for a model.
+    fn check_worker_conflicts(&self, model_id: &str, workers: &[Arc<dyn Worker>]) {
+        let prefill_workers: Vec<_> = workers
+            .iter()
+            .filter(|w| {
+                w.metadata()
+                    .labels
+                    .get("disaggregation_mode")
+                    .map(|s| s.as_str())
+                    == Some("prefill")
+            })
+            .collect();
+
+        let decode_workers: Vec<_> = workers
+            .iter()
+            .filter(|w| {
+                w.metadata()
+                    .labels
+                    .get("disaggregation_mode")
+                    .map(|s| s.as_str())
+                    == Some("decode")
+            })
+            .collect();
+
+        if prefill_workers.is_empty() || decode_workers.is_empty() {
+            return;
+        }
+
+        // Compare configurations of prefill vs decode workers
+        if let (Some(pw), Some(dw)) = (prefill_workers.first(), decode_workers.first()) {
+            let pl = &pw.metadata().labels;
+            let dl = &dw.metadata().labels;
+
+            // Define keys to check for equality
+            let keys_to_check = ["tp_size", "dp_size", "load_balance_method"];
+
+            for key in keys_to_check {
+                let p_val = pl.get(key);
+                let d_val = dl.get(key);
+                if p_val != d_val {
+                    warn!(
+                        "Model {} has conflicting {}: prefill={:?}, decode={:?}",
+                        model_id, key, p_val, d_val
+                    );
+                }
+            }
+
+            // Specific check for Data-Parallel consistency
+            if let Some(dp_size) = pl.get("dp_size").and_then(|s| s.parse::<usize>().ok()) {
+                if dp_size > 1 {
+                    let plb = pl.get("load_balance_method").map(|s| s.as_str());
+                    if plb != Some("follow_bootstrap_room") {
+                        warn!(
+                            "Model {} has dp_size > 1 but load_balance_method is not 'follow_bootstrap_room' on prefill workers. This may cause rank mismatch in disaggregated mode.",
+                            model_id
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl<D: WorkerRegistrationData + WorkflowData> StepExecutor<D> for UpdatePoliciesStep {
+    async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .get_app_context()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+
+        let workers = context
+            .data
+            .get_actual_workers()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
+
+        let labels = context
+            .data
+            .get_labels()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?;
+
+        let policy_hint = labels.get("policy").map(|s| s.as_str());
+
+        // Track unique model IDs we've updated policies for
+        let mut updated_models = Vec::new();
+
+        for worker in workers.iter() {
+            let model_id = worker.model_id().to_string();
+
+            // Notify policy registry
+            app_context
+                .policy_registry
+                .on_worker_added(&model_id, policy_hint);
+
+            // Initialize cache-aware policy if configured
+            let all_workers = app_context.worker_registry.get_by_model(&model_id);
+
+            // Check for configuration conflicts between prefill and decode
+            self.check_worker_conflicts(&model_id, &all_workers);
+            if let Some(policy) = app_context.policy_registry.get_policy(&model_id) {
+                if policy.name() == "cache_aware" {
+                    app_context
+                        .policy_registry
+                        .init_cache_aware_policy(&model_id, &all_workers);
+                }
+            }
+
+            if !updated_models.contains(&model_id) {
+                updated_models.push(model_id);
+            }
+        }
+
+        // Initialize bucket policies for prefill workers (local workers only)
+        let prefill_workers = app_context.worker_registry.get_prefill_workers();
+        if !prefill_workers.is_empty() {
+            let policy = app_context.policy_registry.get_prefill_policy();
+            if policy.name() == "bucket" {
+                app_context
+                    .policy_registry
+                    .init_pd_bucket_policies(&prefill_workers);
+            }
+        }
+
+        debug!(
+            "Updated policies for {} workers across {} models",
+            workers.len(),
+            updated_models.len()
+        );
+
+        Ok(StepResult::Success)
+    }
+
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
+        false
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/workflow_data.rs

@@ -0,0 +1,375 @@
+//! Typed workflow data structures
+//!
+//! This module defines the typed data structures for all workflows, enabling
+//! compile-time type safety and state persistence. Each workflow has its own
+//! strongly-typed data structure, and steps are typed to their specific workflow.
+//!
+//! # Shared Step Trait
+//!
+//! For steps that are shared between local and external worker workflows,
+//! we use the `WorkerRegistrationData` trait. This trait provides a common
+//! interface for accessing worker data while maintaining full type safety.
+
+use std::{collections::HashMap, sync::Arc};
+
+use serde::{Deserialize, Serialize};
+use wfaas::{WorkflowData, WorkflowError};
+
+use super::{
+    mcp_registration::McpServerConfigRequest, tokenizer_registration::TokenizerConfigRequest,
+    wasm_module_registration::WasmModuleConfigRequest,
+    wasm_module_removal::WasmModuleRemovalRequest, worker::local::WorkerRemovalRequest,
+};
+/// Re-export the protocol types for convenience
+pub use crate::protocols::worker_spec::{
+    WorkerConfigRequest, WorkerUpdateRequest as ProtocolUpdateRequest,
+};
+use crate::{
+    app_context::AppContext,
+    core::{model_card::ModelCard, Worker},
+    protocols::worker_spec::{
+        WorkerConfigRequest as ProtocolWorkerConfigRequest,
+        WorkerUpdateRequest as ProtocolWorkerUpdateRequest,
+    },
+};
+
+// ============================================================================
+// Shared trait for worker registration workflows
+// ============================================================================
+
+/// Trait for workflow data that supports worker registration operations.
+///
+/// This trait is implemented by both `LocalWorkerWorkflowData` and
+/// `ExternalWorkerWorkflowData`, allowing shared steps to work with either
+/// workflow type while maintaining full type safety.
+pub trait WorkerRegistrationData: WorkflowData {
+    /// Get the application context (transient, not serialized).
+    fn get_app_context(&self) -> Option<&Arc<AppContext>>;
+
+    /// Get the actual worker objects (transient, not serialized).
+    fn get_actual_workers(&self) -> Option<&Vec<Arc<dyn Worker>>>;
+
+    /// Get the labels for policy registration.
+    fn get_labels(&self) -> Option<&HashMap<String, String>>;
+}
+
+/// Wrapper for worker list that can be serialized
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct WorkerList {
+    /// Worker URLs (we can't serialize Arc<dyn Worker>, so we store URLs)
+    pub worker_urls: Vec<String>,
+}
+
+impl WorkerList {
+    pub fn new() -> Self {
+        Self {
+            worker_urls: Vec::new(),
+        }
+    }
+
+    pub fn from_workers(workers: &[Arc<dyn Worker>]) -> Self {
+        Self {
+            worker_urls: workers.iter().map(|w| w.url().to_string()).collect(),
+        }
+    }
+}
+
+// ============================================================================
+// Workflow-specific data types
+// ============================================================================
+
+/// Data for tokenizer registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TokenizerWorkflowData {
+    pub config: TokenizerConfigRequest,
+    pub vocab_size: Option<usize>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+}
+
+impl WorkflowData for TokenizerWorkflowData {
+    fn workflow_type() -> &'static str {
+        "tokenizer_registration"
+    }
+}
+
+impl TokenizerWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    ///
+    /// Call this after deserializing workflow state to ensure runtime fields
+    /// have been repopulated.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for local worker registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LocalWorkerWorkflowData {
+    pub config: ProtocolWorkerConfigRequest,
+    pub connection_mode: Option<crate::core::ConnectionMode>,
+    pub discovered_labels: HashMap<String, String>,
+    pub dp_info: Option<super::worker::local::DpInfo>,
+    pub workers: Option<WorkerList>,
+    pub final_labels: HashMap<String, String>,
+    /// Detected runtime type (for gRPC workers)
+    pub detected_runtime_type: Option<String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Actual worker objects (transient, not serialized)
+    #[serde(skip, default)]
+    pub actual_workers: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for LocalWorkerWorkflowData {
+    fn workflow_type() -> &'static str {
+        "local_worker_registration"
+    }
+}
+
+impl LocalWorkerWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+impl WorkerRegistrationData for LocalWorkerWorkflowData {
+    fn get_app_context(&self) -> Option<&Arc<AppContext>> {
+        self.app_context.as_ref()
+    }
+
+    fn get_actual_workers(&self) -> Option<&Vec<Arc<dyn Worker>>> {
+        self.actual_workers.as_ref()
+    }
+
+    fn get_labels(&self) -> Option<&HashMap<String, String>> {
+        Some(&self.final_labels)
+    }
+}
+
+/// Data for external worker registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExternalWorkerWorkflowData {
+    pub config: ProtocolWorkerConfigRequest,
+    /// Discovered model cards from /v1/models endpoint
+    pub model_cards: Vec<ModelCard>,
+    pub workers: Option<WorkerList>,
+    /// Labels for policies (derived from config)
+    pub labels: HashMap<String, String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Actual worker objects (transient, not serialized)
+    #[serde(skip, default)]
+    pub actual_workers: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for ExternalWorkerWorkflowData {
+    fn workflow_type() -> &'static str {
+        "external_worker_registration"
+    }
+}
+
+impl ExternalWorkerWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+impl WorkerRegistrationData for ExternalWorkerWorkflowData {
+    fn get_app_context(&self) -> Option<&Arc<AppContext>> {
+        self.app_context.as_ref()
+    }
+
+    fn get_actual_workers(&self) -> Option<&Vec<Arc<dyn Worker>>> {
+        self.actual_workers.as_ref()
+    }
+
+    fn get_labels(&self) -> Option<&HashMap<String, String>> {
+        Some(&self.labels)
+    }
+}
+
+/// Data for worker removal workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WorkerRemovalWorkflowData {
+    pub config: WorkerRemovalRequest,
+    pub workers_to_remove: Option<WorkerList>,
+    /// URLs of workers being removed
+    pub worker_urls: Vec<String>,
+    /// Model IDs affected by the removal
+    pub affected_models: std::collections::HashSet<String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Actual worker objects to remove (transient, not serialized)
+    #[serde(skip, default)]
+    pub actual_workers_to_remove: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for WorkerRemovalWorkflowData {
+    fn workflow_type() -> &'static str {
+        "worker_removal"
+    }
+}
+
+impl WorkerRemovalWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for worker update workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WorkerUpdateWorkflowData {
+    pub config: ProtocolWorkerUpdateRequest,
+    /// URL of worker(s) to update
+    pub worker_url: String,
+    /// Whether to update all DP-aware workers with matching prefix
+    pub dp_aware: bool,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Workers to update (transient, not serialized)
+    #[serde(skip, default)]
+    pub workers_to_update: Option<Vec<Arc<dyn Worker>>>,
+    /// Updated worker objects (transient, not serialized)
+    #[serde(skip, default)]
+    pub updated_workers: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for WorkerUpdateWorkflowData {
+    fn workflow_type() -> &'static str {
+        "worker_update"
+    }
+}
+
+impl WorkerUpdateWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for MCP server registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct McpWorkflowData {
+    pub config: McpServerConfigRequest,
+    pub validated: bool,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Connected MCP client (transient, not serialized)
+    #[serde(skip, default)]
+    pub mcp_client: Option<Arc<rmcp::service::RunningService<rmcp::RoleClient, ()>>>,
+}
+
+impl WorkflowData for McpWorkflowData {
+    fn workflow_type() -> &'static str {
+        "mcp_registration"
+    }
+}
+
+impl McpWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for WASM module registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WasmRegistrationWorkflowData {
+    pub config: WasmModuleConfigRequest,
+    pub wasm_bytes: Option<Vec<u8>>,
+    /// SHA256 hash of the module file (32 bytes)
+    pub sha256_hash: Option<[u8; 32]>,
+    /// File size in bytes
+    pub file_size_bytes: Option<u64>,
+    /// UUID assigned to the registered module
+    pub module_uuid: Option<uuid::Uuid>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+}
+
+impl WorkflowData for WasmRegistrationWorkflowData {
+    fn workflow_type() -> &'static str {
+        "wasm_module_registration"
+    }
+}
+
+impl WasmRegistrationWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for WASM module removal workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WasmRemovalWorkflowData {
+    pub config: WasmModuleRemovalRequest,
+    pub module_id: Option<String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+}
+
+impl WorkflowData for WasmRemovalWorkflowData {
+    fn workflow_type() -> &'static str {
+        "wasm_module_removal"
+    }
+}
+
+impl WasmRemovalWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/steps/workflow_engines.rs

@@ -0,0 +1,166 @@
+//! Typed workflow engines collection
+//!
+//! This module provides a collection of typed workflow engines for different workflow types.
+//! Each workflow type has its own engine with compile-time type safety.
+
+use std::sync::Arc;
+
+use wfaas::{EventSubscriber, InMemoryStore, WorkflowEngine};
+
+use super::{
+    create_external_worker_workflow, create_local_worker_workflow,
+    create_mcp_registration_workflow, create_tokenizer_registration_workflow,
+    create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
+    create_worker_removal_workflow, create_worker_update_workflow, ExternalWorkerWorkflowData,
+    LocalWorkerWorkflowData, McpWorkflowData, TokenizerWorkflowData, WasmRegistrationWorkflowData,
+    WasmRemovalWorkflowData, WorkerRemovalWorkflowData, WorkerUpdateWorkflowData,
+};
+use crate::config::RouterConfig;
+
+/// Type alias for local worker workflow engine
+pub type LocalWorkerEngine =
+    WorkflowEngine<LocalWorkerWorkflowData, InMemoryStore<LocalWorkerWorkflowData>>;
+
+/// Type alias for external worker workflow engine
+pub type ExternalWorkerEngine =
+    WorkflowEngine<ExternalWorkerWorkflowData, InMemoryStore<ExternalWorkerWorkflowData>>;
+
+/// Type alias for worker removal workflow engine
+pub type WorkerRemovalEngine =
+    WorkflowEngine<WorkerRemovalWorkflowData, InMemoryStore<WorkerRemovalWorkflowData>>;
+
+/// Type alias for worker update workflow engine
+pub type WorkerUpdateEngine =
+    WorkflowEngine<WorkerUpdateWorkflowData, InMemoryStore<WorkerUpdateWorkflowData>>;
+
+/// Type alias for MCP registration workflow engine
+pub type McpEngine = WorkflowEngine<McpWorkflowData, InMemoryStore<McpWorkflowData>>;
+
+/// Type alias for tokenizer registration workflow engine
+pub type TokenizerEngine =
+    WorkflowEngine<TokenizerWorkflowData, InMemoryStore<TokenizerWorkflowData>>;
+
+/// Type alias for WASM registration workflow engine
+pub type WasmRegistrationEngine =
+    WorkflowEngine<WasmRegistrationWorkflowData, InMemoryStore<WasmRegistrationWorkflowData>>;
+
+/// Type alias for WASM removal workflow engine
+pub type WasmRemovalEngine =
+    WorkflowEngine<WasmRemovalWorkflowData, InMemoryStore<WasmRemovalWorkflowData>>;
+
+/// Collection of typed workflow engines
+///
+/// Each workflow type has its own engine with compile-time type safety.
+/// This replaces the old `WorkflowEngine<AnyWorkflowData, ...>` approach.
+#[derive(Clone, Debug)]
+pub struct WorkflowEngines {
+    /// Engine for local worker registration workflows
+    pub local_worker: Arc<LocalWorkerEngine>,
+    /// Engine for external worker registration workflows
+    pub external_worker: Arc<ExternalWorkerEngine>,
+    /// Engine for worker removal workflows
+    pub worker_removal: Arc<WorkerRemovalEngine>,
+    /// Engine for worker update workflows
+    pub worker_update: Arc<WorkerUpdateEngine>,
+    /// Engine for MCP server registration workflows
+    pub mcp: Arc<McpEngine>,
+    /// Engine for tokenizer registration workflows
+    pub tokenizer: Arc<TokenizerEngine>,
+    /// Engine for WASM module registration workflows
+    pub wasm_registration: Arc<WasmRegistrationEngine>,
+    /// Engine for WASM module removal workflows
+    pub wasm_removal: Arc<WasmRemovalEngine>,
+}
+
+impl WorkflowEngines {
+    /// Create and initialize all workflow engines with their workflow definitions
+    pub fn new(router_config: &RouterConfig) -> Self {
+        // Create local worker engine
+        let local_worker = WorkflowEngine::new();
+        local_worker
+            .register_workflow(create_local_worker_workflow(router_config))
+            .expect("local_worker_registration workflow should be valid");
+
+        // Create external worker engine
+        let external_worker = WorkflowEngine::new();
+        external_worker
+            .register_workflow(create_external_worker_workflow())
+            .expect("external_worker_registration workflow should be valid");
+
+        // Create worker removal engine
+        let worker_removal = WorkflowEngine::new();
+        worker_removal
+            .register_workflow(create_worker_removal_workflow())
+            .expect("worker_removal workflow should be valid");
+
+        // Create worker update engine
+        let worker_update = WorkflowEngine::new();
+        worker_update
+            .register_workflow(create_worker_update_workflow())
+            .expect("worker_update workflow should be valid");
+
+        // Create MCP engine
+        let mcp = WorkflowEngine::new();
+        mcp.register_workflow(create_mcp_registration_workflow())
+            .expect("mcp_registration workflow should be valid");
+
+        // Create tokenizer engine
+        let tokenizer = WorkflowEngine::new();
+        tokenizer
+            .register_workflow(create_tokenizer_registration_workflow())
+            .expect("tokenizer_registration workflow should be valid");
+
+        // Create WASM registration engine
+        let wasm_registration = WorkflowEngine::new();
+        wasm_registration
+            .register_workflow(create_wasm_module_registration_workflow())
+            .expect("wasm_module_registration workflow should be valid");
+
+        // Create WASM removal engine
+        let wasm_removal = WorkflowEngine::new();
+        wasm_removal
+            .register_workflow(create_wasm_module_removal_workflow())
+            .expect("wasm_module_removal workflow should be valid");
+
+        Self {
+            local_worker: Arc::new(local_worker),
+            external_worker: Arc::new(external_worker),
+            worker_removal: Arc::new(worker_removal),
+            worker_update: Arc::new(worker_update),
+            mcp: Arc::new(mcp),
+            tokenizer: Arc::new(tokenizer),
+            wasm_registration: Arc::new(wasm_registration),
+            wasm_removal: Arc::new(wasm_removal),
+        }
+    }
+
+    /// Subscribe an event subscriber to all workflow engines
+    pub async fn subscribe_all<S: EventSubscriber + 'static>(&self, subscriber: Arc<S>) {
+        self.local_worker
+            .event_bus()
+            .subscribe(subscriber.clone())
+            .await;
+        self.external_worker
+            .event_bus()
+            .subscribe(subscriber.clone())
+            .await;
+        self.worker_removal
+            .event_bus()
+            .subscribe(subscriber.clone())
+            .await;
+        self.worker_update
+            .event_bus()
+            .subscribe(subscriber.clone())
+            .await;
+        self.mcp.event_bus().subscribe(subscriber.clone()).await;
+        self.tokenizer
+            .event_bus()
+            .subscribe(subscriber.clone())
+            .await;
+        self.wasm_registration
+            .event_bus()
+            .subscribe(subscriber.clone())
+            .await;
+        self.wasm_removal.event_bus().subscribe(subscriber).await;
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/token_bucket.rs

@@ -0,0 +1,279 @@
+use std::{
+    sync::Arc,
+    time::{Duration, Instant},
+};
+
+use parking_lot::Mutex;
+use tokio::sync::Notify;
+use tracing::{debug, trace};
+
+/// Token bucket for rate limiting.
+///
+/// This implementation provides:
+/// - Smooth rate limiting with configurable refill rate
+/// - Burst capacity handling
+/// - Fair queuing for waiting requests via Notify
+/// - Sync token return for Drop handlers (via `return_tokens_sync`)
+///
+/// Uses `parking_lot::Mutex` for sync-compatible locking (no async required).
+#[derive(Clone)]
+pub struct TokenBucket {
+    inner: Arc<Mutex<TokenBucketInner>>,
+    notify: Arc<Notify>,
+    capacity: f64,
+    refill_rate: f64, // tokens per second
+}
+
+struct TokenBucketInner {
+    tokens: f64,
+    last_refill: Instant,
+}
+
+impl TokenBucket {
+    /// Create a new token bucket
+    ///
+    /// # Arguments
+    /// * `capacity` - Maximum number of tokens (burst capacity)
+    /// * `refill_rate` - Tokens added per second (0 for pure concurrency limiting)
+    pub fn new(capacity: usize, refill_rate: usize) -> Self {
+        let capacity = capacity as f64;
+        // Allow refill_rate=0 for pure concurrency limiting (semaphore behavior)
+        // When refill_rate=0, tokens are only returned via return_tokens()
+        let refill_rate = refill_rate as f64;
+
+        Self {
+            inner: Arc::new(Mutex::new(TokenBucketInner {
+                tokens: capacity,
+                last_refill: Instant::now(),
+            })),
+            notify: Arc::new(Notify::new()),
+            capacity,
+            refill_rate,
+        }
+    }
+
+    /// Try to acquire tokens immediately.
+    ///
+    /// Returns `Ok(())` if tokens were acquired, `Err(())` if insufficient tokens.
+    pub async fn try_acquire(&self, tokens: f64) -> Result<(), ()> {
+        self.try_acquire_sync(tokens)
+    }
+
+    /// Sync version of try_acquire (for internal use).
+    fn try_acquire_sync(&self, tokens: f64) -> Result<(), ()> {
+        let mut inner = self.inner.lock();
+
+        let now = Instant::now();
+        let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
+        let refill_amount = elapsed * self.refill_rate;
+
+        inner.tokens = (inner.tokens + refill_amount).min(self.capacity);
+        inner.last_refill = now;
+
+        trace!(
+            "Token bucket: {} tokens available, requesting {}",
+            inner.tokens,
+            tokens
+        );
+
+        if inner.tokens >= tokens {
+            inner.tokens -= tokens;
+            debug!(
+                "Token bucket: acquired {} tokens, {} remaining",
+                tokens, inner.tokens
+            );
+            Ok(())
+        } else {
+            Err(())
+        }
+    }
+
+    /// Acquire tokens, waiting if necessary.
+    ///
+    /// When `refill_rate=0`, waits indefinitely for tokens to be returned via `return_tokens()`.
+    /// Use `acquire_timeout()` to set an appropriate timeout.
+    pub async fn acquire(&self, tokens: f64) -> Result<(), tokio::time::error::Elapsed> {
+        if self.try_acquire(tokens).await.is_ok() {
+            return Ok(());
+        }
+
+        // When refill_rate=0 (pure concurrency limiting), tokens only come back
+        // via return_tokens(), so we wait on notify signal only.
+        if self.refill_rate == 0.0 {
+            debug!(
+                "Token bucket: waiting indefinitely for {} tokens (refill_rate=0)",
+                tokens
+            );
+
+            loop {
+                // Wait for notify signal from return_tokens()
+                self.notify.notified().await;
+
+                if self.try_acquire(tokens).await.is_ok() {
+                    return Ok(());
+                }
+            }
+        }
+
+        let wait_time = {
+            let inner = self.inner.lock();
+            let tokens_needed = tokens - inner.tokens;
+            let wait_secs = (tokens_needed / self.refill_rate).max(0.0);
+            Duration::from_secs_f64(wait_secs)
+        };
+
+        debug!(
+            "Token bucket: waiting {:?} for {} tokens",
+            wait_time, tokens
+        );
+
+        tokio::time::timeout(wait_time, async {
+            loop {
+                if self.try_acquire(tokens).await.is_ok() {
+                    return;
+                }
+                tokio::select! {
+                    _ = self.notify.notified() => {},
+                    _ = tokio::time::sleep(Duration::from_millis(10)) => {},
+                }
+            }
+        })
+        .await?;
+
+        Ok(())
+    }
+
+    /// Acquire tokens with custom timeout.
+    pub async fn acquire_timeout(
+        &self,
+        tokens: f64,
+        timeout: Duration,
+    ) -> Result<(), tokio::time::error::Elapsed> {
+        tokio::time::timeout(timeout, self.acquire(tokens)).await?
+    }
+
+    /// Return tokens to the bucket (sync version).
+    ///
+    /// This is safe to call from sync contexts (e.g., Drop handlers).
+    /// Uses `parking_lot::Mutex` which never blocks indefinitely.
+    pub fn return_tokens_sync(&self, tokens: f64) {
+        {
+            let mut inner = self.inner.lock();
+            inner.tokens = (inner.tokens + tokens).min(self.capacity);
+            debug!(
+                "Token bucket: returned {} tokens, {} available",
+                tokens, inner.tokens
+            );
+        } // Release lock before notify
+        self.notify.notify_waiters();
+    }
+
+    /// Return tokens to the bucket (async version for API compatibility).
+    pub async fn return_tokens(&self, tokens: f64) {
+        self.return_tokens_sync(tokens);
+    }
+
+    /// Get current available tokens (for monitoring).
+    pub async fn available_tokens(&self) -> f64 {
+        let mut inner = self.inner.lock();
+
+        let now = Instant::now();
+        let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
+        let refill_amount = elapsed * self.refill_rate;
+
+        inner.tokens = (inner.tokens + refill_amount).min(self.capacity);
+        inner.last_refill = now;
+
+        inner.tokens
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_token_bucket_basic() {
+        let bucket = TokenBucket::new(10, 5);
+
+        assert!(bucket.try_acquire(5.0).await.is_ok());
+        assert!(bucket.try_acquire(5.0).await.is_ok());
+
+        assert!(bucket.try_acquire(1.0).await.is_err());
+
+        tokio::time::sleep(Duration::from_millis(300)).await;
+
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_token_bucket_refill() {
+        let bucket = TokenBucket::new(10, 10);
+
+        assert!(bucket.try_acquire(10.0).await.is_ok());
+
+        tokio::time::sleep(Duration::from_millis(500)).await;
+
+        let available = bucket.available_tokens().await;
+        assert!((4.0..=6.0).contains(&available));
+    }
+
+    #[tokio::test]
+    async fn test_token_bucket_zero_refill_rate() {
+        // With refill_rate=0, tokens should only come back via return_tokens()
+        let bucket = TokenBucket::new(2, 0);
+
+        // Acquire both tokens
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+
+        // No more tokens available
+        assert!(bucket.try_acquire(1.0).await.is_err());
+
+        // Wait - should NOT refill automatically
+        tokio::time::sleep(Duration::from_millis(500)).await;
+        assert!(bucket.try_acquire(1.0).await.is_err());
+
+        // Return a token - now we should be able to acquire
+        bucket.return_tokens(1.0).await;
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+
+        // No more tokens again
+        assert!(bucket.try_acquire(1.0).await.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_token_bucket_zero_refill_with_notify() {
+        // Test that acquire wakes up when tokens are returned
+        let bucket = Arc::new(TokenBucket::new(1, 0));
+
+        // Acquire the only token
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+
+        let bucket_clone = bucket.clone();
+
+        // Spawn a task that will return the token after a delay
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            bucket_clone.return_tokens(1.0).await;
+        });
+
+        // This should wait and then succeed when token is returned
+        let result = bucket.acquire_timeout(1.0, Duration::from_secs(1)).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_return_tokens_sync() {
+        // Test that sync return works correctly
+        let bucket = TokenBucket::new(2, 0);
+
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+        assert!(bucket.try_acquire(1.0).await.is_err());
+
+        // Use sync return
+        bucket.return_tokens_sync(1.0);
+        assert!(bucket.try_acquire(1.0).await.is_ok());
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/worker_builder.rs

@@ -0,0 +1,553 @@
+use std::collections::HashMap;
+
+use super::{
+    circuit_breaker::{CircuitBreaker, CircuitBreakerConfig},
+    model_card::ModelCard,
+    model_type::ModelType,
+    worker::{
+        BasicWorker, ConnectionMode, DPAwareWorker, HealthConfig, RuntimeType, WorkerMetadata,
+        WorkerRoutingKeyLoad, WorkerType,
+    },
+};
+use crate::{observability::metrics::Metrics, routers::grpc::client::GrpcClient};
+
+/// Builder for creating BasicWorker instances with fluent API
+pub struct BasicWorkerBuilder {
+    url: String,
+    api_key: Option<String>,
+    worker_type: WorkerType,
+    connection_mode: ConnectionMode,
+    runtime_type: RuntimeType,
+    labels: HashMap<String, String>,
+    models: Vec<ModelCard>,
+    health_config: HealthConfig,
+    circuit_breaker_config: CircuitBreakerConfig,
+    grpc_client: Option<GrpcClient>,
+}
+
+impl BasicWorkerBuilder {
+    /// Create a new builder with only the URL
+    pub fn new(url: impl Into<String>) -> Self {
+        Self {
+            url: url.into(),
+            api_key: None,
+            worker_type: WorkerType::Regular,
+            connection_mode: ConnectionMode::Http,
+            runtime_type: RuntimeType::default(),
+            labels: HashMap::new(),
+            models: Vec::new(),
+            health_config: HealthConfig::default(),
+            circuit_breaker_config: CircuitBreakerConfig::default(),
+            grpc_client: None,
+        }
+    }
+
+    /// Create a new builder with URL and worker type (for backwards compatibility)
+    pub fn new_with_type(url: impl Into<String>, worker_type: WorkerType) -> Self {
+        Self {
+            url: url.into(),
+            api_key: None,
+            worker_type,
+            connection_mode: ConnectionMode::Http,
+            runtime_type: RuntimeType::default(),
+            labels: HashMap::new(),
+            models: Vec::new(),
+            health_config: HealthConfig::default(),
+            circuit_breaker_config: CircuitBreakerConfig::default(),
+            grpc_client: None,
+        }
+    }
+
+    /// Set the API key
+    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
+        self.api_key = Some(api_key.into());
+        self
+    }
+
+    /// Set the worker type (Regular, Prefill, or Decode)
+    pub fn worker_type(mut self, worker_type: WorkerType) -> Self {
+        self.worker_type = worker_type;
+        self
+    }
+
+    /// Set the connection mode (HTTP or gRPC)
+    pub fn connection_mode(mut self, mode: ConnectionMode) -> Self {
+        self.connection_mode = mode;
+        self
+    }
+
+    /// Set the runtime type (SGLang or vLLM)
+    pub fn runtime_type(mut self, runtime_type: RuntimeType) -> Self {
+        self.runtime_type = runtime_type;
+        self
+    }
+
+    /// Set labels for worker identification
+    pub fn labels(mut self, labels: HashMap<String, String>) -> Self {
+        self.labels = labels;
+        self
+    }
+
+    /// Add a single label
+    pub fn label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
+        self.labels.insert(key.into(), value.into());
+        self
+    }
+
+    /// Set health check configuration
+    pub fn health_config(mut self, config: HealthConfig) -> Self {
+        self.health_config = config;
+        self
+    }
+
+    /// Set circuit breaker configuration
+    pub fn circuit_breaker_config(mut self, config: CircuitBreakerConfig) -> Self {
+        self.circuit_breaker_config = config;
+        self
+    }
+
+    /// Set gRPC client for gRPC workers
+    pub fn grpc_client(mut self, client: GrpcClient) -> Self {
+        self.grpc_client = Some(client);
+        self
+    }
+
+    /// Set models this worker can serve
+    pub fn models(mut self, models: Vec<ModelCard>) -> Self {
+        self.models = models;
+        self
+    }
+
+    /// Add a single model this worker can serve
+    pub fn model(mut self, model: ModelCard) -> Self {
+        self.models.push(model);
+        self
+    }
+
+    /// Build the BasicWorker instance
+    pub fn build(self) -> BasicWorker {
+        use std::sync::{
+            atomic::{AtomicBool, AtomicUsize},
+            Arc, RwLock as StdRwLock,
+        };
+
+        use tokio::sync::OnceCell;
+
+        let bootstrap_host = match url::Url::parse(&self.url) {
+            Ok(parsed) => parsed.host_str().unwrap_or("localhost").to_string(),
+            Err(_) if !self.url.contains("://") => {
+                match url::Url::parse(&format!("http://{}", self.url)) {
+                    Ok(parsed) => parsed.host_str().unwrap_or("localhost").to_string(),
+                    Err(_) => {
+                        tracing::warn!(
+                            "Failed to parse URL '{}', defaulting to localhost",
+                            self.url
+                        );
+                        "localhost".to_string()
+                    }
+                }
+            }
+            Err(_) => {
+                tracing::warn!(
+                    "Failed to parse URL '{}', defaulting to localhost",
+                    self.url
+                );
+                "localhost".to_string()
+            }
+        };
+
+        let bootstrap_port = match self.worker_type {
+            WorkerType::Prefill { bootstrap_port } => bootstrap_port,
+            _ => None,
+        };
+
+        let metadata = WorkerMetadata {
+            url: self.url.clone(),
+            api_key: self.api_key,
+            worker_type: self.worker_type,
+            connection_mode: self.connection_mode,
+            runtime_type: self.runtime_type,
+            labels: self.labels,
+            health_config: self.health_config,
+            bootstrap_host,
+            bootstrap_port,
+            models: self.models,                // Empty = accepts any model
+            default_provider: None,             // Native/passthrough
+            default_model_type: ModelType::LLM, // Standard LLM capabilities
+        };
+
+        // Use OnceCell for lock-free gRPC client access after initialization
+        let grpc_client = Arc::new(match self.grpc_client {
+            Some(client) => {
+                let cell = OnceCell::new();
+                // Pre-set the client if provided (blocking set is fine during construction)
+                cell.set(Arc::new(client)).ok();
+                cell
+            }
+            None => OnceCell::new(),
+        });
+
+        let healthy = true;
+        Metrics::set_worker_health(&self.url, healthy);
+
+        BasicWorker {
+            metadata,
+            load_counter: Arc::new(AtomicUsize::new(0)),
+            worker_routing_key_load: Arc::new(WorkerRoutingKeyLoad::new(&self.url)),
+            processed_counter: Arc::new(AtomicUsize::new(0)),
+            healthy: Arc::new(AtomicBool::new(healthy)),
+            consecutive_failures: Arc::new(AtomicUsize::new(0)),
+            consecutive_successes: Arc::new(AtomicUsize::new(0)),
+            circuit_breaker: CircuitBreaker::with_config_and_label(
+                self.circuit_breaker_config,
+                self.url.clone(),
+            ),
+            grpc_client,
+            models_override: Arc::new(StdRwLock::new(None)),
+        }
+    }
+}
+
+/// Builder for creating DPAwareWorker instances with fluent API
+pub struct DPAwareWorkerBuilder {
+    base_url: String,
+    api_key: Option<String>,
+    dp_rank: usize,
+    dp_size: usize,
+    worker_type: WorkerType,
+    connection_mode: ConnectionMode,
+    runtime_type: RuntimeType,
+    labels: HashMap<String, String>,
+    models: Vec<ModelCard>,
+    health_config: HealthConfig,
+    circuit_breaker_config: CircuitBreakerConfig,
+    grpc_client: Option<GrpcClient>,
+}
+
+impl DPAwareWorkerBuilder {
+    /// Create a new DP-aware worker builder
+    pub fn new(base_url: impl Into<String>, dp_rank: usize, dp_size: usize) -> Self {
+        Self {
+            base_url: base_url.into(),
+            api_key: None,
+            dp_rank,
+            dp_size,
+            worker_type: WorkerType::Regular,
+            connection_mode: ConnectionMode::Http,
+            runtime_type: RuntimeType::default(),
+            labels: HashMap::new(),
+            models: Vec::new(),
+            health_config: HealthConfig::default(),
+            circuit_breaker_config: CircuitBreakerConfig::default(),
+            grpc_client: None,
+        }
+    }
+
+    /// Create a new DP-aware worker builder with worker type (for backwards compatibility)
+    pub fn new_with_type(
+        base_url: impl Into<String>,
+        dp_rank: usize,
+        dp_size: usize,
+        worker_type: WorkerType,
+    ) -> Self {
+        Self {
+            base_url: base_url.into(),
+            api_key: None,
+            dp_rank,
+            dp_size,
+            worker_type,
+            connection_mode: ConnectionMode::Http,
+            runtime_type: RuntimeType::default(),
+            labels: HashMap::new(),
+            models: Vec::new(),
+            health_config: HealthConfig::default(),
+            circuit_breaker_config: CircuitBreakerConfig::default(),
+            grpc_client: None,
+        }
+    }
+
+    /// Set the API key
+    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
+        self.api_key = Some(api_key.into());
+        self
+    }
+
+    /// Set the worker type (Regular, Prefill, or Decode)
+    pub fn worker_type(mut self, worker_type: WorkerType) -> Self {
+        self.worker_type = worker_type;
+        self
+    }
+
+    /// Set the connection mode (HTTP or gRPC)
+    pub fn connection_mode(mut self, mode: ConnectionMode) -> Self {
+        self.connection_mode = mode;
+        self
+    }
+
+    /// Set the runtime type (SGLang or vLLM)
+    pub fn runtime_type(mut self, runtime_type: RuntimeType) -> Self {
+        self.runtime_type = runtime_type;
+        self
+    }
+
+    /// Set labels for worker identification
+    pub fn labels(mut self, labels: HashMap<String, String>) -> Self {
+        self.labels = labels;
+        self
+    }
+
+    /// Add a single label
+    pub fn label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
+        self.labels.insert(key.into(), value.into());
+        self
+    }
+
+    /// Set health check configuration
+    pub fn health_config(mut self, config: HealthConfig) -> Self {
+        self.health_config = config;
+        self
+    }
+
+    /// Set circuit breaker configuration
+    pub fn circuit_breaker_config(mut self, config: CircuitBreakerConfig) -> Self {
+        self.circuit_breaker_config = config;
+        self
+    }
+
+    /// Set gRPC client for gRPC workers
+    pub fn grpc_client(mut self, client: GrpcClient) -> Self {
+        self.grpc_client = Some(client);
+        self
+    }
+
+    /// Set models this worker can serve
+    pub fn models(mut self, models: Vec<ModelCard>) -> Self {
+        self.models = models;
+        self
+    }
+
+    /// Add a single model this worker can serve
+    pub fn model(mut self, model: ModelCard) -> Self {
+        self.models.push(model);
+        self
+    }
+
+    /// Build the DPAwareWorker instance
+    pub fn build(self) -> DPAwareWorker {
+        let worker_url = format!("{}@{}", self.base_url, self.dp_rank);
+        let mut builder = BasicWorkerBuilder::new(worker_url)
+            .models(self.models)
+            .worker_type(self.worker_type)
+            .connection_mode(self.connection_mode)
+            .runtime_type(self.runtime_type)
+            .labels(self.labels)
+            .health_config(self.health_config)
+            .circuit_breaker_config(self.circuit_breaker_config);
+
+        if let Some(client) = self.grpc_client {
+            builder = builder.grpc_client(client);
+        }
+        if let Some(api_key) = self.api_key {
+            builder = builder.api_key(api_key);
+        }
+
+        let base_worker = builder.build();
+        DPAwareWorker::with_base_worker(base_worker, self.base_url, self.dp_rank, self.dp_size)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::*;
+    use crate::core::worker::Worker;
+
+    #[test]
+    fn test_basic_worker_builder_minimal() {
+        let worker = BasicWorkerBuilder::new("http://localhost:8080").build();
+
+        assert_eq!(worker.url(), "http://localhost:8080");
+        assert_eq!(worker.worker_type(), &WorkerType::Regular);
+        assert_eq!(worker.connection_mode(), &ConnectionMode::Http);
+        assert!(worker.is_healthy());
+    }
+
+    #[test]
+    fn test_basic_worker_builder_with_type() {
+        let worker = BasicWorkerBuilder::new("http://localhost:8080")
+            .worker_type(WorkerType::Decode)
+            .build();
+
+        assert_eq!(worker.url(), "http://localhost:8080");
+        assert_eq!(worker.worker_type(), &WorkerType::Decode);
+        assert_eq!(worker.connection_mode(), &ConnectionMode::Http);
+        assert!(worker.is_healthy());
+    }
+
+    #[test]
+    fn test_basic_worker_builder_full() {
+        let mut labels = HashMap::new();
+        labels.insert("env".to_string(), "prod".to_string());
+        labels.insert("region".to_string(), "us-east".to_string());
+
+        let health_config = HealthConfig {
+            endpoint: "/health".to_string(),
+            timeout_secs: 30,
+            check_interval_secs: 60,
+            failure_threshold: 3,
+            success_threshold: 2,
+            disable_health_check: false,
+        };
+
+        let cb_config = CircuitBreakerConfig {
+            failure_threshold: 10,
+            success_threshold: 5,
+            timeout_duration: Duration::from_millis(2000),
+            window_duration: Duration::from_millis(30000),
+        };
+
+        let worker = BasicWorkerBuilder::new("http://localhost:8080")
+            .worker_type(WorkerType::Prefill {
+                bootstrap_port: None,
+            })
+            .connection_mode(ConnectionMode::Grpc { port: Some(50051) })
+            .labels(labels.clone())
+            .health_config(health_config.clone())
+            .circuit_breaker_config(cb_config)
+            .build();
+
+        assert_eq!(worker.url(), "http://localhost:8080");
+        assert_eq!(
+            worker.worker_type(),
+            &WorkerType::Prefill {
+                bootstrap_port: None
+            }
+        );
+        assert_eq!(
+            worker.connection_mode(),
+            &ConnectionMode::Grpc { port: Some(50051) }
+        );
+        assert_eq!(worker.metadata().labels, labels);
+        assert_eq!(
+            worker.metadata().health_config.endpoint,
+            health_config.endpoint
+        );
+        assert_eq!(
+            worker.metadata().health_config.timeout_secs,
+            health_config.timeout_secs
+        );
+        assert_eq!(
+            worker.metadata().health_config.check_interval_secs,
+            health_config.check_interval_secs
+        );
+        assert_eq!(
+            worker.metadata().health_config.failure_threshold,
+            health_config.failure_threshold
+        );
+        assert_eq!(
+            worker.metadata().health_config.success_threshold,
+            health_config.success_threshold
+        );
+    }
+
+    #[test]
+    fn test_basic_worker_builder_with_single_label() {
+        let worker = BasicWorkerBuilder::new("http://localhost:8080")
+            .worker_type(WorkerType::Decode)
+            .label("env", "staging")
+            .label("version", "v1.2.3")
+            .build();
+
+        assert_eq!(
+            worker.metadata().labels.get("env"),
+            Some(&"staging".to_string())
+        );
+        assert_eq!(
+            worker.metadata().labels.get("version"),
+            Some(&"v1.2.3".to_string())
+        );
+    }
+
+    #[test]
+    fn test_dp_aware_worker_builder_minimal() {
+        let worker = DPAwareWorkerBuilder::new("http://localhost:8080", 2, 8).build();
+
+        assert_eq!(worker.url(), "http://localhost:8080@2");
+        assert_eq!(worker.dp_rank(), Some(2));
+        assert_eq!(worker.dp_size(), Some(8));
+        assert_eq!(worker.worker_type(), &WorkerType::Regular);
+    }
+
+    #[test]
+    fn test_dp_aware_worker_builder_full() {
+        let mut labels = HashMap::new();
+        labels.insert("cluster".to_string(), "main".to_string());
+
+        let health_config = HealthConfig {
+            endpoint: "/status".to_string(),
+            timeout_secs: 20,
+            check_interval_secs: 45,
+            failure_threshold: 5,
+            success_threshold: 3,
+            disable_health_check: false,
+        };
+
+        let worker = DPAwareWorkerBuilder::new("http://localhost:8080", 3, 16)
+            .worker_type(WorkerType::Prefill {
+                bootstrap_port: Some(9090),
+            })
+            .connection_mode(ConnectionMode::Http)
+            .labels(labels.clone())
+            .health_config(health_config.clone())
+            .api_key("test_api_key")
+            .build();
+
+        assert_eq!(worker.url(), "http://localhost:8080@3");
+        assert_eq!(worker.dp_rank(), Some(3));
+        assert_eq!(worker.dp_size(), Some(16));
+        assert_eq!(worker.metadata().labels, labels);
+        assert_eq!(
+            worker.metadata().health_config.endpoint,
+            health_config.endpoint
+        );
+        assert_eq!(
+            worker.metadata().health_config.timeout_secs,
+            health_config.timeout_secs
+        );
+        assert_eq!(
+            worker.metadata().health_config.check_interval_secs,
+            health_config.check_interval_secs
+        );
+        assert_eq!(
+            worker.metadata().health_config.failure_threshold,
+            health_config.failure_threshold
+        );
+        assert_eq!(
+            worker.metadata().health_config.success_threshold,
+            health_config.success_threshold
+        );
+    }
+
+    #[test]
+    fn test_dp_aware_worker_with_grpc() {
+        let worker = DPAwareWorkerBuilder::new("grpc://cluster.local", 1, 4)
+            .worker_type(WorkerType::Decode)
+            .connection_mode(ConnectionMode::Grpc { port: Some(50051) })
+            .label("transport", "grpc")
+            .build();
+
+        assert_eq!(worker.url(), "grpc://cluster.local@1");
+        assert_eq!(worker.dp_rank(), Some(1));
+        assert_eq!(worker.dp_size(), Some(4));
+        assert_eq!(worker.worker_type(), &WorkerType::Decode);
+        assert_eq!(
+            worker.connection_mode(),
+            &ConnectionMode::Grpc { port: Some(50051) }
+        );
+        assert_eq!(
+            worker.metadata().labels.get("transport"),
+            Some(&"grpc".to_string())
+        );
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/worker_manager.rs

@@ -0,0 +1,394 @@
+//! Worker Management Module
+//!
+//! Provides worker lifecycle operations and fan-out request utilities.
+
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use axum::response::{IntoResponse, Response};
+use futures::{
+    future,
+    stream::{self, StreamExt},
+};
+use http::StatusCode;
+use serde_json::Value;
+use tokio::{
+    sync::{watch, Mutex},
+    task::JoinHandle,
+};
+use tracing::{debug, info, warn};
+
+use crate::{
+    core::{metrics_aggregator::MetricPack, ConnectionMode, Worker, WorkerRegistry, WorkerType},
+    policies::PolicyRegistry,
+    protocols::worker_spec::{FlushCacheResult, WorkerLoadInfo, WorkerLoadsResult},
+};
+
+const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
+const MAX_CONCURRENT: usize = 32;
+
+/// Result of a fan-out request to a single worker
+struct WorkerResponse {
+    url: String,
+    result: Result<reqwest::Response, reqwest::Error>,
+}
+
+/// Fan out requests to workers in parallel
+async fn fan_out(
+    workers: &[Arc<dyn Worker>],
+    client: &reqwest::Client,
+    endpoint: &str,
+    method: reqwest::Method,
+) -> Vec<WorkerResponse> {
+    let futures: Vec<_> = workers
+        .iter()
+        .map(|worker| {
+            let client = client.clone();
+            let url = worker.url().to_string();
+            let full_url = format!("{}/{}", url, endpoint);
+            let api_key = worker.api_key().clone();
+            let method = method.clone();
+
+            async move {
+                let mut req = client.request(method, &full_url).timeout(REQUEST_TIMEOUT);
+                if let Some(key) = api_key {
+                    req = req.bearer_auth(key);
+                }
+                WorkerResponse {
+                    url,
+                    result: req.send().await,
+                }
+            }
+        })
+        .collect();
+
+    stream::iter(futures)
+        .buffer_unordered(MAX_CONCURRENT)
+        .collect()
+        .await
+}
+
+pub enum EngineMetricsResult {
+    Ok(String),
+    Err(String),
+}
+
+impl IntoResponse for EngineMetricsResult {
+    fn into_response(self) -> Response {
+        match self {
+            Self::Ok(text) => (StatusCode::OK, text).into_response(),
+            Self::Err(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response(),
+        }
+    }
+}
+
+pub struct WorkerManager;
+
+impl WorkerManager {
+    pub fn get_worker_urls(registry: &Arc<WorkerRegistry>) -> Vec<String> {
+        registry
+            .get_all()
+            .iter()
+            .map(|w| w.url().to_string())
+            .collect()
+    }
+
+    pub async fn flush_cache_all(
+        worker_registry: &WorkerRegistry,
+        client: &reqwest::Client,
+    ) -> FlushCacheResult {
+        let workers = worker_registry.get_all();
+        let total_workers = workers.len();
+
+        let http_workers: Vec<_> = workers
+            .into_iter()
+            .filter(|w| matches!(w.connection_mode(), ConnectionMode::Http))
+            .collect();
+
+        if http_workers.is_empty() {
+            return FlushCacheResult {
+                successful: vec![],
+                failed: vec![],
+                total_workers,
+                http_workers: 0,
+                message: "No HTTP workers available for cache flush".to_string(),
+            };
+        }
+
+        info!(
+            "Flushing cache on {} HTTP workers (out of {} total)",
+            http_workers.len(),
+            total_workers
+        );
+
+        let responses = fan_out(&http_workers, client, "flush_cache", reqwest::Method::POST).await;
+
+        let mut successful = Vec::new();
+        let mut failed = Vec::new();
+
+        for resp in responses {
+            match resp.result {
+                Ok(r) if r.status().is_success() => successful.push(resp.url),
+                Ok(r) => failed.push((resp.url, format!("HTTP {}", r.status()))),
+                Err(e) => failed.push((resp.url, e.to_string())),
+            }
+        }
+
+        let message = if failed.is_empty() {
+            format!(
+                "Successfully flushed cache on all {} HTTP workers",
+                successful.len()
+            )
+        } else {
+            format!(
+                "Cache flush: {} succeeded, {} failed",
+                successful.len(),
+                failed.len()
+            )
+        };
+
+        info!("{}", message);
+
+        FlushCacheResult {
+            successful,
+            failed,
+            total_workers,
+            http_workers: http_workers.len(),
+            message,
+        }
+    }
+
+    pub async fn get_all_worker_loads(
+        worker_registry: &WorkerRegistry,
+        client: &reqwest::Client,
+    ) -> WorkerLoadsResult {
+        let workers = worker_registry.get_all();
+        let total_workers = workers.len();
+
+        let futures: Vec<_> = workers
+            .iter()
+            .map(|worker| {
+                let url = worker.url().to_string();
+                let api_key = worker.api_key().clone();
+                let worker_type = match worker.worker_type() {
+                    WorkerType::Regular => None,
+                    WorkerType::Prefill { .. } => Some("prefill".to_string()),
+                    WorkerType::Decode => Some("decode".to_string()),
+                };
+                let is_http = matches!(worker.connection_mode(), ConnectionMode::Http);
+                let client = client.clone();
+
+                async move {
+                    let load = if is_http {
+                        Self::parse_load_response(&client, &url, api_key.as_deref()).await
+                    } else {
+                        -1
+                    };
+                    WorkerLoadInfo {
+                        worker: url,
+                        worker_type,
+                        load,
+                    }
+                }
+            })
+            .collect();
+
+        let loads = future::join_all(futures).await;
+        let successful = loads.iter().filter(|l| l.load >= 0).count();
+        let failed = loads.iter().filter(|l| l.load < 0).count();
+
+        WorkerLoadsResult {
+            loads,
+            total_workers,
+            successful,
+            failed,
+        }
+    }
+
+    async fn parse_load_response(
+        client: &reqwest::Client,
+        url: &str,
+        api_key: Option<&str>,
+    ) -> isize {
+        let load_url = format!("{}/get_load", url);
+        let mut req = client.get(&load_url).timeout(REQUEST_TIMEOUT);
+        if let Some(key) = api_key {
+            req = req.bearer_auth(key);
+        }
+
+        match req.send().await {
+            Ok(r) if r.status().is_success() => match r.json::<Value>().await {
+                Ok(json) if json.is_array() => json
+                    .as_array()
+                    .unwrap()
+                    .iter()
+                    .filter_map(|e| e.get("num_tokens").and_then(|v| v.as_i64()))
+                    .sum::<i64>() as isize,
+                _ => -1,
+            },
+            _ => -1,
+        }
+    }
+
+    pub async fn get_engine_metrics(
+        worker_registry: &WorkerRegistry,
+        client: &reqwest::Client,
+    ) -> EngineMetricsResult {
+        let workers = worker_registry.get_all();
+
+        if workers.is_empty() {
+            return EngineMetricsResult::Err("No available workers".to_string());
+        }
+
+        let responses = fan_out(&workers, client, "metrics", reqwest::Method::GET).await;
+
+        let mut metric_packs = Vec::new();
+        for resp in responses {
+            if let Ok(r) = resp.result {
+                if r.status().is_success() {
+                    if let Ok(text) = r.text().await {
+                        metric_packs.push(MetricPack {
+                            labels: vec![("worker_addr".into(), resp.url)],
+                            metrics_text: text,
+                        });
+                    }
+                }
+            }
+        }
+
+        if metric_packs.is_empty() {
+            return EngineMetricsResult::Err("All backend requests failed".to_string());
+        }
+
+        match crate::core::metrics_aggregator::aggregate_metrics(metric_packs) {
+            Ok(text) => EngineMetricsResult::Ok(text),
+            Err(e) => EngineMetricsResult::Err(format!("Failed to aggregate metrics: {}", e)),
+        }
+    }
+}
+
+/// Load monitoring service that periodically fetches worker loads
+pub struct LoadMonitor {
+    worker_registry: Arc<WorkerRegistry>,
+    policy_registry: Arc<PolicyRegistry>,
+    client: reqwest::Client,
+    interval: Duration,
+    tx: watch::Sender<HashMap<String, isize>>,
+    rx: watch::Receiver<HashMap<String, isize>>,
+    monitor_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl LoadMonitor {
+    pub fn new(
+        worker_registry: Arc<WorkerRegistry>,
+        policy_registry: Arc<PolicyRegistry>,
+        client: reqwest::Client,
+        interval_secs: u64,
+    ) -> Self {
+        let (tx, rx) = watch::channel(HashMap::new());
+
+        Self {
+            worker_registry,
+            policy_registry,
+            client,
+            interval: Duration::from_secs(interval_secs),
+            tx,
+            rx,
+            monitor_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    pub async fn start(&self) {
+        let mut handle_guard = self.monitor_handle.lock().await;
+        if handle_guard.is_some() {
+            debug!("Load monitoring already running");
+            return;
+        }
+
+        info!(
+            "Starting load monitoring with interval: {:?}",
+            self.interval
+        );
+
+        let worker_registry = Arc::clone(&self.worker_registry);
+        let policy_registry = Arc::clone(&self.policy_registry);
+        let client = self.client.clone();
+        let interval = self.interval;
+        let tx = self.tx.clone();
+
+        let handle = tokio::spawn(async move {
+            Self::monitor_loop(worker_registry, policy_registry, client, interval, tx).await;
+        });
+
+        *handle_guard = Some(handle);
+    }
+
+    pub async fn stop(&self) {
+        let mut handle_guard = self.monitor_handle.lock().await;
+        if let Some(handle) = handle_guard.take() {
+            info!("Stopping load monitoring");
+            handle.abort();
+            let _ = handle.await; // Wait for task to finish
+        }
+    }
+
+    pub fn subscribe(&self) -> watch::Receiver<HashMap<String, isize>> {
+        self.rx.clone()
+    }
+
+    async fn monitor_loop(
+        worker_registry: Arc<WorkerRegistry>,
+        policy_registry: Arc<PolicyRegistry>,
+        client: reqwest::Client,
+        interval: Duration,
+        tx: watch::Sender<HashMap<String, isize>>,
+    ) {
+        let mut interval_timer = tokio::time::interval(interval);
+
+        loop {
+            interval_timer.tick().await;
+
+            let power_of_two_policies = policy_registry.get_all_power_of_two_policies();
+
+            if power_of_two_policies.is_empty() {
+                debug!("No PowerOfTwo policies found, skipping load fetch");
+                continue;
+            }
+
+            let result = WorkerManager::get_all_worker_loads(&worker_registry, &client).await;
+
+            let mut loads = HashMap::new();
+            for load_info in result.loads {
+                loads.insert(load_info.worker, load_info.load);
+            }
+
+            if !loads.is_empty() {
+                debug!(
+                    "Fetched loads from {} workers, updating {} PowerOfTwo policies",
+                    loads.len(),
+                    power_of_two_policies.len()
+                );
+                for policy in &power_of_two_policies {
+                    policy.update_loads(&loads);
+                }
+                let _ = tx.send(loads);
+            } else {
+                warn!("No loads fetched from workers");
+            }
+        }
+    }
+
+    pub async fn is_running(&self) -> bool {
+        let handle_guard = self.monitor_handle.lock().await;
+        handle_guard.is_some()
+    }
+}
+
+impl Drop for LoadMonitor {
+    fn drop(&mut self) {
+        if let Ok(mut handle_guard) = self.monitor_handle.try_lock() {
+            if let Some(handle) = handle_guard.take() {
+                handle.abort();
+            }
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/worker_registry.rs

@@ -0,0 +1,835 @@
+//! Worker Registry for multi-router support
+//!
+//! Provides centralized registry for workers with model-based indexing
+//!
+//! # Performance Optimizations
+//! The model index uses immutable Arc snapshots instead of RwLock for lock-free reads.
+//! This is critical for high-concurrency scenarios where many requests query the same model.
+//!
+//! # Consistent Hash Ring
+//! The registry maintains a pre-computed hash ring per model for O(log n) consistent hashing.
+//! The ring is rebuilt only when workers are added/removed, not per-request.
+//! Uses virtual nodes (150 per worker) for even distribution and blake3 for stable hashing.
+
+use std::sync::{Arc, RwLock};
+
+use dashmap::DashMap;
+use smg_mesh::OptionalMeshSyncManager;
+use uuid::Uuid;
+
+use crate::{
+    core::{
+        circuit_breaker::CircuitState,
+        worker::{HealthChecker, RuntimeType, WorkerType},
+        ConnectionMode, Worker,
+    },
+    observability::metrics::Metrics,
+};
+
+/// Number of virtual nodes per physical worker for even distribution.
+/// 150 is a common choice that provides good balance between memory and distribution.
+const VIRTUAL_NODES_PER_WORKER: usize = 150;
+
+/// Consistent hash ring for O(log n) worker selection.
+///
+/// Each worker is placed at multiple positions (virtual nodes) on the ring
+/// based on hash(worker_url + vnode_index). This provides:
+/// - Even key distribution across workers
+/// - Minimal key redistribution when workers are added/removed (~1/N keys move)
+/// - O(log n) lookup via binary search
+///
+/// Uses blake3 for stable, fast hashing that's consistent across Rust versions.
+#[derive(Debug, Clone)]
+pub struct HashRing {
+    /// Sorted list of (ring_position, worker_url)
+    /// Multiple entries per worker (virtual nodes) for even distribution.
+    /// Uses Arc<str> to share URL across all virtual nodes (150 refs vs 150 copies).
+    entries: Arc<[(u64, Arc<str>)]>,
+}
+
+impl HashRing {
+    /// Build a hash ring from a list of workers.
+    /// Creates VIRTUAL_NODES_PER_WORKER entries per worker for even distribution.
+    pub fn new(workers: &[Arc<dyn Worker>]) -> Self {
+        let mut entries: Vec<(u64, Arc<str>)> =
+            Vec::with_capacity(workers.len() * VIRTUAL_NODES_PER_WORKER);
+
+        for worker in workers {
+            // Create Arc<str> once per worker, share across all virtual nodes
+            let url: Arc<str> = Arc::from(worker.url());
+            let url_bytes = url.as_bytes();
+
+            // Create multiple virtual nodes per worker
+            for vnode in 0..VIRTUAL_NODES_PER_WORKER {
+                let mut hasher = blake3::Hasher::new();
+                hasher.update(url_bytes);
+                hasher.update(b"#");
+                hasher.update(&(vnode as u64).to_le_bytes());
+                let hash = hasher.finalize();
+                let pos = u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap());
+                entries.push((pos, Arc::clone(&url)));
+            }
+        }
+
+        // Sort by ring position for binary search
+        entries.sort_unstable_by_key(|(pos, _)| *pos);
+
+        Self {
+            entries: Arc::from(entries.into_boxed_slice()),
+        }
+    }
+
+    /// Hash a string to a ring position using blake3 (stable across versions).
+    #[inline]
+    fn hash_position(s: &str) -> u64 {
+        let hash = blake3::hash(s.as_bytes());
+        // Take first 8 bytes as u64
+        u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap())
+    }
+
+    /// Find worker URL for a key using consistent hashing.
+    /// Returns the first healthy worker URL at or after the key's position (clockwise).
+    ///
+    /// - `key`: The routing key to hash
+    /// - `is_healthy`: Function to check if a worker URL is healthy
+    pub fn find_healthy_url<F>(&self, key: &str, is_healthy: F) -> Option<&str>
+    where
+        F: Fn(&str) -> bool,
+    {
+        if self.entries.is_empty() {
+            return None;
+        }
+
+        let key_pos = Self::hash_position(key);
+
+        // Binary search to find first entry at or after key_pos
+        let start = self.entries.partition_point(|(pos, _)| *pos < key_pos);
+
+        // Walk clockwise from start, wrapping around
+        // Track visited URLs to avoid checking same worker multiple times (virtual nodes)
+        let mut checked_urls =
+            std::collections::HashSet::with_capacity(self.worker_count().min(16));
+
+        for i in 0..self.entries.len() {
+            let (_, url) = &self.entries[(start + i) % self.entries.len()];
+            let url_str: &str = url;
+
+            // Skip if we already checked this worker (from another virtual node)
+            if !checked_urls.insert(url_str) {
+                continue;
+            }
+
+            if is_healthy(url_str) {
+                return Some(url_str);
+            }
+        }
+
+        None
+    }
+
+    /// Check if the ring is empty
+    pub fn is_empty(&self) -> bool {
+        self.entries.is_empty()
+    }
+
+    /// Get the number of entries in the ring (including virtual nodes)
+    pub fn len(&self) -> usize {
+        self.entries.len()
+    }
+
+    /// Get the number of unique workers in the ring
+    pub fn worker_count(&self) -> usize {
+        self.entries.len() / VIRTUAL_NODES_PER_WORKER.max(1)
+    }
+}
+
+/// Unique identifier for a worker
+#[derive(Debug, Clone, Hash, Eq, PartialEq)]
+pub struct WorkerId(String);
+
+impl WorkerId {
+    /// Create a new worker ID
+    pub fn new() -> Self {
+        Self(Uuid::new_v4().to_string())
+    }
+
+    /// Create a worker ID from a string
+    pub fn from_string(s: String) -> Self {
+        Self(s)
+    }
+
+    /// Get the ID as a string
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl Default for WorkerId {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Model index using immutable snapshots for lock-free reads.
+/// Each model maps to an Arc'd slice of workers that can be read without locking.
+/// Updates create new snapshots (copy-on-write semantics).
+type ModelIndex = Arc<DashMap<String, Arc<[Arc<dyn Worker>]>>>;
+
+/// Worker registry with model-based indexing
+#[derive(Debug)]
+pub struct WorkerRegistry {
+    /// All workers indexed by ID
+    workers: Arc<DashMap<WorkerId, Arc<dyn Worker>>>,
+
+    /// Model index for O(1) lookups using immutable snapshots.
+    /// Uses Arc<[T]> instead of Arc<RwLock<Vec<T>>> for lock-free reads.
+    model_index: ModelIndex,
+
+    /// Consistent hash rings per model for O(log n) routing.
+    /// Rebuilt on worker add/remove (copy-on-write).
+    hash_rings: Arc<DashMap<String, Arc<HashRing>>>,
+
+    /// Workers indexed by worker type
+    type_workers: Arc<DashMap<WorkerType, Vec<WorkerId>>>,
+
+    /// Workers indexed by connection mode
+    connection_workers: Arc<DashMap<ConnectionMode, Vec<WorkerId>>>,
+
+    /// URL to worker ID mapping
+    url_to_id: Arc<DashMap<String, WorkerId>>,
+    /// Optional mesh sync manager for state synchronization
+    /// When None, the registry works independently without mesh synchronization
+    /// Uses RwLock for thread-safe access when setting mesh_sync after initialization
+    mesh_sync: Arc<RwLock<OptionalMeshSyncManager>>,
+}
+
+impl WorkerRegistry {
+    /// Create a new worker registry
+    pub fn new() -> Self {
+        Self {
+            workers: Arc::new(DashMap::new()),
+            model_index: Arc::new(DashMap::new()),
+            hash_rings: Arc::new(DashMap::new()),
+            type_workers: Arc::new(DashMap::new()),
+            connection_workers: Arc::new(DashMap::new()),
+            url_to_id: Arc::new(DashMap::new()),
+            mesh_sync: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    /// Rebuild the hash ring for a model based on current workers in the model index
+    fn rebuild_hash_ring(&self, model_id: &str) {
+        if let Some(workers) = self.model_index.get(model_id) {
+            let ring = HashRing::new(&workers);
+            self.hash_rings.insert(model_id.to_string(), Arc::new(ring));
+        } else {
+            // No workers for this model, remove the ring
+            self.hash_rings.remove(model_id);
+        }
+    }
+
+    /// Get the hash ring for a model (O(1) lookup)
+    pub fn get_hash_ring(&self, model_id: &str) -> Option<Arc<HashRing>> {
+        self.hash_rings.get(model_id).map(|r| Arc::clone(&r))
+    }
+
+    /// Set mesh sync manager (thread-safe, can be called after initialization)
+    pub fn set_mesh_sync(&self, mesh_sync: OptionalMeshSyncManager) {
+        *self.mesh_sync.write().unwrap() = mesh_sync;
+    }
+
+    /// Register a new worker
+    pub fn register(&self, worker: Arc<dyn Worker>) -> WorkerId {
+        let worker_id = if let Some(existing_id) = self.url_to_id.get(worker.url()) {
+            // Worker with this URL already exists, update it
+            existing_id.clone()
+        } else {
+            WorkerId::new()
+        };
+
+        // Store worker
+        self.workers.insert(worker_id.clone(), worker.clone());
+
+        // Update URL mapping
+        self.url_to_id
+            .insert(worker.url().to_string(), worker_id.clone());
+
+        // Update model index for O(1) lookups using copy-on-write
+        // This creates a new immutable snapshot with the added worker
+        let model_id = worker.model_id().to_string();
+        self.model_index
+            .entry(model_id.clone())
+            .and_modify(|existing| {
+                // Create new snapshot with the additional worker
+                let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect();
+                new_workers.push(worker.clone());
+                *existing = Arc::from(new_workers.into_boxed_slice());
+            })
+            .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice()));
+
+        // Rebuild hash ring for this model
+        self.rebuild_hash_ring(&model_id);
+
+        // Update type index (clone needed for DashMap key ownership)
+        self.type_workers
+            .entry(worker.worker_type().clone())
+            .or_default()
+            .push(worker_id.clone());
+
+        // Update connection mode index (clone needed for DashMap key ownership)
+        self.connection_workers
+            .entry(worker.connection_mode().clone())
+            .or_default()
+            .push(worker_id.clone());
+
+        // Sync to mesh if enabled (no-op if mesh is not enabled)
+        if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+            mesh_sync.sync_worker_state(
+                worker_id.as_str().to_string(),
+                worker.model_id().to_string(),
+                worker.url().to_string(),
+                worker.is_healthy(),
+                0.0, // TODO: Get actual load
+            );
+        }
+
+        worker_id
+    }
+
+    /// Reserve (or retrieve) a stable UUID for a worker URL.
+    /// Uses atomic entry API to avoid race conditions between check and insert.
+    pub fn reserve_id_for_url(&self, url: &str) -> WorkerId {
+        self.url_to_id.entry(url.to_string()).or_default().clone()
+    }
+
+    /// Best-effort lookup of the URL for a given worker ID.
+    pub fn get_url_by_id(&self, worker_id: &WorkerId) -> Option<String> {
+        if let Some(worker) = self.get(worker_id) {
+            return Some(worker.url().to_string());
+        }
+        self.url_to_id
+            .iter()
+            .find_map(|entry| (entry.value() == worker_id).then(|| entry.key().clone()))
+    }
+
+    /// Remove a worker by ID
+    pub fn remove(&self, worker_id: &WorkerId) -> Option<Arc<dyn Worker>> {
+        if let Some((_, worker)) = self.workers.remove(worker_id) {
+            // Remove from URL mapping
+            self.url_to_id.remove(worker.url());
+
+            // Remove from model index using copy-on-write
+            // Create new snapshot without the removed worker
+            let worker_url = worker.url();
+            let model_id = worker.model_id().to_string();
+            if let Some(mut entry) = self.model_index.get_mut(&model_id) {
+                let new_workers: Vec<Arc<dyn Worker>> = entry
+                    .iter()
+                    .filter(|w| w.url() != worker_url)
+                    .cloned()
+                    .collect();
+                *entry = Arc::from(new_workers.into_boxed_slice());
+            }
+
+            // Rebuild hash ring for this model
+            self.rebuild_hash_ring(&model_id);
+
+            // Remove from type index
+            if let Some(mut type_workers) = self.type_workers.get_mut(worker.worker_type()) {
+                type_workers.retain(|id| id != worker_id);
+            }
+
+            // Remove from connection mode index
+            if let Some(mut conn_workers) =
+                self.connection_workers.get_mut(worker.connection_mode())
+            {
+                conn_workers.retain(|id| id != worker_id);
+            }
+
+            worker.set_healthy(false);
+            Metrics::remove_worker_metrics(worker.url());
+
+            // Sync removal to mesh if enabled (no-op if mesh is not enabled)
+            if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+                mesh_sync.remove_worker_state(worker_id.as_str());
+            }
+
+            Some(worker)
+        } else {
+            None
+        }
+    }
+
+    /// Remove a worker by URL
+    pub fn remove_by_url(&self, url: &str) -> Option<Arc<dyn Worker>> {
+        if let Some((_, worker_id)) = self.url_to_id.remove(url) {
+            self.remove(&worker_id)
+        } else {
+            None
+        }
+    }
+
+    /// Get a worker by ID
+    pub fn get(&self, worker_id: &WorkerId) -> Option<Arc<dyn Worker>> {
+        self.workers.get(worker_id).map(|entry| entry.clone())
+    }
+
+    /// Get a worker by URL
+    pub fn get_by_url(&self, url: &str) -> Option<Arc<dyn Worker>> {
+        self.url_to_id.get(url).and_then(|id| self.get(&id))
+    }
+
+    /// Empty worker slice constant for returning when no workers found
+    const EMPTY_WORKERS: &'static [Arc<dyn Worker>] = &[];
+
+    /// Get all workers for a model (O(1) optimized, lock-free)
+    /// Returns an Arc to the immutable worker slice - just an atomic refcount bump.
+    /// This is the fastest possible read path with zero contention.
+    pub fn get_by_model(&self, model_id: &str) -> Arc<[Arc<dyn Worker>]> {
+        self.model_index
+            .get(model_id)
+            .map(|workers| Arc::clone(&workers))
+            .unwrap_or_else(|| Arc::from(Self::EMPTY_WORKERS))
+    }
+
+    /// Get all workers by worker type
+    pub fn get_by_type(&self, worker_type: &WorkerType) -> Vec<Arc<dyn Worker>> {
+        self.type_workers
+            .get(worker_type)
+            .map(|ids| ids.iter().filter_map(|id| self.get(id)).collect())
+            .unwrap_or_default()
+    }
+
+    /// Update worker health status and sync to mesh
+    pub fn update_worker_health(&self, worker_id: &WorkerId, is_healthy: bool) {
+        if let Some(worker) = self.workers.get(worker_id) {
+            // Update worker health (if Worker trait has a method for this)
+            // For now, we'll just sync to mesh
+
+            // Sync to mesh if enabled (no-op if mesh is not enabled)
+            if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+                mesh_sync.sync_worker_state(
+                    worker_id.as_str().to_string(),
+                    worker.model_id().to_string(),
+                    worker.url().to_string(),
+                    is_healthy,
+                    0.0, // TODO: Get actual load
+                );
+            }
+        }
+    }
+
+    /// Get all prefill workers (regardless of bootstrap_port)
+    pub fn get_prefill_workers(&self) -> Vec<Arc<dyn Worker>> {
+        self.workers
+            .iter()
+            .filter_map(|entry| {
+                let worker = entry.value();
+                match worker.worker_type() {
+                    WorkerType::Prefill { .. } => Some(worker.clone()),
+                    _ => None,
+                }
+            })
+            .collect()
+    }
+
+    /// Get all decode workers
+    pub fn get_decode_workers(&self) -> Vec<Arc<dyn Worker>> {
+        self.get_by_type(&WorkerType::Decode)
+    }
+
+    /// Get all workers by connection mode
+    pub fn get_by_connection(&self, connection_mode: &ConnectionMode) -> Vec<Arc<dyn Worker>> {
+        self.connection_workers
+            .get(connection_mode)
+            .map(|ids| ids.iter().filter_map(|id| self.get(id)).collect())
+            .unwrap_or_default()
+    }
+
+    /// Get the number of workers in the registry
+    pub fn len(&self) -> usize {
+        self.workers.len()
+    }
+
+    /// Check if the registry is empty
+    pub fn is_empty(&self) -> bool {
+        self.workers.is_empty()
+    }
+
+    /// Get all workers
+    pub fn get_all(&self) -> Vec<Arc<dyn Worker>> {
+        self.workers
+            .iter()
+            .map(|entry| entry.value().clone())
+            .collect()
+    }
+
+    /// Get all workers with their IDs
+    pub fn get_all_with_ids(&self) -> Vec<(WorkerId, Arc<dyn Worker>)> {
+        self.workers
+            .iter()
+            .map(|entry| (entry.key().clone(), entry.value().clone()))
+            .collect()
+    }
+
+    /// Get all worker URLs
+    pub fn get_all_urls(&self) -> Vec<String> {
+        self.workers
+            .iter()
+            .map(|entry| entry.value().url().to_string())
+            .collect()
+    }
+
+    pub fn get_all_urls_with_api_key(&self) -> Vec<(String, Option<String>)> {
+        self.workers
+            .iter()
+            .map(|entry| {
+                (
+                    entry.value().url().to_string(),
+                    entry.value().api_key().clone(),
+                )
+            })
+            .collect()
+    }
+
+    /// Get all model IDs with workers (lock-free)
+    pub fn get_models(&self) -> Vec<String> {
+        self.model_index
+            .iter()
+            .filter(|entry| !entry.value().is_empty())
+            .map(|entry| entry.key().clone())
+            .collect()
+    }
+
+    /// Get workers filtered by multiple criteria
+    ///
+    /// This method allows flexible filtering of workers based on:
+    /// - model_id: Filter by specific model
+    /// - worker_type: Filter by worker type (Regular, Prefill, Decode)
+    /// - connection_mode: Filter by connection mode (Http, Grpc)
+    /// - runtime_type: Filter by runtime type (Sglang, Vllm, External)
+    /// - healthy_only: Only return healthy workers
+    pub fn get_workers_filtered(
+        &self,
+        model_id: Option<&str>,
+        worker_type: Option<WorkerType>,
+        connection_mode: Option<ConnectionMode>,
+        runtime_type: Option<RuntimeType>,
+        healthy_only: bool,
+    ) -> Vec<Arc<dyn Worker>> {
+        // Start with the most efficient collection based on filters
+        // Use model index when possible as it's O(1) lookup
+        let workers: Vec<Arc<dyn Worker>> = if let Some(model) = model_id {
+            self.get_by_model(model).to_vec()
+        } else {
+            self.get_all()
+        };
+
+        // Apply remaining filters
+        workers
+            .into_iter()
+            .filter(|w| {
+                // Check worker_type if specified
+                if let Some(ref wtype) = worker_type {
+                    if *w.worker_type() != *wtype {
+                        return false;
+                    }
+                }
+
+                // Check connection_mode if specified (using matches for flexible gRPC matching)
+                if let Some(ref conn) = connection_mode {
+                    if !w.connection_mode().matches(conn) {
+                        return false;
+                    }
+                }
+
+                // Check runtime_type if specified
+                if let Some(ref rt) = runtime_type {
+                    if w.metadata().runtime_type != *rt {
+                        return false;
+                    }
+                }
+
+                // Check health if required
+                if healthy_only && !w.is_healthy() {
+                    return false;
+                }
+
+                true
+            })
+            .collect()
+    }
+
+    /// Get worker statistics (lock-free)
+    pub fn stats(&self) -> WorkerRegistryStats {
+        let total_workers = self.workers.len();
+        // Count models directly instead of allocating Vec via get_models() (lock-free)
+        let total_models = self
+            .model_index
+            .iter()
+            .filter(|entry| !entry.value().is_empty())
+            .count();
+
+        let mut healthy_count = 0;
+        let mut total_load = 0;
+        let mut regular_count = 0;
+        let mut prefill_count = 0;
+        let mut decode_count = 0;
+        let mut http_count = 0;
+        let mut grpc_count = 0;
+        let mut cb_open_count = 0;
+        let mut cb_half_open_count = 0;
+
+        // Iterate DashMap directly to avoid cloning all workers via get_all()
+        for entry in self.workers.iter() {
+            let worker = entry.value();
+            if worker.is_healthy() {
+                healthy_count += 1;
+            }
+            total_load += worker.load();
+
+            match worker.worker_type() {
+                WorkerType::Regular => regular_count += 1,
+                WorkerType::Prefill { .. } => prefill_count += 1,
+                WorkerType::Decode => decode_count += 1,
+            }
+
+            match worker.connection_mode() {
+                ConnectionMode::Http => http_count += 1,
+                ConnectionMode::Grpc { .. } => grpc_count += 1,
+            }
+
+            match worker.circuit_breaker().state() {
+                CircuitState::Open => cb_open_count += 1,
+                CircuitState::HalfOpen => cb_half_open_count += 1,
+                CircuitState::Closed => {}
+            }
+        }
+
+        WorkerRegistryStats {
+            total_workers,
+            total_models,
+            healthy_workers: healthy_count,
+            unhealthy_workers: total_workers.saturating_sub(healthy_count),
+            total_load,
+            regular_workers: regular_count,
+            prefill_workers: prefill_count,
+            decode_workers: decode_count,
+            http_workers: http_count,
+            grpc_workers: grpc_count,
+            circuit_breaker_open: cb_open_count,
+            circuit_breaker_half_open: cb_half_open_count,
+        }
+    }
+
+    /// Get counts of regular and PD workers efficiently (O(1))
+    /// This avoids the overhead of get_all() which allocates memory and iterates all workers
+    pub fn get_worker_distribution(&self) -> (usize, usize) {
+        // Use the existing type_workers index for O(1) lookup
+        let regular_count = self
+            .type_workers
+            .get(&WorkerType::Regular)
+            .map(|v| v.len())
+            .unwrap_or(0);
+
+        // Get total workers count efficiently from DashMap
+        let total_workers = self.workers.len();
+
+        // PD workers are any workers that are not Regular
+        let pd_count = total_workers.saturating_sub(regular_count);
+
+        (regular_count, pd_count)
+    }
+
+    /// Start a health checker for all workers in the registry
+    /// This should be called once after the registry is populated with workers
+    pub(crate) fn start_health_checker(&self, check_interval_secs: u64) -> HealthChecker {
+        use std::sync::{
+            atomic::{AtomicBool, Ordering},
+            Arc,
+        };
+
+        let shutdown = Arc::new(AtomicBool::new(false));
+        let shutdown_clone = shutdown.clone();
+        let workers_ref = self.workers.clone();
+
+        let handle = tokio::spawn(async move {
+            let mut interval =
+                tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs));
+
+            loop {
+                interval.tick().await;
+
+                // Check for shutdown signal
+                if shutdown_clone.load(Ordering::Acquire) {
+                    tracing::debug!("Registry health checker shutting down");
+                    break;
+                }
+
+                // Get all workers from registry
+                let workers: Vec<Arc<dyn Worker>> = workers_ref
+                    .iter()
+                    .map(|entry| entry.value().clone())
+                    .collect();
+
+                // Perform health checks in parallel for better performance
+                // This is especially important when there are many workers
+                let health_futures: Vec<_> = workers
+                    .iter()
+                    .filter(|worker| !worker.metadata().health_config.disable_health_check)
+                    .map(|worker| {
+                        let worker = worker.clone();
+                        async move {
+                            let _ = worker.check_health_async().await;
+                        }
+                    })
+                    .collect();
+                futures::future::join_all(health_futures).await;
+            }
+        });
+
+        HealthChecker::new(handle, shutdown)
+    }
+}
+
+impl Default for WorkerRegistry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Statistics for the worker registry
+#[derive(Debug, Clone)]
+pub struct WorkerRegistryStats {
+    /// Total number of registered workers
+    pub total_workers: usize,
+    /// Number of unique models served
+    pub total_models: usize,
+    /// Number of workers passing health checks
+    pub healthy_workers: usize,
+    /// Number of workers failing health checks
+    pub unhealthy_workers: usize,
+    /// Sum of current load across all workers
+    pub total_load: usize,
+    /// Number of regular (non-PD) workers
+    pub regular_workers: usize,
+    /// Number of prefill workers (PD mode)
+    pub prefill_workers: usize,
+    /// Number of decode workers (PD mode)
+    pub decode_workers: usize,
+    /// Number of HTTP-connected workers
+    pub http_workers: usize,
+    /// Number of gRPC-connected workers
+    pub grpc_workers: usize,
+    /// Number of workers with circuit breaker in Open state (not accepting requests)
+    pub circuit_breaker_open: usize,
+    /// Number of workers with circuit breaker in HalfOpen state (testing recovery)
+    pub circuit_breaker_half_open: usize,
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{circuit_breaker::CircuitBreakerConfig, BasicWorkerBuilder};
+
+    #[test]
+    fn test_worker_registry() {
+        let registry = WorkerRegistry::new();
+
+        // Create a worker with labels
+        let mut labels = HashMap::new();
+        labels.insert("model_id".to_string(), "llama-3-8b".to_string());
+        labels.insert("priority".to_string(), "50".to_string());
+        labels.insert("cost".to_string(), "0.8".to_string());
+
+        let worker: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://worker1:8080")
+                .worker_type(WorkerType::Regular)
+                .labels(labels)
+                .circuit_breaker_config(CircuitBreakerConfig::default())
+                .api_key("test_api_key")
+                .build(),
+        );
+
+        // Register worker
+        let worker_id = registry.register(Arc::from(worker));
+
+        assert!(registry.get(&worker_id).is_some());
+        assert!(registry.get_by_url("http://worker1:8080").is_some());
+        assert_eq!(registry.get_by_model("llama-3-8b").len(), 1);
+        assert_eq!(registry.get_by_type(&WorkerType::Regular).len(), 1);
+        assert_eq!(registry.get_by_connection(&ConnectionMode::Http).len(), 1);
+
+        let stats = registry.stats();
+        assert_eq!(stats.total_workers, 1);
+        assert_eq!(stats.total_models, 1);
+
+        // Remove worker
+        registry.remove(&worker_id);
+        assert!(registry.get(&worker_id).is_none());
+    }
+
+    #[test]
+    fn test_model_index_fast_lookup() {
+        let registry = WorkerRegistry::new();
+
+        // Create workers for different models
+        let mut labels1 = HashMap::new();
+        labels1.insert("model_id".to_string(), "llama-3".to_string());
+        let worker1: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://worker1:8080")
+                .worker_type(WorkerType::Regular)
+                .labels(labels1)
+                .circuit_breaker_config(CircuitBreakerConfig::default())
+                .api_key("test_api_key")
+                .build(),
+        );
+
+        let mut labels2 = HashMap::new();
+        labels2.insert("model_id".to_string(), "llama-3".to_string());
+        let worker2: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://worker2:8080")
+                .worker_type(WorkerType::Regular)
+                .labels(labels2)
+                .circuit_breaker_config(CircuitBreakerConfig::default())
+                .api_key("test_api_key")
+                .build(),
+        );
+
+        let mut labels3 = HashMap::new();
+        labels3.insert("model_id".to_string(), "gpt-4".to_string());
+        let worker3: Box<dyn Worker> = Box::new(
+            BasicWorkerBuilder::new("http://worker3:8080")
+                .worker_type(WorkerType::Regular)
+                .labels(labels3)
+                .circuit_breaker_config(CircuitBreakerConfig::default())
+                .api_key("test_api_key")
+                .build(),
+        );
+
+        // Register workers
+        registry.register(Arc::from(worker1));
+        registry.register(Arc::from(worker2));
+        registry.register(Arc::from(worker3));
+
+        let llama_workers = registry.get_by_model("llama-3");
+        assert_eq!(llama_workers.len(), 2);
+        let urls: Vec<String> = llama_workers.iter().map(|w| w.url().to_string()).collect();
+        assert!(urls.contains(&"http://worker1:8080".to_string()));
+        assert!(urls.contains(&"http://worker2:8080".to_string()));
+
+        let gpt_workers = registry.get_by_model("gpt-4");
+        assert_eq!(gpt_workers.len(), 1);
+        assert_eq!(gpt_workers[0].url(), "http://worker3:8080");
+
+        let unknown_workers = registry.get_by_model("unknown-model");
+        assert_eq!(unknown_workers.len(), 0);
+
+        registry.remove_by_url("http://worker1:8080");
+        let llama_workers_after = registry.get_by_model("llama-3");
+        assert_eq!(llama_workers_after.len(), 1);
+        assert_eq!(llama_workers_after[0].url(), "http://worker2:8080");
+    }
+}

third_party/sglang/sgl-model-gateway/src/core/worker_service.rs

@@ -0,0 +1,366 @@
+//! Worker Service - Business logic layer for worker operations
+//!
+//! This module provides a clean separation between HTTP concerns (in routers)
+//! and business logic for worker management. The service orchestrates
+//! WorkerRegistry and JobQueue operations.
+
+use std::sync::Arc;
+
+use axum::{
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    Json,
+};
+use serde_json::json;
+use tracing::warn;
+
+use crate::{
+    config::RouterConfig,
+    core::{worker::worker_to_info, worker_registry::WorkerId, Job, JobQueue, WorkerRegistry},
+    protocols::worker_spec::{
+        WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest,
+    },
+};
+
+/// Error types for worker service operations
+#[derive(Debug)]
+pub enum WorkerServiceError {
+    /// Worker with given ID was not found
+    NotFound { worker_id: String },
+    /// Invalid worker ID format (expected UUID)
+    InvalidId { raw: String, message: String },
+    /// Job queue not initialized
+    QueueNotInitialized,
+    /// Failed to submit job to queue
+    QueueSubmitFailed { message: String },
+}
+
+impl WorkerServiceError {
+    pub fn error_code(&self) -> &'static str {
+        match self {
+            Self::NotFound { .. } => "WORKER_NOT_FOUND",
+            Self::InvalidId { .. } => "BAD_REQUEST",
+            Self::QueueNotInitialized => "INTERNAL_SERVER_ERROR",
+            Self::QueueSubmitFailed { .. } => "INTERNAL_SERVER_ERROR",
+        }
+    }
+
+    pub fn status_code(&self) -> StatusCode {
+        match self {
+            Self::NotFound { .. } => StatusCode::NOT_FOUND,
+            Self::InvalidId { .. } => StatusCode::BAD_REQUEST,
+            Self::QueueNotInitialized => StatusCode::INTERNAL_SERVER_ERROR,
+            Self::QueueSubmitFailed { .. } => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+}
+
+impl std::fmt::Display for WorkerServiceError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::NotFound { worker_id } => write!(f, "Worker {} not found", worker_id),
+            Self::InvalidId { raw, message } => {
+                write!(
+                    f,
+                    "Invalid worker_id '{}' (expected UUID). Error: {}",
+                    raw, message
+                )
+            }
+            Self::QueueNotInitialized => write!(f, "Job queue not initialized"),
+            Self::QueueSubmitFailed { message } => write!(f, "{}", message),
+        }
+    }
+}
+
+impl std::error::Error for WorkerServiceError {}
+
+impl IntoResponse for WorkerServiceError {
+    fn into_response(self) -> Response {
+        let error = WorkerErrorResponse {
+            error: self.to_string(),
+            code: self.error_code().to_string(),
+        };
+        (self.status_code(), Json(error)).into_response()
+    }
+}
+
+/// Result of creating a worker (async job submission)
+#[derive(Debug)]
+pub struct CreateWorkerResult {
+    pub worker_id: WorkerId,
+    pub url: String,
+    pub location: String,
+}
+
+impl IntoResponse for CreateWorkerResult {
+    fn into_response(self) -> Response {
+        let response = json!({
+            "status": "accepted",
+            "worker_id": self.worker_id.as_str(),
+            "url": self.url,
+            "location": self.location,
+            "message": "Worker addition queued for background processing"
+        });
+        (
+            StatusCode::ACCEPTED,
+            [(http::header::LOCATION, self.location)],
+            Json(response),
+        )
+            .into_response()
+    }
+}
+
+/// Result of deleting a worker (async job submission)
+#[derive(Debug)]
+pub struct DeleteWorkerResult {
+    pub worker_id: WorkerId,
+    pub url: String,
+}
+
+impl IntoResponse for DeleteWorkerResult {
+    fn into_response(self) -> Response {
+        let response = json!({
+            "status": "accepted",
+            "worker_id": self.worker_id.as_str(),
+            "message": "Worker removal queued for background processing"
+        });
+        (StatusCode::ACCEPTED, Json(response)).into_response()
+    }
+}
+
+/// Result of updating a worker (async job submission)
+#[derive(Debug)]
+pub struct UpdateWorkerResult {
+    pub worker_id: WorkerId,
+    pub url: String,
+}
+
+impl IntoResponse for UpdateWorkerResult {
+    fn into_response(self) -> Response {
+        let response = json!({
+            "status": "accepted",
+            "worker_id": self.worker_id.as_str(),
+            "message": "Worker update queued for background processing"
+        });
+        (StatusCode::ACCEPTED, Json(response)).into_response()
+    }
+}
+
+/// Result of listing workers
+#[derive(Debug)]
+pub struct ListWorkersResult {
+    pub workers: Vec<WorkerInfo>,
+    pub total: usize,
+    pub prefill_count: usize,
+    pub decode_count: usize,
+    pub regular_count: usize,
+}
+
+impl IntoResponse for ListWorkersResult {
+    fn into_response(self) -> Response {
+        let response = json!({
+            "workers": self.workers,
+            "total": self.total,
+            "stats": {
+                "prefill_count": self.prefill_count,
+                "decode_count": self.decode_count,
+                "regular_count": self.regular_count,
+            }
+        });
+        Json(response).into_response()
+    }
+}
+
+/// Wrapper for WorkerInfo to implement IntoResponse
+pub struct GetWorkerResponse(pub WorkerInfo);
+
+impl IntoResponse for GetWorkerResponse {
+    fn into_response(self) -> Response {
+        Json(self.0).into_response()
+    }
+}
+
+/// Worker Service - Orchestrates worker business logic
+///
+/// This service provides a clean API for worker operations, separating
+/// business logic from HTTP concerns. Handlers in server.rs become thin
+/// wrappers that translate between HTTP and this service.
+pub struct WorkerService {
+    worker_registry: Arc<WorkerRegistry>,
+    job_queue: Arc<std::sync::OnceLock<Arc<JobQueue>>>,
+    router_config: RouterConfig,
+}
+
+impl WorkerService {
+    /// Create a new WorkerService
+    pub fn new(
+        worker_registry: Arc<WorkerRegistry>,
+        job_queue: Arc<std::sync::OnceLock<Arc<JobQueue>>>,
+        router_config: RouterConfig,
+    ) -> Self {
+        Self {
+            worker_registry,
+            job_queue,
+            router_config,
+        }
+    }
+
+    /// Parse and validate a worker ID string
+    pub fn parse_worker_id(&self, raw: &str) -> Result<WorkerId, WorkerServiceError> {
+        uuid::Uuid::parse_str(raw)
+            .map(|_| WorkerId::from_string(raw.to_string()))
+            .map_err(|e| WorkerServiceError::InvalidId {
+                raw: raw.to_string(),
+                message: e.to_string(),
+            })
+    }
+
+    /// Get the job queue, returning an error if not initialized
+    fn get_job_queue(&self) -> Result<&Arc<JobQueue>, WorkerServiceError> {
+        self.job_queue
+            .get()
+            .ok_or(WorkerServiceError::QueueNotInitialized)
+    }
+
+    pub async fn create_worker(
+        &self,
+        mut config: WorkerConfigRequest,
+    ) -> Result<CreateWorkerResult, WorkerServiceError> {
+        if self.router_config.api_key.is_some() && config.api_key.is_none() {
+            warn!(
+                "Adding worker {} without API key while router has API key configured. \
+                Worker will be accessible without authentication. \
+                If the worker requires the same API key as the router, please specify it explicitly.",
+                config.url
+            );
+        }
+
+        config.dp_aware = self.router_config.dp_aware;
+
+        let worker_url = config.url.clone();
+        let worker_id = self.worker_registry.reserve_id_for_url(&worker_url);
+
+        let job = Job::AddWorker {
+            config: Box::new(config),
+        };
+
+        self.get_job_queue()?
+            .submit(job)
+            .await
+            .map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?;
+
+        let location = format!("/workers/{}", worker_id.as_str());
+
+        Ok(CreateWorkerResult {
+            worker_id,
+            url: worker_url,
+            location,
+        })
+    }
+
+    /// List all workers with their info
+    pub fn list_workers(&self) -> ListWorkersResult {
+        let workers = self.worker_registry.get_all_with_ids();
+        let worker_infos: Vec<WorkerInfo> = workers
+            .iter()
+            .map(|(worker_id, worker)| {
+                let mut info = worker_to_info(worker);
+                info.id = worker_id.as_str().to_string();
+                info
+            })
+            .collect();
+
+        let stats = self.worker_registry.stats();
+
+        ListWorkersResult {
+            workers: worker_infos,
+            total: stats.total_workers,
+            prefill_count: stats.prefill_workers,
+            decode_count: stats.decode_workers,
+            regular_count: stats.regular_workers,
+        }
+    }
+
+    pub fn get_worker(&self, worker_id_raw: &str) -> Result<GetWorkerResponse, WorkerServiceError> {
+        let worker_id = self.parse_worker_id(worker_id_raw)?;
+        let job_queue = self.get_job_queue()?;
+
+        if let Some(worker) = self.worker_registry.get(&worker_id) {
+            let worker_url = worker.url().to_string();
+            let mut worker_info = worker_to_info(&worker);
+            worker_info.id = worker_id.as_str().to_string();
+            if let Some(status) = job_queue.get_status(&worker_url) {
+                worker_info.job_status = Some(status);
+            }
+            return Ok(GetWorkerResponse(worker_info));
+        }
+
+        if let Some(worker_url) = self.worker_registry.get_url_by_id(&worker_id) {
+            if let Some(status) = job_queue.get_status(&worker_url) {
+                return Ok(GetWorkerResponse(WorkerInfo::pending(
+                    worker_id.as_str(),
+                    worker_url,
+                    Some(status),
+                )));
+            }
+        }
+
+        Err(WorkerServiceError::NotFound {
+            worker_id: worker_id_raw.to_string(),
+        })
+    }
+
+    /// Delete a worker by ID (submits async job)
+    pub async fn delete_worker(
+        &self,
+        worker_id_raw: &str,
+    ) -> Result<DeleteWorkerResult, WorkerServiceError> {
+        let worker_id = self.parse_worker_id(worker_id_raw)?;
+
+        let url = self
+            .worker_registry
+            .get_url_by_id(&worker_id)
+            .ok_or_else(|| WorkerServiceError::NotFound {
+                worker_id: worker_id_raw.to_string(),
+            })?;
+
+        let job = Job::RemoveWorker { url: url.clone() };
+
+        let job_queue = self.get_job_queue()?;
+        job_queue
+            .submit(job)
+            .await
+            .map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?;
+
+        Ok(DeleteWorkerResult { worker_id, url })
+    }
+
+    /// Update a worker by ID (submits async job)
+    pub async fn update_worker(
+        &self,
+        worker_id_raw: &str,
+        update: WorkerUpdateRequest,
+    ) -> Result<UpdateWorkerResult, WorkerServiceError> {
+        let worker_id = self.parse_worker_id(worker_id_raw)?;
+
+        let url = self
+            .worker_registry
+            .get_url_by_id(&worker_id)
+            .ok_or_else(|| WorkerServiceError::NotFound {
+                worker_id: worker_id_raw.to_string(),
+            })?;
+
+        let job = Job::UpdateWorker {
+            url: url.clone(),
+            update: Box::new(update),
+        };
+
+        let job_queue = self.get_job_queue()?;
+        job_queue
+            .submit(job)
+            .await
+            .map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?;
+
+        Ok(UpdateWorkerResult { worker_id, url })
+    }
+}

third_party/sglang/sgl-model-gateway/src/lib.rs

@@ -0,0 +1,16 @@
+pub mod app_context;
+pub use smg_auth as auth;
+pub mod config;
+pub mod core;
+pub mod middleware;
+pub mod observability;
+pub mod policies;
+pub use openai_protocol as protocols;
+pub use reasoning_parser;
+pub mod routers;
+pub mod server;
+pub mod service_discovery;
+pub use llm_tokenizer as tokenizer;
+pub use tool_parser;
+pub mod version;
+pub mod wasm;

third_party/sglang/sgl-model-gateway/src/observability/events.rs

@@ -0,0 +1,90 @@
+//! Request events for observability and monitoring.
+//!
+//! Events use DEBUG level when OTEL is disabled, INFO when enabled.
+
+use tracing::{debug, event, Level};
+
+use super::otel_trace::is_otel_enabled;
+
+/// Module path used by CustomOtelFilter to identify events for OTEL export.
+#[inline]
+pub const fn get_module_path() -> &'static str {
+    "smg::observability::events"
+}
+
+pub trait Event {
+    fn emit(&self);
+}
+
+/// Event emitted when a prefill-decode request pair is sent.
+#[derive(Debug, Clone, Copy)]
+pub struct RequestPDSentEvent<'a> {
+    pub prefill_url: &'a str,
+    pub decode_url: &'a str,
+}
+
+impl Event for RequestPDSentEvent<'_> {
+    #[inline]
+    fn emit(&self) {
+        if is_otel_enabled() {
+            event!(
+                Level::INFO,
+                prefill_url = %self.prefill_url,
+                decode_url = %self.decode_url,
+                "Sending concurrent requests"
+            );
+        } else {
+            debug!(
+                prefill_url = %self.prefill_url,
+                decode_url = %self.decode_url,
+                "Sending concurrent requests"
+            );
+        }
+    }
+}
+
+/// Event emitted when a request is sent to a worker.
+#[derive(Debug, Clone, Copy)]
+pub struct RequestSentEvent<'a> {
+    pub url: &'a str,
+}
+
+impl Event for RequestSentEvent<'_> {
+    #[inline]
+    fn emit(&self) {
+        if is_otel_enabled() {
+            event!(Level::INFO, url = %self.url, "Sending request");
+        } else {
+            debug!(url = %self.url, "Sending request");
+        }
+    }
+}
+
+/// Event emitted when concurrent requests are received.
+#[derive(Debug, Clone, Copy)]
+pub struct RequestReceivedEvent;
+
+impl Event for RequestReceivedEvent {
+    #[inline]
+    fn emit(&self) {
+        if is_otel_enabled() {
+            event!(Level::INFO, "Received concurrent requests");
+        } else {
+            debug!("Received concurrent requests");
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::mem::size_of;
+
+    use super::*;
+
+    #[test]
+    fn test_event_sizes() {
+        assert_eq!(size_of::<RequestReceivedEvent>(), 0);
+        assert_eq!(size_of::<RequestSentEvent>(), 16);
+        assert_eq!(size_of::<RequestPDSentEvent>(), 32);
+    }
+}

third_party/sglang/sgl-model-gateway/src/observability/gauge_histogram.rs

@@ -0,0 +1,583 @@
+//! Non-cumulative gauge histogram for Grafana heatmap visualization.
+//!
+//! Unlike Prometheus Histogram which uses cumulative `le` buckets, this emits
+//! non-cumulative bucket counts with `(gt, le]` ranges suitable for heatmaps.
+//!
+//! # Design: True Zero-Allocation Hot Path
+//!
+//! The key insight is that `gauge!` returns a `Gauge` handle that can be stored.
+//! By pre-registering all gauge handles at startup, the hot path becomes just
+//! N+1 atomic `gauge.set()` calls with zero allocations.
+//!
+//! # Performance Characteristics
+//!
+//! Setup (once per label combination):
+//! - `register()`: N+1 gauge registrations, N+1 String allocations for gt/le
+//!
+//! Hot path (`set_counts()`):
+//! - **Zero heap allocations**
+//! - **Zero key lookups** (handles are pre-registered)
+//! - N+1 atomic `gauge.set()` calls
+//!
+//! # Example
+//!
+//! ```ignore
+//! use crate::observability::gauge_histogram::{BucketBounds, GaugeHistogramVec};
+//!
+//! // Define at module level
+//! static BOUNDS: BucketBounds<10> = BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
+//! static HISTOGRAM: GaugeHistogramVec<10> = GaugeHistogramVec::new("smg_request_dist", &BOUNDS);
+//!
+//! // At startup: register for each label combination
+//! let handle = HISTOGRAM.register(&[("router", "round_robin"), ("model", "llama")]);
+//!
+//! // Pre-allocate counts buffer
+//! let mut counts = vec![0usize; BOUNDS.bucket_count()];
+//!
+//! // Hot path: TRUE zero allocation
+//! fn update(handle: &GaugeHistogramHandle, counts: &mut [usize], observations: &[u64]) {
+//!     BOUNDS.compute_counts_into(counts, observations);
+//!     handle.set_counts(counts);  // Just N+1 atomic gauge.set() calls!
+//! }
+//! ```
+
+use std::sync::Arc;
+
+use dashmap::DashMap;
+use metrics::{gauge, Label};
+
+// =============================================================================
+// BUCKET BOUNDS
+// =============================================================================
+
+/// Static bucket boundary configuration.
+///
+/// Uses const generics to define bucket bounds at compile time with validation.
+/// The bounds define `N` upper limits, creating `N + 1` buckets:
+/// `(0, b[0]], (b[0], b[1]], ..., (b[N-1], +Inf]`.
+#[derive(Debug)]
+pub struct BucketBounds<const N: usize> {
+    bounds: [u64; N],
+}
+
+impl<const N: usize> BucketBounds<N> {
+    /// Create new bucket bounds from a sorted array of upper limits.
+    ///
+    /// # Panics
+    ///
+    /// Panics at compile time (in const context) or runtime if bounds are not
+    /// strictly ascending.
+    #[must_use]
+    pub const fn new(bounds: [u64; N]) -> Self {
+        let mut i = 1;
+        while i < N {
+            assert!(
+                bounds[i] > bounds[i - 1],
+                "bucket bounds must be strictly ascending"
+            );
+            i += 1;
+        }
+        Self { bounds }
+    }
+
+    /// Returns the number of buckets (one more than the number of bounds).
+    #[inline]
+    #[must_use]
+    pub const fn bucket_count(&self) -> usize {
+        N + 1
+    }
+
+    /// Returns the number of bounds.
+    #[inline]
+    #[must_use]
+    pub const fn bound_count(&self) -> usize {
+        N
+    }
+
+    /// Get the bounds array.
+    #[inline]
+    #[must_use]
+    pub const fn bounds(&self) -> &[u64; N] {
+        &self.bounds
+    }
+
+    /// Find the bucket index for a value. O(log N).
+    #[inline]
+    #[must_use]
+    pub fn bucket_index(&self, value: u64) -> usize {
+        self.bounds.partition_point(|&bound| bound < value)
+    }
+
+    /// Get the upper bound for a bucket index, or None for the +Inf bucket.
+    #[inline]
+    #[must_use]
+    pub const fn upper_bound(&self, idx: usize) -> Option<u64> {
+        if idx < N {
+            Some(self.bounds[idx])
+        } else {
+            None
+        }
+    }
+
+    /// Get the lower bound for a bucket index (0 for the first bucket).
+    #[inline]
+    #[must_use]
+    pub const fn lower_bound(&self, idx: usize) -> u64 {
+        if idx == 0 {
+            0
+        } else {
+            self.bounds[idx - 1]
+        }
+    }
+
+    /// Compute bucket counts into a pre-allocated buffer. **Zero allocation.**
+    ///
+    /// # Panics
+    ///
+    /// Panics if `counts.len() < bucket_count()`.
+    #[inline]
+    pub fn compute_counts_into(&self, counts: &mut [usize], observations: &[u64]) {
+        debug_assert!(
+            counts.len() >= self.bucket_count(),
+            "counts buffer too small"
+        );
+        counts[..self.bucket_count()].fill(0);
+        for &value in observations {
+            let idx = self.bucket_index(value);
+            counts[idx] += 1;
+        }
+    }
+
+    /// Compute bucket counts, allocating a new Vec.
+    ///
+    /// Prefer `compute_counts_into` in hot paths to avoid allocation.
+    #[must_use]
+    pub fn compute_counts(&self, observations: &[u64]) -> Vec<usize> {
+        let mut counts = vec![0usize; self.bucket_count()];
+        self.compute_counts_into(&mut counts, observations);
+        counts
+    }
+}
+
+// =============================================================================
+// GAUGE HISTOGRAM HANDLE (pre-registered, zero-alloc hot path)
+// =============================================================================
+
+/// Pre-registered gauge handles for a histogram with specific labels.
+///
+/// This is what you use in the hot path. Calling `set_counts()` does only
+/// N+1 atomic `gauge.set()` operations - zero allocations, zero lookups.
+#[derive(Clone)]
+pub struct GaugeHistogramHandle {
+    gauges: Vec<metrics::Gauge>,
+}
+
+impl GaugeHistogramHandle {
+    /// Set bucket counts. **TRUE zero allocation.**
+    ///
+    /// Just N+1 atomic `gauge.set()` calls - no key lookup, no allocation.
+    #[inline]
+    pub fn set_counts(&self, counts: &[usize]) {
+        debug_assert_eq!(
+            counts.len(),
+            self.gauges.len(),
+            "counts length must match bucket count"
+        );
+        for (gauge, &count) in self.gauges.iter().zip(counts.iter()) {
+            gauge.set(count as f64);
+        }
+    }
+
+    /// Number of buckets.
+    #[inline]
+    pub fn bucket_count(&self) -> usize {
+        self.gauges.len()
+    }
+
+    /// Zero out all gauges. **Zero allocation.**
+    #[inline]
+    pub fn zero_counts(&self) {
+        for gauge in &self.gauges {
+            gauge.set(0.0);
+        }
+    }
+}
+
+// =============================================================================
+// GAUGE HISTOGRAM VEC (factory for registering handles)
+// =============================================================================
+
+/// Factory for creating pre-registered histogram handles.
+///
+/// Define as a static constant, then call `register()` for each label combination
+/// you need. The returned `GaugeHistogramHandle` provides zero-allocation updates.
+#[derive(Debug)]
+pub struct GaugeHistogramVec<const N: usize> {
+    name: &'static str,
+    bounds: &'static BucketBounds<N>,
+}
+
+impl<const N: usize> GaugeHistogramVec<N> {
+    /// Create a new gauge histogram factory.
+    ///
+    /// This just stores the name and bounds - no allocation or registration yet.
+    #[must_use]
+    pub const fn new(name: &'static str, bounds: &'static BucketBounds<N>) -> Self {
+        Self { name, bounds }
+    }
+
+    #[inline]
+    pub const fn name(&self) -> &'static str {
+        self.name
+    }
+
+    #[inline]
+    pub const fn bounds(&self) -> &BucketBounds<N> {
+        self.bounds
+    }
+
+    /// Register gauges for a specific label combination.
+    ///
+    /// Call this once per unique label combination (at startup or when first seen).
+    /// The returned handle can be cloned cheaply (just Arc clones internally).
+    ///
+    /// # Arguments
+    ///
+    /// - `labels`: Static key-value label pairs for this histogram instance
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// let handle = HISTOGRAM.register(&[("router", "round_robin"), ("model", "llama")]);
+    /// ```
+    pub fn register(&self, labels: &[(&'static str, &str)]) -> GaugeHistogramHandle {
+        let bucket_count = self.bounds.bucket_count();
+        let mut gauges = Vec::with_capacity(bucket_count);
+
+        for i in 0..bucket_count {
+            // Build gt/le labels for this bucket
+            let gt_str = if i == 0 {
+                "0".to_string()
+            } else {
+                self.bounds.bounds[i - 1].to_string()
+            };
+
+            let le_str = if i < N {
+                self.bounds.bounds[i].to_string()
+            } else {
+                "+Inf".to_string()
+            };
+
+            // Build complete label set
+            let mut all_labels: Vec<Label> = Vec::with_capacity(labels.len() + 2);
+            for &(k, v) in labels {
+                all_labels.push(Label::new(k, v.to_string()));
+            }
+            all_labels.push(Label::new("gt", gt_str));
+            all_labels.push(Label::new("le", le_str));
+
+            // Register and store the gauge handle
+            let g = gauge!(self.name, all_labels);
+            gauges.push(g);
+        }
+
+        GaugeHistogramHandle { gauges }
+    }
+
+    /// Register with no additional labels (just gt/le).
+    pub fn register_no_labels(&self) -> GaugeHistogramHandle {
+        self.register(&[])
+    }
+}
+
+// =============================================================================
+// CACHED GAUGE HISTOGRAM (for dynamic labels discovered at runtime)
+// =============================================================================
+
+/// A gauge histogram with automatic handle caching for dynamic labels.
+///
+/// Use this when label values (like worker names) are discovered at runtime.
+/// Handles are registered on first use and cached for subsequent calls.
+///
+/// # Example
+///
+/// ```ignore
+/// static BOUNDS: BucketBounds<10> = BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
+/// static HISTOGRAM: GaugeHistogramVec<10> = GaugeHistogramVec::new("smg_worker_dist", &BOUNDS);
+///
+/// // Create cached wrapper (do this once, store in your router state)
+/// let cached = CachedGaugeHistogram::new(&HISTOGRAM);
+///
+/// // Hot path - first call registers, subsequent calls use cached handle
+/// cached.observe("worker-1", &request_counts);
+/// cached.observe("worker-2", &request_counts);
+/// cached.observe("worker-1", &request_counts);  // Uses cached handle
+/// ```
+pub struct CachedGaugeHistogram<const N: usize> {
+    histogram: &'static GaugeHistogramVec<N>,
+    /// Cache of label value -> (handle, counts_buffer)
+    cache: DashMap<Arc<str>, (GaugeHistogramHandle, Vec<usize>)>,
+    /// Static label key (e.g., "worker", "model")
+    label_key: &'static str,
+}
+
+impl<const N: usize> CachedGaugeHistogram<N> {
+    /// Create a new cached histogram for a single dynamic label.
+    ///
+    /// # Arguments
+    ///
+    /// - `histogram`: The static histogram factory
+    /// - `label_key`: The label key for the dynamic value (e.g., "worker")
+    pub fn new(histogram: &'static GaugeHistogramVec<N>, label_key: &'static str) -> Self {
+        Self {
+            histogram,
+            cache: DashMap::new(),
+            label_key,
+        }
+    }
+
+    /// Get or create a handle for the given label value.
+    ///
+    /// First call for a label value registers the gauges (allocates).
+    /// Subsequent calls return the cached handle (no allocation).
+    /// Thread-safe: uses entry API to avoid race conditions.
+    pub fn get_or_register(
+        &self,
+        label_value: &str,
+    ) -> dashmap::mapref::one::Ref<'_, Arc<str>, (GaugeHistogramHandle, Vec<usize>)> {
+        // Fast path: already cached
+        if let Some(entry) = self.cache.get(label_value) {
+            return entry;
+        }
+
+        // Slow path: use entry API to handle concurrent inserts atomically
+        self.cache.entry(Arc::from(label_value)).or_insert_with(|| {
+            let handle = self.histogram.register(&[(self.label_key, label_value)]);
+            let counts_buf = vec![0usize; self.histogram.bounds.bucket_count()];
+            (handle, counts_buf)
+        });
+
+        self.cache.get(label_value).unwrap()
+    }
+
+    /// Observe a distribution for a label value. **Zero allocation after first call.**
+    ///
+    /// First call for a new label value registers gauges (allocates).
+    /// All subsequent calls are zero-allocation.
+    /// Thread-safe: uses entry API to avoid race conditions.
+    pub fn observe(&self, label_value: &str, observations: &[u64]) {
+        // Fast path: existing entry
+        if let Some(mut entry) = self.cache.get_mut(label_value) {
+            let (ref handle, ref mut counts_buf) = entry.value_mut();
+            self.histogram
+                .bounds
+                .compute_counts_into(counts_buf, observations);
+            handle.set_counts(counts_buf);
+            return;
+        }
+
+        // Slow path: use entry API to handle concurrent inserts atomically
+        let mut entry = self.cache.entry(Arc::from(label_value)).or_insert_with(|| {
+            let handle = self.histogram.register(&[(self.label_key, label_value)]);
+            let counts_buf = vec![0usize; self.histogram.bounds.bucket_count()];
+            (handle, counts_buf)
+        });
+
+        let (ref handle, ref mut counts_buf) = entry.value_mut();
+        self.histogram
+            .bounds
+            .compute_counts_into(counts_buf, observations);
+        handle.set_counts(counts_buf);
+    }
+
+    /// Number of cached label combinations.
+    pub fn cache_size(&self) -> usize {
+        self.cache.len()
+    }
+
+    /// Remove a worker and zero out its metrics. **Zero allocation.**
+    ///
+    /// Call this when a worker is removed from the pool.
+    /// Sets all bucket counts to 0 (so Grafana shows it as empty).
+    ///
+    /// Note: The gauge handles remain in the Prometheus registry (the `metrics`
+    /// crate doesn't support unregistering). But memory in our cache is freed.
+    pub fn remove(&self, label_value: &str) {
+        if let Some((_, (handle, _))) = self.cache.remove(label_value) {
+            handle.zero_counts();
+        }
+    }
+
+    /// Remove workers not in the provided set.
+    ///
+    /// Call this periodically with your current active workers to clean up stale entries.
+    /// Uses `DashMap::retain` for atomic operation without intermediate allocation.
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// let active: HashSet<&str> = workers.iter().map(|w| w.name.as_str()).collect();
+    /// cached.retain_only(&active);
+    /// ```
+    pub fn retain_only<S: std::borrow::Borrow<str> + std::hash::Hash + Eq>(
+        &self,
+        active_labels: &std::collections::HashSet<S>,
+    ) {
+        self.cache.retain(|key, (handle, _)| {
+            if active_labels.contains(key.as_ref()) {
+                true
+            } else {
+                handle.zero_counts();
+                false
+            }
+        });
+    }
+
+    /// Get all currently tracked label values.
+    pub fn tracked_labels(&self) -> Vec<Arc<str>> {
+        self.cache.iter().map(|e| Arc::clone(e.key())).collect()
+    }
+}
+
+// =============================================================================
+// CONVENIENCE CONSTANTS
+// =============================================================================
+
+/// Common bucket bounds for request counts.
+pub static REQUEST_COUNT_BOUNDS: BucketBounds<10> =
+    BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
+
+// =============================================================================
+// TESTS
+// =============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_bucket_bounds_creation() {
+        let bounds = BucketBounds::new([10, 30, 60]);
+        assert_eq!(bounds.bucket_count(), 4);
+        assert_eq!(bounds.bound_count(), 3);
+    }
+
+    #[test]
+    fn test_bucket_bounds_const_creation() {
+        static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
+        assert_eq!(BOUNDS.bucket_count(), 4);
+    }
+
+    #[test]
+    #[should_panic(expected = "bucket bounds must be strictly ascending")]
+    fn test_bucket_bounds_not_ascending_panics() {
+        let _ = BucketBounds::new([10, 5, 60]);
+    }
+
+    #[test]
+    fn test_bucket_index() {
+        let bounds = BucketBounds::new([10, 30, 60]);
+        assert_eq!(bounds.bucket_index(0), 0);
+        assert_eq!(bounds.bucket_index(10), 0);
+        assert_eq!(bounds.bucket_index(11), 1);
+        assert_eq!(bounds.bucket_index(30), 1);
+        assert_eq!(bounds.bucket_index(31), 2);
+        assert_eq!(bounds.bucket_index(60), 2);
+        assert_eq!(bounds.bucket_index(61), 3);
+    }
+
+    #[test]
+    fn test_compute_counts() {
+        let bounds = BucketBounds::new([10, 30, 60]);
+        assert_eq!(
+            bounds.compute_counts(&[5, 10, 15, 40, 100]),
+            vec![2, 1, 1, 1]
+        );
+    }
+
+    #[test]
+    fn test_compute_counts_into() {
+        let bounds = BucketBounds::new([10, 30, 60]);
+        let mut counts = [0usize; 4];
+        bounds.compute_counts_into(&mut counts, &[5, 10, 15, 40, 100]);
+        assert_eq!(counts, [2, 1, 1, 1]);
+    }
+
+    #[test]
+    fn test_gauge_histogram_vec_creation() {
+        static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
+        static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_metric", &BOUNDS);
+
+        assert_eq!(HISTOGRAM.name(), "test_metric");
+        assert_eq!(HISTOGRAM.bounds().bucket_count(), 4);
+    }
+
+    #[test]
+    fn test_gauge_histogram_handle_registration() {
+        static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
+        static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_hist", &BOUNDS);
+
+        let handle = HISTOGRAM.register(&[("router", "rr")]);
+        assert_eq!(handle.bucket_count(), 4);
+
+        // This should be zero-allocation
+        handle.set_counts(&[1, 2, 3, 4]);
+    }
+
+    #[test]
+    fn test_request_count_bounds() {
+        assert_eq!(REQUEST_COUNT_BOUNDS.bucket_count(), 11);
+        assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(1), 0);
+        assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(2), 1);
+        assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(201), 10);
+    }
+
+    #[test]
+    fn test_cached_histogram() {
+        static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
+        static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_cached", &BOUNDS);
+
+        let cached = CachedGaugeHistogram::new(&HISTOGRAM, "worker");
+
+        // First call registers
+        cached.observe("worker-1", &[5, 15, 45, 100]);
+        assert_eq!(cached.cache_size(), 1);
+
+        // Second call uses cache
+        cached.observe("worker-1", &[1, 2, 3]);
+        assert_eq!(cached.cache_size(), 1);
+
+        // New worker registers
+        cached.observe("worker-2", &[10, 20, 30]);
+        assert_eq!(cached.cache_size(), 2);
+    }
+
+    #[test]
+    fn test_cached_histogram_removal() {
+        static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
+        static HISTOGRAM: GaugeHistogramVec<3> =
+            GaugeHistogramVec::new("test_cached_remove", &BOUNDS);
+
+        let cached = CachedGaugeHistogram::new(&HISTOGRAM, "worker");
+
+        // Add some workers
+        cached.observe("worker-1", &[5, 15]);
+        cached.observe("worker-2", &[10, 20]);
+        cached.observe("worker-3", &[1, 2]);
+        assert_eq!(cached.cache_size(), 3);
+
+        // Remove one
+        cached.remove("worker-2");
+        assert_eq!(cached.cache_size(), 2);
+
+        // retain_only
+        let active: std::collections::HashSet<&str> = ["worker-1"].into_iter().collect();
+        cached.retain_only(&active);
+        assert_eq!(cached.cache_size(), 1);
+
+        // Check tracked labels
+        let labels = cached.tracked_labels();
+        assert_eq!(labels.len(), 1);
+        assert_eq!(&*labels[0], "worker-1");
+    }
+}

third_party/sglang/sgl-model-gateway/src/observability/inflight_tracker.rs

@@ -0,0 +1,195 @@
+use std::{
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc, LazyLock, OnceLock,
+    },
+    time::Instant,
+};
+
+use dashmap::DashMap;
+
+use super::gauge_histogram::{BucketBounds, GaugeHistogramHandle, GaugeHistogramVec};
+use crate::policies::utils::PeriodicTask;
+
+static INFLIGHT_AGE_BOUNDS: BucketBounds<11> =
+    BucketBounds::new([30, 60, 180, 300, 600, 1200, 3600, 7200, 14400, 28800, 86400]);
+static INFLIGHT_AGE_HISTOGRAM: GaugeHistogramVec<11> =
+    GaugeHistogramVec::new("smg_http_inflight_request_age_count", &INFLIGHT_AGE_BOUNDS);
+static INFLIGHT_AGE_HANDLE: LazyLock<GaugeHistogramHandle> =
+    LazyLock::new(|| INFLIGHT_AGE_HISTOGRAM.register_no_labels());
+
+pub struct InFlightRequestTracker {
+    requests: DashMap<u64, Instant>,
+    next_id: AtomicU64,
+    sampler: OnceLock<PeriodicTask>,
+}
+
+impl InFlightRequestTracker {
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            requests: DashMap::new(),
+            next_id: AtomicU64::new(0),
+            sampler: OnceLock::new(),
+        })
+    }
+
+    pub fn start_sampler(self: &Arc<Self>, interval_secs: u64) {
+        let tracker = self.clone();
+        let task = PeriodicTask::spawn(interval_secs, "InFlightRequestSampler", move || {
+            tracker.sample_and_record();
+        });
+        self.sampler.set(task).unwrap();
+    }
+
+    pub fn track(self: &Arc<Self>) -> InFlightGuard {
+        let request_id = self.next_id.fetch_add(1, Ordering::Relaxed);
+        self.requests.insert(request_id, Instant::now());
+        InFlightGuard {
+            tracker: self.clone(),
+            request_id,
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.requests.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.requests.is_empty()
+    }
+
+    pub fn compute_bucket_counts(&self) -> Vec<usize> {
+        let ages = self.collect_ages();
+        INFLIGHT_AGE_BOUNDS.compute_counts(&ages)
+    }
+
+    fn collect_ages(&self) -> Vec<u64> {
+        let now = Instant::now();
+        self.requests
+            .iter()
+            .map(|entry| now.duration_since(*entry.value()).as_secs())
+            .collect()
+    }
+
+    fn sample_and_record(&self) {
+        let counts = self.compute_bucket_counts();
+        INFLIGHT_AGE_HANDLE.set_counts(&counts);
+    }
+}
+
+pub struct InFlightGuard {
+    tracker: Arc<InFlightRequestTracker>,
+    request_id: u64,
+}
+
+impl Drop for InFlightGuard {
+    fn drop(&mut self) {
+        self.tracker.requests.remove(&self.request_id);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::*;
+
+    impl InFlightRequestTracker {
+        fn insert_with_time(&self, request_id: u64, start_time: Instant) {
+            self.requests.insert(request_id, start_time);
+        }
+    }
+
+    #[test]
+    fn test_track_and_drop() {
+        let tracker = InFlightRequestTracker::new();
+        {
+            let _guard1 = tracker.track();
+            let _guard2 = tracker.track();
+            assert_eq!(tracker.len(), 2);
+        }
+        assert_eq!(tracker.len(), 0);
+    }
+
+    #[test]
+    fn test_guard_auto_deregister() {
+        let tracker = InFlightRequestTracker::new();
+        let guard = tracker.track();
+        assert_eq!(tracker.len(), 1);
+        drop(guard);
+        assert_eq!(tracker.len(), 0);
+    }
+
+    #[test]
+    fn test_request_age_tracking() {
+        let tracker = InFlightRequestTracker::new();
+        let _guard = tracker.track();
+        std::thread::sleep(Duration::from_millis(100));
+
+        let entry = tracker.requests.iter().next().unwrap();
+        let age = entry.value().elapsed();
+        assert!(age >= Duration::from_millis(100));
+    }
+
+    #[test]
+    fn test_collect_ages_empty() {
+        let tracker = InFlightRequestTracker::new();
+        let ages = tracker.collect_ages();
+        assert!(ages.is_empty());
+    }
+
+    #[test]
+    fn test_collect_ages() {
+        let tracker = InFlightRequestTracker::new();
+        let now = Instant::now();
+
+        tracker.insert_with_time(1, now);
+        tracker.insert_with_time(2, now - Duration::from_secs(45));
+        tracker.insert_with_time(3, now - Duration::from_secs(100));
+
+        let ages = tracker.collect_ages();
+        assert_eq!(ages.len(), 3);
+        // Ages should be approximately 0, 45, 100 (order may vary due to DashMap)
+        let mut sorted_ages = ages.clone();
+        sorted_ages.sort();
+        assert!(sorted_ages[0] <= 1); // ~0s
+        assert!((44..=46).contains(&sorted_ages[1])); // ~45s
+        assert!((99..=101).contains(&sorted_ages[2])); // ~100s
+    }
+
+    #[test]
+    fn test_concurrent_tracking() {
+        use std::thread;
+
+        let tracker = InFlightRequestTracker::new();
+        let mut handles = vec![];
+
+        for _ in 0..10 {
+            let t = tracker.clone();
+            handles.push(thread::spawn(move || {
+                (0..100).map(|_| t.track()).collect::<Vec<_>>()
+            }));
+        }
+
+        let all_guards: Vec<_> = handles
+            .into_iter()
+            .flat_map(|h| h.join().unwrap())
+            .collect();
+
+        assert_eq!(tracker.len(), 1000);
+        drop(all_guards);
+        assert_eq!(tracker.len(), 0);
+    }
+
+    #[test]
+    fn test_unique_ids() {
+        let tracker = InFlightRequestTracker::new();
+        let g1 = tracker.track();
+        let g2 = tracker.track();
+        let g3 = tracker.track();
+
+        assert_ne!(g1.request_id, g2.request_id);
+        assert_ne!(g2.request_id, g3.request_id);
+        assert_eq!(tracker.len(), 3);
+    }
+}

third_party/sglang/sgl-model-gateway/src/observability/logging.rs

@@ -0,0 +1,174 @@
+//! Logging infrastructure with non-blocking file I/O.
+
+use std::path::PathBuf;
+
+use tracing::Level;
+use tracing_appender::{
+    non_blocking::WorkerGuard,
+    rolling::{RollingFileAppender, Rotation},
+};
+use tracing_log::LogTracer;
+use tracing_subscriber::{
+    fmt::time::ChronoUtc, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer,
+};
+
+use super::otel_trace::get_otel_layer;
+use crate::config::TraceConfig;
+
+const TIME_FORMAT: &str = "%Y-%m-%d %H:%M:%S";
+const DEFAULT_LOG_TARGET: &str = "smg";
+
+#[derive(Debug, Clone)]
+pub struct LoggingConfig {
+    pub level: Level,
+    pub json_format: bool,
+    pub log_dir: Option<String>,
+    pub colorize: bool,
+    pub log_file_name: String,
+    pub log_targets: Option<Vec<String>>,
+}
+
+impl Default for LoggingConfig {
+    #[inline]
+    fn default() -> Self {
+        Self {
+            level: Level::INFO,
+            json_format: false,
+            log_dir: None,
+            colorize: true,
+            log_file_name: "smg".to_string(),
+            log_targets: Some(vec![DEFAULT_LOG_TARGET.to_string()]),
+        }
+    }
+}
+
+/// Guard that keeps the file appender thread alive.
+#[allow(dead_code)]
+pub struct LogGuard {
+    _file_guard: Option<WorkerGuard>,
+}
+
+#[inline]
+const fn level_to_str(level: Level) -> &'static str {
+    match level {
+        Level::TRACE => "trace",
+        Level::DEBUG => "debug",
+        Level::INFO => "info",
+        Level::WARN => "warn",
+        Level::ERROR => "error",
+    }
+}
+
+#[inline]
+fn build_filter_string(targets: &[String], level_filter: &str) -> String {
+    // Exact capacity: sum of target lengths + "=" and level per target + commas between
+    let capacity = targets.iter().map(String::len).sum::<usize>()
+        + targets.len() * (1 + level_filter.len())
+        + targets.len().saturating_sub(1);
+    let mut filter_string = String::with_capacity(capacity);
+
+    for (i, target) in targets.iter().enumerate() {
+        if i > 0 {
+            filter_string.push(',');
+        }
+        filter_string.push_str(target);
+        filter_string.push('=');
+        filter_string.push_str(level_filter);
+    }
+
+    filter_string
+}
+
+pub fn init_logging(config: LoggingConfig, otel_layer_config: Option<TraceConfig>) -> LogGuard {
+    let _ = LogTracer::init();
+
+    let level_filter = level_to_str(config.level);
+
+    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
+        let filter_string = match &config.log_targets {
+            Some(targets) if !targets.is_empty() => build_filter_string(targets, level_filter),
+            _ => {
+                let mut s =
+                    String::with_capacity(DEFAULT_LOG_TARGET.len() + 1 + level_filter.len());
+                s.push_str(DEFAULT_LOG_TARGET);
+                s.push('=');
+                s.push_str(level_filter);
+                s
+            }
+        };
+
+        EnvFilter::new(filter_string)
+    });
+
+    let mut layers = Vec::with_capacity(3);
+
+    let stdout_layer = tracing_subscriber::fmt::layer()
+        .with_ansi(config.colorize)
+        .with_file(true)
+        .with_line_number(true)
+        .with_timer(ChronoUtc::new(TIME_FORMAT.to_string()));
+
+    let stdout_layer = if config.json_format {
+        stdout_layer.json().flatten_event(true).boxed()
+    } else {
+        stdout_layer.boxed()
+    };
+
+    layers.push(stdout_layer);
+
+    let mut file_guard = None;
+
+    if let Some(log_dir) = &config.log_dir {
+        let log_dir = PathBuf::from(log_dir);
+
+        if !log_dir.exists() {
+            if let Err(e) = std::fs::create_dir_all(&log_dir) {
+                eprintln!("Failed to create log directory: {}", e);
+                return LogGuard { _file_guard: None };
+            }
+        }
+
+        let file_appender =
+            RollingFileAppender::new(Rotation::DAILY, log_dir, &config.log_file_name);
+
+        let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
+        file_guard = Some(guard);
+
+        let file_layer = tracing_subscriber::fmt::layer()
+            .with_ansi(false)
+            .with_file(true)
+            .with_line_number(true)
+            .with_timer(ChronoUtc::new(TIME_FORMAT.to_string()))
+            .with_writer(non_blocking);
+
+        let file_layer = if config.json_format {
+            file_layer.json().flatten_event(true).boxed()
+        } else {
+            file_layer.boxed()
+        };
+
+        layers.push(file_layer);
+    }
+
+    if let Some(otel_layer_config) = &otel_layer_config {
+        if otel_layer_config.enable_trace {
+            match get_otel_layer() {
+                Ok(otel_layer) => {
+                    layers.push(otel_layer);
+                }
+                Err(e) => {
+                    eprintln!("Failed to initialize OpenTelemetry: {}", e);
+                }
+            }
+        }
+    }
+
+    let _ = tracing_subscriber::registry()
+        .with(env_filter)
+        .with(layers)
+        .try_init();
+
+    LogGuard {
+        _file_guard: file_guard,
+    }
+}

third_party/sglang/sgl-model-gateway/src/observability/mod.rs

@@ -0,0 +1,8 @@
+//! Observability utilities for logging, metrics, and tracing.
+
+pub mod events;
+pub mod gauge_histogram;
+pub mod inflight_tracker;
+pub mod logging;
+pub mod metrics;
+pub mod otel_trace;

third_party/sglang/sgl-model-gateway/src/observability/otel_trace.rs

@@ -0,0 +1,279 @@
+//! OpenTelemetry tracing integration.
+
+use std::{
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        OnceLock,
+    },
+    time::Duration,
+};
+
+use anyhow::Result;
+use axum::http::{HeaderMap, HeaderName, HeaderValue};
+use opentelemetry::{global, trace::TracerProvider as _, KeyValue};
+use opentelemetry_otlp::WithExportConfig;
+use opentelemetry_sdk::{
+    propagation::TraceContextPropagator,
+    runtime,
+    trace::{BatchConfigBuilder, BatchSpanProcessor, Tracer as SdkTracer, TracerProvider},
+    Resource,
+};
+use tokio::task::spawn_blocking;
+use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
+use tracing::{Metadata, Subscriber};
+use tracing_opentelemetry::{self, OpenTelemetrySpanExt};
+use tracing_subscriber::{
+    layer::{Context, Filter},
+    Layer,
+};
+
+use super::events::get_module_path as events_module_path;
+
+/// Whether OpenTelemetry tracing is enabled.
+///
+/// This flag guards access to TRACER and PROVIDER. We use Release/Acquire
+/// ordering to ensure proper synchronization: writes to TRACER/PROVIDER
+/// happen-before the Release store, and Acquire loads happen-before reads.
+static ENABLED: AtomicBool = AtomicBool::new(false);
+static TRACER: OnceLock<SdkTracer> = OnceLock::new();
+static PROVIDER: OnceLock<TracerProvider> = OnceLock::new();
+static ALLOWED_TARGETS: OnceLock<[&'static str; 3]> = OnceLock::new();
+
+#[inline]
+fn get_allowed_targets() -> &'static [&'static str; 3] {
+    ALLOWED_TARGETS.get_or_init(|| {
+        [
+            "smg::otel-trace",
+            "smg::observability::otel_trace",
+            events_module_path(),
+        ]
+    })
+}
+
+/// Filter that only allows specific module targets to be exported to OTEL.
+#[derive(Clone, Copy, Default)]
+pub(crate) struct CustomOtelFilter;
+
+impl CustomOtelFilter {
+    #[inline]
+    pub const fn new() -> Self {
+        Self
+    }
+
+    #[inline]
+    fn is_allowed(target: &str) -> bool {
+        get_allowed_targets()
+            .iter()
+            .any(|allowed| target.starts_with(allowed))
+    }
+}
+
+impl<S> Filter<S> for CustomOtelFilter
+where
+    S: Subscriber,
+{
+    #[inline]
+    fn enabled(&self, meta: &Metadata<'_>, _cx: &Context<'_, S>) -> bool {
+        Self::is_allowed(meta.target())
+    }
+
+    #[inline]
+    fn callsite_enabled(&self, meta: &'static Metadata<'static>) -> tracing::subscriber::Interest {
+        if Self::is_allowed(meta.target()) {
+            tracing::subscriber::Interest::always()
+        } else {
+            tracing::subscriber::Interest::never()
+        }
+    }
+}
+
+pub fn otel_tracing_init(enable: bool, otlp_endpoint: Option<&str>) -> Result<()> {
+    if !enable {
+        // Use Release to ensure any prior OTEL state changes are visible
+        ENABLED.store(false, Ordering::Release);
+        return Ok(());
+    }
+
+    let endpoint = otlp_endpoint.unwrap_or("localhost:4317");
+    let endpoint = if !endpoint.starts_with("http://") && !endpoint.starts_with("https://") {
+        format!("http://{}", endpoint)
+    } else {
+        endpoint.to_string()
+    };
+
+    global::set_text_map_propagator(TraceContextPropagator::new());
+
+    let exporter = opentelemetry_otlp::SpanExporter::builder()
+        .with_tonic()
+        .with_endpoint(&endpoint)
+        .with_protocol(opentelemetry_otlp::Protocol::Grpc)
+        .build()
+        .map_err(|e| {
+            eprintln!("[tracing] Failed to create OTLP exporter: {}", e);
+            anyhow::anyhow!("Failed to create OTLP exporter: {}", e)
+        })?;
+
+    let batch_config = BatchConfigBuilder::default()
+        .with_scheduled_delay(Duration::from_millis(500))
+        .with_max_export_batch_size(64)
+        .build();
+
+    let span_processor = BatchSpanProcessor::builder(exporter, runtime::Tokio)
+        .with_batch_config(batch_config)
+        .build();
+
+    let resource =
+        Resource::default().merge(&Resource::new(vec![KeyValue::new("service.name", "smg")]));
+
+    let provider = TracerProvider::builder()
+        .with_span_processor(span_processor)
+        .with_resource(resource)
+        .build();
+
+    PROVIDER
+        .set(provider.clone())
+        .map_err(|_| anyhow::anyhow!("Provider already initialized"))?;
+
+    let tracer = provider.tracer("smg");
+
+    TRACER
+        .set(tracer)
+        .map_err(|_| anyhow::anyhow!("Tracer already initialized"))?;
+
+    let _ = global::set_tracer_provider(provider);
+
+    // Use Release ordering: all writes to TRACER/PROVIDER happen-before this store,
+    // so any thread that loads ENABLED with Acquire will see the initialized state.
+    ENABLED.store(true, Ordering::Release);
+
+    eprintln!("[tracing] OpenTelemetry initialized successfully");
+    Ok(())
+}
+
+/// Get the OpenTelemetry tracing layer. Must be called after `otel_tracing_init`.
+pub fn get_otel_layer<S>() -> Result<Box<dyn Layer<S> + Send + Sync + 'static>>
+where
+    S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a> + Send + Sync,
+{
+    if !is_otel_enabled() {
+        anyhow::bail!("OpenTelemetry is not enabled");
+    }
+
+    let tracer = TRACER
+        .get()
+        .ok_or_else(|| anyhow::anyhow!("Tracer not initialized. Call otel_tracing_init first."))?
+        .clone();
+
+    let layer = tracing_opentelemetry::layer()
+        .with_tracer(tracer)
+        .with_filter(CustomOtelFilter::new());
+
+    Ok(Box::new(layer))
+}
+
+/// Check if OpenTelemetry tracing is enabled.
+///
+/// Uses Acquire ordering to synchronize with the Release store in `otel_tracing_init`,
+/// ensuring that if this returns true, TRACER and PROVIDER are fully initialized.
+#[inline]
+pub fn is_otel_enabled() -> bool {
+    ENABLED.load(Ordering::Acquire)
+}
+
+pub async fn flush_spans_async() -> Result<()> {
+    if !is_otel_enabled() {
+        return Ok(());
+    }
+
+    let provider = PROVIDER
+        .get()
+        .ok_or_else(|| anyhow::anyhow!("Provider not initialized"))?
+        .clone();
+
+    spawn_blocking(move || provider.force_flush())
+        .await
+        .map_err(|e| anyhow::anyhow!("Failed to flush spans: {}", e))?;
+
+    Ok(())
+}
+
+pub fn shutdown_otel() {
+    // Use Acquire to ensure we see any prior OTEL operations
+    if ENABLED.load(Ordering::Acquire) {
+        global::shutdown_tracer_provider();
+        // Use Release to ensure shutdown completes before flag is cleared
+        ENABLED.store(false, Ordering::Release);
+        eprintln!("[tracing] OpenTelemetry shut down");
+    }
+}
+
+/// Inject W3C trace context headers into an HTTP request.
+#[inline]
+pub fn inject_trace_context_http(headers: &mut HeaderMap) {
+    if !is_otel_enabled() {
+        return;
+    }
+
+    let context = tracing::Span::current().context();
+
+    struct HeaderInjector<'a>(&'a mut HeaderMap);
+
+    impl opentelemetry::propagation::Injector for HeaderInjector<'_> {
+        #[inline]
+        fn set(&mut self, key: &str, value: String) {
+            if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
+                if let Ok(header_value) = HeaderValue::from_str(&value) {
+                    self.0.insert(header_name, header_value);
+                }
+            }
+        }
+    }
+
+    global::get_text_map_propagator(|propagator| {
+        propagator.inject_context(&context, &mut HeaderInjector(headers));
+    });
+}
+
+/// Inject W3C trace context into gRPC metadata.
+#[inline]
+pub fn inject_trace_context_grpc(metadata: &mut MetadataMap) {
+    if !is_otel_enabled() {
+        return;
+    }
+
+    let context = tracing::Span::current().context();
+
+    struct MetadataInjector<'a>(&'a mut MetadataMap);
+
+    impl opentelemetry::propagation::Injector for MetadataInjector<'_> {
+        #[inline]
+        fn set(&mut self, key: &str, value: String) {
+            if let Ok(metadata_key) = MetadataKey::from_bytes(key.as_bytes()) {
+                if let Ok(metadata_value) = MetadataValue::try_from(&value) {
+                    self.0.insert(metadata_key, metadata_value);
+                }
+            }
+        }
+    }
+
+    global::get_text_map_propagator(|propagator| {
+        propagator.inject_context(&context, &mut MetadataInjector(metadata));
+    });
+}
+
+/// OpenTelemetry trace injector implementing the `smg_grpc_client::TraceInjector` trait.
+///
+/// This bridges sglang's OTel integration with the `smg-grpc-client` crate's
+/// trace injection interface, enabling distributed tracing across gRPC calls.
+#[derive(Clone, Default)]
+pub struct OtelTraceInjector;
+
+impl smg_grpc_client::TraceInjector for OtelTraceInjector {
+    fn inject(
+        &self,
+        metadata: &mut MetadataMap,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        inject_trace_context_grpc(metadata);
+        Ok(())
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/cache_aware.rs

@@ -0,0 +1,887 @@
+/*
+    Cache-Aware Load Balancing Router
+
+    This router combines two strategies to optimize both cache utilization and request distribution:
+
+    1. Cache-Aware Routing (Approximate Tree)
+    2. Load Balancing (Shortest Queue with Balance Thresholds)
+
+    The router dynamically switches between these strategies based on load conditions:
+    - Uses load balancing when the system is imbalanced
+    - Uses cache-aware routing when the system is balanced
+
+    A system is considered imbalanced if both conditions are met:
+    1. (max - min) > abs_threshold
+    2. max > rel_threshold * min
+
+    Strategy Details:
+
+    1. Cache-Aware Routing (Approximate Tree)
+    -------------------------------------------
+    This strategy maintains an approximate radix tree for each worker based on request history,
+    eliminating the need for direct cache state queries. The tree stores raw text characters
+    instead of token IDs to avoid tokenization overhead.
+
+    Process:
+    a. For each request, find the worker with the highest prefix match
+    b. If match rate > cache_threshold:
+    Route to the worker with highest match (likely has relevant data cached)
+    c. If match rate ≤ cache_threshold:
+    Route to the worker with smallest tree size (most available cache capacity)
+    d. Background maintenance:
+    Periodically evict least recently used leaf nodes to prevent memory overflow
+
+    2. Load Balancing (Shortest Queue)
+    -------------------------------------------
+    This strategy tracks pending request counts per worker and routes new requests
+    to the least busy worker when the system is detected to be imbalanced.
+
+    Configuration Parameters:
+    ------------------------
+    1. cache_threshold: (float, 0.0 to 1.0)
+    Minimum prefix match ratio to use highest-match routing.
+    Below this threshold, routes to worker with most available cache space.
+
+    2. balance_abs_threshold: (integer)
+    Absolute difference threshold for load imbalance detection.
+    System is potentially imbalanced if (max_load - min_load) > abs_threshold
+
+    3. balance_rel_threshold: (float)
+    Relative ratio threshold for load imbalance detection.
+    System is potentially imbalanced if max_load > min_load * rel_threshold
+    Used in conjunction with abs_threshold to determine final imbalance state.
+
+    4. eviction_interval_secs: (integer)
+    Interval between LRU eviction cycles for the approximate trees.
+
+    5. max_tree_size: (integer)
+    Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted
+    during the next eviction cycle.
+*/
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use dashmap::DashMap;
+use rand::Rng;
+use smg_mesh::{tree_ops::TreeOperation, OptionalMeshSyncManager};
+use tracing::{debug, warn};
+
+use super::{
+    get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask,
+    CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo,
+};
+use crate::core::{Worker, UNKNOWN_MODEL_ID};
+
+/// Cache-aware routing policy
+///
+/// Routes requests based on cache affinity when load is balanced,
+/// switches to shortest-queue routing when load is imbalanced.
+/// Maintains separate trees per model for multi-model support.
+/// Supports mesh synchronization of tree operations across cluster nodes.
+/// When mesh is not enabled, the policy works independently without synchronization.
+#[derive(Debug)]
+pub struct CacheAwarePolicy {
+    config: CacheAwareConfig,
+    trees: Arc<DashMap<String, Arc<Tree>>>,
+    mesh_sync: OptionalMeshSyncManager,
+    _eviction_task: Option<PeriodicTask>,
+}
+
+impl CacheAwarePolicy {
+    pub fn new() -> Self {
+        Self::with_config(CacheAwareConfig::default())
+    }
+
+    pub fn with_config(config: CacheAwareConfig) -> Self {
+        let trees = Arc::new(DashMap::<String, Arc<Tree>>::new());
+
+        // Start background eviction thread if configured
+        let eviction_task = if config.eviction_interval_secs > 0 {
+            let trees_clone = Arc::clone(&trees);
+            let max_tree_size = config.max_tree_size;
+
+            Some(PeriodicTask::spawn(
+                config.eviction_interval_secs,
+                "Eviction",
+                move || {
+                    for tree_ref in trees_clone.iter() {
+                        let model_id = tree_ref.key();
+                        let tree = tree_ref.value();
+                        tree.evict_tenant_by_size(max_tree_size);
+
+                        debug!(
+                            "Cache eviction completed for model {}, max_size: {}",
+                            model_id, max_tree_size
+                        );
+                    }
+                },
+            ))
+        } else {
+            None
+        };
+
+        Self {
+            config,
+            trees,
+            mesh_sync: None,
+            _eviction_task: eviction_task,
+        }
+    }
+
+    /// Set mesh sync manager (can be called after construction)
+    pub fn set_mesh_sync(&mut self, mesh_sync: OptionalMeshSyncManager) {
+        self.mesh_sync = mesh_sync.clone();
+        if mesh_sync.is_some() {
+            self.restore_tree_state_from_mesh();
+        }
+    }
+
+    /// Initialize the tree with worker URLs (used only during initial setup)
+    pub fn init_workers(&self, workers: &[Arc<dyn Worker>]) {
+        // Group workers by model
+        let mut model_workers: std::collections::HashMap<String, Vec<&Arc<dyn Worker>>> =
+            std::collections::HashMap::new();
+        for worker in workers {
+            let tree_key = normalize_model_key(worker.model_id());
+            model_workers
+                .entry(tree_key.to_string())
+                .or_default()
+                .push(worker);
+        }
+
+        // Initialize tree for each model
+        for (tree_key, model_workers) in model_workers {
+            let tree = self
+                .trees
+                .entry(tree_key)
+                .or_insert_with(|| Arc::new(Tree::new()));
+            for worker in model_workers {
+                tree.insert("", worker.url());
+            }
+        }
+    }
+
+    /// Add a single worker to the tree (incremental update)
+    pub fn add_worker(&self, worker: &dyn Worker) {
+        let tree_key = normalize_model_key(worker.model_id());
+        let tree = self
+            .trees
+            .entry(tree_key.to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+        tree.insert("", worker.url());
+    }
+
+    /// Add a worker by URL and model (for backward compatibility)
+    pub fn add_worker_by_url(&self, url: &str, model_id: &str) {
+        let tree = self
+            .trees
+            .entry(model_id.to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+        tree.insert("", url);
+    }
+
+    /// Remove a worker from the tree
+    pub fn remove_worker(&self, worker: &dyn Worker) {
+        let tree_key = normalize_model_key(worker.model_id());
+        if let Some(tree) = self.trees.get(tree_key) {
+            tree.remove_tenant(worker.url());
+        }
+    }
+
+    /// Remove a worker by URL (removes from all model trees for backward compatibility)
+    pub fn remove_worker_by_url(&self, url: &str) {
+        // Remove from all trees since we don't know which model it belongs to
+        for tree_ref in self.trees.iter() {
+            tree_ref.value().remove_tenant(url);
+        }
+    }
+
+    /// Restore tree state from mesh store
+    /// This is called during initialization to rebuild trees from synchronized state
+    fn restore_tree_state_from_mesh(&self) {
+        if let Some(ref mesh_sync) = self.mesh_sync {
+            // Get all tree states from mesh
+            // We need to iterate through all models that have tree states
+            // For now, we'll restore trees for models that are already in our trees map
+            // In a full implementation, we might want to query mesh for all tree states
+
+            for tree_ref in self.trees.iter() {
+                let model_id = tree_ref.key();
+                if let Some(tree_state) = mesh_sync.get_tree_state(model_id) {
+                    debug!(
+                        "Restoring tree state for model {} with {} operations",
+                        model_id,
+                        tree_state.operations.len()
+                    );
+
+                    let tree = tree_ref.value();
+                    // Apply all operations to rebuild the tree
+                    for operation in &tree_state.operations {
+                        match operation {
+                            TreeOperation::Insert(insert_op) => {
+                                tree.insert(&insert_op.text, &insert_op.tenant);
+                            }
+                            TreeOperation::Remove(remove_op) => {
+                                tree.remove_tenant(&remove_op.tenant);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    /// Normalize model_id for mesh synchronization
+    /// Converts empty model_id to UNKNOWN_MODEL_ID for consistency
+    fn normalize_mesh_model_id(model_id: &str) -> &str {
+        if model_id.is_empty() {
+            UNKNOWN_MODEL_ID
+        } else {
+            model_id
+        }
+    }
+
+    /// Apply remote tree operation from mesh
+    /// This is called when receiving tree state updates from other nodes
+    pub fn apply_remote_tree_operation(&self, model_id: &str, operation: &TreeOperation) {
+        let tree_key = Self::normalize_mesh_model_id(model_id);
+
+        let tree = self
+            .trees
+            .entry(tree_key.to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+
+        match operation {
+            TreeOperation::Insert(insert_op) => {
+                tree.insert(&insert_op.text, &insert_op.tenant);
+                debug!(
+                    "Applied remote tree insert: model={}, text={}, tenant={}",
+                    model_id, insert_op.text, insert_op.tenant
+                );
+            }
+            TreeOperation::Remove(remove_op) => {
+                tree.remove_tenant(&remove_op.tenant);
+                debug!(
+                    "Applied remote tree remove: model={}, tenant={}",
+                    model_id, remove_op.tenant
+                );
+            }
+        }
+    }
+
+    /// Run cache eviction to prevent unbounded growth
+    pub fn evict_cache(&self, max_size: usize) {
+        for tree_ref in self.trees.iter() {
+            let model_id = tree_ref.key();
+            let tree = tree_ref.value();
+            tree.evict_tenant_by_size(max_size);
+            debug!(
+                "Cache eviction for model {}, max_size: {}",
+                model_id, max_size
+            );
+        }
+    }
+
+    fn select_worker_min_load(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        request_text: &Option<&str>,
+        healthy_indices: &[usize],
+        model_id: &str,
+        max_load: usize,
+        min_load: usize,
+    ) -> Option<usize> {
+        // Log load balancing trigger (only compute worker loads if debug enabled)
+        if tracing::enabled!(tracing::Level::DEBUG) {
+            let worker_loads: Vec<(&str, usize)> =
+                workers.iter().map(|w| (w.url(), w.load())).collect();
+            debug!(
+                "Load balancing triggered | max: {} | min: {} | workers: {:?}",
+                max_load, min_load, worker_loads
+            );
+        }
+
+        // Use shortest queue when imbalanced
+        let min_load_idx = healthy_indices
+            .iter()
+            .min_by_key(|&&idx| workers[idx].load())
+            .copied()?;
+
+        // Even in imbalanced mode, update the tree to maintain cache state
+        if let Some(text) = request_text {
+            // Get the tree reference without locking the entire HashMap
+            // DashMap only locks the specific shard containing this key
+            let tree = self.trees.get(model_id).map(|entry| entry.value().clone());
+
+            if let Some(tree) = tree {
+                let worker_url = workers[min_load_idx].url();
+                // Now we can work with the tree without holding the HashMap lock
+                tree.insert(text, worker_url);
+
+                // Sync insert operation to mesh if enabled (no-op if mesh is not enabled)
+                if let Some(ref mesh_sync) = self.mesh_sync {
+                    use smg_mesh::tree_ops::TreeInsertOp;
+                    let op = TreeOperation::Insert(TreeInsertOp {
+                        text: text.to_string(),
+                        tenant: worker_url.to_string(),
+                    });
+                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                        warn!("Failed to sync tree insert operation to mesh: {}", e);
+                    }
+                }
+            } else {
+                debug!(
+                    "Warning: No tree found for model '{}', skipping cache update",
+                    model_id
+                );
+            }
+        }
+
+        // Increment processed counter
+        workers[min_load_idx].increment_processed();
+
+        Some(min_load_idx)
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for CacheAwarePolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let request_text = info.request_text;
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        // Determine the model for this set of workers (router pre-filters by model)
+        // All workers should be from the same model
+        let model_id = normalize_model_key(workers[healthy_indices[0]].model_id());
+
+        // Get current load statistics - compute min/max in single pass without allocation
+        let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| {
+            let load = w.load();
+            (min.min(load), max.max(load))
+        });
+        let min_load = if min_load == usize::MAX { 0 } else { min_load };
+
+        // Check if load is imbalanced
+        let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold
+            && (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold);
+
+        if is_imbalanced {
+            return self.select_worker_min_load(
+                workers,
+                &request_text,
+                &healthy_indices,
+                model_id,
+                max_load,
+                min_load,
+            );
+        }
+
+        // Use cache-aware routing when balanced
+        let text = request_text.unwrap_or("");
+
+        // Get the tree reference without locking the entire HashMap
+        // DashMap only locks the specific shard containing this key
+        let tree = self.trees.get(model_id).map(|entry| entry.value().clone());
+
+        if let Some(tree) = tree {
+            // Now we work with the tree without holding the HashMap lock
+            // Use prefix_match_with_counts to avoid redundant chars().count() calls
+            let result = tree.prefix_match_with_counts(text);
+            let match_rate = if result.input_char_count == 0 {
+                0.0
+            } else {
+                result.matched_char_count as f32 / result.input_char_count as f32
+            };
+
+            // Select worker without String allocation
+            let selected_idx = if match_rate > self.config.cache_threshold {
+                // Cache hit path: find worker by URL (compare &str directly, no allocation)
+                let tenant_url: &str = &result.tenant;
+                workers
+                    .iter()
+                    .position(|w| w.url() == tenant_url)
+                    .filter(|&idx| workers[idx].is_healthy())
+            } else {
+                // Low cache match: use worker with minimum load
+                healthy_indices
+                    .iter()
+                    .min_by_key(|&&idx| workers[idx].load())
+                    .copied()
+            };
+
+            if let Some(idx) = selected_idx {
+                // Update the tree with this request (use worker URL directly, no allocation)
+                tree.insert(text, workers[idx].url());
+
+                // Sync insert operation to mesh if enabled (no-op if mesh is not enabled)
+                if let Some(ref mesh_sync) = self.mesh_sync {
+                    use smg_mesh::tree_ops::TreeInsertOp;
+                    let op = TreeOperation::Insert(TreeInsertOp {
+                        text: text.to_string(),
+                        tenant: workers[idx].url().to_string(),
+                    });
+                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                        warn!("Failed to sync tree insert operation to mesh: {}", e);
+                    }
+                }
+
+                // Increment processed counter
+                workers[idx].increment_processed();
+
+                return Some(idx);
+            }
+
+            // Selected worker no longer exists or unhealthy, remove stale tenant from tree
+            if match_rate > self.config.cache_threshold {
+                let tenant_url: &str = &result.tenant;
+                tree.remove_tenant(tenant_url);
+                debug!("Removed stale worker {} from cache tree", tenant_url);
+
+                // Sync removal to mesh if enabled (no-op if mesh is not enabled)
+                if let Some(ref mesh_sync) = self.mesh_sync {
+                    use smg_mesh::tree_ops::TreeRemoveOp;
+                    let op = TreeOperation::Remove(TreeRemoveOp {
+                        tenant: tenant_url.to_string(),
+                    });
+                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                        warn!("Failed to sync tree remove operation to mesh: {}", e);
+                    }
+                }
+            }
+
+            // Fallback to first healthy worker
+            healthy_indices.first().copied()
+        } else {
+            // No tree for this model, log warning and use random selection
+            debug!(
+                "Warning: No tree found for model '{}', using random worker selection",
+                model_id
+            );
+            // Return a random healthy worker
+            let mut rng = rand::rng();
+            let random_idx = rng.random_range(0..healthy_indices.len());
+            Some(healthy_indices[random_idx])
+        }
+    }
+
+    fn on_request_complete(&self, worker_url: &str, success: bool) {
+        // Could track success rates per worker for more intelligent routing
+        if !success {
+            // Optionally reduce affinity for failed requests
+            tracing::debug!(
+                "Request to {} completed with success={}",
+                worker_url,
+                success
+            );
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        "cache_aware"
+    }
+
+    fn needs_request_text(&self) -> bool {
+        true // Cache-aware policy needs request text for cache affinity
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+impl Default for CacheAwarePolicy {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_cache_aware_with_balanced_load() {
+        // Create policy without eviction thread for testing
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0, // Disable eviction thread
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+        ];
+
+        // Initialize the policy with workers
+        policy.init_workers(&workers);
+
+        // First request should be distributed
+        let idx1 = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("hello world"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+        // Same request should go to same worker (cache hit)
+        let idx2 = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("hello world"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx1, idx2);
+
+        // Similar request should also go to same worker
+        let idx3 = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("hello"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx1, idx3);
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_with_imbalanced_load() {
+        let policy = CacheAwarePolicy::with_config(CacheAwareConfig {
+            cache_threshold: 0.5,
+            balance_abs_threshold: 5,
+            balance_rel_threshold: 2.0,
+            eviction_interval_secs: 0, // Disable eviction thread
+            max_tree_size: 10000,
+        });
+
+        let worker1 = BasicWorkerBuilder::new("http://w1:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+        let worker2 = BasicWorkerBuilder::new("http://w2:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // Create significant load imbalance
+        for _ in 0..20 {
+            worker1.increment_load();
+        }
+        // worker2 has load 0
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker1), Arc::new(worker2)];
+        policy.init_workers(&workers);
+
+        // Should select worker2 (lower load) despite cache affinity
+        let info = SelectWorkerInfo {
+            request_text: Some("test"),
+            ..Default::default()
+        };
+        for _ in 0..5 {
+            let idx = policy.select_worker(&workers, &info).await.unwrap();
+            assert_eq!(idx, 1); // Should always pick worker2
+        }
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_worker_removal() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0, // Disable eviction thread
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        policy.init_workers(&workers);
+
+        // Route some requests
+        policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test1"),
+                    ..Default::default()
+                },
+            )
+            .await;
+        policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test2"),
+                    ..Default::default()
+                },
+            )
+            .await;
+
+        // Remove a worker
+        policy.remove_worker_by_url("http://w1:8000");
+        workers[0].set_healthy(false);
+
+        // All requests should now go to worker2
+        let idx = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test1"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx, 1);
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_sync_tree_operation_to_mesh() {
+        use std::sync::Arc;
+
+        use smg_mesh::{stores::StateStores, sync::MeshSyncManager};
+
+        let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let mut policy = CacheAwarePolicy::with_config(config);
+        policy.set_mesh_sync(Some(mesh_sync.clone()));
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        )];
+
+        policy.init_workers(&workers);
+
+        // Select worker with a request - should sync to mesh
+        let _idx = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test request"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+        // Verify tree operation was synced to mesh (under UNKNOWN_MODEL_ID since no model was specified)
+        let tree_state = mesh_sync.get_tree_state(UNKNOWN_MODEL_ID);
+        assert!(tree_state.is_some());
+        let tree = tree_state.unwrap();
+        assert!(!tree.operations.is_empty());
+    }
+
+    #[test]
+    fn test_cache_aware_restore_tree_state_from_mesh() {
+        use std::sync::Arc;
+
+        use smg_mesh::{
+            stores::StateStores,
+            sync::MeshSyncManager,
+            tree_ops::{TreeInsertOp, TreeOperation},
+        };
+
+        let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
+
+        // Pre-populate mesh with tree state
+        let op1 = TreeOperation::Insert(TreeInsertOp {
+            text: "test_text_1".to_string(),
+            tenant: "http://w1:8000".to_string(),
+        });
+        mesh_sync
+            .sync_tree_operation("model1".to_string(), op1)
+            .unwrap();
+
+        let op2 = TreeOperation::Insert(TreeInsertOp {
+            text: "test_text_2".to_string(),
+            tenant: "http://w2:8000".to_string(),
+        });
+        mesh_sync
+            .sync_tree_operation("model1".to_string(), op2)
+            .unwrap();
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let mut policy = CacheAwarePolicy::with_config(config);
+        policy.set_mesh_sync(Some(mesh_sync.clone()));
+
+        // Initialize with a model to trigger restore
+        let _workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        )];
+
+        // Create a tree entry for model1 to trigger restore
+        let _tree = policy
+            .trees
+            .entry("model1".to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+
+        // Manually trigger restore (normally done in constructor)
+        // For testing, we'll verify the tree state exists in mesh
+        let tree_state = mesh_sync.get_tree_state("model1");
+        assert!(tree_state.is_some());
+        let state = tree_state.unwrap();
+        assert_eq!(state.operations.len(), 2);
+    }
+
+    #[test]
+    fn test_cache_aware_apply_remote_tree_operation() {
+        use std::sync::Arc;
+
+        use smg_mesh::{
+            stores::StateStores,
+            sync::MeshSyncManager,
+            tree_ops::{TreeInsertOp, TreeOperation},
+        };
+
+        let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let mut policy = CacheAwarePolicy::with_config(config);
+        policy.set_mesh_sync(Some(mesh_sync.clone()));
+
+        // Apply remote tree operation
+        let remote_op = TreeOperation::Insert(TreeInsertOp {
+            text: "remote_text".to_string(),
+            tenant: "http://remote:8000".to_string(),
+        });
+
+        policy.apply_remote_tree_operation("model1", &remote_op);
+
+        // Verify the tree was updated
+        let tree = policy.trees.get("model1");
+        assert!(tree.is_some());
+    }
+
+    #[test]
+    fn test_cache_aware_multi_node_consistency() {
+        use std::sync::Arc;
+
+        use smg_mesh::{
+            stores::StateStores,
+            sync::MeshSyncManager,
+            tree_ops::{TreeInsertOp, TreeOperation},
+        };
+
+        // Simulate two nodes
+        let stores1 = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync1 = Arc::new(MeshSyncManager::new(stores1.clone(), "node1".to_string()));
+
+        let stores2 = Arc::new(StateStores::with_self_name("node2".to_string()));
+        let mesh_sync2 = Arc::new(MeshSyncManager::new(stores2.clone(), "node2".to_string()));
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+
+        let mut _policy1 = CacheAwarePolicy::with_config(config.clone());
+        _policy1.set_mesh_sync(Some(mesh_sync1.clone()));
+        let mut _policy2 = CacheAwarePolicy::with_config(config);
+        _policy2.set_mesh_sync(Some(mesh_sync2.clone()));
+
+        // Node1 syncs a tree operation
+        let op = TreeOperation::Insert(TreeInsertOp {
+            text: "shared_text".to_string(),
+            tenant: "http://shared:8000".to_string(),
+        });
+        mesh_sync1
+            .sync_tree_operation("model1".to_string(), op.clone())
+            .unwrap();
+
+        // Node2 should be able to get the tree state
+        let tree_state = mesh_sync2.get_tree_state("model1");
+        // Note: In a real scenario, this would be synced via gossip protocol
+        // For unit test, we verify the sync mechanism works
+        // Tree state may or may not exist depending on sync timing
+        let _ = tree_state;
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_without_mesh() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        )];
+
+        policy.init_workers(&workers);
+
+        // Should work without mesh
+        let idx = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test request"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx, 0);
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/consistent_hashing.rs

@@ -0,0 +1,521 @@
+//! Consistent hashing routing policy with header-based routing support
+//!
+//! Supports two routing mechanisms via HTTP headers:
+//! - `X-SMG-Target-Worker`: Direct routing by worker index (0-based), returns None if unavailable
+//! - `X-SMG-Routing-Key`: Consistent hash routing for session affinity
+//!
+//! ## Consistent Hashing
+//!
+//! Uses a pre-computed hash ring from WorkerRegistry where:
+//! 1. Each worker is placed at a fixed position based on hash(worker_url)
+//! 2. Keys are hashed to the ring, then walk clockwise to find first healthy worker
+//! 3. When workers scale up/down, only keys in the affected range redistribute (~1/N keys move)
+//!
+//! The ring is built once when workers are added/removed, not per-request.
+//! This ensures O(log n) lookup performance.
+//!
+//! Complexity: O(log n) binary search + O(k) walk where k = consecutive unhealthy workers.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use rand::Rng as _;
+
+use super::{LoadBalancingPolicy, SelectWorkerInfo};
+use crate::{
+    core::Worker,
+    observability::metrics::Metrics,
+    routers::header_utils::{extract_routing_key, extract_target_worker},
+};
+
+/// Execution branch for metrics
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Branch {
+    NoHealthyWorkers,
+    TargetWorkerHit,
+    TargetWorkerMiss,
+    RoutingKeyHit,
+    RandomFallback,
+}
+
+impl Branch {
+    #[inline]
+    const fn as_str(&self) -> &'static str {
+        match self {
+            Self::NoHealthyWorkers => "no_healthy_workers",
+            Self::TargetWorkerHit => "target_worker_hit",
+            Self::TargetWorkerMiss => "target_worker_miss",
+            Self::RoutingKeyHit => "routing_key_hit",
+            Self::RandomFallback => "random_fallback",
+        }
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct ConsistentHashingPolicy;
+
+impl ConsistentHashingPolicy {
+    pub fn new() -> Self {
+        Self
+    }
+
+    /// Use consistent hashing to find a worker for the given key.
+    /// Uses pre-computed ring from SelectWorkerInfo if available.
+    ///
+    /// The ring returns a worker URL, which we then map to an index in the workers array.
+    /// This correctly handles filtered worker arrays since we match by URL, not by index.
+    ///
+    /// Complexity: O(n) to build healthy URL map + O(log n) ring lookup + O(k) walk
+    fn find_by_consistent_hash(
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+        key: &str,
+    ) -> Option<usize> {
+        // Build URL→index map for healthy workers: O(n) once, O(1) lookups
+        let healthy_url_to_idx: std::collections::HashMap<&str, usize> = workers
+            .iter()
+            .enumerate()
+            .filter(|(_, w)| w.is_healthy())
+            .map(|(i, w)| (w.url(), i))
+            .collect();
+
+        if healthy_url_to_idx.is_empty() {
+            return None;
+        }
+
+        // Use pre-computed ring if available
+        if let Some(ref ring) = info.hash_ring {
+            // O(1) lookup per URL checked instead of O(n)
+            let url = ring.find_healthy_url(key, |url| healthy_url_to_idx.contains_key(url))?;
+            return healthy_url_to_idx.get(url).copied();
+        }
+
+        // Fallback: no ring provided, use simple modulo (less optimal but functional)
+        // This shouldn't happen in normal operation as WorkerSelectionStage provides the ring
+        let mut healthy_indices: Vec<usize> = healthy_url_to_idx.values().copied().collect();
+        healthy_indices.sort_unstable(); // Ensure deterministic order
+
+        // Use blake3 for consistent hashing in fallback too
+        let hash = blake3::hash(key.as_bytes());
+        let hash_val = u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap());
+        let idx = (hash_val as usize) % healthy_indices.len();
+        Some(healthy_indices[idx])
+    }
+
+    fn select_worker_impl(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+    ) -> (Option<usize>, Branch) {
+        if workers.is_empty() {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        let target_worker = extract_target_worker(info.headers);
+        let routing_key = extract_routing_key(info.headers);
+
+        // Priority 1: X-SMG-Target-Worker - direct routing by worker index
+        // O(1) parse + O(1) bounds check + O(1) health check
+        if let Some(idx_str) = target_worker {
+            if let Ok(idx) = idx_str.parse::<usize>() {
+                if idx < workers.len() && workers[idx].is_healthy() {
+                    return (Some(idx), Branch::TargetWorkerHit);
+                }
+            }
+            return (None, Branch::TargetWorkerMiss);
+        }
+
+        // Priority 2: X-SMG-Routing-Key - consistent hash routing (O(log n))
+        if let Some(key) = routing_key {
+            return match Self::find_by_consistent_hash(workers, info, key) {
+                Some(idx) => (Some(idx), Branch::RoutingKeyHit),
+                None => (None, Branch::NoHealthyWorkers),
+            };
+        }
+
+        // Priority 3: Implicit routing key from stable headers (session affinity)
+        let implicit_key = info.headers.and_then(|h| {
+            h.get("authorization")
+                .or_else(|| h.get("x-forwarded-for"))
+                .or_else(|| h.get("cookie"))
+                .and_then(|v| v.to_str().ok())
+                .filter(|s| !s.is_empty())
+        });
+
+        if let Some(key) = implicit_key {
+            return match Self::find_by_consistent_hash(workers, info, key) {
+                Some(idx) => (Some(idx), Branch::RoutingKeyHit),
+                None => (None, Branch::NoHealthyWorkers),
+            };
+        }
+
+        // Fallback: random selection (truly anonymous client)
+        let healthy_count = workers.iter().filter(|w| w.is_healthy()).count();
+        if healthy_count == 0 {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        let random_healthy_idx = rand::rng().random_range(0..healthy_count);
+        let idx = workers
+            .iter()
+            .enumerate()
+            .filter(|(_, w)| w.is_healthy())
+            .nth(random_healthy_idx)
+            .map(|(i, _)| i)
+            .unwrap();
+
+        (Some(idx), Branch::RandomFallback)
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for ConsistentHashingPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let (result, branch) = self.select_worker_impl(workers, info);
+        Metrics::record_worker_consistent_hashing_policy_branch(branch.as_str());
+        result
+    }
+
+    fn name(&self) -> &'static str {
+        "consistent_hashing"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, HashRing, WorkerType};
+
+    fn headers_with_routing_key(key: &str) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-routing-key", key.parse().unwrap());
+        headers
+    }
+
+    fn headers_with_target_worker(idx: usize) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-target-worker", idx.to_string().parse().unwrap());
+        headers
+    }
+
+    fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .collect()
+    }
+
+    #[tokio::test]
+    async fn test_consistent_routing() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let headers = headers_with_routing_key("user-123");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, _) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+
+        // Same key should always route to same worker
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(first_idx));
+            assert_eq!(branch, Branch::RoutingKeyHit);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_different_keys_distribute() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let mut distribution = HashMap::new();
+        for i in 0..100 {
+            let headers = headers_with_routing_key(&format!("user-{}", i));
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        assert!(distribution.len() > 1, "Should distribute across workers");
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_hit() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_target_worker(1);
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(1));
+        assert_eq!(branch, Branch::TargetWorkerHit);
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_miss_out_of_bounds() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_target_worker(5); // Out of bounds
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::TargetWorkerMiss);
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_miss_unhealthy() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+        workers[1].set_healthy(false);
+
+        let headers = headers_with_target_worker(1);
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::TargetWorkerMiss);
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_priority_over_routing_key() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-target-worker", "1".parse().unwrap());
+        headers.insert("x-smg-routing-key", "some-key".parse().unwrap());
+
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(1));
+        assert_eq!(branch, Branch::TargetWorkerHit);
+    }
+
+    #[tokio::test]
+    async fn test_fallback_random_distribution() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        // Without routing headers, should distribute randomly across workers
+        let mut distribution = HashMap::new();
+        for _ in 0..100 {
+            let info = SelectWorkerInfo::default();
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert!(result.is_some());
+            assert_eq!(branch, Branch::RandomFallback);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        // Should distribute across multiple workers (not always same one)
+        assert!(
+            distribution.len() > 1,
+            "Random fallback should distribute across workers"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_no_healthy_workers() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000"]);
+        workers[0].set_healthy(false);
+
+        let headers = headers_with_routing_key("test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoHealthyWorkers);
+    }
+
+    #[tokio::test]
+    async fn test_empty_workers() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![];
+
+        let info = SelectWorkerInfo::default();
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoHealthyWorkers);
+    }
+
+    #[tokio::test]
+    async fn test_consistent_hash_minimal_redistribution() {
+        // Test that consistent hashing moves fewer keys than random redistribution
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&[
+            "http://w0:8000",
+            "http://w1:8000",
+            "http://w2:8000",
+            "http://w3:8000",
+        ]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Record which worker each key routes to with all workers healthy
+        let mut key_to_worker_before: HashMap<String, usize> = HashMap::new();
+        for i in 0..100 {
+            let key = format!("user-{}", i);
+            let headers = headers_with_routing_key(&key);
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                hash_ring: Some(ring.clone()),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            key_to_worker_before.insert(key, result.unwrap());
+        }
+
+        // Mark worker 1 as unhealthy
+        workers[1].set_healthy(false);
+
+        // Record new routing and count how many keys moved
+        let mut moved_count = 0;
+        for i in 0..100 {
+            let key = format!("user-{}", i);
+            let headers = headers_with_routing_key(&key);
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                hash_ring: Some(ring.clone()),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            let new_worker = result.unwrap();
+            let old_worker = key_to_worker_before[&key];
+
+            if new_worker != old_worker {
+                moved_count += 1;
+            }
+        }
+
+        // With consistent hashing, approximately 1/N keys should move (N = worker count)
+        // Random redistribution would move approximately (N-1)/N = 75% of keys
+        // Verify we're significantly better than random (< 50% moved)
+        let keys_on_failed_worker = key_to_worker_before.values().filter(|&&w| w == 1).count();
+        assert!(
+            moved_count <= keys_on_failed_worker + 5,
+            "Consistent hashing should only move keys from failed worker (+small variance). \
+             Expected ~{}, got {}",
+            keys_on_failed_worker,
+            moved_count
+        );
+        assert!(
+            moved_count < 50,
+            "Consistent hashing should move fewer than 50% of keys (random would move ~75%), got {}%",
+            moved_count
+        );
+    }
+
+    #[tokio::test]
+    async fn test_routing_key_failover_and_recovery() {
+        // Test that when a worker fails, keys move to another worker,
+        // and when it recovers, keys return to the original worker
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w0:8000", "http://w1:8000", "http://w2:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Find which worker a key routes to when all are healthy
+        let test_key = "session-abc-123";
+        let headers = headers_with_routing_key(test_key);
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+        let (result, _) = policy.select_worker_impl(&workers, &info);
+        let original_idx = result.unwrap();
+
+        // Mark that worker unhealthy
+        workers[original_idx].set_healthy(false);
+
+        // Key should now route to a different healthy worker
+        let (failover_result, _) = policy.select_worker_impl(&workers, &info);
+        let failover_idx = failover_result.unwrap();
+        assert_ne!(
+            failover_idx, original_idx,
+            "Should failover to different worker"
+        );
+        assert!(
+            workers[failover_idx].is_healthy(),
+            "Failover target should be healthy"
+        );
+
+        // Failover should be consistent
+        for _ in 0..5 {
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(failover_idx), "Failover should be consistent");
+        }
+
+        // Recover the original worker
+        workers[original_idx].set_healthy(true);
+
+        // Key should route back to original worker
+        let (recovered_result, _) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(
+            recovered_result,
+            Some(original_idx),
+            "Should return to original worker after recovery"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_empty_routing_key_uses_fallback() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert!(result.is_some());
+        assert_eq!(branch, Branch::RandomFallback);
+    }
+
+    #[tokio::test]
+    async fn test_policy_name() {
+        let policy = ConsistentHashingPolicy::new();
+        assert_eq!(policy.name(), "consistent_hashing");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/factory.rs

@@ -0,0 +1,156 @@
+//! Factory for creating load balancing policies
+
+use std::sync::Arc;
+
+use super::{
+    BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy,
+    LoadBalancingPolicy, ManualConfig, ManualPolicy, PowerOfTwoPolicy, PrefixHashConfig,
+    PrefixHashPolicy, RandomPolicy, RoundRobinPolicy,
+};
+use crate::config::PolicyConfig;
+
+/// Factory for creating policy instances
+pub struct PolicyFactory;
+
+impl PolicyFactory {
+    /// Create a policy from configuration
+    pub fn create_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
+        match config {
+            PolicyConfig::Random => Arc::new(RandomPolicy::new()),
+            PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
+            PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
+            PolicyConfig::CacheAware {
+                cache_threshold,
+                balance_abs_threshold,
+                balance_rel_threshold,
+                eviction_interval_secs,
+                max_tree_size,
+            } => {
+                let config = CacheAwareConfig {
+                    cache_threshold: *cache_threshold,
+                    balance_abs_threshold: *balance_abs_threshold,
+                    balance_rel_threshold: *balance_rel_threshold,
+                    eviction_interval_secs: *eviction_interval_secs,
+                    max_tree_size: *max_tree_size,
+                };
+                Arc::new(CacheAwarePolicy::with_config(config))
+            }
+            PolicyConfig::Bucket {
+                balance_abs_threshold,
+                balance_rel_threshold,
+                bucket_adjust_interval_secs,
+            } => {
+                let config = BucketConfig {
+                    balance_abs_threshold: *balance_abs_threshold,
+                    balance_rel_threshold: *balance_rel_threshold,
+                    bucket_adjust_interval_secs: *bucket_adjust_interval_secs,
+                };
+                Arc::new(BucketPolicy::with_config(config))
+            }
+            PolicyConfig::Manual {
+                eviction_interval_secs,
+                max_idle_secs,
+                assignment_mode,
+            } => {
+                let config = ManualConfig {
+                    eviction_interval_secs: *eviction_interval_secs,
+                    max_idle_secs: *max_idle_secs,
+                    assignment_mode: *assignment_mode,
+                };
+                Arc::new(ManualPolicy::with_config(config))
+            }
+            PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()),
+            PolicyConfig::PrefixHash {
+                prefix_token_count,
+                load_factor,
+            } => {
+                let config = PrefixHashConfig {
+                    prefix_token_count: *prefix_token_count,
+                    load_factor: *load_factor,
+                };
+                Arc::new(PrefixHashPolicy::new(config))
+            }
+        }
+    }
+
+    /// Create a policy by name (for dynamic loading)
+    pub fn create_by_name(name: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
+        match name.to_lowercase().as_str() {
+            "random" => Some(Arc::new(RandomPolicy::new())),
+            "round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())),
+            "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())),
+            "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
+            "bucket" => Some(Arc::new(BucketPolicy::new())),
+            "manual" => Some(Arc::new(ManualPolicy::new())),
+            "consistent_hashing" | "consistenthashing" => {
+                Some(Arc::new(ConsistentHashingPolicy::new()))
+            }
+            "prefix_hash" | "prefixhash" => Some(Arc::new(PrefixHashPolicy::with_defaults())),
+            _ => None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_create_from_config() {
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::Random);
+        assert_eq!(policy.name(), "random");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::RoundRobin);
+        assert_eq!(policy.name(), "round_robin");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::PowerOfTwo {
+            load_check_interval_secs: 60,
+        });
+        assert_eq!(policy.name(), "power_of_two");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAware {
+            cache_threshold: 0.7,
+            balance_abs_threshold: 10,
+            balance_rel_threshold: 1.5,
+            eviction_interval_secs: 30,
+            max_tree_size: 1000,
+        });
+        assert_eq!(policy.name(), "cache_aware");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::Bucket {
+            balance_abs_threshold: 10,
+            balance_rel_threshold: 1.5,
+            bucket_adjust_interval_secs: 5,
+        });
+        assert_eq!(policy.name(), "bucket");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::Manual {
+            eviction_interval_secs: 60,
+            max_idle_secs: 4 * 3600,
+            assignment_mode: Default::default(),
+        });
+        assert_eq!(policy.name(), "manual");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::ConsistentHashing);
+        assert_eq!(policy.name(), "consistent_hashing");
+    }
+
+    #[tokio::test]
+    async fn test_create_by_name() {
+        assert!(PolicyFactory::create_by_name("random").is_some());
+        assert!(PolicyFactory::create_by_name("RANDOM").is_some());
+        assert!(PolicyFactory::create_by_name("round_robin").is_some());
+        assert!(PolicyFactory::create_by_name("RoundRobin").is_some());
+        assert!(PolicyFactory::create_by_name("power_of_two").is_some());
+        assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some());
+        assert!(PolicyFactory::create_by_name("cache_aware").is_some());
+        assert!(PolicyFactory::create_by_name("CacheAware").is_some());
+        assert!(PolicyFactory::create_by_name("bucket").is_some());
+        assert!(PolicyFactory::create_by_name("Bucket").is_some());
+        assert!(PolicyFactory::create_by_name("manual").is_some());
+        assert!(PolicyFactory::create_by_name("Manual").is_some());
+        assert!(PolicyFactory::create_by_name("consistent_hashing").is_some());
+        assert!(PolicyFactory::create_by_name("ConsistentHashing").is_some());
+        assert!(PolicyFactory::create_by_name("unknown").is_none());
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/manual.rs

@@ -0,0 +1,981 @@
+//! Manual routing policy based on routing key header
+//!
+//! This policy provides sticky session routing where each unique routing key
+//! is consistently mapped to the same worker. Unlike consistent hashing,
+//! this policy:
+//! - Does NOT redistribute any sessions when workers are added
+//! - Only remaps sessions when their assigned worker becomes unhealthy
+//! - Maintains up to 2 candidate workers per routing key for fast failover
+//!
+//! Use this when you need stronger stickiness guarantees than consistent hashing,
+//! for example with stateful chat sessions where context is stored on the worker.
+//!
+//! ## Header
+//! - `X-SMG-Routing-Key`: The routing key for sticky session routing
+
+use std::{sync::Arc, time::Instant};
+
+use async_trait::async_trait;
+use dashmap::{mapref::entry::Entry, DashMap};
+use rand::Rng;
+use tracing::info;
+
+use super::{
+    get_healthy_worker_indices, utils::PeriodicTask, LoadBalancingPolicy, SelectWorkerInfo,
+};
+use crate::{
+    config::ManualAssignmentMode, core::Worker, observability::metrics::Metrics,
+    routers::header_utils::extract_routing_key,
+};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum ExecutionBranch {
+    NoHealthyWorkers,
+    OccupiedHit,
+    OccupiedMiss,
+    Vacant,
+    NoRoutingId,
+}
+
+impl ExecutionBranch {
+    fn as_str(&self) -> &'static str {
+        match self {
+            Self::NoHealthyWorkers => "no_healthy_workers",
+            Self::OccupiedHit => "occupied_hit",
+            Self::OccupiedMiss => "occupied_miss",
+            Self::Vacant => "vacant",
+            Self::NoRoutingId => "no_routing_id",
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct RoutingId(String);
+
+impl RoutingId {
+    fn new(id: impl Into<String>) -> Self {
+        Self(id.into())
+    }
+}
+
+const MAX_CANDIDATE_WORKERS: usize = 2;
+
+#[derive(Debug, Clone)]
+pub struct ManualConfig {
+    pub eviction_interval_secs: u64,
+    pub max_idle_secs: u64,
+    pub assignment_mode: ManualAssignmentMode,
+}
+
+impl Default for ManualConfig {
+    fn default() -> Self {
+        Self {
+            eviction_interval_secs: 60,
+            max_idle_secs: 4 * 3600,
+            assignment_mode: ManualAssignmentMode::Random,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Node {
+    candi_worker_urls: Vec<String>,
+    last_access: Instant,
+}
+
+impl Node {
+    fn push_bounded(&mut self, url: String) {
+        while self.candi_worker_urls.len() >= MAX_CANDIDATE_WORKERS {
+            self.candi_worker_urls.remove(0);
+        }
+        self.candi_worker_urls.push(url);
+    }
+}
+
+#[derive(Debug)]
+pub struct ManualPolicy {
+    routing_map: Arc<DashMap<RoutingId, Node>>,
+    assignment_mode: ManualAssignmentMode,
+    _eviction_task: Option<PeriodicTask>,
+}
+
+impl Default for ManualPolicy {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ManualPolicy {
+    pub fn new() -> Self {
+        Self::with_config(ManualConfig::default())
+    }
+
+    pub fn with_config(config: ManualConfig) -> Self {
+        use std::time::Duration;
+
+        let routing_map = Arc::new(DashMap::<RoutingId, Node>::new());
+
+        let eviction_task = if config.eviction_interval_secs > 0 && config.max_idle_secs > 0 {
+            let routing_map_clone = Arc::clone(&routing_map);
+            let max_idle = Duration::from_secs(config.max_idle_secs);
+
+            Some(PeriodicTask::spawn(
+                config.eviction_interval_secs,
+                "ManualPolicyEviction",
+                move || {
+                    let now = Instant::now();
+                    let before_size = routing_map_clone.len();
+
+                    routing_map_clone
+                        .retain(|_, node| now.duration_since(node.last_access) < max_idle);
+
+                    let evicted_count = before_size - routing_map_clone.len();
+                    if evicted_count > 0 {
+                        info!(
+                            "ManualPolicy TTL eviction: evicted {} entries, remaining {} (max_idle: {}s)",
+                            evicted_count,
+                            routing_map_clone.len(),
+                            max_idle.as_secs()
+                        );
+                    }
+                },
+            ))
+        } else {
+            None
+        };
+
+        Self {
+            routing_map,
+            assignment_mode: config.assignment_mode,
+            _eviction_task: eviction_task,
+        }
+    }
+
+    fn select_new_worker(&self, workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> usize {
+        match self.assignment_mode {
+            ManualAssignmentMode::Random => random_select(healthy_indices),
+            ManualAssignmentMode::MinLoad => min_load_select(workers, healthy_indices),
+            ManualAssignmentMode::MinGroup => min_group_select(workers, healthy_indices),
+        }
+    }
+
+    fn select_by_routing_id(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        routing_id: &str,
+        healthy_indices: &[usize],
+    ) -> (usize, ExecutionBranch) {
+        let routing_id = RoutingId::new(routing_id);
+
+        match self.routing_map.entry(routing_id) {
+            Entry::Occupied(mut entry) => {
+                let node = entry.get_mut();
+                node.last_access = Instant::now();
+                if let Some(idx) =
+                    find_healthy_worker(&node.candi_worker_urls, workers, healthy_indices)
+                {
+                    (idx, ExecutionBranch::OccupiedHit)
+                } else {
+                    let selected_idx = self.select_new_worker(workers, healthy_indices);
+                    node.push_bounded(workers[selected_idx].url().to_string());
+                    (selected_idx, ExecutionBranch::OccupiedMiss)
+                }
+            }
+            Entry::Vacant(entry) => {
+                let selected_idx = self.select_new_worker(workers, healthy_indices);
+                entry.insert(Node {
+                    candi_worker_urls: vec![workers[selected_idx].url().to_string()],
+                    last_access: Instant::now(),
+                });
+                (selected_idx, ExecutionBranch::Vacant)
+            }
+        }
+    }
+
+    fn select_worker_impl(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> (Option<usize>, ExecutionBranch) {
+        let healthy_indices = get_healthy_worker_indices(workers);
+        if healthy_indices.is_empty() {
+            return (None, ExecutionBranch::NoHealthyWorkers);
+        }
+
+        if let Some(routing_id) = extract_routing_key(info.headers) {
+            let (idx, branch) = self.select_by_routing_id(workers, routing_id, &healthy_indices);
+            return (Some(idx), branch);
+        }
+
+        (
+            Some(random_select(&healthy_indices)),
+            ExecutionBranch::NoRoutingId,
+        )
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for ManualPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let (result, branch) = self.select_worker_impl(workers, info);
+        Metrics::record_worker_manual_policy_branch(branch.as_str());
+        Metrics::set_manual_policy_cache_entries(self.routing_map.len());
+        result
+    }
+
+    fn name(&self) -> &'static str {
+        "manual"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+fn find_healthy_worker(
+    urls: &[String],
+    workers: &[Arc<dyn Worker>],
+    healthy_indices: &[usize],
+) -> Option<usize> {
+    for url in urls {
+        if let Some(idx) = find_worker_index_by_url(workers, url) {
+            if healthy_indices.contains(&idx) {
+                return Some(idx);
+            }
+        }
+    }
+    None
+}
+
+fn find_worker_index_by_url(workers: &[Arc<dyn Worker>], url: &str) -> Option<usize> {
+    workers.iter().position(|w| w.url() == url)
+}
+
+fn random_select(healthy_indices: &[usize]) -> usize {
+    let mut rng = rand::rng();
+    let random_idx = rng.random_range(0..healthy_indices.len());
+    healthy_indices[random_idx]
+}
+
+fn select_min_by<K, V, F>(indices: &[K], get_value: F) -> K
+where
+    K: Copy,
+    V: Ord,
+    F: Fn(K) -> V,
+{
+    let mut min_val: Option<V> = None;
+    let mut candidates = Vec::new();
+
+    for &idx in indices {
+        let val = get_value(idx);
+        match min_val.as_ref().map(|m| val.cmp(m)) {
+            None | Some(std::cmp::Ordering::Less) => {
+                min_val = Some(val);
+                candidates.clear();
+                candidates.push(idx);
+            }
+            Some(std::cmp::Ordering::Equal) => {
+                candidates.push(idx);
+            }
+            Some(std::cmp::Ordering::Greater) => {}
+        }
+    }
+
+    if candidates.len() == 1 {
+        candidates[0]
+    } else {
+        let mut rng = rand::rng();
+        candidates[rng.random_range(0..candidates.len())]
+    }
+}
+
+fn min_load_select(workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> usize {
+    select_min_by(healthy_indices, |idx| workers[idx].load())
+}
+
+fn min_group_select(workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> usize {
+    select_min_by(healthy_indices, |idx| {
+        workers[idx].worker_routing_key_load().value()
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .collect()
+    }
+
+    fn headers_with_routing_key(key: &str) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-routing-key", key.parse().unwrap());
+        headers
+    }
+
+    #[test]
+    fn test_manual_consistent_routing() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let headers = headers_with_routing_key("user-123");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(
+                result,
+                Some(first_idx),
+                "Same routing_id should route to same worker"
+            );
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_different_routing_ids() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let mut distribution = HashMap::new();
+        for i in 0..100 {
+            let headers = headers_with_routing_key(&format!("user-{}", i));
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::Vacant);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        assert!(
+            distribution.len() > 1,
+            "Should distribute across multiple workers"
+        );
+    }
+
+    #[test]
+    fn test_manual_fallback_random() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let mut counts = HashMap::new();
+        for _ in 0..100 {
+            let info = SelectWorkerInfo::default();
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::NoRoutingId);
+            if let Some(idx) = result {
+                *counts.entry(idx).or_insert(0) += 1;
+            }
+        }
+
+        assert_eq!(counts.len(), 2, "Random fallback should use all workers");
+    }
+
+    #[test]
+    fn test_manual_with_unhealthy_workers() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        workers[0].set_healthy(false);
+
+        let headers = headers_with_routing_key("test-routing-id");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(1), "Should only select healthy worker");
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(1), "Should only select healthy worker");
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_no_healthy_workers() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000"]);
+
+        workers[0].set_healthy(false);
+        let headers = headers_with_routing_key("test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, ExecutionBranch::NoHealthyWorkers);
+    }
+
+    #[test]
+    fn test_manual_empty_routing_id() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let mut counts = HashMap::new();
+        for _ in 0..100 {
+            let headers = headers_with_routing_key("");
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::NoRoutingId);
+            if let Some(idx) = result {
+                *counts.entry(idx).or_insert(0) += 1;
+            }
+        }
+
+        assert_eq!(
+            counts.len(),
+            2,
+            "Empty routing_id should use random fallback"
+        );
+    }
+
+    #[test]
+    fn test_manual_remaps_when_worker_becomes_unhealthy() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("sticky-user");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        workers[first_idx].set_healthy(false);
+
+        let (new_result, branch) = policy.select_worker_impl(&workers, &info);
+        let new_idx = new_result.unwrap();
+        assert_ne!(new_idx, first_idx, "Should remap to healthy worker");
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(
+                result,
+                Some(new_idx),
+                "Should consistently route to new worker"
+            );
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_empty_workers() {
+        let policy = ManualPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![];
+        let headers = headers_with_routing_key("test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, ExecutionBranch::NoHealthyWorkers);
+    }
+
+    #[test]
+    fn test_manual_single_worker() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000"]);
+
+        let headers = headers_with_routing_key("single-test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(0));
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(0));
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_worker_recovery() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("recovery-test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        workers[first_idx].set_healthy(false);
+
+        let (second_result, branch) = policy.select_worker_impl(&workers, &info);
+        let second_idx = second_result.unwrap();
+        assert_ne!(second_idx, first_idx);
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        workers[first_idx].set_healthy(true);
+
+        let (after_recovery, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(
+            after_recovery,
+            Some(first_idx),
+            "Should return to original worker after recovery since it's first in candidate list"
+        );
+        assert_eq!(branch, ExecutionBranch::OccupiedHit);
+    }
+
+    #[test]
+    fn test_manual_max_candidate_workers_eviction() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let headers = headers_with_routing_key("eviction-test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        workers[first_idx].set_healthy(false);
+
+        let (second_result, branch) = policy.select_worker_impl(&workers, &info);
+        let second_idx = second_result.unwrap();
+        assert_ne!(second_idx, first_idx);
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        workers[second_idx].set_healthy(false);
+
+        let remaining_idx = (0..3).find(|&i| i != first_idx && i != second_idx).unwrap();
+        let (third_result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(
+            third_result,
+            Some(remaining_idx),
+            "Should select the only remaining healthy worker"
+        );
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        workers[first_idx].set_healthy(true);
+
+        let (idx_after_restore, branch) = policy.select_worker_impl(&workers, &info);
+        assert_ne!(
+            idx_after_restore,
+            Some(first_idx),
+            "First worker should be evicted from candidates due to MAX_CANDIDATE_WORKERS=2"
+        );
+        assert_eq!(branch, ExecutionBranch::OccupiedHit);
+    }
+
+    #[test]
+    fn test_manual_policy_name() {
+        let policy = ManualPolicy::new();
+        assert_eq!(policy.name(), "manual");
+    }
+
+    #[test]
+    fn test_manual_routing_info_push_bounded() {
+        let mut info = Node {
+            candi_worker_urls: vec!["http://w1:8000".to_string()],
+            last_access: Instant::now(),
+        };
+
+        info.push_bounded("http://w2:8000".to_string());
+        assert_eq!(info.candi_worker_urls.len(), 2);
+        assert_eq!(info.candi_worker_urls[0], "http://w1:8000");
+        assert_eq!(info.candi_worker_urls[1], "http://w2:8000");
+
+        info.push_bounded("http://w3:8000".to_string());
+        assert_eq!(info.candi_worker_urls.len(), 2);
+        assert_eq!(
+            info.candi_worker_urls[0], "http://w2:8000",
+            "Oldest entry should be removed"
+        );
+        assert_eq!(info.candi_worker_urls[1], "http://w3:8000");
+    }
+
+    #[test]
+    fn test_manual_find_healthy_worker_priority() {
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let urls = vec![
+            "http://w1:8000".to_string(),
+            "http://w2:8000".to_string(),
+            "http://w3:8000".to_string(),
+        ];
+        let healthy_indices = vec![0, 1, 2];
+
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(
+            result,
+            Some(0),
+            "Should return first healthy worker in urls"
+        );
+
+        workers[0].set_healthy(false);
+        let healthy_indices = vec![1, 2];
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(result, Some(1), "Should skip unhealthy and return next");
+
+        workers[1].set_healthy(false);
+        let healthy_indices = vec![2];
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(result, Some(2), "Should return last healthy worker");
+
+        workers[2].set_healthy(false);
+        let healthy_indices: Vec<usize> = vec![];
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(result, None, "Should return None when no healthy workers");
+    }
+
+    #[test]
+    fn test_manual_find_worker_index_by_url() {
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        assert_eq!(
+            find_worker_index_by_url(&workers, "http://w1:8000"),
+            Some(0)
+        );
+        assert_eq!(
+            find_worker_index_by_url(&workers, "http://w2:8000"),
+            Some(1)
+        );
+        assert_eq!(
+            find_worker_index_by_url(&workers, "http://w3:8000"),
+            None,
+            "Should return None for unknown URL"
+        );
+    }
+
+    #[test]
+    fn test_manual_config_default() {
+        let config = ManualConfig::default();
+        assert_eq!(config.eviction_interval_secs, 60);
+        assert_eq!(config.max_idle_secs, 4 * 3600);
+    }
+
+    #[test]
+    fn test_manual_with_disabled_eviction() {
+        let config = ManualConfig {
+            eviction_interval_secs: 0,
+            max_idle_secs: 3600,
+            assignment_mode: ManualAssignmentMode::Random,
+        };
+        let policy = ManualPolicy::with_config(config);
+        assert!(policy._eviction_task.is_none());
+    }
+
+    #[test]
+    fn test_manual_last_access_updates() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+        let headers = headers_with_routing_key("test-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let routing_id = RoutingId::new("test-key");
+
+        // Vacant: first access
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(branch, ExecutionBranch::Vacant);
+        let first_idx = result.unwrap();
+        let access_after_vacant = policy.routing_map.get(&routing_id).unwrap().last_access;
+        assert!(access_after_vacant.elapsed().as_millis() < 100);
+
+        std::thread::sleep(std::time::Duration::from_millis(10));
+
+        // OccupiedHit: same worker still healthy
+        let (_, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        let access_after_hit = policy.routing_map.get(&routing_id).unwrap().last_access;
+        assert!(access_after_hit > access_after_vacant);
+
+        std::thread::sleep(std::time::Duration::from_millis(10));
+
+        // OccupiedMiss: worker becomes unhealthy
+        workers[first_idx].set_healthy(false);
+        let (_, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+        let access_after_miss = policy.routing_map.get(&routing_id).unwrap().last_access;
+        assert!(access_after_miss > access_after_hit);
+    }
+
+    #[test]
+    fn test_manual_ttl_eviction_logic() {
+        use std::time::Duration;
+
+        let config = ManualConfig {
+            eviction_interval_secs: 2,
+            max_idle_secs: 2,
+            assignment_mode: ManualAssignmentMode::Random,
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("key-0");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        policy.select_worker_impl(&workers, &info);
+
+        assert_eq!(policy.routing_map.len(), 1);
+
+        std::thread::sleep(Duration::from_secs(4));
+
+        assert_eq!(policy.routing_map.len(), 0);
+    }
+
+    #[test]
+    fn test_min_group_select_distributes_evenly() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinGroup,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        for i in 0..9 {
+            let routing_key = format!("key-{}", i);
+            let headers = headers_with_routing_key(&routing_key);
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert!(result.is_some());
+            assert_eq!(branch, ExecutionBranch::Vacant);
+
+            let selected_idx = result.unwrap();
+            workers[selected_idx]
+                .worker_routing_key_load()
+                .increment(&routing_key);
+        }
+
+        let distribution: HashMap<_, usize> = policy
+            .routing_map
+            .iter()
+            .map(|e| e.candi_worker_urls.first().unwrap().clone())
+            .fold(HashMap::new(), |mut acc, url| {
+                *acc.entry(url).or_default() += 1;
+                acc
+            });
+
+        assert_eq!(distribution.len(), 3, "Should use all 3 workers");
+        for count in distribution.values() {
+            assert_eq!(*count, 3, "Each worker should have exactly 3 routing keys");
+        }
+    }
+
+    #[test]
+    fn test_min_group_select_prefers_worker_with_fewer_routing_keys() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinGroup,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        workers[0].worker_routing_key_load().increment("existing-1");
+        workers[0].worker_routing_key_load().increment("existing-2");
+        workers[1].worker_routing_key_load().increment("existing-3");
+
+        assert_eq!(workers[0].worker_routing_key_load().value(), 2);
+        assert_eq!(workers[1].worker_routing_key_load().value(), 1);
+        assert_eq!(workers[2].worker_routing_key_load().value(), 0);
+
+        let headers = headers_with_routing_key("new-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, _) = policy.select_worker_impl(&workers, &info);
+        let selected_idx = result.unwrap();
+
+        assert_eq!(selected_idx, 2, "Should select worker with 0 routing keys");
+    }
+
+    #[test]
+    fn test_min_load_select_prefers_worker_with_fewer_requests() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinLoad,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        workers[0].increment_load();
+        workers[0].increment_load();
+        workers[1].increment_load();
+
+        assert_eq!(workers[0].load(), 2);
+        assert_eq!(workers[1].load(), 1);
+        assert_eq!(workers[2].load(), 0);
+
+        let headers = headers_with_routing_key("new-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, _) = policy.select_worker_impl(&workers, &info);
+        let selected_idx = result.unwrap();
+
+        assert_eq!(selected_idx, 2, "Should select worker with 0 load");
+    }
+
+    #[test]
+    fn test_min_group_sticky_after_assignment() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinGroup,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        workers[0].worker_routing_key_load().increment("key-0");
+        workers[1].worker_routing_key_load().increment("key-1");
+        workers[1].worker_routing_key_load().increment("key-2");
+
+        let headers = headers_with_routing_key("new-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+        assert_eq!(
+            first_idx, 0,
+            "Should select worker 0 (has 1 routing key vs 2)"
+        );
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(
+                result,
+                Some(first_idx),
+                "Same routing key should route to same worker"
+            );
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_random_mode_does_not_consider_load() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::Random,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        workers[0].worker_routing_key_load().increment("key-1");
+        workers[0].worker_routing_key_load().increment("key-2");
+        workers[0].worker_routing_key_load().increment("key-3");
+
+        let mut selected_worker_0 = false;
+        for i in 0..50 {
+            let headers = headers_with_routing_key(&format!("test-{}", i));
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            if result == Some(0) {
+                selected_worker_0 = true;
+                break;
+            }
+        }
+        assert!(
+            selected_worker_0,
+            "Random mode should sometimes select worker 0 despite higher load"
+        );
+    }
+
+    fn assert_no_routing_key_uses_random(
+        assignment_mode: ManualAssignmentMode,
+        setup_load: impl Fn(&[Arc<dyn Worker>]),
+    ) {
+        let config = ManualConfig {
+            assignment_mode,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+        setup_load(&workers);
+
+        let mut selected_worker_0 = false;
+        for _ in 0..50 {
+            let info = SelectWorkerInfo::default();
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::NoRoutingId);
+            if result == Some(0) {
+                selected_worker_0 = true;
+                break;
+            }
+        }
+        assert!(
+            selected_worker_0,
+            "Should randomly select worker 0 despite higher load"
+        );
+    }
+
+    #[test]
+    fn test_no_routing_key_uses_random_even_with_min_load_mode() {
+        assert_no_routing_key_uses_random(ManualAssignmentMode::MinLoad, |workers| {
+            workers[0].increment_load();
+            workers[0].increment_load();
+        });
+    }
+
+    #[test]
+    fn test_no_routing_key_uses_random_even_with_min_group_mode() {
+        assert_no_routing_key_uses_random(ManualAssignmentMode::MinGroup, |workers| {
+            workers[0].worker_routing_key_load().increment("k1");
+            workers[0].worker_routing_key_load().increment("k2");
+        });
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/mod.rs

@@ -0,0 +1,213 @@
+//! Load balancing policies for SGLang router
+//!
+//! This module provides a unified abstraction for routing policies that work
+//! across both regular and prefill-decode (PD) routing modes.
+
+use std::{fmt::Debug, sync::Arc};
+
+use async_trait::async_trait;
+use smg_mesh::OptionalMeshSyncManager;
+
+use crate::core::{HashRing, Worker};
+
+mod bucket;
+mod cache_aware;
+mod consistent_hashing;
+mod factory;
+mod manual;
+mod power_of_two;
+mod prefix_hash;
+mod random;
+mod registry;
+mod round_robin;
+pub mod tree;
+pub(crate) mod utils;
+pub use bucket::BucketPolicy;
+pub use cache_aware::CacheAwarePolicy;
+pub use consistent_hashing::ConsistentHashingPolicy;
+pub use factory::PolicyFactory;
+pub use manual::{ManualConfig, ManualPolicy};
+pub use power_of_two::PowerOfTwoPolicy;
+pub use prefix_hash::{PrefixHashConfig, PrefixHashPolicy};
+pub use random::RandomPolicy;
+pub use registry::PolicyRegistry;
+pub use round_robin::RoundRobinPolicy;
+pub use tree::PrefixMatchResult;
+
+/// Core trait for load balancing policies
+///
+/// This trait provides a unified interface for implementing routing algorithms
+/// that can work with both regular single-worker selection and PD dual-worker selection.
+#[async_trait]
+pub trait LoadBalancingPolicy: Send + Sync + Debug {
+    /// Select a single worker from the available workers
+    ///
+    /// This is used for regular routing mode where requests go to a single worker.
+    /// Now uses Arc<dyn Worker> for better performance and to avoid unnecessary cloning.
+    ///
+    /// # Arguments
+    /// * `workers` - Available workers to select from
+    /// * `info` - Additional information for routing decisions
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize>;
+
+    /// Update policy state after request completion
+    ///
+    /// This is called when a request completes (successfully or not) to allow
+    /// policies to update their internal state.
+    fn on_request_complete(&self, _worker_url: &str, _success: bool) {
+        // Default: no-op for stateless policies
+    }
+
+    /// Get policy name for metrics and debugging
+    fn name(&self) -> &'static str;
+
+    /// Check if this policy needs request text for routing decisions
+    fn needs_request_text(&self) -> bool {
+        false // Default: most policies don't need request text
+    }
+
+    /// Update worker load information
+    ///
+    /// This is called periodically with current load information for load-aware policies.
+    fn update_loads(&self, _loads: &std::collections::HashMap<String, isize>) {
+        // Default: no-op for policies that don't use load information
+    }
+
+    /// Set mesh sync manager
+    fn set_mesh_sync(&mut self, _mesh_sync: OptionalMeshSyncManager) {
+        // Default: no-op for policies that don't use mesh sync
+    }
+
+    /// Reset any internal state
+    ///
+    /// This is useful for policies that maintain state (e.g., round-robin counters).
+    fn reset(&self) {
+        // Default: no-op for stateless policies
+    }
+
+    /// Get as Any for downcasting
+    fn as_any(&self) -> &dyn std::any::Any;
+}
+
+/// Configuration for cache-aware policy
+#[derive(Debug, Clone)]
+pub struct CacheAwareConfig {
+    pub cache_threshold: f32,
+    pub balance_abs_threshold: usize,
+    pub balance_rel_threshold: f32,
+    pub eviction_interval_secs: u64,
+    pub max_tree_size: usize,
+}
+
+impl Default for CacheAwareConfig {
+    fn default() -> Self {
+        Self {
+            cache_threshold: 0.5,
+            balance_abs_threshold: 32,
+            balance_rel_threshold: 1.1,
+            eviction_interval_secs: 30,
+            max_tree_size: 10000,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct BucketConfig {
+    pub balance_abs_threshold: usize,
+    pub balance_rel_threshold: f32,
+    pub bucket_adjust_interval_secs: usize,
+}
+
+impl Default for BucketConfig {
+    fn default() -> Self {
+        Self {
+            balance_abs_threshold: 32,
+            balance_rel_threshold: 1.0001,
+            bucket_adjust_interval_secs: 5,
+        }
+    }
+}
+
+/// Helper function to filter healthy workers and return their indices
+pub(crate) fn get_healthy_worker_indices(workers: &[Arc<dyn Worker>]) -> Vec<usize> {
+    workers
+        .iter()
+        .enumerate()
+        .filter(|(_, w)| w.is_healthy() && w.circuit_breaker().can_execute())
+        .map(|(idx, _)| idx)
+        .collect()
+}
+
+/// Helper function to normalize model_id to a key for policy lookups.
+///
+/// Returns UNKNOWN_MODEL_ID for empty model_ids to ensure consistent behavior
+/// across single-model and multi-model deployments.
+#[inline]
+pub(crate) fn normalize_model_key(model_id: &str) -> &str {
+    if model_id.is_empty() {
+        crate::core::UNKNOWN_MODEL_ID
+    } else {
+        model_id
+    }
+}
+
+/// Information passed to policy for worker selection
+#[derive(Debug, Clone, Default)]
+pub struct SelectWorkerInfo<'a> {
+    /// Request text for cache-aware routing
+    pub request_text: Option<&'a str>,
+    /// Tokenized request for prefix-hash routing
+    /// Used by PrefixHashPolicy for token-based prefix hashing
+    pub tokens: Option<&'a [u32]>,
+    /// HTTP headers for header-based routing policies
+    /// Policies can extract routing information from headers like:
+    /// - X-SMG-Target-Worker: Direct routing to a specific worker by index
+    /// - X-SMG-Routing-Key: Consistent hash routing for session affinity
+    pub headers: Option<&'a http::HeaderMap>,
+    /// Pre-computed hash ring for O(log n) consistent hashing
+    /// Built and cached by WorkerRegistry, passed through to avoid per-request rebuilds
+    pub hash_ring: Option<Arc<HashRing>>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_get_healthy_worker_indices() {
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key2")
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+        ];
+
+        // All healthy initially
+        let indices = get_healthy_worker_indices(&workers);
+        assert_eq!(indices, vec![0, 1, 2]);
+
+        // Mark one unhealthy
+        workers[1].set_healthy(false);
+        let indices = get_healthy_worker_indices(&workers);
+        assert_eq!(indices, vec![0, 2]);
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/power_of_two.rs

@@ -0,0 +1,389 @@
+//! Power-of-two choices load balancing policy
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
+
+use async_trait::async_trait;
+use rand::Rng;
+use tracing::debug;
+
+use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
+use crate::core::Worker;
+
+/// Power-of-two choices policy
+///
+/// Randomly selects two workers and routes to the one with lower load.
+/// This provides good load distribution with minimal coordination overhead.
+#[derive(Debug)]
+pub struct PowerOfTwoPolicy {
+    /// Cached load information from external monitoring
+    cached_loads: RwLock<HashMap<String, isize>>,
+}
+
+impl PowerOfTwoPolicy {
+    pub fn new() -> Self {
+        Self {
+            cached_loads: RwLock::new(HashMap::new()),
+        }
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for PowerOfTwoPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        _info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        if healthy_indices.len() == 1 {
+            return Some(healthy_indices[0]);
+        }
+
+        // Select two random workers - use offset to guarantee different selection in O(1)
+        let mut rng = rand::rng();
+        let idx1 = rng.random_range(0..healthy_indices.len());
+        // Pick idx2 from remaining indices: offset by 1 + random from (len-1) to guarantee different
+        let idx2 =
+            (idx1 + 1 + rng.random_range(0..healthy_indices.len() - 1)) % healthy_indices.len();
+
+        let worker_idx1 = healthy_indices[idx1];
+        let worker_idx2 = healthy_indices[idx2];
+        let worker1 = &workers[worker_idx1];
+        let worker2 = &workers[worker_idx2];
+
+        // Access cached loads safely
+        let loads_guard = self.cached_loads.read().ok();
+
+        // Try to get high-fidelity token loads for BOTH workers
+        let load1_tokens = loads_guard
+            .as_ref()
+            .and_then(|m| m.get(worker1.url()).copied());
+        let load2_tokens = loads_guard
+            .as_ref()
+            .and_then(|m| m.get(worker2.url()).copied());
+
+        // If either worker is missing token data (e.g. monitor failure),
+        // we must degrade BOTH to request counts to ensure fairness.
+        let (load1, load2) = match (load1_tokens, load2_tokens) {
+            (Some(t1), Some(t2)) => {
+                // Both have token data. Compare Tokens.
+                (t1, t2)
+            }
+            _ => {
+                // If One or both are missing token data.
+                // Fallback to local request counts for BOTH.
+                (worker1.load() as isize, worker2.load() as isize)
+            }
+        };
+
+        // Select worker with lower load
+        let selected_idx = if load1 <= load2 {
+            worker_idx1
+        } else {
+            worker_idx2
+        };
+
+        debug!(
+            "Power-of-two selection: {}={} vs {}={} -> selected {}",
+            worker1.url(),
+            load1,
+            worker2.url(),
+            load2,
+            workers[selected_idx].url()
+        );
+
+        // Increment processed counter
+        workers[selected_idx].increment_processed();
+
+        Some(selected_idx)
+    }
+
+    fn name(&self) -> &'static str {
+        "power_of_two"
+    }
+
+    fn update_loads(&self, loads: &HashMap<String, isize>) {
+        if let Ok(mut cached) = self.cached_loads.write() {
+            *cached = loads.clone();
+        }
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+impl Default for PowerOfTwoPolicy {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_power_of_two_selection() {
+        let policy = PowerOfTwoPolicy::new();
+        let worker1 = BasicWorkerBuilder::new("http://w1:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+        let worker2 = BasicWorkerBuilder::new("http://w2:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+        let worker3 = BasicWorkerBuilder::new("http://w3:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // Set different loads
+        for _ in 0..10 {
+            worker1.increment_load();
+        }
+        for _ in 0..5 {
+            worker2.increment_load();
+        }
+        // worker3 has load 0
+
+        let workers: Vec<Arc<dyn Worker>> =
+            vec![Arc::new(worker1), Arc::new(worker2), Arc::new(worker3)];
+
+        // Run multiple selections
+        let mut selected_counts = [0; 3];
+        let info = SelectWorkerInfo::default();
+        for _ in 0..100 {
+            if let Some(idx) = policy.select_worker(&workers, &info).await {
+                selected_counts[idx] += 1;
+            }
+        }
+
+        // Worker with lowest load (worker3) should be selected most often
+        assert!(selected_counts[2] > selected_counts[1]);
+        assert!(selected_counts[1] > selected_counts[0]);
+    }
+
+    #[tokio::test]
+    async fn test_power_of_two_with_cached_loads() {
+        let policy = PowerOfTwoPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        // Update cached loads
+        let mut loads = HashMap::new();
+        loads.insert("http://w1:8000".to_string(), 100);
+        loads.insert("http://w2:8000".to_string(), 10);
+        policy.update_loads(&loads);
+
+        // Should prefer worker2 with lower cached load
+        let mut w2_selected = 0;
+        let info = SelectWorkerInfo::default();
+        for _ in 0..50 {
+            if let Some(idx) = policy.select_worker(&workers, &info).await {
+                if idx == 1 {
+                    w2_selected += 1;
+                }
+            }
+        }
+
+        // Worker2 should be selected significantly more often
+        assert!(w2_selected > 35); // Should win most of the time
+    }
+
+    #[tokio::test]
+    async fn test_power_of_two_single_worker() {
+        let policy = PowerOfTwoPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        )];
+
+        // With single worker, should always select it
+        assert_eq!(
+            policy
+                .select_worker(&workers, &SelectWorkerInfo::default())
+                .await,
+            Some(0)
+        );
+    }
+
+    #[tokio::test]
+    async fn test_reproduce_incompatible_metric_bug() {
+        use std::{collections::HashMap, sync::Arc};
+
+        use crate::core::{BasicWorkerBuilder, WorkerType};
+
+        // 1. Setup the policy
+        let policy = PowerOfTwoPolicy::new();
+
+        // 2. Create Worker A: Idle (0 reqs), but has high token usage in cache
+        let worker_a = BasicWorkerBuilder::new("http://worker_a:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // 3. Create Worker B: Busy (5 reqs), but missing from cache
+        let worker_b = BasicWorkerBuilder::new("http://worker_b:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // Manually increment load on Worker B to simulate active requests
+        for _ in 0..5 {
+            worker_b.increment_load();
+        }
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker_a), Arc::new(worker_b)];
+
+        // 4. Simulate LoadMonitor update:
+        // Only Worker A gets a token report. Worker B is missing (e.g. monitor failure).
+        let mut loads = HashMap::new();
+        loads.insert("http://worker_a:8000".to_string(), 50_000); // 50k tokens load
+        policy.update_loads(&loads);
+
+        // 5. Run selection
+        let selected_idx = policy
+            .select_worker(&workers, &SelectWorkerInfo::default())
+            .await
+            .expect("Should select a worker");
+
+        // 6. Verify the Fix
+        // Logic:
+        // - Worker A has token load (50k) but Worker B has NO token load.
+        // - Policy should fallback to request counts for BOTH.
+        // - A has 0 requests, B has 5 requests.
+        // - 0 <= 5, so A should be selected.
+
+        if selected_idx == 0 {
+            println!("Bug Fixed: System correctly fell back to request counts and selected idle Worker A.");
+        } else {
+            println!(
+                "Bug PERSISTS: Selected Worker B (Load: 5 reqs) over Worker A (Load: 50k tokens)"
+            );
+        }
+
+        // Assert that the CORRECT worker (A, index 0) is selected
+        assert_eq!(
+            selected_idx, 0,
+            "The policy failed to handle incompatible metrics. Should select idle Worker A."
+        );
+    }
+    #[tokio::test]
+    async fn test_power_of_two_edge_cases() {
+        use std::{collections::HashMap, sync::Arc};
+
+        use crate::core::{BasicWorkerBuilder, WorkerType};
+
+        let policy = PowerOfTwoPolicy::new();
+
+        // Helper to create a worker with specific request load
+        let create_worker = |url: &str, reqs: usize| {
+            let w = BasicWorkerBuilder::new(url)
+                .worker_type(WorkerType::Regular)
+                .build();
+            for _ in 0..reqs {
+                w.increment_load();
+            }
+            Arc::new(w)
+        };
+
+        //  Scenario 1: Happy Path (Both have Token Data)
+        // Worker A: 10 requests, but only 1,000 tokens (Light usage) -> Should be CHOSEN
+        // Worker B:  2 requests, but 100,000 tokens (Heavy usage) -> Should be AVOIDED
+        // This proves we use high-fidelity metrics when available, ignoring request counts.
+        let w_a = create_worker("http://a:8000", 10);
+        let w_b = create_worker("http://b:8000", 2);
+        let workers_1: Vec<Arc<dyn Worker>> = vec![w_a.clone(), w_b.clone()];
+
+        let mut loads_1 = HashMap::new();
+        loads_1.insert("http://a:8000".to_string(), 1_000);
+        loads_1.insert("http://b:8000".to_string(), 100_000);
+        policy.update_loads(&loads_1);
+
+        let idx_1 = policy
+            .select_worker(&workers_1, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(
+            idx_1, 0,
+            "Happy Path Failed: Should select Worker A (fewer tokens) despite higher request count"
+        );
+
+        // Scenario 2: Partial Failure (Worker A has tokens, Worker B is missing)
+        // Worker A: 10 requests, 1,000 tokens (Cached)
+        // Worker B:  2 requests, MISSING cache
+        // Logic: Fallback to requests -> Compare 10 (A) vs 2 (B) -> Select B
+        let w_c = create_worker("http://c:8000", 10);
+        let w_d = create_worker("http://d:8000", 2);
+        let workers_2: Vec<Arc<dyn Worker>> = vec![w_c.clone(), w_d.clone()];
+
+        let mut loads_2 = HashMap::new();
+        loads_2.insert("http://c:8000".to_string(), 1_000);
+        // http://d:8000 is MISSING
+        policy.update_loads(&loads_2);
+
+        let idx_2 = policy
+            .select_worker(&workers_2, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(idx_2, 1, "Partial Fail 1 Failed: Should fallback to requests and select Worker B (fewer requests)");
+
+        // Scenario 3: Partial Failure (Worker A is missing, Worker B has tokens)
+        // Worker A:  2 requests, MISSING cache
+        // Worker B: 10 requests, 1,000 tokens (Cached)
+        // Logic: Fallback to requests -> Compare 2 (A) vs 10 (B) -> Select A
+        let w_e = create_worker("http://e:8000", 2);
+        let w_f = create_worker("http://f:8000", 10);
+        let workers_3: Vec<Arc<dyn Worker>> = vec![w_e.clone(), w_f.clone()];
+
+        let mut loads_3 = HashMap::new();
+        // http://e:8000 is MISSING
+        loads_3.insert("http://f:8000".to_string(), 1_000);
+        policy.update_loads(&loads_3);
+
+        let idx_3 = policy
+            .select_worker(&workers_3, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(idx_3, 0, "Partial Fail 2 Failed: Should fallback to requests and select Worker A (fewer requests)");
+
+        // Scenario 4: Total Failure (Both missing)
+        // Worker A: 5 requests
+        // Worker B: 3 requests
+        // Logic: Requests vs Requests -> Select B
+        let w_g = create_worker("http://g:8000", 5);
+        let w_h = create_worker("http://h:8000", 3);
+        let workers_4: Vec<Arc<dyn Worker>> = vec![w_g.clone(), w_h.clone()];
+
+        let loads_4 = HashMap::new();
+        policy.update_loads(&loads_4);
+
+        let idx_4 = policy
+            .select_worker(&workers_4, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(
+            idx_4, 1,
+            "Total Fail Failed: Should select Worker B based on request count"
+        );
+
+        println!("All edge case tests passed successfully.");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/prefix_hash.rs

@@ -0,0 +1,414 @@
+//! Prefix Hash routing policy for KV cache-aware load balancing
+//!
+//! A lightweight alternative to the full radix tree cache_aware policy.
+//! Routes requests based on a hash of their prefix tokens to maximize
+//! KV cache hits across workers.
+//!
+//! ## Algorithm
+//!
+//! 1. Extract first N tokens from the request (configurable prefix length)
+//! 2. Hash the token sequence using xxhash for fast, stable hashing
+//! 3. Use consistent hash ring to find the target worker
+//! 4. If worker is overloaded (load > avg * load_factor), find least loaded
+//! 5. Return least loaded worker that passes load check, or initial if all overloaded
+//!
+//! ## Complexity
+//!
+//! - Hash computation: O(prefix_length)
+//! - Ring lookup: O(log n) binary search
+//! - Load balance fallback: O(n) scan for least loaded
+//!
+//! ## Comparison with cache_aware
+//!
+//! | Aspect          | prefix_hash       | cache_aware (radix) |
+//! |-----------------|-------------------|---------------------|
+//! | Lookup          | O(log n)          | O(prefix_len)       |
+//! | Memory          | O(workers × vn)   | O(total_tokens)     |
+//! | Update          | O(1)              | O(prefix_len)       |
+//! | Precision       | Prefix grouping   | Exact matching      |
+//!
+//! prefix_hash trades optimal cache utilization for predictable O(log n) performance.
+
+use std::sync::Arc;
+
+use super::{LoadBalancingPolicy, SelectWorkerInfo};
+use crate::{core::Worker, observability::metrics::Metrics};
+
+/// Configuration for the PrefixHash load balancing policy
+#[derive(Debug, Clone)]
+pub struct PrefixHashConfig {
+    /// Number of prefix tokens to use for hashing.
+    /// Longer prefixes = more precise routing but less grouping.
+    /// Shorter prefixes = more requests grouped together.
+    /// Default: 256 tokens (~1 paragraph of text)
+    pub prefix_token_count: usize,
+
+    /// Load factor threshold for walking the ring.
+    /// If a worker's load > (total_load / num_workers) * load_factor,
+    /// walk clockwise to the next worker.
+    /// Default: 1.25 (125% of average load)
+    pub load_factor: f64,
+}
+
+impl Default for PrefixHashConfig {
+    fn default() -> Self {
+        Self {
+            prefix_token_count: 256,
+            load_factor: 1.25,
+        }
+    }
+}
+
+/// Execution branch for metrics
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Branch {
+    NoHealthyWorkers,
+    NoTokens,
+    RingHit,
+    LoadBalanceWalk,
+    FallbackLeastLoad,
+}
+
+impl Branch {
+    #[inline]
+    const fn as_str(&self) -> &'static str {
+        match self {
+            Self::NoHealthyWorkers => "no_healthy_workers",
+            Self::NoTokens => "no_tokens",
+            Self::RingHit => "ring_hit",
+            Self::LoadBalanceWalk => "load_balance_walk",
+            Self::FallbackLeastLoad => "fallback_least_load",
+        }
+    }
+}
+
+/// Prefix Hash load balancing policy
+///
+/// Routes requests based on prefix token hash for KV cache locality.
+/// Uses consistent hashing with bounded load balancing.
+#[derive(Debug)]
+pub struct PrefixHashPolicy {
+    config: PrefixHashConfig,
+}
+
+impl PrefixHashPolicy {
+    /// Create a new PrefixHashPolicy with the given configuration
+    pub fn new(config: PrefixHashConfig) -> Self {
+        Self { config }
+    }
+
+    /// Create a new PrefixHashPolicy with default configuration
+    pub fn with_defaults() -> Self {
+        Self::new(PrefixHashConfig::default())
+    }
+
+    /// Compute hash of prefix tokens using xxhash
+    #[inline]
+    fn compute_prefix_hash(&self, tokens: &[u32]) -> u64 {
+        let prefix_len = tokens.len().min(self.config.prefix_token_count);
+        let prefix = &tokens[..prefix_len];
+
+        let bytes: &[u8] = bytemuck::cast_slice(prefix);
+        xxhash_rust::xxh3::xxh3_64(bytes)
+    }
+
+    /// Check if a worker's load is acceptable
+    #[inline]
+    fn load_ok(&self, worker_load: usize, total_load: usize, num_workers: usize) -> bool {
+        if total_load == 0 || num_workers == 0 {
+            return true;
+        }
+
+        // Average load per worker (with +1 to simulate incoming request)
+        let avg_load = (total_load + 1) as f64 / num_workers as f64;
+        let threshold = avg_load * self.config.load_factor;
+
+        (worker_load as f64) <= threshold
+    }
+
+    /// Find worker using consistent hash ring with load balancing
+    fn find_worker_with_load_balance(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+        prefix_hash: u64,
+    ) -> (Option<usize>, Branch) {
+        // Build healthy worker URL to index map
+        let healthy_workers: Vec<(usize, &Arc<dyn Worker>)> = workers
+            .iter()
+            .enumerate()
+            .filter(|(_, w)| w.is_healthy())
+            .collect();
+
+        if healthy_workers.is_empty() {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        // Calculate total load for load balancing
+        let total_load: usize = healthy_workers.iter().map(|(_, w)| w.load()).sum();
+        let num_workers = healthy_workers.len();
+
+        // Use pre-computed ring if available
+        if let Some(ref ring) = info.hash_ring {
+            // Convert prefix hash to a ring key string for lookup
+            let key = format!("{:016x}", prefix_hash);
+
+            // Build URL to (index, worker) map for healthy workers
+            let healthy_url_map: std::collections::HashMap<&str, (usize, &Arc<dyn Worker>)> =
+                healthy_workers
+                    .iter()
+                    .map(|(idx, w)| (w.url(), (*idx, *w)))
+                    .collect();
+
+            // Find initial worker from ring
+            if let Some(initial_url) =
+                ring.find_healthy_url(&key, |url| healthy_url_map.contains_key(url))
+            {
+                if let Some(&(idx, worker)) = healthy_url_map.get(initial_url) {
+                    let worker_load = worker.load();
+
+                    // Check if initial worker has acceptable load
+                    if self.load_ok(worker_load, total_load, num_workers) {
+                        return (Some(idx), Branch::RingHit);
+                    }
+
+                    // Initial worker overloaded, find least loaded healthy worker
+                    // This is a simpler approach than walking the ring
+                    let least_loaded = healthy_workers
+                        .iter()
+                        .filter(|(_, w)| self.load_ok(w.load(), total_load, num_workers))
+                        .min_by_key(|(_, w)| w.load());
+
+                    if let Some(&(idx, _)) = least_loaded {
+                        return (Some(idx), Branch::LoadBalanceWalk);
+                    }
+
+                    // All workers overloaded, use initial worker anyway
+                    return (Some(idx), Branch::LoadBalanceWalk);
+                }
+            }
+        }
+
+        // Fallback: no ring or ring lookup failed, use least loaded worker
+        let least_loaded = healthy_workers
+            .iter()
+            .min_by_key(|(_, w)| w.load())
+            .map(|(idx, _)| *idx);
+
+        (least_loaded, Branch::FallbackLeastLoad)
+    }
+
+    fn select_worker_impl(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+    ) -> (Option<usize>, Branch) {
+        if workers.is_empty() {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        // Get tokens from SelectWorkerInfo
+        let tokens = match info.tokens {
+            Some(t) if !t.is_empty() => t,
+            _ => return (None, Branch::NoTokens),
+        };
+
+        // Compute prefix hash
+        let prefix_hash = self.compute_prefix_hash(tokens);
+
+        // Find worker using ring with load balancing
+        self.find_worker_with_load_balance(workers, info, prefix_hash)
+    }
+}
+
+#[async_trait::async_trait]
+impl LoadBalancingPolicy for PrefixHashPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let (result, branch) = self.select_worker_impl(workers, info);
+        Metrics::record_worker_prefix_hash_policy_branch(branch.as_str());
+        result
+    }
+
+    fn name(&self) -> &'static str {
+        "prefix_hash"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, HashRing, WorkerType};
+
+    fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .collect()
+    }
+
+    #[test]
+    fn test_prefix_hash_consistent_routing() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Same tokens should always route to same worker
+        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        let info = SelectWorkerInfo {
+            tokens: Some(&tokens),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+
+        let (first_result, _) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+
+        // Verify consistency
+        for _ in 0..10 {
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(first_idx));
+        }
+    }
+
+    #[test]
+    fn test_different_prefixes_distribute() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        let mut distribution = std::collections::HashMap::new();
+
+        // Different token sequences should distribute across workers
+        for i in 0..100 {
+            let tokens: Vec<u32> = vec![i, i + 1, i + 2, i + 3];
+            let info = SelectWorkerInfo {
+                tokens: Some(&tokens),
+                hash_ring: Some(ring.clone()),
+                ..Default::default()
+            };
+
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        assert!(
+            distribution.len() > 1,
+            "Should distribute across workers, got {:?}",
+            distribution
+        );
+    }
+
+    #[test]
+    fn test_shared_prefix_routes_same() {
+        let policy = PrefixHashPolicy::new(PrefixHashConfig {
+            prefix_token_count: 5, // Only look at first 5 tokens
+            ..Default::default()
+        });
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Two sequences with same first 5 tokens should route to same worker
+        let tokens1: Vec<u32> = vec![1, 2, 3, 4, 5, 100, 200, 300];
+        let tokens2: Vec<u32> = vec![1, 2, 3, 4, 5, 999, 888, 777];
+
+        let info1 = SelectWorkerInfo {
+            tokens: Some(&tokens1),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+        let info2 = SelectWorkerInfo {
+            tokens: Some(&tokens2),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+
+        let (result1, _) = policy.select_worker_impl(&workers, &info1);
+        let (result2, _) = policy.select_worker_impl(&workers, &info2);
+
+        assert_eq!(result1, result2, "Same prefix should route to same worker");
+    }
+
+    #[test]
+    fn test_no_tokens_returns_none() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Empty tokens
+        let tokens: Vec<u32> = vec![];
+        let info = SelectWorkerInfo {
+            tokens: Some(&tokens),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoTokens);
+
+        // No tokens field
+        let info_no_tokens = SelectWorkerInfo {
+            tokens: None,
+            hash_ring: Some(ring),
+            ..Default::default()
+        };
+
+        let (result2, branch2) = policy.select_worker_impl(&workers, &info_no_tokens);
+        assert_eq!(result2, None);
+        assert_eq!(branch2, Branch::NoTokens);
+    }
+
+    #[test]
+    fn test_no_healthy_workers() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000"]);
+        workers[0].set_healthy(false);
+
+        let ring = Arc::new(HashRing::new(&workers));
+        let tokens: Vec<u32> = vec![1, 2, 3];
+        let info = SelectWorkerInfo {
+            tokens: Some(&tokens),
+            hash_ring: Some(ring),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoHealthyWorkers);
+    }
+
+    #[test]
+    fn test_load_ok_calculation() {
+        let policy = PrefixHashPolicy::new(PrefixHashConfig {
+            load_factor: 1.25,
+            ..Default::default()
+        });
+
+        // Total load 100, 4 workers -> avg 25, threshold 31.25
+        assert!(policy.load_ok(30, 100, 4)); // 30 <= 31.25
+        assert!(!policy.load_ok(35, 100, 4)); // 35 > 31.25
+
+        // Edge cases
+        assert!(policy.load_ok(0, 0, 4)); // No load = OK
+        assert!(policy.load_ok(100, 0, 0)); // No workers = OK (shouldn't happen)
+    }
+
+    #[test]
+    fn test_policy_name() {
+        let policy = PrefixHashPolicy::with_defaults();
+        assert_eq!(policy.name(), "prefix_hash");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/random.rs

@@ -0,0 +1,138 @@
+//! Random load balancing policy
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use rand::Rng;
+
+use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
+use crate::core::Worker;
+
+/// Random selection policy
+///
+/// Selects workers randomly with uniform distribution among healthy workers.
+#[derive(Debug, Default)]
+pub struct RandomPolicy;
+
+impl RandomPolicy {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for RandomPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        _info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        let mut rng = rand::rng();
+        let random_idx = rng.random_range(0..healthy_indices.len());
+
+        Some(healthy_indices[random_idx])
+    }
+
+    fn name(&self) -> &'static str {
+        "random"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_random_selection() {
+        let policy = RandomPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        let mut counts = HashMap::new();
+        for _ in 0..100 {
+            if let Some(idx) = policy
+                .select_worker(&workers, &SelectWorkerInfo::default())
+                .await
+            {
+                *counts.entry(idx).or_insert(0) += 1;
+            }
+        }
+
+        assert_eq!(counts.len(), 3);
+        assert!(counts.values().all(|&count| count > 0));
+    }
+
+    #[tokio::test]
+    async fn test_random_with_unhealthy_workers() {
+        let policy = RandomPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        workers[0].set_healthy(false);
+
+        for _ in 0..10 {
+            assert_eq!(
+                policy
+                    .select_worker(&workers, &SelectWorkerInfo::default())
+                    .await,
+                Some(1)
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn test_random_no_healthy_workers() {
+        let policy = RandomPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        )];
+
+        workers[0].set_healthy(false);
+        assert_eq!(
+            policy
+                .select_worker(&workers, &SelectWorkerInfo::default())
+                .await,
+            None
+        );
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/registry.rs

@@ -0,0 +1,512 @@
+use std::sync::{Arc, OnceLock, RwLock};
+
+use dashmap::DashMap;
+use serde_json;
+use smg_mesh::OptionalMeshSyncManager;
+use tracing::{debug, info, warn};
+
+/// Policy Registry for managing model-to-policy mappings
+///
+/// This registry manages the dynamic assignment of load balancing policies to models.
+/// When the first worker of a new model is added, it determines the policy for that model.
+/// All subsequent workers of the same model use the established policy.
+/// When the last worker of a model is removed, the policy mapping is cleaned up.
+use super::{BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PolicyFactory};
+use crate::{config::types::PolicyConfig, core::Worker};
+
+/// Registry for managing model-to-policy mappings
+#[derive(Clone)]
+pub struct PolicyRegistry {
+    /// Model ID -> Policy instance mapping (lock-free reads via DashMap)
+    model_policies: Arc<DashMap<String, Arc<dyn LoadBalancingPolicy>>>,
+
+    /// Model ID -> Worker count for cleanup tracking (lock-free reads via DashMap)
+    model_worker_counts: Arc<DashMap<String, usize>>,
+
+    /// Default policy instance (cached, immutable after creation)
+    default_policy: Arc<dyn LoadBalancingPolicy>,
+
+    /// Prefill policy for PD mode (set once at startup, lock-free reads via OnceLock)
+    prefill_policy: Arc<OnceLock<Arc<dyn LoadBalancingPolicy>>>,
+
+    /// Decode policy for PD mode (set once at startup, lock-free reads via OnceLock)
+    decode_policy: Arc<OnceLock<Arc<dyn LoadBalancingPolicy>>>,
+
+    /// Optional mesh sync manager for state synchronization
+    /// When None, the registry works independently without mesh synchronization
+    /// Uses RwLock for thread-safe access when setting mesh_sync after initialization
+    mesh_sync: Arc<RwLock<OptionalMeshSyncManager>>,
+}
+
+impl PolicyRegistry {
+    /// Create a new PolicyRegistry with a default policy
+    pub fn new(default_policy_config: PolicyConfig) -> Self {
+        let default_policy = Self::create_policy_from_config(&default_policy_config);
+
+        Self {
+            model_policies: Arc::new(DashMap::new()),
+            model_worker_counts: Arc::new(DashMap::new()),
+            default_policy,
+            prefill_policy: Arc::new(OnceLock::new()),
+            decode_policy: Arc::new(OnceLock::new()),
+            mesh_sync: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    /// Set mesh sync manager (thread-safe, can be called after initialization)
+    pub fn set_mesh_sync(&self, mesh_sync: OptionalMeshSyncManager) {
+        *self.mesh_sync.write().unwrap() = mesh_sync;
+    }
+
+    /// Called when a worker is added
+    /// Returns the policy that should be used for this worker's model
+    pub fn on_worker_added(
+        &self,
+        model_id: &str,
+        policy_hint: Option<&str>,
+    ) -> Arc<dyn LoadBalancingPolicy> {
+        // Increment worker count using DashMap entry API
+        let count = self
+            .model_worker_counts
+            .entry(model_id.to_string())
+            .and_modify(|c| *c += 1)
+            .or_insert(1);
+        debug!("Worker added for model {}, count: {}", model_id, *count);
+        drop(count); // Release the entry lock
+
+        // Check if model already has a policy (lock-free read via DashMap)
+        if let Some(existing_policy) = self.model_policies.get(model_id) {
+            debug!(
+                "Model {} already has policy: {}",
+                model_id,
+                existing_policy.name()
+            );
+            return Arc::clone(&existing_policy);
+        }
+
+        // New model - determine policy
+        let policy = self.determine_policy_for_model(model_id, policy_hint);
+
+        info!(
+            "Assigning policy {} to new model {}",
+            policy.name(),
+            model_id
+        );
+
+        // Store policy for this model (DashMap handles concurrent inserts)
+        self.model_policies
+            .insert(model_id.to_string(), Arc::clone(&policy));
+
+        // Sync to mesh if enabled (no-op if mesh is not enabled)
+        if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+            // Serialize policy config (simplified - just store policy name for now)
+            let config = serde_json::to_vec(&policy.name()).unwrap_or_default();
+            mesh_sync.sync_policy_state(model_id.to_string(), policy.name().to_string(), config);
+        }
+
+        policy
+    }
+
+    /// Called when a worker is removed
+    pub fn on_worker_removed(&self, model_id: &str) {
+        // Decrement worker count and check if cleanup needed
+        let should_cleanup = if let Some(mut count_ref) = self.model_worker_counts.get_mut(model_id)
+        {
+            *count_ref = count_ref.saturating_sub(1);
+            debug!(
+                "Worker removed for model {}, count: {}",
+                model_id, *count_ref
+            );
+            if *count_ref == 0 {
+                drop(count_ref); // Release before remove
+                self.model_worker_counts.remove(model_id);
+                true
+            } else {
+                false
+            }
+        } else {
+            warn!(
+                "Attempted to remove worker for model {} with no registered workers",
+                model_id
+            );
+            false
+        };
+
+        // Clean up policy if this was the last worker
+        if should_cleanup {
+            if let Some((_, policy)) = self.model_policies.remove(model_id) {
+                info!(
+                    "Removed policy {} for model {} (last worker removed)",
+                    policy.name(),
+                    model_id
+                );
+            }
+
+            // Sync removal to mesh if enabled (no-op if mesh is not enabled)
+            if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+                mesh_sync.remove_policy_state(model_id);
+            }
+        }
+    }
+
+    /// Get the policy for a model (lock-free via DashMap)
+    pub fn get_policy(&self, model_id: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
+        self.model_policies.get(model_id).map(|r| Arc::clone(&r))
+    }
+
+    /// Get the default policy
+    pub fn get_default_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
+        Arc::clone(&self.default_policy)
+    }
+
+    /// Get policy for a model, or default if not found
+    pub fn get_policy_or_default(&self, model_id: &str) -> Arc<dyn LoadBalancingPolicy> {
+        self.get_policy(model_id)
+            .unwrap_or_else(|| self.get_default_policy())
+    }
+
+    /// Determine policy for a new model
+    fn determine_policy_for_model(
+        &self,
+        model_id: &str,
+        policy_hint: Option<&str>,
+    ) -> Arc<dyn LoadBalancingPolicy> {
+        // 1. Check policy hint from worker
+        if let Some(policy_type) = policy_hint {
+            debug!("Using policy hint '{}' for model {}", policy_type, model_id);
+            return self.create_policy_from_type(policy_type);
+        }
+
+        // 2. Use default policy
+        debug!("Using default policy for model {}", model_id);
+        Arc::clone(&self.default_policy)
+    }
+
+    /// Create a policy from a type string (delegates to PolicyFactory)
+    fn create_policy_from_type(&self, policy_type: &str) -> Arc<dyn LoadBalancingPolicy> {
+        if policy_type == "cache_aware" {
+            let mut cache_aware = CacheAwarePolicy::new();
+            let mesh_sync = &*self.mesh_sync.read().unwrap();
+            cache_aware.set_mesh_sync(mesh_sync.clone());
+            Arc::new(cache_aware)
+        } else {
+            PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| {
+                warn!("Unknown policy type '{}', using default", policy_type);
+                Arc::clone(&self.default_policy)
+            })
+        }
+    }
+
+    /// Create a policy from a PolicyConfig (delegates to PolicyFactory)
+    fn create_policy_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
+        PolicyFactory::create_from_config(config)
+    }
+
+    /// Get current model->policy mappings (for debugging/monitoring)
+    pub fn get_all_mappings(&self) -> std::collections::HashMap<String, String> {
+        self.model_policies
+            .iter()
+            .map(|entry| (entry.key().clone(), entry.value().name().to_string()))
+            .collect()
+    }
+
+    /// Get worker counts per model
+    pub fn get_worker_counts(&self) -> std::collections::HashMap<String, usize> {
+        self.model_worker_counts
+            .iter()
+            .map(|entry| (entry.key().clone(), *entry.value()))
+            .collect()
+    }
+
+    /// Clear all policies (useful for testing)
+    pub fn clear(&self) {
+        self.model_policies.clear();
+        self.model_worker_counts.clear();
+    }
+
+    /// Set the prefill policy for PD mode (lock-free, set once at startup)
+    pub fn set_prefill_policy(&self, policy: Arc<dyn LoadBalancingPolicy>) {
+        // OnceLock::set returns Err if already set, which we ignore since
+        // the policy should only be set once at startup
+        let _ = self.prefill_policy.set(policy);
+    }
+
+    /// Set the decode policy for PD mode (lock-free, set once at startup)
+    pub fn set_decode_policy(&self, policy: Arc<dyn LoadBalancingPolicy>) {
+        // OnceLock::set returns Err if already set, which we ignore since
+        // the policy should only be set once at startup
+        let _ = self.decode_policy.set(policy);
+    }
+
+    /// Get the prefill policy for PD mode, or default if not set (lock-free)
+    pub fn get_prefill_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
+        self.prefill_policy
+            .get()
+            .map(Arc::clone)
+            .unwrap_or_else(|| self.get_default_policy())
+    }
+
+    /// Get the decode policy for PD mode, or default if not set (lock-free)
+    pub fn get_decode_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
+        self.decode_policy
+            .get()
+            .map(Arc::clone)
+            .unwrap_or_else(|| self.get_default_policy())
+    }
+
+    /// Get all PowerOfTwo policies that need load updates (lock-free)
+    pub fn get_all_power_of_two_policies(&self) -> Vec<Arc<dyn LoadBalancingPolicy>> {
+        let mut power_of_two_policies = Vec::new();
+
+        if self.default_policy.name() == "power_of_two" {
+            power_of_two_policies.push(Arc::clone(&self.default_policy));
+        }
+
+        // Get prefill and decode policies (lock-free via OnceLock::get)
+        let prefill_policy_opt = self.prefill_policy.get();
+        let decode_policy_opt = self.decode_policy.get();
+
+        if let Some(policy) = prefill_policy_opt {
+            if policy.name() == "power_of_two" && !Arc::ptr_eq(policy, &self.default_policy) {
+                power_of_two_policies.push(Arc::clone(policy));
+            }
+        }
+
+        if let Some(policy) = decode_policy_opt {
+            if policy.name() == "power_of_two"
+                && !Arc::ptr_eq(policy, &self.default_policy)
+                && !prefill_policy_opt.is_some_and(|p| Arc::ptr_eq(p, policy))
+            {
+                power_of_two_policies.push(Arc::clone(policy));
+            }
+        }
+
+        for entry in self.model_policies.iter() {
+            let policy = entry.value();
+            if policy.name() == "power_of_two" {
+                let already_added = power_of_two_policies.iter().any(|p| Arc::ptr_eq(p, policy));
+                if !already_added {
+                    power_of_two_policies.push(Arc::clone(policy));
+                }
+            }
+        }
+
+        power_of_two_policies
+    }
+
+    /// Initialize cache-aware policy with workers if applicable
+    /// This should be called after workers are registered for a model
+    pub fn init_cache_aware_policy(&self, model_id: &str, workers: &[Arc<dyn Worker>]) {
+        // Get the policy for this model
+        if let Some(policy) = self.get_policy(model_id) {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    debug!(
+                        "Initializing cache-aware policy with {} workers for model {}",
+                        workers.len(),
+                        model_id
+                    );
+                    cache_aware.init_workers(workers);
+                }
+            }
+        }
+    }
+
+    /// Remove a worker from cache-aware policy if applicable
+    /// This should be called when a worker is being removed
+    pub fn remove_worker_from_cache_aware(&self, model_id: &str, worker_url: &str) {
+        // Get the policy for this model
+        if let Some(policy) = self.get_policy(model_id) {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    cache_aware.remove_worker_by_url(worker_url);
+                    debug!(
+                        "Removed worker {} from cache-aware policy for model {}",
+                        worker_url, model_id
+                    );
+                }
+            }
+        }
+    }
+
+    /// Initialize cache-aware policies for PD mode (prefill and decode) - lock-free
+    pub fn init_pd_cache_aware_policies(
+        &self,
+        prefill_workers: &[Arc<dyn Worker>],
+        decode_workers: &[Arc<dyn Worker>],
+    ) {
+        // Initialize prefill policy if it's cache-aware (lock-free via OnceLock::get)
+        if let Some(prefill_policy) = self.prefill_policy.get() {
+            if prefill_policy.name() == "cache_aware" {
+                if let Some(cache_aware) =
+                    prefill_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    if !prefill_workers.is_empty() {
+                        debug!(
+                            "Initializing prefill cache-aware policy with {} workers",
+                            prefill_workers.len()
+                        );
+                        cache_aware.init_workers(prefill_workers);
+                    }
+                }
+            }
+        }
+
+        // Initialize decode policy if it's cache-aware (lock-free via OnceLock::get)
+        if let Some(decode_policy) = self.decode_policy.get() {
+            if decode_policy.name() == "cache_aware" {
+                if let Some(cache_aware) = decode_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    if !decode_workers.is_empty() {
+                        debug!(
+                            "Initializing decode cache-aware policy with {} workers",
+                            decode_workers.len()
+                        );
+                        cache_aware.init_workers(decode_workers);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Initialize bucket policies for PD mode - lock-free
+    pub fn init_pd_bucket_policies(&self, prefill_workers: &[Arc<dyn Worker>]) {
+        // Initialize prefill policy if it's bucket (lock-free via OnceLock::get)
+        if let Some(prefill_policy) = self.prefill_policy.get() {
+            if prefill_policy.name() == "bucket" {
+                if let Some(bucket) = prefill_policy.as_any().downcast_ref::<BucketPolicy>() {
+                    if !prefill_workers.is_empty() {
+                        debug!(
+                            "Initializing prefill bucket policy with {} workers",
+                            prefill_workers.len()
+                        );
+                        bucket.init_prefill_worker_urls(prefill_workers);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Apply remote tree operation to cache-aware policy for a model
+    /// This is called when receiving tree state updates from mesh
+    pub fn apply_remote_tree_operation(
+        &self,
+        model_id: &str,
+        operation: &smg_mesh::tree_ops::TreeOperation,
+    ) {
+        // Try to find the policy for this model
+        if let Some(policy) = self.get_policy(model_id) {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    cache_aware.apply_remote_tree_operation(model_id, operation);
+                }
+            }
+        }
+
+        // Also check default policy if it's cache-aware
+        if self.default_policy.name() == "cache_aware" {
+            if let Some(cache_aware) = self
+                .default_policy
+                .as_any()
+                .downcast_ref::<CacheAwarePolicy>()
+            {
+                cache_aware.apply_remote_tree_operation(model_id, operation);
+            }
+        }
+
+        // Check prefill and decode policies for PD mode
+        if let Some(prefill_policy) = self.prefill_policy.get() {
+            if prefill_policy.name() == "cache_aware" {
+                if let Some(cache_aware) =
+                    prefill_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    cache_aware.apply_remote_tree_operation(model_id, operation);
+                }
+            }
+        }
+
+        if let Some(decode_policy) = self.decode_policy.get() {
+            if decode_policy.name() == "cache_aware" {
+                if let Some(cache_aware) = decode_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    cache_aware.apply_remote_tree_operation(model_id, operation);
+                }
+            }
+        }
+    }
+}
+
+impl std::fmt::Debug for PolicyRegistry {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("PolicyRegistry")
+            .field("model_policies", &self.model_policies)
+            .field("model_worker_counts", &self.model_worker_counts)
+            .field("default_policy", &self.default_policy.name())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_policy_registry_basic() {
+        let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+        // First worker of a model sets the policy
+        let policy1 = registry.on_worker_added("llama-3", Some("cache_aware"));
+        assert_eq!(policy1.name(), "cache_aware");
+
+        // Second worker of same model uses existing policy
+        let policy2 = registry.on_worker_added("llama-3", Some("round_robin"));
+        assert_eq!(policy2.name(), "cache_aware"); // Ignores hint, uses existing
+
+        // Different model can have different policy
+        let policy3 = registry.on_worker_added("gpt-4", Some("random"));
+        assert_eq!(policy3.name(), "random");
+
+        // Check mappings
+        let mappings = registry.get_all_mappings();
+        assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware");
+        assert_eq!(mappings.get("gpt-4").unwrap(), "random");
+
+        // Check worker counts
+        let counts = registry.get_worker_counts();
+        assert_eq!(*counts.get("llama-3").unwrap(), 2);
+        assert_eq!(*counts.get("gpt-4").unwrap(), 1);
+    }
+
+    #[tokio::test]
+    async fn test_policy_registry_cleanup() {
+        let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+        // Add workers
+        registry.on_worker_added("llama-3", Some("cache_aware"));
+        registry.on_worker_added("llama-3", None);
+        assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&2));
+
+        // Remove one worker - policy should remain
+        registry.on_worker_removed("llama-3");
+        assert!(registry.get_policy("llama-3").is_some());
+        assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&1));
+
+        // Remove last worker - policy should be cleaned up
+        registry.on_worker_removed("llama-3");
+        assert!(registry.get_policy("llama-3").is_none());
+        assert_eq!(registry.get_worker_counts().get("llama-3"), None);
+    }
+
+    #[tokio::test]
+    async fn test_default_policy() {
+        let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+        // No hint, no template - uses default
+        let policy = registry.on_worker_added("unknown-model", None);
+        assert_eq!(policy.name(), "round_robin");
+
+        // Get default directly
+        let default = registry.get_default_policy();
+        assert_eq!(default.name(), "round_robin");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/round_robin.rs

@@ -0,0 +1,149 @@
+//! Round-robin load balancing policy
+
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use async_trait::async_trait;
+
+use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
+use crate::core::Worker;
+
+/// Round-robin selection policy
+///
+/// Selects workers in sequential order, cycling through all healthy workers.
+#[derive(Debug, Default)]
+pub struct RoundRobinPolicy {
+    counter: AtomicUsize,
+}
+
+impl RoundRobinPolicy {
+    pub fn new() -> Self {
+        Self {
+            counter: AtomicUsize::new(0),
+        }
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for RoundRobinPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        _info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        // Get and increment counter atomically
+        let count = self.counter.fetch_add(1, Ordering::Relaxed);
+        let selected_idx = count % healthy_indices.len();
+
+        Some(healthy_indices[selected_idx])
+    }
+
+    fn name(&self) -> &'static str {
+        "round_robin"
+    }
+
+    fn reset(&self) {
+        self.counter.store(0, Ordering::Relaxed);
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_round_robin_selection() {
+        let policy = RoundRobinPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        let info = SelectWorkerInfo::default();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(1));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(2));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(1));
+    }
+
+    #[tokio::test]
+    async fn test_round_robin_with_unhealthy_workers() {
+        let policy = RoundRobinPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        workers[1].set_healthy(false);
+
+        let info = SelectWorkerInfo::default();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(2));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(2));
+    }
+
+    #[tokio::test]
+    async fn test_round_robin_reset() {
+        let policy = RoundRobinPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        let info = SelectWorkerInfo::default();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(1));
+
+        policy.reset();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/utils.rs

@@ -0,0 +1,110 @@
+use std::{
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    thread::{self, JoinHandle},
+    time::Duration,
+};
+
+use tracing::debug;
+
+#[derive(Debug)]
+pub struct PeriodicTask {
+    debug_name: &'static str,
+    shutdown_flag: Arc<AtomicBool>,
+    handle: Option<JoinHandle<()>>,
+}
+
+impl PeriodicTask {
+    /// Spawn a background thread that periodically executes a task.
+    pub fn spawn<F>(interval_secs: u64, debug_name: &'static str, task: F) -> Self
+    where
+        F: Fn() + Send + 'static,
+    {
+        let shutdown_flag = Arc::new(AtomicBool::new(false));
+        let shutdown_clone = Arc::clone(&shutdown_flag);
+
+        let handle = thread::spawn(move || {
+            let check_interval_ms = 100u64;
+            let total_sleep_ms = interval_secs * 1000;
+
+            loop {
+                // Sleep in small increments, checking shutdown flag periodically
+                let mut slept_ms = 0u64;
+                while slept_ms < total_sleep_ms {
+                    if shutdown_clone.load(Ordering::Relaxed) {
+                        debug!("{} thread received shutdown signal", debug_name);
+                        return;
+                    }
+                    thread::sleep(Duration::from_millis(check_interval_ms));
+                    slept_ms += check_interval_ms;
+                }
+
+                // Check shutdown before starting task
+                if shutdown_clone.load(Ordering::Relaxed) {
+                    debug!("{} thread received shutdown signal", debug_name);
+                    return;
+                }
+
+                task();
+            }
+        });
+
+        Self {
+            debug_name,
+            shutdown_flag,
+            handle: Some(handle),
+        }
+    }
+}
+
+impl Drop for PeriodicTask {
+    fn drop(&mut self) {
+        self.shutdown_flag.store(true, Ordering::Relaxed);
+
+        if let Some(handle) = self.handle.take() {
+            match handle.join() {
+                Ok(()) => debug!("{} thread shut down cleanly", self.debug_name),
+                Err(_) => debug!("{} thread panicked during shutdown", self.debug_name),
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{sync::atomic::AtomicUsize, time::Instant};
+
+    use super::*;
+
+    #[test]
+    fn test_periodic_task_executes() {
+        let counter = Arc::new(AtomicUsize::new(0));
+        let counter_clone = Arc::clone(&counter);
+
+        let _task = PeriodicTask::spawn(1, "test", move || {
+            counter_clone.fetch_add(1, Ordering::SeqCst);
+        });
+
+        // Wait for at least one execution
+        thread::sleep(Duration::from_millis(1200));
+        assert!(counter.load(Ordering::SeqCst) >= 1);
+
+        // Task will be stopped on drop
+    }
+
+    #[test]
+    fn test_periodic_task_responds_to_shutdown() {
+        let task = PeriodicTask::spawn(60, "test", || {
+            // Long interval task
+        });
+
+        let start = Instant::now();
+        drop(task);
+        let elapsed = start.elapsed();
+
+        // Should shutdown within ~200ms (2 check intervals), not 60 seconds
+        assert!(elapsed < Duration::from_millis(500));
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/conversations/handlers.rs

@@ -0,0 +1,633 @@
+//! Conversation CRUD handlers for the /v1/conversations API - shared across routers
+
+use std::sync::Arc;
+
+use axum::{
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    Json,
+};
+use chrono::Utc;
+use data_connector::{
+    Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage,
+    ConversationStorage, ListParams, NewConversation, NewConversationItem, SortOrder,
+};
+use serde_json::{json, Value};
+use tracing::{info, warn};
+
+use crate::routers::persistence_utils::item_to_json;
+
+// ============================================================================
+// Constants
+// ============================================================================
+
+pub const MAX_METADATA_PROPERTIES: usize = 16;
+const MAX_ITEMS_PER_REQUEST: usize = 20;
+
+const SUPPORTED_ITEM_TYPES: &[&str] = &[
+    "message",
+    "reasoning",
+    "mcp_list_tools",
+    "mcp_call",
+    "item_reference",
+    "function_call",
+    "function_call_output",
+    "file_search_call",
+    "computer_call",
+    "computer_call_output",
+    "web_search_call",
+    "image_generation_call",
+    "code_interpreter_call",
+    "local_shell_call",
+    "local_shell_call_output",
+    "mcp_approval_request",
+    "mcp_approval_response",
+    "custom_tool_call",
+    "custom_tool_call_output",
+];
+
+const IMPLEMENTED_ITEM_TYPES: &[&str] = &[
+    "message",
+    "reasoning",
+    "mcp_list_tools",
+    "mcp_call",
+    "item_reference",
+];
+
+// ============================================================================
+// Error Response Helpers
+// ============================================================================
+
+fn bad_request(message: impl Into<String>) -> Response {
+    (
+        StatusCode::BAD_REQUEST,
+        Json(json!({"error": message.into()})),
+    )
+        .into_response()
+}
+
+fn not_found(message: impl Into<String>) -> Response {
+    (
+        StatusCode::NOT_FOUND,
+        Json(json!({"error": message.into()})),
+    )
+        .into_response()
+}
+
+fn internal_error(message: impl Into<String>) -> Response {
+    (
+        StatusCode::INTERNAL_SERVER_ERROR,
+        Json(json!({"error": message.into()})),
+    )
+        .into_response()
+}
+
+fn bad_request_structured(error_obj: Value) -> Response {
+    (StatusCode::BAD_REQUEST, Json(json!({"error": error_obj}))).into_response()
+}
+
+// ============================================================================
+// Storage Helpers
+// ============================================================================
+
+async fn ensure_conversation_exists(
+    storage: &Arc<dyn ConversationStorage>,
+    conv_id: &ConversationId,
+) -> Result<Conversation, Response> {
+    match storage.get_conversation(conv_id).await {
+        Ok(Some(conv)) => Ok(conv),
+        Ok(None) => Err(not_found("Conversation not found")),
+        Err(e) => Err(internal_error(format!("Failed to get conversation: {e}"))),
+    }
+}
+
+// ============================================================================
+// Metadata Operations
+// ============================================================================
+
+fn validate_metadata(value: &Value) -> Result<Option<serde_json::Map<String, Value>>, String> {
+    match value.get("metadata") {
+        Some(Value::Object(map)) => {
+            if map.len() > MAX_METADATA_PROPERTIES {
+                Err(format!(
+                    "metadata cannot have more than {MAX_METADATA_PROPERTIES} properties"
+                ))
+            } else {
+                Ok(Some(map.clone()))
+            }
+        }
+        Some(_) => Err("metadata must be an object".to_string()),
+        None => Ok(None),
+    }
+}
+
+fn apply_metadata_patches(
+    current: Option<serde_json::Map<String, Value>>,
+    body: &Value,
+) -> Result<Option<serde_json::Map<String, Value>>, String> {
+    let patch_map = match body.get("metadata") {
+        Some(Value::Object(map)) => map,
+        Some(_) => return Err("metadata must be an object".to_string()),
+        None => return Ok(current),
+    };
+
+    let mut result = current.unwrap_or_default();
+    for (k, v) in patch_map {
+        if v.is_null() {
+            result.remove(k);
+        } else {
+            result.insert(k.clone(), v.clone());
+        }
+    }
+
+    if result.len() > MAX_METADATA_PROPERTIES {
+        return Err(format!(
+            "metadata cannot have more than {MAX_METADATA_PROPERTIES} properties"
+        ));
+    }
+
+    Ok(if result.is_empty() {
+        None
+    } else {
+        Some(result)
+    })
+}
+
+// ============================================================================
+// Conversation CRUD Handlers
+// ============================================================================
+
+pub async fn create_conversation(storage: &Arc<dyn ConversationStorage>, body: Value) -> Response {
+    let metadata = match validate_metadata(&body) {
+        Ok(m) => m,
+        Err(msg) => return bad_request(msg),
+    };
+
+    let new_conv = NewConversation { id: None, metadata };
+
+    match storage.create_conversation(new_conv).await {
+        Ok(conversation) => {
+            info!(conversation_id = %conversation.id.0, "Created conversation");
+            (StatusCode::OK, Json(conversation_to_json(&conversation))).into_response()
+        }
+        Err(e) => internal_error(format!("Failed to create conversation: {e}")),
+    }
+}
+
+pub async fn get_conversation(storage: &Arc<dyn ConversationStorage>, conv_id: &str) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+
+    match storage.get_conversation(&conversation_id).await {
+        Ok(Some(conversation)) => {
+            (StatusCode::OK, Json(conversation_to_json(&conversation))).into_response()
+        }
+        Ok(None) => not_found("Conversation not found"),
+        Err(e) => internal_error(format!("Failed to get conversation: {e}")),
+    }
+}
+
+pub async fn update_conversation(
+    storage: &Arc<dyn ConversationStorage>,
+    conv_id: &str,
+    body: Value,
+) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+
+    let current = match ensure_conversation_exists(storage, &conversation_id).await {
+        Ok(c) => c,
+        Err(response) => return response,
+    };
+
+    let final_metadata = match apply_metadata_patches(current.metadata.clone(), &body) {
+        Ok(m) => m,
+        Err(msg) => return bad_request(msg),
+    };
+
+    match storage
+        .update_conversation(&conversation_id, final_metadata)
+        .await
+    {
+        Ok(Some(conversation)) => {
+            info!(conversation_id = %conversation_id.0, "Updated conversation");
+            (StatusCode::OK, Json(conversation_to_json(&conversation))).into_response()
+        }
+        Ok(None) => not_found("Conversation not found"),
+        Err(e) => internal_error(format!("Failed to update conversation: {e}")),
+    }
+}
+
+pub async fn delete_conversation(
+    storage: &Arc<dyn ConversationStorage>,
+    conv_id: &str,
+) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+
+    if let Err(response) = ensure_conversation_exists(storage, &conversation_id).await {
+        return response;
+    }
+
+    match storage.delete_conversation(&conversation_id).await {
+        Ok(_) => {
+            info!(conversation_id = %conversation_id.0, "Deleted conversation");
+            (
+                StatusCode::OK,
+                Json(json!({
+                    "id": conversation_id.0,
+                    "object": "conversation.deleted",
+                    "deleted": true
+                })),
+            )
+                .into_response()
+        }
+        Err(e) => internal_error(format!("Failed to delete conversation: {e}")),
+    }
+}
+
+// ============================================================================
+// Conversation Item Handlers
+// ============================================================================
+
+pub async fn list_conversation_items(
+    conversation_storage: &Arc<dyn ConversationStorage>,
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conv_id: &str,
+    limit: Option<usize>,
+    order: Option<&str>,
+    after: Option<&str>,
+) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+
+    if let Err(response) = ensure_conversation_exists(conversation_storage, &conversation_id).await
+    {
+        return response;
+    }
+
+    let limit = limit.unwrap_or(100);
+    let order = match order {
+        Some("asc") => SortOrder::Asc,
+        _ => SortOrder::Desc,
+    };
+
+    let params = ListParams {
+        limit,
+        order,
+        after: after.map(String::from),
+    };
+
+    match item_storage.list_items(&conversation_id, params).await {
+        Ok(items) => {
+            let item_values: Vec<Value> = items
+                .iter()
+                .map(|item| {
+                    let mut item_json = item_to_json(item);
+                    if let Some(obj) = item_json.as_object_mut() {
+                        obj.insert("created_at".to_string(), json!(item.created_at));
+                    }
+                    item_json
+                })
+                .collect();
+
+            (
+                StatusCode::OK,
+                Json(json!({
+                    "object": "list",
+                    "data": item_values,
+                    "has_more": items.len() == limit,
+                    "first_id": items.first().map(|item| &item.id.0),
+                    "last_id": items.last().map(|item| &item.id.0),
+                })),
+            )
+                .into_response()
+        }
+        Err(e) => internal_error(format!("Failed to list items: {e}")),
+    }
+}
+
+pub async fn create_conversation_items(
+    conversation_storage: &Arc<dyn ConversationStorage>,
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conv_id: &str,
+    body: Value,
+) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+
+    if let Err(response) = ensure_conversation_exists(conversation_storage, &conversation_id).await
+    {
+        return response;
+    }
+
+    let items_array = match body.get("items").and_then(|v| v.as_array()) {
+        Some(arr) => arr,
+        None => return bad_request("Missing or invalid 'items' field"),
+    };
+
+    if items_array.len() > MAX_ITEMS_PER_REQUEST {
+        return bad_request(format!(
+            "Cannot add more than {MAX_ITEMS_PER_REQUEST} items at a time"
+        ));
+    }
+
+    let mut created_items = Vec::new();
+    let mut warnings = Vec::new();
+    let added_at = Utc::now();
+
+    for item_val in items_array {
+        match process_item(item_storage, &conversation_id, item_val, added_at).await {
+            Ok((item_json, warning)) => {
+                created_items.push(item_json);
+                if let Some(w) = warning {
+                    warnings.push(w);
+                }
+            }
+            Err(response) => return response,
+        }
+    }
+
+    let mut response = json!({
+        "object": "list",
+        "data": created_items,
+        "first_id": created_items.first().and_then(|v| v.get("id")),
+        "last_id": created_items.last().and_then(|v| v.get("id")),
+        "has_more": false
+    });
+
+    if !warnings.is_empty() {
+        if let Some(obj) = response.as_object_mut() {
+            obj.insert("warnings".to_string(), json!(warnings));
+        }
+    }
+
+    (StatusCode::OK, Json(response)).into_response()
+}
+
+/// Process a single item for creation/linking
+async fn process_item(
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conversation_id: &ConversationId,
+    item_val: &Value,
+    added_at: chrono::DateTime<Utc>,
+) -> Result<(Value, Option<String>), Response> {
+    let item_type = item_val
+        .get("type")
+        .and_then(|v| v.as_str())
+        .unwrap_or("message");
+
+    // Handle item_reference specially - just link existing item
+    if item_type == "item_reference" {
+        return process_item_reference(item_storage, conversation_id, item_val, added_at).await;
+    }
+
+    let user_provided_id = item_val.get("id").and_then(|v| v.as_str());
+
+    let (item, warning) = if let Some(id_str) = user_provided_id {
+        process_item_with_id(item_storage, conversation_id, item_val, id_str).await?
+    } else {
+        process_new_item(item_storage, item_val).await?
+    };
+
+    // Link item to conversation
+    if let Err(e) = item_storage
+        .link_item(conversation_id, &item.id, added_at)
+        .await
+    {
+        warn!("Failed to link item {}: {}", item.id.0, e);
+    }
+
+    Ok((item_to_json(&item), warning))
+}
+
+/// Process an item_reference - link an existing item to the conversation
+async fn process_item_reference(
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conversation_id: &ConversationId,
+    item_val: &Value,
+    added_at: chrono::DateTime<Utc>,
+) -> Result<(Value, Option<String>), Response> {
+    let ref_id = item_val
+        .get("id")
+        .and_then(|v| v.as_str())
+        .ok_or_else(|| bad_request("item_reference requires 'id' field"))?;
+
+    let item_id = ConversationItemId::from(ref_id);
+
+    let existing_item = match item_storage.get_item(&item_id).await {
+        Ok(Some(item)) => item,
+        Ok(None) => return Err(not_found(format!("Referenced item '{ref_id}' not found"))),
+        Err(e) => {
+            return Err(internal_error(format!(
+                "Failed to get referenced item: {e}"
+            )))
+        }
+    };
+
+    if let Err(e) = item_storage
+        .link_item(conversation_id, &existing_item.id, added_at)
+        .await
+    {
+        warn!("Failed to link item {}: {}", existing_item.id.0, e);
+    }
+
+    Ok((item_to_json(&existing_item), None))
+}
+
+/// Process an item with a user-provided ID
+async fn process_item_with_id(
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conversation_id: &ConversationId,
+    item_val: &Value,
+    id_str: &str,
+) -> Result<(ConversationItem, Option<String>), Response> {
+    let item_id = ConversationItemId::from(id_str);
+
+    // Check if already linked
+    let is_linked = item_storage
+        .is_item_linked(conversation_id, &item_id)
+        .await
+        .map_err(|e| internal_error(format!("Failed to check item link: {e}")))?;
+
+    if is_linked {
+        return Err(bad_request_structured(json!({
+            "message": "Item already in conversation",
+            "type": "invalid_request_error",
+            "param": "items",
+            "code": "item_already_in_conversation"
+        })));
+    }
+
+    // Check if item exists globally
+    match item_storage.get_item(&item_id).await {
+        Ok(Some(existing)) => Ok((existing, None)),
+        Ok(None) => {
+            // Create new item with the provided ID
+            let (mut new_item, warning) = parse_item_from_value(item_val).map_err(bad_request)?;
+            new_item.id = Some(item_id);
+
+            let created = item_storage
+                .create_item(new_item)
+                .await
+                .map_err(|e| internal_error(format!("Failed to create item: {e}")))?;
+
+            Ok((created, warning))
+        }
+        Err(e) => Err(internal_error(format!(
+            "Failed to check item existence: {e}"
+        ))),
+    }
+}
+
+/// Process a new item without a user-provided ID
+async fn process_new_item(
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    item_val: &Value,
+) -> Result<(ConversationItem, Option<String>), Response> {
+    let (new_item, warning) = parse_item_from_value(item_val).map_err(bad_request)?;
+
+    let created = item_storage
+        .create_item(new_item)
+        .await
+        .map_err(|e| internal_error(format!("Failed to create item: {e}")))?;
+
+    Ok((created, warning))
+}
+
+pub async fn get_conversation_item(
+    conversation_storage: &Arc<dyn ConversationStorage>,
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conv_id: &str,
+    item_id: &str,
+    _include: Option<Vec<String>>,
+) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+    let item_id = ConversationItemId::from(item_id);
+
+    if let Err(response) = ensure_conversation_exists(conversation_storage, &conversation_id).await
+    {
+        return response;
+    }
+
+    let is_linked = match item_storage
+        .is_item_linked(&conversation_id, &item_id)
+        .await
+    {
+        Ok(linked) => linked,
+        Err(e) => return internal_error(format!("Failed to check item link: {e}")),
+    };
+
+    if !is_linked {
+        return not_found("Item not found in this conversation");
+    }
+
+    match item_storage.get_item(&item_id).await {
+        Ok(Some(item)) => (StatusCode::OK, Json(item_to_json(&item))).into_response(),
+        Ok(None) => not_found("Item not found"),
+        Err(e) => internal_error(format!("Failed to get item: {e}")),
+    }
+}
+
+pub async fn delete_conversation_item(
+    conversation_storage: &Arc<dyn ConversationStorage>,
+    item_storage: &Arc<dyn ConversationItemStorage>,
+    conv_id: &str,
+    item_id: &str,
+) -> Response {
+    let conversation_id = ConversationId::from(conv_id);
+    let item_id = ConversationItemId::from(item_id);
+
+    let conversation =
+        match ensure_conversation_exists(conversation_storage, &conversation_id).await {
+            Ok(conv) => conv,
+            Err(response) => return response,
+        };
+
+    match item_storage.delete_item(&conversation_id, &item_id).await {
+        Ok(_) => {
+            info!(
+                conversation_id = %conversation_id.0,
+                item_id = %item_id.0,
+                "Deleted conversation item"
+            );
+            (StatusCode::OK, Json(conversation_to_json(&conversation))).into_response()
+        }
+        Err(e) => internal_error(format!("Failed to delete item: {e}")),
+    }
+}
+
+// ============================================================================
+// Parsing and Serialization
+// ============================================================================
+
+fn parse_item_from_value(
+    item_val: &Value,
+) -> Result<(NewConversationItem, Option<String>), String> {
+    let item_type = item_val
+        .get("type")
+        .and_then(|v| v.as_str())
+        .unwrap_or("message");
+
+    if !SUPPORTED_ITEM_TYPES.contains(&item_type) {
+        return Err(format!(
+            "Unsupported item type '{}'. Supported types: {}",
+            item_type,
+            SUPPORTED_ITEM_TYPES.join(", ")
+        ));
+    }
+
+    let warning = if !IMPLEMENTED_ITEM_TYPES.contains(&item_type) {
+        Some(format!(
+            "Item type '{}' is accepted but not yet implemented. \
+             The item will be stored but may not function as expected.",
+            item_type
+        ))
+    } else {
+        None
+    };
+
+    let role = item_val
+        .get("role")
+        .and_then(|v| v.as_str())
+        .map(String::from);
+
+    if item_type == "message" && role.is_none() {
+        return Err("Message items require 'role' field".to_string());
+    }
+
+    let status = item_val
+        .get("status")
+        .and_then(|v| v.as_str())
+        .map(String::from)
+        .or_else(|| Some("completed".to_string()));
+
+    let content = if item_type == "message" || item_type == "reasoning" {
+        item_val.get("content").cloned().unwrap_or(json!([]))
+    } else {
+        item_val.clone()
+    };
+
+    Ok((
+        NewConversationItem {
+            id: None,
+            response_id: None,
+            item_type: item_type.to_string(),
+            role,
+            content,
+            status,
+        },
+        warning,
+    ))
+}
+
+pub fn conversation_to_json(conversation: &Conversation) -> Value {
+    let mut obj = json!({
+        "id": conversation.id.0,
+        "object": "conversation",
+        "created_at": conversation.created_at.timestamp()
+    });
+
+    if let Some(metadata) = &conversation.metadata {
+        if !metadata.is_empty() {
+            obj["metadata"] = Value::Object(metadata.clone());
+        }
+    }
+
+    obj
+}

third_party/sglang/sgl-model-gateway/src/routers/conversations/mod.rs

@@ -0,0 +1,8 @@
+//! Shared conversation management module.
+//!
+//! This module provides conversation CRUD operations that can be shared
+//! across different router implementations.
+
+mod handlers;
+
+pub use handlers::*;

third_party/sglang/sgl-model-gateway/src/routers/error.rs

@@ -0,0 +1,95 @@
+use axum::{
+    http::{HeaderMap, HeaderValue, StatusCode},
+    response::{IntoResponse, Response},
+    Json,
+};
+use serde::Serialize;
+
+#[derive(Serialize)]
+struct ErrorResponse<'a> {
+    error: ErrorDetail<'a>,
+}
+
+#[derive(Serialize)]
+struct ErrorDetail<'a> {
+    #[serde(rename = "type")]
+    error_type: &'static str,
+    code: &'a str,
+    message: &'a str,
+}
+
+pub const HEADER_X_SMG_ERROR_CODE: &str = "X-SMG-Error-Code";
+
+pub fn internal_error(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::INTERNAL_SERVER_ERROR, code, message)
+}
+
+pub fn bad_request(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::BAD_REQUEST, code, message)
+}
+
+pub fn not_found(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::NOT_FOUND, code, message)
+}
+
+pub fn service_unavailable(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::SERVICE_UNAVAILABLE, code, message)
+}
+
+pub fn failed_dependency(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::FAILED_DEPENDENCY, code, message)
+}
+
+pub fn not_implemented(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::NOT_IMPLEMENTED, code, message)
+}
+
+pub fn bad_gateway(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::BAD_GATEWAY, code, message)
+}
+
+pub fn method_not_allowed(code: impl Into<String>, message: impl Into<String>) -> Response {
+    create_error(StatusCode::METHOD_NOT_ALLOWED, code, message)
+}
+
+pub fn create_error(
+    status: StatusCode,
+    code: impl Into<String>,
+    message: impl Into<String>,
+) -> Response {
+    let code_str = code.into();
+    let message_str = message.into();
+
+    let mut headers = HeaderMap::with_capacity(1);
+    headers.insert(
+        HEADER_X_SMG_ERROR_CODE,
+        HeaderValue::from_str(&code_str).unwrap(),
+    );
+
+    (
+        status,
+        headers,
+        Json(ErrorResponse {
+            error: ErrorDetail {
+                error_type: status_code_to_str(status),
+                code: &code_str,
+                message: &message_str,
+            },
+        }),
+    )
+        .into_response()
+}
+
+fn status_code_to_str(status_code: StatusCode) -> &'static str {
+    status_code
+        .canonical_reason()
+        .unwrap_or("Unknown Status Code")
+}
+
+pub fn extract_error_code_from_response<B>(response: &Response<B>) -> &str {
+    response
+        .headers()
+        .get(HEADER_X_SMG_ERROR_CODE)
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or_default()
+}

third_party/sglang/sgl-model-gateway/src/routers/factory.rs

@@ -0,0 +1,130 @@
+//! Factory for creating router instances
+
+use std::sync::Arc;
+
+use super::{
+    grpc::{pd_router::GrpcPDRouter, router::GrpcRouter},
+    http::{pd_router::PDRouter, router::Router},
+    openai::OpenAIRouter,
+    RouterTrait,
+};
+use crate::{
+    app_context::AppContext,
+    config::{PolicyConfig, RoutingMode},
+    core::ConnectionMode,
+    policies::PolicyFactory,
+};
+
+/// Factory for creating router instances based on configuration
+pub struct RouterFactory;
+
+impl RouterFactory {
+    /// Create a router instance from application context
+    pub async fn create_router(ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
+        match ctx.router_config.connection_mode {
+            ConnectionMode::Grpc { .. } => match &ctx.router_config.mode {
+                RoutingMode::Regular { .. } => Self::create_grpc_router(ctx).await,
+                RoutingMode::PrefillDecode {
+                    prefill_policy,
+                    decode_policy,
+                    ..
+                } => {
+                    Self::create_grpc_pd_router(
+                        prefill_policy.as_ref(),
+                        decode_policy.as_ref(),
+                        &ctx.router_config.policy,
+                        ctx,
+                    )
+                    .await
+                }
+                RoutingMode::OpenAI { .. } => {
+                    Err("OpenAI mode requires HTTP connection_mode".to_string())
+                }
+            },
+            ConnectionMode::Http => match &ctx.router_config.mode {
+                RoutingMode::Regular { .. } => Self::create_regular_router(ctx).await,
+                RoutingMode::PrefillDecode {
+                    prefill_policy,
+                    decode_policy,
+                    ..
+                } => {
+                    Self::create_pd_router(
+                        prefill_policy.as_ref(),
+                        decode_policy.as_ref(),
+                        &ctx.router_config.policy,
+                        ctx,
+                    )
+                    .await
+                }
+                RoutingMode::OpenAI { .. } => Self::create_openai_router(ctx).await,
+            },
+        }
+    }
+
+    /// Create a regular router
+    pub async fn create_regular_router(
+        ctx: &Arc<AppContext>,
+    ) -> Result<Box<dyn RouterTrait>, String> {
+        let router = Router::new(ctx).await?;
+
+        Ok(Box::new(router))
+    }
+
+    /// Create a PD router with injected policy
+    pub async fn create_pd_router(
+        prefill_policy_config: Option<&PolicyConfig>,
+        decode_policy_config: Option<&PolicyConfig>,
+        main_policy_config: &PolicyConfig,
+        ctx: &Arc<AppContext>,
+    ) -> Result<Box<dyn RouterTrait>, String> {
+        let prefill_policy =
+            PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config));
+        let decode_policy =
+            PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config));
+
+        ctx.policy_registry.set_prefill_policy(prefill_policy);
+        ctx.policy_registry.set_decode_policy(decode_policy);
+
+        let router = PDRouter::new(ctx).await?;
+
+        Ok(Box::new(router))
+    }
+
+    /// Create a gRPC router with injected policy
+    pub async fn create_grpc_router(ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
+        let router = GrpcRouter::new(ctx).await?;
+
+        Ok(Box::new(router))
+    }
+
+    /// Create a gRPC PD router with tokenizer and worker configuration
+    pub async fn create_grpc_pd_router(
+        prefill_policy_config: Option<&PolicyConfig>,
+        decode_policy_config: Option<&PolicyConfig>,
+        main_policy_config: &PolicyConfig,
+        ctx: &Arc<AppContext>,
+    ) -> Result<Box<dyn RouterTrait>, String> {
+        let prefill_policy =
+            PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config));
+        let decode_policy =
+            PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config));
+
+        ctx.policy_registry.set_prefill_policy(prefill_policy);
+        ctx.policy_registry.set_decode_policy(decode_policy);
+        let router = GrpcPDRouter::new(ctx).await?;
+
+        Ok(Box::new(router))
+    }
+
+    /// Create an OpenAI router
+    ///
+    /// Workers should be registered via the external worker registration workflow
+    /// before using this router. The workflow discovers models from the provided
+    /// endpoints and creates external workers in the registry.
+    pub async fn create_openai_router(
+        ctx: &Arc<AppContext>,
+    ) -> Result<Box<dyn RouterTrait>, String> {
+        let router = OpenAIRouter::new(ctx).await?;
+        Ok(Box::new(router))
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/client.rs

@@ -0,0 +1,208 @@
+//! Unified gRPC client wrapper for SGLang and vLLM backends
+
+use std::sync::Arc;
+
+use smg_grpc_client::{SglangSchedulerClient, VllmEngineClient};
+
+use crate::{
+    observability::otel_trace::OtelTraceInjector,
+    routers::grpc::proto_wrapper::{
+        ProtoEmbedRequest, ProtoEmbedResponse, ProtoGenerateRequest, ProtoStream,
+    },
+};
+
+/// Health check response (common across backends)
+#[derive(Debug, Clone)]
+pub struct HealthCheckResponse {
+    pub healthy: bool,
+    pub message: String,
+}
+
+/// Polymorphic gRPC client that wraps either SGLang or vLLM
+#[derive(Clone)]
+pub enum GrpcClient {
+    Sglang(SglangSchedulerClient),
+    Vllm(VllmEngineClient),
+}
+
+impl GrpcClient {
+    /// Get reference to SGLang client (panics if vLLM)
+    pub fn as_sglang(&self) -> &SglangSchedulerClient {
+        match self {
+            Self::Sglang(client) => client,
+            Self::Vllm(_) => panic!("Expected SGLang client, got vLLM"),
+        }
+    }
+
+    /// Get mutable reference to SGLang client (panics if vLLM)
+    pub fn as_sglang_mut(&mut self) -> &mut SglangSchedulerClient {
+        match self {
+            Self::Sglang(client) => client,
+            Self::Vllm(_) => panic!("Expected SGLang client, got vLLM"),
+        }
+    }
+
+    /// Get reference to vLLM client (panics if SGLang)
+    pub fn as_vllm(&self) -> &VllmEngineClient {
+        match self {
+            Self::Vllm(client) => client,
+            Self::Sglang(_) => panic!("Expected vLLM client, got SGLang"),
+        }
+    }
+
+    /// Get mutable reference to vLLM client (panics if SGLang)
+    pub fn as_vllm_mut(&mut self) -> &mut VllmEngineClient {
+        match self {
+            Self::Vllm(client) => client,
+            Self::Sglang(_) => panic!("Expected vLLM client, got SGLang"),
+        }
+    }
+
+    /// Check if this is a SGLang client
+    pub fn is_sglang(&self) -> bool {
+        matches!(self, Self::Sglang(_))
+    }
+
+    /// Check if this is a vLLM client
+    pub fn is_vllm(&self) -> bool {
+        matches!(self, Self::Vllm(_))
+    }
+
+    /// Connect to gRPC server (runtime-aware)
+    pub async fn connect(
+        url: &str,
+        runtime_type: &str,
+    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
+        let trace_injector = Arc::new(OtelTraceInjector);
+        match runtime_type {
+            "sglang" => Ok(Self::Sglang(
+                SglangSchedulerClient::connect_with_trace_injector(url, trace_injector).await?,
+            )),
+            "vllm" => Ok(Self::Vllm(
+                VllmEngineClient::connect_with_trace_injector(url, trace_injector).await?,
+            )),
+            _ => Err(format!("Unknown runtime type: {}", runtime_type).into()),
+        }
+    }
+
+    /// Perform health check (dispatches to appropriate backend)
+    pub async fn health_check(
+        &self,
+    ) -> Result<HealthCheckResponse, Box<dyn std::error::Error + Send + Sync>> {
+        match self {
+            Self::Sglang(client) => {
+                let resp = client.health_check().await?;
+                Ok(HealthCheckResponse {
+                    healthy: resp.healthy,
+                    message: resp.message,
+                })
+            }
+            Self::Vllm(client) => {
+                let resp = client.health_check().await?;
+                Ok(HealthCheckResponse {
+                    healthy: resp.healthy,
+                    message: resp.message,
+                })
+            }
+        }
+    }
+
+    /// Get model info (returns enum wrapping backend-specific response)
+    pub async fn get_model_info(
+        &self,
+    ) -> Result<ModelInfo, Box<dyn std::error::Error + Send + Sync>> {
+        match self {
+            Self::Sglang(client) => {
+                let info = client.get_model_info().await?;
+                Ok(ModelInfo::Sglang(Box::new(info)))
+            }
+            Self::Vllm(client) => {
+                let info = client.get_model_info().await?;
+                Ok(ModelInfo::Vllm(info))
+            }
+        }
+    }
+
+    /// Generate streaming response from request
+    ///
+    /// Dispatches to the appropriate backend client and wraps the result in ProtoStream
+    pub async fn generate(
+        &mut self,
+        req: ProtoGenerateRequest,
+    ) -> Result<ProtoStream, Box<dyn std::error::Error + Send + Sync>> {
+        match (self, req) {
+            (Self::Sglang(client), ProtoGenerateRequest::Sglang(boxed_req)) => {
+                let stream = client.generate(*boxed_req).await?;
+                Ok(ProtoStream::Sglang(stream))
+            }
+            (Self::Vllm(client), ProtoGenerateRequest::Vllm(boxed_req)) => {
+                let stream = client.generate(*boxed_req).await?;
+                Ok(ProtoStream::Vllm(stream))
+            }
+            _ => panic!("Mismatched client and request types"),
+        }
+    }
+
+    /// Submit an embedding request
+    pub async fn embed(
+        &mut self,
+        req: ProtoEmbedRequest,
+    ) -> Result<ProtoEmbedResponse, Box<dyn std::error::Error + Send + Sync>> {
+        match (self, req) {
+            (Self::Sglang(client), ProtoEmbedRequest::Sglang(boxed_req)) => {
+                let resp = client.embed(*boxed_req).await?;
+                Ok(ProtoEmbedResponse::Sglang(resp))
+            }
+            _ => panic!("Mismatched client and request types or unsupported embedding backend"),
+        }
+    }
+}
+
+/// Unified ModelInfo wrapper
+pub enum ModelInfo {
+    Sglang(Box<smg_grpc_client::sglang_proto::GetModelInfoResponse>),
+    Vllm(smg_grpc_client::vllm_proto::GetModelInfoResponse),
+}
+
+impl ModelInfo {
+    /// Convert model info to label map for worker metadata
+    pub fn to_labels(&self) -> std::collections::HashMap<String, String> {
+        let mut labels = std::collections::HashMap::new();
+
+        // Serialize to JSON Value (like pydantic's model_dump)
+        let value = match self {
+            ModelInfo::Sglang(info) => serde_json::to_value(info).ok(),
+            ModelInfo::Vllm(info) => serde_json::to_value(info).ok(),
+        };
+
+        // Convert JSON object to HashMap, filtering out empty/zero/false values
+        if let Some(serde_json::Value::Object(obj)) = value {
+            for (key, val) in obj {
+                match val {
+                    // Insert non-empty strings
+                    serde_json::Value::String(s) if !s.is_empty() => {
+                        labels.insert(key, s);
+                    }
+                    // Insert positive numbers
+                    serde_json::Value::Number(n) if n.as_i64().unwrap_or(0) > 0 => {
+                        labels.insert(key, n.to_string());
+                    }
+                    // Insert true booleans
+                    serde_json::Value::Bool(true) => {
+                        labels.insert(key, "true".to_string());
+                    }
+                    // Insert non-empty arrays as JSON strings (for architectures, etc.)
+                    serde_json::Value::Array(arr) if !arr.is_empty() => {
+                        if let Ok(json_str) = serde_json::to_string(&arr) {
+                            labels.insert(key, json_str);
+                        }
+                    }
+                    // Skip empty strings, zeros, false, nulls, empty arrays, objects
+                    _ => {}
+                }
+            }
+        }
+
+        labels
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/mod.rs

@@ -0,0 +1,6 @@
+//! Shared code for both regular and harmony routers
+
+pub(crate) mod response_collection;
+pub(crate) mod response_formatting;
+pub(crate) mod responses;
+pub(crate) mod stages;

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/response_collection.rs

@@ -0,0 +1,98 @@
+//! Shared response collection logic
+//!
+//! This module contains common logic for collecting responses from execution results.
+//! Both regular and harmony processors use these functions to avoid duplication.
+
+use axum::response::Response;
+
+use crate::routers::{
+    error,
+    grpc::{context::ExecutionResult, proto_wrapper::ProtoGenerateComplete, utils},
+};
+
+/// Collect and merge responses from execution result
+///
+/// Handles both Single and Dual (prefill-decode) execution modes.
+/// For Dual mode, merges prefill input_logprobs into decode responses if requested.
+///
+/// # Arguments
+/// * `execution_result` - The execution result containing stream(s)
+/// * `merge_logprobs` - Whether to merge prefill input_logprobs (for chat with logprobs=true)
+///
+/// # Returns
+/// Vector of GenerateComplete responses, one per index (n parameter)
+pub(crate) async fn collect_responses(
+    execution_result: ExecutionResult,
+    merge_logprobs: bool,
+) -> Result<Vec<ProtoGenerateComplete>, Response> {
+    let all_responses = match execution_result {
+        ExecutionResult::Single { mut stream } => {
+            let responses = utils::collect_stream_responses(&mut stream, "Single").await?;
+            stream.mark_completed();
+            responses
+        }
+        ExecutionResult::Dual {
+            mut prefill,
+            decode,
+        } => {
+            // Collect prefill for input_logprobs (don't mark completed yet)
+            let prefill_responses =
+                utils::collect_stream_responses(&mut prefill, "Prefill").await?;
+
+            // Collect decode for actual output (don't mark completed yet)
+            let mut decode_stream = *decode;
+            let mut decode_responses =
+                utils::collect_stream_responses(&mut decode_stream, "Decode").await?;
+
+            // Mark both streams as completed now that both succeeded
+            prefill.mark_completed();
+            decode_stream.mark_completed();
+
+            // Merge prefill input_logprobs if requested
+            if merge_logprobs {
+                merge_prefill_logprobs(&prefill_responses, &mut decode_responses);
+            }
+
+            decode_responses
+        }
+        ExecutionResult::Embedding { .. } => {
+            // Embeddings do not support this path (no generate complete response)
+            return Err(error::internal_error(
+                "invalid_execution_mode",
+                "Embedding result encountered in response collection",
+            ));
+        }
+    };
+
+    if all_responses.is_empty() {
+        return Err(error::internal_error(
+            "no_responses_from_server",
+            "No responses from server",
+        ));
+    }
+
+    Ok(all_responses)
+}
+
+/// Merge prefill input_logprobs into decode responses
+///
+/// Takes input_logprobs from the first prefill response and copies them
+/// into all decode responses. This is used in PD mode when logprobs are requested.
+/// Only works with SGLang (vLLM doesn't support PD mode).
+fn merge_prefill_logprobs(
+    prefill_responses: &[ProtoGenerateComplete],
+    decode_responses: &mut [ProtoGenerateComplete],
+) {
+    // Only SGLang supports PD mode and has input_logprobs
+    if let Some(ProtoGenerateComplete::Sglang(prefill_first)) = prefill_responses.first() {
+        // Use ref to borrow input_logprobs instead of cloning upfront
+        // This avoids one allocation when the Option is Some
+        if let Some(ref prefill_input_logprobs) = prefill_first.input_logprobs {
+            for response in decode_responses.iter_mut() {
+                if let ProtoGenerateComplete::Sglang(decode_resp) = response {
+                    decode_resp.input_logprobs = Some(prefill_input_logprobs.clone());
+                }
+            }
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/response_formatting.rs

@@ -0,0 +1,29 @@
+//! Shared response formatting logic
+//!
+//! This module contains common logic for formatting responses, including:
+//! - Usage calculation from gRPC responses
+//! - ChatCompletionResponse construction
+
+use crate::{protocols::common::Usage, routers::grpc::proto_wrapper::ProtoGenerateComplete};
+
+/// Build usage information from collected gRPC responses
+///
+/// Sums prompt_tokens and completion_tokens across all responses.
+/// Typically used with n>1 parameter where multiple completions are generated.
+///
+/// # Arguments
+/// * `responses` - Vector of GenerateComplete responses from the backend
+///
+/// # Returns
+/// Usage object with aggregated token counts
+pub(crate) fn build_usage(responses: &[ProtoGenerateComplete]) -> Usage {
+    let total_prompt_tokens: u32 = responses.iter().map(|r| r.prompt_tokens() as u32).sum();
+    let total_completion_tokens: u32 = responses.iter().map(|r| r.completion_tokens() as u32).sum();
+
+    Usage {
+        prompt_tokens: total_prompt_tokens,
+        completion_tokens: total_completion_tokens,
+        total_tokens: total_prompt_tokens + total_completion_tokens,
+        completion_tokens_details: None,
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/responses/context.rs

@@ -0,0 +1,60 @@
+//! Shared context for /v1/responses endpoint handlers
+//!
+//! This context is used by both regular and harmony response implementations.
+
+use std::sync::{Arc, RwLock as StdRwLock};
+
+use data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage};
+use smg_mcp::McpManager;
+
+use crate::routers::grpc::{context::SharedComponents, pipeline::RequestPipeline};
+
+/// Context for /v1/responses endpoint
+///
+/// Used by both regular and harmony implementations.
+/// All fields are Arc/shared references, so cloning this context is cheap.
+#[derive(Clone)]
+pub(crate) struct ResponsesContext {
+    /// Chat pipeline for executing requests
+    pub pipeline: Arc<RequestPipeline>,
+
+    /// Shared components (tokenizer, parsers)
+    pub components: Arc<SharedComponents>,
+
+    /// Response storage backend
+    pub response_storage: Arc<dyn ResponseStorage>,
+
+    /// Conversation storage backend
+    pub conversation_storage: Arc<dyn ConversationStorage>,
+
+    /// Conversation item storage backend
+    pub conversation_item_storage: Arc<dyn ConversationItemStorage>,
+
+    /// MCP manager for tool support
+    pub mcp_manager: Arc<McpManager>,
+
+    /// Server keys for MCP tools requested in this context
+    pub requested_servers: Arc<StdRwLock<Vec<String>>>,
+}
+
+impl ResponsesContext {
+    /// Create a new responses context
+    pub fn new(
+        pipeline: Arc<RequestPipeline>,
+        components: Arc<SharedComponents>,
+        response_storage: Arc<dyn ResponseStorage>,
+        conversation_storage: Arc<dyn ConversationStorage>,
+        conversation_item_storage: Arc<dyn ConversationItemStorage>,
+        mcp_manager: Arc<McpManager>,
+    ) -> Self {
+        Self {
+            pipeline,
+            components,
+            response_storage,
+            conversation_storage,
+            conversation_item_storage,
+            mcp_manager,
+            requested_servers: Arc::new(StdRwLock::new(Vec::new())),
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs

@@ -0,0 +1,74 @@
+//! Shared response handlers for both regular and harmony implementations
+//!
+//! These handlers are used by both pipelines for retrieving and cancelling responses.
+
+use axum::response::{IntoResponse, Response};
+use data_connector::ResponseId;
+
+use super::ResponsesContext;
+use crate::routers::error;
+
+/// Implementation for GET /v1/responses/{response_id}
+///
+/// Retrieves a stored response from the database.
+/// Used by both regular and harmony implementations.
+pub(crate) async fn get_response_impl(ctx: &ResponsesContext, response_id: &str) -> Response {
+    let resp_id = ResponseId::from(response_id);
+
+    // Retrieve response from storage
+    match ctx.response_storage.get_response(&resp_id).await {
+        Ok(Some(stored_response)) => axum::Json(stored_response.raw_response).into_response(),
+        Ok(None) => error::not_found(
+            "response_not_found",
+            format!("Response with id '{}' not found", response_id),
+        ),
+        Err(e) => error::internal_error(
+            "retrieve_response_failed",
+            format!("Failed to retrieve response: {}", e),
+        ),
+    }
+}
+
+/// Implementation for POST /v1/responses/{response_id}/cancel
+///
+/// Background mode is no longer supported, so this endpoint always returns
+/// an error indicating that cancellation is not available.
+pub(crate) async fn cancel_response_impl(ctx: &ResponsesContext, response_id: &str) -> Response {
+    let resp_id = ResponseId::from(response_id);
+
+    // Check if response exists
+    match ctx.response_storage.get_response(&resp_id).await {
+        Ok(Some(stored_response)) => {
+            let current_status = stored_response
+                .raw_response
+                .get("status")
+                .and_then(|v| v.as_str())
+                .unwrap_or("unknown");
+
+            match current_status {
+                "completed" => error::bad_request(
+                    "response_already_completed",
+                    "Cannot cancel completed response",
+                ),
+                "failed" => {
+                    error::bad_request("response_already_failed", "Cannot cancel failed response")
+                }
+                _ => {
+                    // Background mode is no longer supported, so there's nothing to cancel
+                    error::bad_request(
+                        "cancellation_not_supported",
+                        "Background mode is not supported. Synchronous and streaming responses cannot be cancelled.",
+                    )
+                }
+            }
+        }
+        Ok(None) => error::not_found(
+            "response_not_found",
+            format!("Response with id '{}' not found", response_id),
+        ),
+        Err(e) => error::internal_error(
+            "retrieve_response_failed",
+            format!("Failed to retrieve response: {}", e),
+        ),
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/responses/mod.rs

@@ -0,0 +1,11 @@
+//! Shared response functionality used by both regular and harmony implementations
+
+pub(crate) mod context;
+pub(crate) mod handlers;
+pub(crate) mod streaming;
+pub(crate) mod utils;
+
+// Re-export commonly used items
+pub(crate) use context::ResponsesContext;
+pub(crate) use streaming::build_sse_response;
+pub(crate) use utils::{ensure_mcp_connection, persist_response_if_needed};

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/responses/streaming.rs

@@ -0,0 +1,842 @@
+//! Streaming infrastructure for /v1/responses endpoint
+
+use std::collections::HashMap;
+
+use axum::{body::Body, http::StatusCode, response::Response};
+use bytes::Bytes;
+use serde_json::json;
+use smg_mcp as mcp;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::UnboundedReceiverStream;
+use uuid::Uuid;
+
+use crate::{
+    protocols::{
+        chat::ChatCompletionStreamResponse,
+        common::{Usage, UsageInfo},
+        event_types::{
+            ContentPartEvent, FunctionCallEvent, McpEvent, OutputItemEvent, OutputTextEvent,
+            ResponseEvent,
+        },
+        responses::{
+            ResponseOutputItem, ResponseStatus, ResponsesRequest, ResponsesResponse, ResponsesUsage,
+        },
+    },
+    routers::grpc::harmony::responses::ToolResult,
+};
+
+pub(crate) enum OutputItemType {
+    Message,
+    McpListTools,
+    McpCall,
+    FunctionCall,
+    Reasoning,
+}
+
+/// Status of an output item
+#[derive(Debug, Clone, PartialEq)]
+enum ItemStatus {
+    InProgress,
+    Completed,
+}
+
+/// State tracking for a single output item
+#[derive(Debug, Clone)]
+struct OutputItemState {
+    output_index: usize,
+    status: ItemStatus,
+    item_data: Option<serde_json::Value>,
+}
+
+/// OpenAI-compatible event emitter for /v1/responses streaming
+///
+/// Manages state and sequence numbers to emit proper event types:
+/// - response.created
+/// - response.in_progress
+/// - response.output_item.added
+/// - response.content_part.added
+/// - response.output_text.delta (multiple)
+/// - response.output_text.done
+/// - response.content_part.done
+/// - response.output_item.done
+/// - response.completed
+/// - response.mcp_list_tools.in_progress
+/// - response.mcp_list_tools.completed
+/// - response.mcp_call.in_progress
+/// - response.mcp_call_arguments.delta
+/// - response.mcp_call_arguments.done
+/// - response.mcp_call.completed
+/// - response.mcp_call.failed
+pub(crate) struct ResponseStreamEventEmitter {
+    sequence_number: u64,
+    pub response_id: String,
+    model: String,
+    created_at: u64,
+    message_id: String,
+    accumulated_text: String,
+    has_emitted_created: bool,
+    has_emitted_in_progress: bool,
+    has_emitted_output_item_added: bool,
+    has_emitted_content_part_added: bool,
+    // MCP call tracking
+    mcp_call_accumulated_args: HashMap<String, String>,
+    pub(crate) mcp_server_label: Option<String>, // Server label for MCP tools
+    // Output item tracking
+    output_items: Vec<OutputItemState>,
+    next_output_index: usize,
+    current_message_output_index: Option<usize>, // Tracks output_index of current message
+    current_item_id: Option<String>,             // Tracks item_id of current item
+    original_request: Option<ResponsesRequest>,
+}
+
+impl ResponseStreamEventEmitter {
+    pub fn new(response_id: String, model: String, created_at: u64) -> Self {
+        let message_id = format!("msg_{}", Uuid::new_v4());
+
+        Self {
+            sequence_number: 0,
+            response_id,
+            model,
+            created_at,
+            message_id,
+            accumulated_text: String::new(),
+            has_emitted_created: false,
+            has_emitted_in_progress: false,
+            has_emitted_output_item_added: false,
+            has_emitted_content_part_added: false,
+            mcp_call_accumulated_args: HashMap::new(),
+            mcp_server_label: None,
+            output_items: Vec::new(),
+            next_output_index: 0,
+            current_message_output_index: None,
+            current_item_id: None,
+            original_request: None,
+        }
+    }
+
+    /// Set the original request for including all fields in response.completed
+    pub fn set_original_request(&mut self, request: ResponsesRequest) {
+        self.original_request = Some(request);
+    }
+
+    /// Set the MCP server label for MCP tool calls
+    pub fn set_mcp_server_label(&mut self, server_label: String) {
+        self.mcp_server_label = Some(server_label);
+    }
+
+    /// Update mcp_call output items with tool execution results
+    ///
+    /// After MCP tools are executed, this updates the stored output items
+    /// to include the output field from the tool results.
+    pub(crate) fn update_mcp_call_outputs(&mut self, tool_results: &[ToolResult]) {
+        for tool_result in tool_results {
+            // Find the output item with matching call_id
+            for item_state in self.output_items.iter_mut() {
+                if let Some(ref mut item_data) = item_state.item_data {
+                    // Check if this is an mcp_call item with matching call_id
+                    if item_data.get("type").and_then(|t| t.as_str()) == Some("mcp_call")
+                        && item_data.get("call_id").and_then(|c| c.as_str())
+                            == Some(&tool_result.call_id)
+                    {
+                        // Add output field
+                        let output_str = serde_json::to_string(&tool_result.output)
+                            .unwrap_or_else(|_| "{}".to_string());
+                        item_data["output"] = json!(output_str);
+
+                        // Update status based on success
+                        if tool_result.is_error {
+                            item_data["status"] = json!("failed");
+                        }
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    fn next_sequence(&mut self) -> u64 {
+        let seq = self.sequence_number;
+        self.sequence_number += 1;
+        seq
+    }
+
+    pub fn emit_created(&mut self) -> serde_json::Value {
+        self.has_emitted_created = true;
+        json!({
+            "type": ResponseEvent::CREATED,
+            "sequence_number": self.next_sequence(),
+            "response": {
+                "id": self.response_id,
+                "object": "response",
+                "created_at": self.created_at,
+                "status": "in_progress",
+                "model": self.model,
+                "output": []
+            }
+        })
+    }
+
+    pub fn emit_in_progress(&mut self) -> serde_json::Value {
+        self.has_emitted_in_progress = true;
+        json!({
+            "type": ResponseEvent::IN_PROGRESS,
+            "sequence_number": self.next_sequence(),
+            "response": {
+                "id": self.response_id,
+                "object": "response",
+                "status": "in_progress"
+            }
+        })
+    }
+
+    pub fn emit_content_part_added(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        content_index: usize,
+    ) -> serde_json::Value {
+        self.has_emitted_content_part_added = true;
+        json!({
+            "type": ContentPartEvent::ADDED,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "content_index": content_index,
+            "part": {
+                "type": "text",
+                "text": ""
+            }
+        })
+    }
+
+    pub fn emit_text_delta(
+        &mut self,
+        delta: &str,
+        output_index: usize,
+        item_id: &str,
+        content_index: usize,
+    ) -> serde_json::Value {
+        self.accumulated_text.push_str(delta);
+        json!({
+            "type": OutputTextEvent::DELTA,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "content_index": content_index,
+            "delta": delta
+        })
+    }
+
+    pub fn emit_text_done(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        content_index: usize,
+    ) -> serde_json::Value {
+        json!({
+            "type": OutputTextEvent::DONE,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "content_index": content_index,
+            "text": self.accumulated_text.clone()
+        })
+    }
+
+    pub fn emit_content_part_done(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        content_index: usize,
+    ) -> serde_json::Value {
+        json!({
+            "type": ContentPartEvent::DONE,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "content_index": content_index,
+            "part": {
+                "type": "text",
+                "text": self.accumulated_text.clone()
+            }
+        })
+    }
+
+    pub fn emit_completed(&mut self, usage: Option<&serde_json::Value>) -> serde_json::Value {
+        // Build output array from tracked items
+        let output: Vec<serde_json::Value> = self
+            .output_items
+            .iter()
+            .filter_map(|item| {
+                if item.status == ItemStatus::Completed {
+                    item.item_data.clone()
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        // If no items were tracked (legacy path), fall back to generic message
+        let output = if output.is_empty() {
+            vec![json!({
+                "id": self.message_id.clone(),
+                "type": "message",
+                "role": "assistant",
+                "content": [{
+                    "type": "text",
+                    "text": self.accumulated_text.clone()
+                }]
+            })]
+        } else {
+            output
+        };
+
+        // Build base response object
+        let mut response_obj = json!({
+            "id": self.response_id,
+            "object": "response",
+            "created_at": self.created_at,
+            "status": "completed",
+            "model": self.model,
+            "output": output
+        });
+
+        // Add usage if provided
+        if let Some(usage_val) = usage {
+            response_obj["usage"] = usage_val.clone();
+        }
+
+        // Add all original request fields if available
+        if let Some(ref req) = self.original_request {
+            Self::add_optional_field(&mut response_obj, "instructions", &req.instructions);
+            Self::add_optional_field(
+                &mut response_obj,
+                "max_output_tokens",
+                &req.max_output_tokens,
+            );
+            Self::add_optional_field(&mut response_obj, "max_tool_calls", &req.max_tool_calls);
+            Self::add_optional_field(
+                &mut response_obj,
+                "previous_response_id",
+                &req.previous_response_id,
+            );
+            Self::add_optional_field(&mut response_obj, "reasoning", &req.reasoning);
+            Self::add_optional_field(&mut response_obj, "temperature", &req.temperature);
+            Self::add_optional_field(&mut response_obj, "top_p", &req.top_p);
+            Self::add_optional_field(&mut response_obj, "truncation", &req.truncation);
+            Self::add_optional_field(&mut response_obj, "user", &req.user);
+
+            response_obj["parallel_tool_calls"] = json!(req.parallel_tool_calls.unwrap_or(true));
+            response_obj["store"] = json!(req.store.unwrap_or(true));
+            response_obj["tools"] = json!(req.tools.as_ref().unwrap_or(&vec![]));
+            response_obj["metadata"] = json!(req.metadata.as_ref().unwrap_or(&Default::default()));
+
+            // tool_choice: serialize if present, otherwise use "auto"
+            if let Some(ref tc) = req.tool_choice {
+                response_obj["tool_choice"] = json!(tc);
+            } else {
+                response_obj["tool_choice"] = json!("auto");
+            }
+        }
+
+        json!({
+            "type": ResponseEvent::COMPLETED,
+            "sequence_number": self.next_sequence(),
+            "response": response_obj
+        })
+    }
+
+    /// Helper to add optional fields to JSON object
+    fn add_optional_field<T: serde::Serialize>(
+        obj: &mut serde_json::Value,
+        key: &str,
+        value: &Option<T>,
+    ) {
+        if let Some(val) = value {
+            obj[key] = json!(val);
+        }
+    }
+
+    // ========================================================================
+    // MCP Event Emission Methods
+    // ========================================================================
+
+    pub fn emit_mcp_list_tools_in_progress(&mut self, output_index: usize) -> serde_json::Value {
+        json!({
+            "type": McpEvent::LIST_TOOLS_IN_PROGRESS,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index
+        })
+    }
+
+    pub fn emit_mcp_list_tools_completed(
+        &mut self,
+        output_index: usize,
+        tools: &[mcp::Tool],
+    ) -> serde_json::Value {
+        let tool_items: Vec<_> = tools
+            .iter()
+            .map(|t| {
+                json!({
+                    "name": &t.name,
+                    "description": &t.description,
+                    "input_schema": t.input_schema.clone()
+                })
+            })
+            .collect();
+
+        json!({
+            "type": McpEvent::LIST_TOOLS_COMPLETED,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "tools": tool_items
+        })
+    }
+
+    pub fn emit_mcp_call_in_progress(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+    ) -> serde_json::Value {
+        json!({
+            "type": McpEvent::CALL_IN_PROGRESS,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id
+        })
+    }
+
+    pub fn emit_mcp_call_arguments_delta(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        delta: &str,
+    ) -> serde_json::Value {
+        // Accumulate arguments for this call
+        self.mcp_call_accumulated_args
+            .entry(item_id.to_string())
+            .or_default()
+            .push_str(delta);
+
+        json!({
+            "type": McpEvent::CALL_ARGUMENTS_DELTA,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "delta": delta
+        })
+    }
+
+    pub fn emit_mcp_call_arguments_done(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        arguments: &str,
+    ) -> serde_json::Value {
+        json!({
+            "type": McpEvent::CALL_ARGUMENTS_DONE,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "arguments": arguments
+        })
+    }
+
+    pub fn emit_mcp_call_completed(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+    ) -> serde_json::Value {
+        json!({
+            "type": McpEvent::CALL_COMPLETED,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id
+        })
+    }
+
+    pub fn emit_mcp_call_failed(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        error: &str,
+    ) -> serde_json::Value {
+        json!({
+            "type": McpEvent::CALL_FAILED,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "error": error
+        })
+    }
+
+    // ========================================================================
+    // Function Call Event Emission Methods
+    // ========================================================================
+
+    pub fn emit_function_call_arguments_delta(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        delta: &str,
+    ) -> serde_json::Value {
+        json!({
+            "type": FunctionCallEvent::ARGUMENTS_DELTA,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "delta": delta
+        })
+    }
+
+    pub fn emit_function_call_arguments_done(
+        &mut self,
+        output_index: usize,
+        item_id: &str,
+        arguments: &str,
+    ) -> serde_json::Value {
+        json!({
+            "type": FunctionCallEvent::ARGUMENTS_DONE,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item_id": item_id,
+            "arguments": arguments
+        })
+    }
+
+    // ========================================================================
+    // Output Item Wrapper Events
+    // ========================================================================
+
+    /// Emit response.output_item.added event
+    pub fn emit_output_item_added(
+        &mut self,
+        output_index: usize,
+        item: &serde_json::Value,
+    ) -> serde_json::Value {
+        json!({
+            "type": OutputItemEvent::ADDED,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item": item
+        })
+    }
+
+    /// Emit response.output_item.done event
+    pub fn emit_output_item_done(
+        &mut self,
+        output_index: usize,
+        item: &serde_json::Value,
+    ) -> serde_json::Value {
+        // Store the item data for later use in emit_completed
+        self.store_output_item_data(output_index, item.clone());
+
+        json!({
+            "type": OutputItemEvent::DONE,
+            "sequence_number": self.next_sequence(),
+            "output_index": output_index,
+            "item": item
+        })
+    }
+
+    /// Generate unique ID for item type
+    fn generate_item_id(prefix: &str) -> String {
+        format!("{}_{}", prefix, Uuid::new_v4().to_string().replace("-", ""))
+    }
+
+    /// Allocate next output index and track item
+    pub fn allocate_output_index(&mut self, item_type: OutputItemType) -> (usize, String) {
+        let index = self.next_output_index;
+        self.next_output_index += 1;
+
+        let id_prefix = match &item_type {
+            OutputItemType::McpListTools => "mcpl",
+            OutputItemType::McpCall => "mcp",
+            OutputItemType::FunctionCall => "fc",
+            OutputItemType::Message => "msg",
+            OutputItemType::Reasoning => "rs",
+        };
+
+        let id = Self::generate_item_id(id_prefix);
+
+        self.output_items.push(OutputItemState {
+            output_index: index,
+            status: ItemStatus::InProgress,
+            item_data: None,
+        });
+
+        (index, id)
+    }
+
+    /// Mark output item as completed and store its data
+    pub fn complete_output_item(&mut self, output_index: usize) {
+        if let Some(item) = self
+            .output_items
+            .iter_mut()
+            .find(|i| i.output_index == output_index)
+        {
+            item.status = ItemStatus::Completed;
+        }
+    }
+
+    /// Store output item data when emitting output_item.done
+    pub fn store_output_item_data(&mut self, output_index: usize, item_data: serde_json::Value) {
+        if let Some(item) = self
+            .output_items
+            .iter_mut()
+            .find(|i| i.output_index == output_index)
+        {
+            item.item_data = Some(item_data);
+        }
+    }
+
+    /// Finalize and return the complete ResponsesResponse
+    ///
+    /// This constructs the final ResponsesResponse from all accumulated output items
+    /// for persistence. Should be called after streaming is complete.
+    pub fn finalize(&self, usage: Option<Usage>) -> ResponsesResponse {
+        // Build output array from tracked items
+        let output: Vec<ResponseOutputItem> = self
+            .output_items
+            .iter()
+            .filter_map(|item| {
+                item.item_data
+                    .as_ref()
+                    .and_then(|data| serde_json::from_value(data.clone()).ok())
+            })
+            .collect();
+
+        // Convert Usage to ResponsesUsage
+        let responses_usage = usage.map(|u| {
+            let usage_info = UsageInfo {
+                prompt_tokens: u.prompt_tokens,
+                completion_tokens: u.completion_tokens,
+                total_tokens: u.total_tokens,
+                reasoning_tokens: u
+                    .completion_tokens_details
+                    .as_ref()
+                    .and_then(|d| d.reasoning_tokens),
+                prompt_tokens_details: None,
+            };
+            ResponsesUsage::Classic(usage_info)
+        });
+
+        // Build response using builder
+        ResponsesResponse::builder(&self.response_id, &self.model)
+            .created_at(self.created_at as i64)
+            .status(ResponseStatus::Completed)
+            .output(output)
+            .maybe_copy_from_request(self.original_request.as_ref())
+            .maybe_usage(responses_usage)
+            .build()
+    }
+
+    /// Emit reasoning item wrapper events (added + done)
+    ///
+    /// Reasoning items in OpenAI format are simple placeholders emitted between tool iterations.
+    /// They don't have streaming content - just wrapper events with empty/null content.
+    pub fn emit_reasoning_item(
+        &mut self,
+        tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+        reasoning_content: Option<String>,
+    ) -> Result<(), String> {
+        // Allocate output index and generate ID
+        let (output_index, item_id) = self.allocate_output_index(OutputItemType::Reasoning);
+
+        // Build reasoning item structure
+        let item = json!({
+            "id": item_id,
+            "type": "reasoning",
+            "summary": [],
+            "content": reasoning_content,
+            "encrypted_content": null,
+            "status": null
+        });
+
+        // Emit output_item.added
+        let added_event = self.emit_output_item_added(output_index, &item);
+        self.send_event(&added_event, tx)?;
+
+        // Immediately emit output_item.done (no streaming for reasoning)
+        let done_event = self.emit_output_item_done(output_index, &item);
+        self.send_event(&done_event, tx)?;
+
+        // Mark as completed
+        self.complete_output_item(output_index);
+
+        Ok(())
+    }
+
+    /// Process a chunk and emit appropriate events
+    pub fn process_chunk(
+        &mut self,
+        chunk: &ChatCompletionStreamResponse,
+        tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+    ) -> Result<(), String> {
+        // Process content if present
+        if let Some(choice) = chunk.choices.first() {
+            if let Some(content) = &choice.delta.content {
+                if !content.is_empty() {
+                    // Allocate output_index and item_id for this message item (once per message)
+                    if self.current_item_id.is_none() {
+                        let (output_index, item_id) =
+                            self.allocate_output_index(OutputItemType::Message);
+
+                        // Build message item structure
+                        let item = json!({
+                            "id": item_id,
+                            "type": "message",
+                            "role": "assistant",
+                            "content": []
+                        });
+
+                        // Emit output_item.added
+                        let event = self.emit_output_item_added(output_index, &item);
+                        self.send_event(&event, tx)?;
+                        self.has_emitted_output_item_added = true;
+
+                        // Store for subsequent events
+                        self.current_item_id = Some(item_id);
+                        self.current_message_output_index = Some(output_index);
+                    }
+
+                    let output_index = self.current_message_output_index.unwrap();
+                    let item_id = self.current_item_id.clone().unwrap(); // Clone to avoid borrow checker issues
+                    let content_index = 0; // Single content part for now
+
+                    // Emit content_part.added before first delta
+                    if !self.has_emitted_content_part_added {
+                        let event =
+                            self.emit_content_part_added(output_index, &item_id, content_index);
+                        self.send_event(&event, tx)?;
+                        self.has_emitted_content_part_added = true;
+                    }
+
+                    // Emit text delta
+                    let event =
+                        self.emit_text_delta(content, output_index, &item_id, content_index);
+                    self.send_event(&event, tx)?;
+                }
+            }
+
+            // Check for finish_reason to emit completion events
+            if let Some(reason) = &choice.finish_reason {
+                if reason == "stop" || reason == "length" {
+                    let output_index = self.current_message_output_index.unwrap();
+                    let item_id = self.current_item_id.clone().unwrap(); // Clone to avoid borrow checker issues
+                    let content_index = 0;
+
+                    // Emit closing events
+                    if self.has_emitted_content_part_added {
+                        let event = self.emit_text_done(output_index, &item_id, content_index);
+                        self.send_event(&event, tx)?;
+                        let event =
+                            self.emit_content_part_done(output_index, &item_id, content_index);
+                        self.send_event(&event, tx)?;
+                    }
+
+                    if self.has_emitted_output_item_added {
+                        // Build complete message item for output_item.done
+                        let item = json!({
+                            "id": item_id,
+                            "type": "message",
+                            "role": "assistant",
+                            "content": [{
+                                "type": "text",
+                                "text": self.accumulated_text.clone()
+                            }]
+                        });
+                        let event = self.emit_output_item_done(output_index, &item);
+                        self.send_event(&event, tx)?;
+                    }
+
+                    // Mark item as completed
+                    self.complete_output_item(output_index);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn send_event(
+        &self,
+        event: &serde_json::Value,
+        tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+    ) -> Result<(), String> {
+        let event_json = serde_json::to_string(event)
+            .map_err(|e| format!("Failed to serialize event: {}", e))?;
+
+        // Extract event type from the JSON for SSE event field
+        let event_type = event
+            .get("type")
+            .and_then(|v| v.as_str())
+            .unwrap_or("message");
+
+        // Format as SSE with event: field
+        let sse_message = format!("event: {}\ndata: {}\n\n", event_type, event_json);
+
+        if tx.send(Ok(Bytes::from(sse_message))).is_err() {
+            return Err("Client disconnected".to_string());
+        }
+
+        Ok(())
+    }
+
+    /// Send event and log any errors (typically client disconnect)
+    ///
+    /// This is a convenience method for streaming scenarios where client
+    /// disconnection is expected and should be logged but not fail the operation.
+    /// Returns true if sent successfully, false if client disconnected.
+    pub fn send_event_best_effort(
+        &self,
+        event: &serde_json::Value,
+        tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+    ) -> bool {
+        match self.send_event(event, tx) {
+            Ok(()) => true,
+            Err(e) => {
+                tracing::debug!("Failed to send event (likely client disconnect): {}", e);
+                false
+            }
+        }
+    }
+
+    /// Emit an error event
+    ///
+    /// Creates and sends an error event with the given error message.
+    /// Uses OpenAI's error event format.
+    /// Use this for terminal errors that should abort the streaming response.
+    pub fn emit_error(
+        &mut self,
+        error_msg: &str,
+        error_code: Option<&str>,
+        tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+    ) {
+        let event = json!({
+            "type": "error",
+            "code": error_code.unwrap_or("internal_error"),
+            "message": error_msg,
+            "param": null,
+            "sequence_number": self.next_sequence()
+        });
+        let sse_data = format!("data: {}\n\n", serde_json::to_string(&event).unwrap());
+        let _ = tx.send(Ok(Bytes::from(sse_data)));
+    }
+}
+
+/// Build a Server-Sent Events (SSE) response
+///
+/// Creates a Response with proper SSE headers and streaming body.
+pub(crate) fn build_sse_response(
+    rx: mpsc::UnboundedReceiver<Result<Bytes, std::io::Error>>,
+) -> Response {
+    let stream = UnboundedReceiverStream::new(rx);
+    Response::builder()
+        .status(StatusCode::OK)
+        .header("Content-Type", "text/event-stream")
+        .header("Cache-Control", "no-cache")
+        .header("Connection", "keep-alive")
+        .body(Body::from_stream(stream))
+        .unwrap()
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs

@@ -0,0 +1,156 @@
+//! Utility functions for /v1/responses endpoint
+
+use std::sync::Arc;
+
+use axum::response::Response;
+use data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage};
+use serde_json::to_value;
+use smg_mcp::McpManager;
+use tracing::{debug, error, warn};
+
+use crate::{
+    core::WorkerRegistry,
+    protocols::{
+        common::Tool,
+        responses::{ResponseTool, ResponseToolType, ResponsesRequest, ResponsesResponse},
+    },
+    routers::{
+        error, mcp_utils::ensure_request_mcp_client, persistence_utils::persist_conversation_items,
+    },
+};
+
+/// Ensure MCP connection succeeds if MCP tools are declared
+///
+/// Checks if request declares MCP tools, and if so, validates that
+/// the MCP clients can be created and connected.
+/// Returns Ok((has_mcp_tools, server_keys)) on success.
+pub(crate) async fn ensure_mcp_connection(
+    mcp_manager: &Arc<McpManager>,
+    tools: Option<&[ResponseTool]>,
+) -> Result<(bool, Vec<String>), Response> {
+    let has_mcp_tools = tools
+        .map(|t| {
+            t.iter()
+                .any(|tool| matches!(tool.r#type, ResponseToolType::Mcp))
+        })
+        .unwrap_or(false);
+
+    if has_mcp_tools {
+        if let Some(tools) = tools {
+            match ensure_request_mcp_client(mcp_manager, tools).await {
+                Some((_manager, server_keys)) => {
+                    return Ok((true, server_keys));
+                }
+                None => {
+                    error!(
+                        function = "ensure_mcp_connection",
+                        "Failed to connect to MCP servers"
+                    );
+                    return Err(error::failed_dependency(
+                        "connect_mcp_server_failed",
+                        "Failed to connect to MCP servers. Check server_url and authorization.",
+                    ));
+                }
+            }
+        }
+    }
+
+    Ok((false, Vec::new()))
+}
+
+/// Validate that workers are available for the requested model
+pub(crate) fn validate_worker_availability(
+    worker_registry: &Arc<WorkerRegistry>,
+    model: &str,
+) -> Option<Response> {
+    let available_models = worker_registry.get_models();
+
+    if !available_models.contains(&model.to_string()) {
+        return Some(error::service_unavailable(
+            "no_available_workers",
+            format!(
+                "No workers available for model '{}'. Available models: {}",
+                model,
+                available_models.join(", ")
+            ),
+        ));
+    }
+
+    None
+}
+
+/// Extract function tools (and optionally MCP tools) from ResponseTools
+///
+/// This utility consolidates the logic for extracting tools with schemas from ResponseTools.
+/// It's used by both Harmony and Regular routers for different purposes:
+///
+/// - **Harmony router**: Extracts both Function and MCP tools (with `include_mcp: true`)
+///   because MCP schemas are populated by convert_mcp_tools_to_response_tools() before the
+///   pipeline runs. These tools are used to generate structural constraints in the
+///   Harmony preparation stage.
+///
+/// - **Regular router**: Extracts only Function tools (with `include_mcp: false`) during
+///   the initial conversion from ResponsesRequest to ChatCompletionRequest. MCP tools
+///   are merged later by the tool loop before being sent to the chat pipeline, where
+///   tool_choice constraints are generated for ALL tools (function + MCP combined).
+pub(crate) fn extract_tools_from_response_tools(
+    response_tools: Option<&[ResponseTool]>,
+    include_mcp: bool,
+) -> Vec<Tool> {
+    let Some(tools) = response_tools else {
+        return Vec::new();
+    };
+
+    tools
+        .iter()
+        .filter_map(|rt| {
+            match rt.r#type {
+                // Function tools: Schema in request
+                ResponseToolType::Function => rt.function.as_ref().map(|f| Tool {
+                    tool_type: "function".to_string(),
+                    function: f.clone(),
+                }),
+                // MCP tools: Schema populated by convert_mcp_tools_to_response_tools()
+                // Only include if requested (Harmony case)
+                ResponseToolType::Mcp if include_mcp => rt.function.as_ref().map(|f| Tool {
+                    tool_type: "function".to_string(),
+                    function: f.clone(),
+                }),
+                // Hosted tools: No schema available, skip
+                _ => None,
+            }
+        })
+        .collect()
+}
+
+/// Persist response to storage if store=true
+///
+/// Common helper function to avoid duplication across sync and streaming paths
+/// in both harmony and regular responses implementations.
+pub(crate) async fn persist_response_if_needed(
+    conversation_storage: Arc<dyn ConversationStorage>,
+    conversation_item_storage: Arc<dyn ConversationItemStorage>,
+    response_storage: Arc<dyn ResponseStorage>,
+    response: &ResponsesResponse,
+    original_request: &ResponsesRequest,
+) {
+    if !original_request.store.unwrap_or(true) {
+        return;
+    }
+
+    if let Ok(response_json) = to_value(response) {
+        if let Err(e) = persist_conversation_items(
+            conversation_storage,
+            conversation_item_storage,
+            response_storage,
+            &response_json,
+            original_request,
+        )
+        .await
+        {
+            warn!("Failed to persist response: {}", e);
+        } else {
+            debug!("Persisted response: {}", response.id);
+        }
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/stages/client_acquisition.rs

@@ -0,0 +1,69 @@
+//! Client acquisition stage: Get gRPC clients from selected workers
+
+use async_trait::async_trait;
+use axum::response::Response;
+use tracing::error;
+
+use super::PipelineStage;
+use crate::routers::{
+    error,
+    grpc::{
+        context::{ClientSelection, RequestContext, WorkerSelection},
+        utils,
+    },
+};
+
+/// Client acquisition stage: Get gRPC clients from selected workers
+pub(crate) struct ClientAcquisitionStage;
+
+#[async_trait]
+impl PipelineStage for ClientAcquisitionStage {
+    async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, Response> {
+        let workers = ctx.state.workers.as_ref().ok_or_else(|| {
+            error!(
+                function = "ClientAcquisitionStage::execute",
+                "Worker selection stage not completed"
+            );
+            error::internal_error(
+                "worker_selection_not_completed",
+                "Worker selection not completed",
+            )
+        })?;
+
+        let clients = match workers {
+            WorkerSelection::Single { worker } => {
+                let client = utils::get_grpc_client_from_worker(worker).await?;
+                ClientSelection::Single { client }
+            }
+            WorkerSelection::Dual { prefill, decode } => {
+                let prefill_client = utils::get_grpc_client_from_worker(prefill).await?;
+                let decode_client = utils::get_grpc_client_from_worker(decode).await?;
+
+                // vLLM does not support dual (PD disaggregated) mode
+                if prefill_client.is_vllm() || decode_client.is_vllm() {
+                    error!(
+                        function = "ClientAcquisitionStage::execute",
+                        "vLLM backend does not support dual (PD disaggregated) mode"
+                    );
+                    return Err(error::bad_request(
+                        "vllm_pd_mode_not_supported",
+                        "vLLM backend does not support prefill/decode disaggregated mode. \
+                         Please use runtime_type: sglang for PD mode, or use a regular (non-PD) worker configuration."
+                    ));
+                }
+
+                ClientSelection::Dual {
+                    prefill: prefill_client,
+                    decode: decode_client,
+                }
+            }
+        };
+
+        ctx.state.clients = Some(clients);
+        Ok(None)
+    }
+
+    fn name(&self) -> &'static str {
+        "ClientAcquisition"
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/stages/dispatch_metadata.rs

@@ -0,0 +1,77 @@
+//! Dispatch metadata stage: Prepare metadata for dispatch
+
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use async_trait::async_trait;
+use axum::response::Response;
+use tracing::error;
+
+use super::PipelineStage;
+use crate::{
+    core::UNKNOWN_MODEL_ID,
+    routers::{
+        error,
+        grpc::context::{DispatchMetadata, RequestContext, RequestType, WorkerSelection},
+    },
+};
+
+/// Dispatch metadata stage: Prepare metadata for dispatch
+pub(crate) struct DispatchMetadataStage;
+
+#[async_trait]
+impl PipelineStage for DispatchMetadataStage {
+    async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, Response> {
+        let proto_request = ctx.state.proto_request.as_ref().ok_or_else(|| {
+            error!(
+                function = "DispatchMetadataStage::execute",
+                "Proto request not built"
+            );
+            error::internal_error("proto_request_not_built", "Proto request not built")
+        })?;
+
+        let request_id = proto_request.request_id().to_string();
+        let model = match &ctx.input.request_type {
+            RequestType::Chat(req) => req.model.clone(),
+            RequestType::Generate(_req) => {
+                // Generate requests don't have a model field
+                // Use model_id from input or UNKNOWN_MODEL_ID
+                ctx.input
+                    .model_id
+                    .clone()
+                    .unwrap_or_else(|| UNKNOWN_MODEL_ID.to_string())
+            }
+            RequestType::Responses(req) => req.model.clone(),
+            RequestType::Embedding(req) => req.model.clone(),
+            RequestType::Classify(req) => req.model.clone(),
+        };
+
+        let weight_version = ctx
+            .state
+            .workers
+            .as_ref()
+            .map(|w| match w {
+                WorkerSelection::Single { worker } => worker,
+                WorkerSelection::Dual { decode, .. } => decode,
+            })
+            .and_then(|w| w.metadata().labels.get("weight_version").cloned())
+            .unwrap_or_else(|| "default".to_string());
+
+        let created = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
+
+        ctx.state.dispatch = Some(DispatchMetadata {
+            request_id,
+            model,
+            created,
+            weight_version: Some(weight_version),
+        });
+
+        Ok(None)
+    }
+
+    fn name(&self) -> &'static str {
+        "DispatchMetadata"
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/stages/helpers.rs

@@ -0,0 +1,41 @@
+//! Common helper functions shared across stages
+
+use std::sync::Arc;
+
+use rand::Rng;
+use smg_grpc_client::sglang_proto::DisaggregatedParams;
+use tracing::debug;
+
+use crate::{core::Worker, routers::grpc::proto_wrapper::ProtoGenerateRequest};
+
+/// Inject PD bootstrap metadata into a gRPC request
+///
+/// Used by both chat and generate request building stages when in PD mode.
+/// Only SGLang supports PD (prefill/decode) disaggregated mode.
+pub(crate) fn inject_bootstrap_metadata(
+    request: &mut ProtoGenerateRequest,
+    prefill_worker: &Arc<dyn Worker>,
+) {
+    let hostname = prefill_worker.bootstrap_host();
+    let bootstrap_port = prefill_worker.bootstrap_port().unwrap_or(8998);
+
+    // Generate room ID for bootstrap
+    let room_id = rand::rng().random_range(0..i32::MAX);
+
+    // Create DisaggregatedParams
+    let disagg_params = DisaggregatedParams {
+        bootstrap_host: hostname.to_string(),
+        bootstrap_port: bootstrap_port as i32,
+        bootstrap_room: room_id,
+    };
+
+    // Inject metadata directly into SGLang request
+    // (vLLM doesn't support PD mode, so this will panic if called with vLLM)
+    let sglang_request = request.as_sglang_mut();
+    sglang_request.disaggregated_params = Some(disagg_params);
+
+    debug!(
+        "Injected bootstrap metadata: host={}, port={}, room={}",
+        hostname, bootstrap_port, room_id
+    );
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/stages/mod.rs

@@ -0,0 +1,39 @@
+//! Common pipeline stages shared across all endpoints and model types
+//!
+//! These stages are endpoint-agnostic and model-agnostic:
+//! - Worker selection
+//! - Client acquisition
+//! - Dispatch metadata generation
+//! - Request execution
+
+use async_trait::async_trait;
+use axum::response::Response;
+
+use crate::routers::grpc::context::RequestContext;
+
+/// Trait for pipeline stages that process requests
+#[async_trait]
+pub trait PipelineStage: Send + Sync {
+    /// Execute this stage, mutating the context
+    ///
+    /// Returns:
+    /// - `Ok(None)` - Continue to next stage
+    /// - `Ok(Some(response))` - Pipeline complete, return this response (e.g., streaming)
+    /// - `Err(response)` - Error occurred, return this error response
+    async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, Response>;
+
+    /// Stage name for logging
+    fn name(&self) -> &'static str;
+}
+
+mod client_acquisition;
+mod dispatch_metadata;
+pub(crate) mod helpers;
+mod request_execution;
+mod worker_selection;
+
+// Export stage implementations
+pub(crate) use client_acquisition::ClientAcquisitionStage;
+pub(crate) use dispatch_metadata::DispatchMetadataStage;
+pub(crate) use request_execution::{ExecutionMode, RequestExecutionStage};
+pub(crate) use worker_selection::{WorkerSelectionMode, WorkerSelectionStage};

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -0,0 +1,283 @@
+//! Request execution stage: Execute gRPC requests (single or dual dispatch)
+
+use async_trait::async_trait;
+use axum::response::Response;
+use tracing::{error, info_span, Instrument};
+
+use super::PipelineStage;
+use crate::routers::{
+    error,
+    grpc::{
+        context::{ClientSelection, ExecutionResult, LoadGuards, RequestContext, WorkerSelection},
+        proto_wrapper::{
+            ProtoEmbedRequest, ProtoEmbedResponseVariant, ProtoGenerateRequest, ProtoRequest,
+            ProtoStream,
+        },
+    },
+};
+
+type StreamResult = Result<ProtoStream, Box<dyn std::error::Error + Send + Sync>>;
+
+/// Request execution stage: Execute gRPC requests (single or dual dispatch)
+pub(crate) struct RequestExecutionStage {
+    mode: ExecutionMode,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum ExecutionMode {
+    /// Regular mode: single worker execution
+    Single,
+    /// PD mode: dual dispatch to prefill + decode workers
+    DualDispatch,
+}
+
+impl RequestExecutionStage {
+    pub fn new(mode: ExecutionMode) -> Self {
+        Self { mode }
+    }
+}
+
+#[async_trait]
+impl PipelineStage for RequestExecutionStage {
+    async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, Response> {
+        let proto_request = ctx.state.proto_request.take().ok_or_else(|| {
+            error!(
+                function = "RequestExecutionStage::execute",
+                "Proto request not built"
+            );
+            error::internal_error("proto_request_not_built", "Proto request not built")
+        })?;
+
+        let clients = ctx.state.clients.as_mut().ok_or_else(|| {
+            error!(
+                function = "RequestExecutionStage::execute",
+                "Client acquisition not completed"
+            );
+            error::internal_error(
+                "client_acquisition_not_completed",
+                "Client acquisition not completed",
+            )
+        })?;
+
+        // Create load guards for worker load tracking (increment load when created)
+        // They will be automatically dropped (and decrement load) when RequestContext is dropped
+        let workers = ctx.state.workers.as_ref().ok_or_else(|| {
+            error!(
+                function = "RequestExecutionStage::execute",
+                "Worker selection not completed"
+            );
+            error::internal_error(
+                "worker_selection_not_completed",
+                "Worker selection not completed",
+            )
+        })?;
+
+        ctx.state.load_guards = Some(LoadGuards::new(workers, ctx.input.headers.as_ref()));
+
+        // Extract dispatch metadata for tracing span
+        let request_id = ctx
+            .state
+            .dispatch
+            .as_ref()
+            .map(|d| d.request_id.as_str())
+            .unwrap_or("unknown");
+        let model = ctx
+            .state
+            .dispatch
+            .as_ref()
+            .map(|d| d.model.as_str())
+            .unwrap_or("unknown");
+
+        // Create OTEL span for gRPC request execution
+        let span = info_span!(
+            target: "smg::otel-trace",
+            "grpc_generate",
+            request_id = %request_id,
+            model = %model,
+            mode = ?self.mode,
+        );
+
+        let result = async {
+            match proto_request {
+                ProtoRequest::Generate(req) => match self.mode {
+                    ExecutionMode::Single => self.execute_single(req, clients, workers).await,
+                    ExecutionMode::DualDispatch => {
+                        self.execute_dual_dispatch(req, clients, workers).await
+                    }
+                },
+                ProtoRequest::Embed(req) => self.execute_single_embed(req, clients).await,
+            }
+        }
+        .instrument(span)
+        .await?;
+
+        // Store result in context for ResponseProcessingStage
+        ctx.state.response.execution_result = Some(result);
+        Ok(None)
+    }
+
+    fn name(&self) -> &'static str {
+        "RequestExecution"
+    }
+}
+
+impl RequestExecutionStage {
+    async fn execute_single(
+        &self,
+        proto_request: ProtoGenerateRequest,
+        clients: &mut ClientSelection,
+        workers: &WorkerSelection,
+    ) -> Result<ExecutionResult, Response> {
+        let client = clients.single_mut().ok_or_else(|| {
+            error!(
+                function = "execute_single",
+                "Expected single client but got dual"
+            );
+            error::internal_error(
+                "expected_single_client_got_dual",
+                "Expected single client but got dual",
+            )
+        })?;
+
+        let result = client.generate(proto_request).await;
+
+        // Record circuit breaker outcome
+        workers.record_outcome(result.is_ok());
+
+        let stream = result.map_err(|e| {
+            error!(
+                function = "execute_single",
+                error = %e,
+                "Failed to start generation"
+            );
+            error::internal_error(
+                "start_generation_failed",
+                format!("Failed to start generation: {}", e),
+            )
+        })?;
+
+        Ok(ExecutionResult::Single { stream })
+    }
+
+    async fn execute_single_embed(
+        &self,
+        proto_request: ProtoEmbedRequest,
+        clients: &mut ClientSelection,
+    ) -> Result<ExecutionResult, Response> {
+        let client = clients.single_mut().ok_or_else(|| {
+            error!(
+                function = "execute_single_embed",
+                "Expected single client but got dual"
+            );
+            error::internal_error(
+                "expected_single_client_got_dual",
+                "Expected single client but got dual",
+            )
+        })?;
+
+        let response = client.embed(proto_request).await.map_err(|e| {
+            error!(
+                function = "execute_single_embed",
+                error = %e,
+                "Failed to start embedding"
+            );
+            error::internal_error(
+                "start_embedding_failed",
+                format!("Failed to start embedding: {}", e),
+            )
+        })?;
+
+        match response.into_response() {
+            ProtoEmbedResponseVariant::Complete(complete) => {
+                Ok(ExecutionResult::Embedding { response: complete })
+            }
+            ProtoEmbedResponseVariant::Error(e) => {
+                error!(
+                    function = "execute_single_embed",
+                    error = %e.message(),
+                    "Embedding execution failed"
+                );
+                Err(error::internal_error(
+                    "embedding_execution_failed",
+                    e.message().to_string(),
+                ))
+            }
+            ProtoEmbedResponseVariant::None => {
+                error!(
+                    function = "execute_single_embed",
+                    "Embedding execution returned no response"
+                );
+                Err(error::internal_error(
+                    "embedding_no_response",
+                    "Embedding execution returned no response",
+                ))
+            }
+        }
+    }
+
+    async fn execute_dual_dispatch(
+        &self,
+        proto_request: ProtoGenerateRequest,
+        clients: &mut ClientSelection,
+        workers: &WorkerSelection,
+    ) -> Result<ExecutionResult, Response> {
+        let (prefill_client, decode_client) = clients.dual_mut().ok_or_else(|| {
+            error!(
+                function = "execute_dual_dispatch",
+                "Expected dual clients but got single"
+            );
+            error::internal_error(
+                "expected_dual_clients_got_single",
+                "Expected dual clients but got single",
+            )
+        })?;
+
+        let prefill_request = proto_request.clone_inner();
+        let decode_request = proto_request;
+
+        let (prefill_result, decode_result): (StreamResult, StreamResult) = tokio::join!(
+            prefill_client.generate(prefill_request),
+            decode_client.generate(decode_request)
+        );
+
+        // Record circuit breaker outcomes for each worker individually
+        workers.record_dual_outcomes(prefill_result.is_ok(), decode_result.is_ok());
+
+        // Handle prefill result
+        let prefill_stream = match prefill_result {
+            Ok(s) => s,
+            Err(e) => {
+                error!(
+                    function = "execute_dual_dispatch",
+                    error = %e,
+                    "Prefill worker failed to start"
+                );
+                return Err(error::internal_error(
+                    "prefill_worker_failed_to_start",
+                    format!("Prefill worker failed to start: {}", e),
+                ));
+            }
+        };
+
+        // Handle decode result
+        let decode_stream = match decode_result {
+            Ok(s) => s,
+            Err(e) => {
+                error!(
+                    function = "execute_dual_dispatch",
+                    error = %e,
+                    "Decode worker failed to start"
+                );
+                return Err(error::internal_error(
+                    "decode_worker_failed_to_start",
+                    format!("Decode worker failed to start: {}", e),
+                ));
+            }
+        };
+
+        Ok(ExecutionResult::Dual {
+            prefill: prefill_stream,
+            decode: Box::new(decode_stream),
+        })
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs

@@ -0,0 +1,275 @@
+//! Worker selection stage: Select appropriate worker(s) based on routing mode
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use axum::response::Response;
+use tracing::{error, warn};
+
+use super::PipelineStage;
+use crate::{
+    core::{ConnectionMode, Worker, WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID},
+    observability::metrics::{metrics_labels, Metrics},
+    policies::{PolicyRegistry, SelectWorkerInfo},
+    routers::{
+        error,
+        grpc::context::{RequestContext, WorkerSelection},
+    },
+};
+
+/// Worker selection stage: Select appropriate worker(s) based on routing mode
+pub(crate) struct WorkerSelectionStage {
+    worker_registry: Arc<WorkerRegistry>,
+    policy_registry: Arc<PolicyRegistry>,
+    mode: WorkerSelectionMode,
+}
+
+pub(crate) enum WorkerSelectionMode {
+    /// Regular mode: select single worker
+    Regular,
+    /// PD mode: select prefill + decode workers
+    PrefillDecode,
+}
+
+impl WorkerSelectionStage {
+    pub fn new(
+        worker_registry: Arc<WorkerRegistry>,
+        policy_registry: Arc<PolicyRegistry>,
+        mode: WorkerSelectionMode,
+    ) -> Self {
+        Self {
+            worker_registry,
+            policy_registry,
+            mode,
+        }
+    }
+}
+
+#[async_trait]
+impl PipelineStage for WorkerSelectionStage {
+    async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, Response> {
+        let prep = ctx.state.preparation.as_ref().ok_or_else(|| {
+            error!(
+                function = "WorkerSelectionStage::execute",
+                "Preparation stage not completed"
+            );
+            error::internal_error(
+                "preparation_stage_not_completed",
+                "Preparation stage not completed",
+            )
+        })?;
+
+        // For Harmony, use selection_text produced during Harmony encoding
+        // Otherwise, use original_text from regular preparation
+        let text = if prep.harmony_mode {
+            prep.selection_text.as_deref()
+        } else {
+            prep.original_text.as_deref()
+        };
+
+        // Get tokens for PrefixHash policy support
+        let tokens = if prep.token_ids.is_empty() {
+            None
+        } else {
+            Some(prep.token_ids.as_slice())
+        };
+
+        let headers = ctx.input.headers.as_ref();
+
+        let workers = match self.mode {
+            WorkerSelectionMode::Regular => {
+                match self
+                    .select_single_worker(ctx.input.model_id.as_deref(), text, tokens, headers)
+                    .await
+                {
+                    Some(w) => WorkerSelection::Single { worker: w },
+                    None => {
+                        let model = ctx.input.model_id.as_deref().unwrap_or(UNKNOWN_MODEL_ID);
+                        error!(
+                            function = "WorkerSelectionStage::execute",
+                            mode = "Regular",
+                            model_id = %model,
+                            "No available workers for model"
+                        );
+                        return Err(error::service_unavailable(
+                            "no_available_workers",
+                            format!("No available workers for model: {}", model),
+                        ));
+                    }
+                }
+            }
+            WorkerSelectionMode::PrefillDecode => {
+                match self
+                    .select_pd_pair(ctx.input.model_id.as_deref(), text, tokens, headers)
+                    .await
+                {
+                    Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode },
+                    None => {
+                        let model = ctx.input.model_id.as_deref().unwrap_or(UNKNOWN_MODEL_ID);
+                        error!(
+                            function = "WorkerSelectionStage::execute",
+                            mode = "PrefillDecode",
+                            model_id = %model,
+                            "No available PD worker pairs for model"
+                        );
+                        return Err(error::service_unavailable(
+                            "no_available_pd_worker_pairs",
+                            format!("No available PD worker pairs for model: {}", model),
+                        ));
+                    }
+                }
+            }
+        };
+
+        ctx.state.workers = Some(workers);
+        Ok(None)
+    }
+
+    fn name(&self) -> &'static str {
+        "WorkerSelection"
+    }
+}
+
+impl WorkerSelectionStage {
+    async fn select_single_worker(
+        &self,
+        model_id: Option<&str>,
+        text: Option<&str>,
+        tokens: Option<&[u32]>,
+        headers: Option<&http::HeaderMap>,
+    ) -> Option<Arc<dyn Worker>> {
+        // Get workers for the specified model, filtered by connection mode
+        let workers = self.worker_registry.get_workers_filtered(
+            model_id,
+            Some(WorkerType::Regular),
+            Some(ConnectionMode::Grpc { port: None }),
+            None,  // any runtime type
+            false, // get all workers, we'll filter by is_available() next
+        );
+
+        // Use into_iter() to take ownership of Arcs without cloning (avoids atomic inc/dec)
+        let available: Vec<Arc<dyn Worker>> =
+            workers.into_iter().filter(|w| w.is_available()).collect();
+
+        if available.is_empty() {
+            return None;
+        }
+
+        // Get the appropriate policy for this model
+        let policy = match model_id {
+            Some(model) => self.policy_registry.get_policy_or_default(model),
+            None => self.policy_registry.get_default_policy(),
+        };
+
+        // Get cached hash ring for consistent hashing (O(log n) lookup)
+        let hash_ring = self
+            .worker_registry
+            .get_hash_ring(model_id.unwrap_or(UNKNOWN_MODEL_ID));
+
+        // Select worker using the policy
+        let idx = policy
+            .select_worker(
+                &available,
+                &SelectWorkerInfo {
+                    request_text: text,
+                    tokens,
+                    headers,
+                    hash_ring,
+                },
+            )
+            .await?;
+        let selected = available[idx].clone();
+
+        // Record worker selection metric
+        Metrics::record_worker_selection(
+            metrics_labels::WORKER_REGULAR,
+            metrics_labels::CONNECTION_GRPC,
+            model_id.unwrap_or(UNKNOWN_MODEL_ID),
+            policy.name(),
+        );
+
+        Some(selected)
+    }
+
+    async fn select_pd_pair(
+        &self,
+        model_id: Option<&str>,
+        text: Option<&str>,
+        tokens: Option<&[u32]>,
+        headers: Option<&http::HeaderMap>,
+    ) -> Option<(Arc<dyn Worker>, Arc<dyn Worker>)> {
+        let all_workers = self.worker_registry.get_workers_filtered(
+            model_id,
+            None,
+            Some(ConnectionMode::Grpc { port: None }), // Match any gRPC worker
+            None,                                      // any runtime type
+            false,
+        );
+
+        let (available_prefill, available_decode): (Vec<_>, Vec<_>) =
+            all_workers
+                .into_iter()
+                .fold((Vec::new(), Vec::new()), |mut acc, w| {
+                    if w.is_available() {
+                        match w.metadata().worker_type {
+                            WorkerType::Prefill { .. } => acc.0.push(w),
+                            WorkerType::Decode => acc.1.push(w),
+                            _ => {}
+                        }
+                    }
+                    acc
+                });
+
+        if available_prefill.is_empty() {
+            warn!("No available prefill workers");
+            return None;
+        }
+
+        if available_decode.is_empty() {
+            warn!("No available decode workers");
+            return None;
+        }
+
+        // Select using policies
+        let policy = match model_id {
+            Some(model) => self.policy_registry.get_policy_or_default(model),
+            None => self.policy_registry.get_default_policy(),
+        };
+
+        // Get cached hash ring for consistent hashing (O(log n) lookup)
+        let hash_ring = self
+            .worker_registry
+            .get_hash_ring(model_id.unwrap_or(UNKNOWN_MODEL_ID));
+
+        let info = SelectWorkerInfo {
+            request_text: text,
+            tokens,
+            headers,
+            hash_ring,
+        };
+        let prefill_idx = policy.select_worker(&available_prefill, &info).await?;
+        let decode_idx = policy.select_worker(&available_decode, &info).await?;
+
+        let model = model_id.unwrap_or(UNKNOWN_MODEL_ID);
+        let policy_name = policy.name();
+
+        // Record worker selection metrics for both prefill and decode
+        Metrics::record_worker_selection(
+            metrics_labels::WORKER_PREFILL,
+            metrics_labels::CONNECTION_GRPC,
+            model,
+            policy_name,
+        );
+        Metrics::record_worker_selection(
+            metrics_labels::WORKER_DECODE,
+            metrics_labels::CONNECTION_GRPC,
+            model,
+            policy_name,
+        );
+
+        Some((
+            available_prefill[prefill_idx].clone(),
+            available_decode[decode_idx].clone(),
+        ))
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/context.rs

@@ -0,0 +1,486 @@
+//! Request context types for gRPC router pipeline
+//!
+//! This module provides the core context types that flow through the router pipeline,
+//! eliminating deep parameter passing chains and providing a single source of truth
+//! for request state.
+
+use std::sync::Arc;
+
+use axum::http::HeaderMap;
+
+use super::{
+    client::GrpcClient,
+    proto_wrapper::{ProtoEmbedComplete, ProtoRequest, ProtoStream},
+};
+use crate::{
+    core::{Worker, WorkerLoadGuard},
+    protocols::{
+        chat::{ChatCompletionRequest, ChatCompletionResponse},
+        classify::{ClassifyRequest, ClassifyResponse},
+        embedding::{EmbeddingRequest, EmbeddingResponse},
+        generate::{GenerateRequest, GenerateResponse},
+        responses::ResponsesRequest,
+    },
+    reasoning_parser::ParserFactory as ReasoningParserFactory,
+    tokenizer::{stop::StopSequenceDecoder, traits::Tokenizer, TokenizerRegistry},
+    tool_parser::ParserFactory as ToolParserFactory,
+};
+
+/// Main request processing context
+///
+/// This is the single source of truth for all request state as it flows
+/// through the pipeline stages. Uses Rust's type system to enforce proper
+/// stage ordering at compile time.
+pub(crate) struct RequestContext {
+    pub input: RequestInput,
+    pub components: Arc<SharedComponents>,
+    pub state: ProcessingState,
+}
+
+/// Immutable request input
+pub(crate) struct RequestInput {
+    pub request_type: RequestType,
+    pub headers: Option<HeaderMap>,
+    pub model_id: Option<String>,
+}
+
+/// Request type variants
+/// Using Arc instead of Box to enable cheap cloning for background tasks
+pub(crate) enum RequestType {
+    Chat(Arc<ChatCompletionRequest>),
+    Generate(Arc<GenerateRequest>),
+    Responses(Arc<ResponsesRequest>),
+    Embedding(Arc<EmbeddingRequest>),
+    Classify(Arc<ClassifyRequest>),
+}
+
+/// Shared components (injected once at creation)
+pub(crate) struct SharedComponents {
+    pub tokenizer_registry: Arc<TokenizerRegistry>,
+    #[allow(dead_code)]
+    pub tool_parser_factory: ToolParserFactory,
+    #[allow(dead_code)]
+    pub reasoning_parser_factory: ReasoningParserFactory,
+}
+
+/// Mutable processing state (evolves through pipeline stages)
+#[derive(Default)]
+pub(crate) struct ProcessingState {
+    // Stage 1: Preparation outputs
+    pub preparation: Option<PreparationOutput>,
+
+    /// Resolved tokenizer (set once in preparation, reused in response processing)
+    /// This avoids redundant registry lookups across pipeline stages.
+    pub tokenizer: Option<Arc<dyn Tokenizer>>,
+
+    // Stage 2: Worker selection outputs
+    pub workers: Option<WorkerSelection>,
+
+    // Stage 3: Client acquisition outputs
+    pub clients: Option<ClientSelection>,
+
+    // Stage 4: Request building outputs
+    pub proto_request: Option<ProtoRequest>,
+
+    // Stage 5: Dispatch metadata
+    pub dispatch: Option<DispatchMetadata>,
+
+    // Load guard for worker load tracking (created at execution stage)
+    pub load_guards: Option<LoadGuards>,
+
+    // Stage 6: Response processing state
+    pub response: ResponseState,
+}
+
+/// Output from preparation stage (Step 1)
+pub(crate) struct PreparationOutput {
+    /// Original text (for chat) or resolved text (for generate)
+    pub original_text: Option<String>,
+
+    /// Tokenized input
+    pub token_ids: Vec<u32>,
+
+    /// Processed messages (chat only)
+    pub processed_messages: Option<super::ProcessedMessages>,
+
+    /// Tool call constraints (if applicable)
+    pub tool_constraints: Option<(String, String)>,
+
+    /// Filtered request (if tools were filtered)
+    pub filtered_request: Option<ChatCompletionRequest>,
+
+    // Harmony-specific fields
+    /// Whether this is a Harmony request (default: false)
+    pub harmony_mode: bool,
+
+    /// Selection text for worker routing (Harmony only)
+    pub selection_text: Option<String>,
+
+    /// Harmony messages for history tracking (Harmony only)
+    #[allow(dead_code)]
+    pub harmony_messages: Option<Vec<super::harmony::HarmonyMessage>>,
+
+    /// Stop token IDs for Harmony models
+    pub harmony_stop_ids: Option<Vec<u32>>,
+}
+
+/// Worker selection (Step 2)
+pub(crate) enum WorkerSelection {
+    Single {
+        worker: Arc<dyn Worker>,
+    },
+    Dual {
+        prefill: Arc<dyn Worker>,
+        decode: Arc<dyn Worker>,
+    },
+}
+
+/// Client selection (Step 3)
+pub(crate) enum ClientSelection {
+    Single {
+        client: GrpcClient,
+    },
+    Dual {
+        prefill: GrpcClient,
+        decode: GrpcClient,
+    },
+}
+
+/// Dispatch metadata (Step 5)
+#[derive(Clone)]
+pub(crate) struct DispatchMetadata {
+    pub request_id: String,
+    pub model: String,
+    pub created: u64,
+    pub weight_version: Option<String>,
+}
+
+/// Load guards for worker load tracking
+/// Automatically decrements load when dropped
+pub(crate) enum LoadGuards {
+    Single {
+        _guard: WorkerLoadGuard,
+    },
+    Dual {
+        _prefill: WorkerLoadGuard,
+        _decode: WorkerLoadGuard,
+    },
+}
+
+impl LoadGuards {
+    pub fn new(selection: &WorkerSelection, headers: Option<&HeaderMap>) -> Self {
+        match selection {
+            WorkerSelection::Single { worker } => LoadGuards::Single {
+                _guard: WorkerLoadGuard::new(worker.clone(), headers),
+            },
+            WorkerSelection::Dual { prefill, decode } => LoadGuards::Dual {
+                _prefill: WorkerLoadGuard::new(prefill.clone(), headers),
+                _decode: WorkerLoadGuard::new(decode.clone(), headers),
+            },
+        }
+    }
+}
+
+/// Response processing state (Step 6)
+#[derive(Default)]
+pub(crate) struct ResponseState {
+    /// Stop sequence decoder
+    pub stop_decoder: Option<StopSequenceDecoder>,
+
+    /// Execution result (streams from workers)
+    pub execution_result: Option<ExecutionResult>,
+
+    /// Final processed response
+    pub final_response: Option<FinalResponse>,
+
+    /// Responses API iteration result (Harmony only, for tool loop orchestration)
+    pub responses_iteration_result: Option<super::harmony::ResponsesIterationResult>,
+}
+
+impl RequestContext {
+    /// Create context for chat completion request
+    pub fn for_chat(
+        request: Arc<ChatCompletionRequest>,
+        headers: Option<HeaderMap>,
+        model_id: Option<String>,
+        components: Arc<SharedComponents>,
+    ) -> Self {
+        Self {
+            input: RequestInput {
+                request_type: RequestType::Chat(request),
+                headers,
+                model_id,
+            },
+            components,
+            state: ProcessingState::default(),
+        }
+    }
+
+    /// Create context for generate request
+    pub fn for_generate(
+        request: Arc<GenerateRequest>,
+        headers: Option<HeaderMap>,
+        model_id: Option<String>,
+        components: Arc<SharedComponents>,
+    ) -> Self {
+        Self {
+            input: RequestInput {
+                request_type: RequestType::Generate(request),
+                headers,
+                model_id,
+            },
+            components,
+            state: ProcessingState::default(),
+        }
+    }
+
+    /// Create context for Responses API request
+    pub fn for_responses(
+        request: Arc<ResponsesRequest>,
+        headers: Option<HeaderMap>,
+        model_id: Option<String>,
+        components: Arc<SharedComponents>,
+    ) -> Self {
+        Self {
+            input: RequestInput {
+                request_type: RequestType::Responses(request),
+                headers,
+                model_id,
+            },
+            components,
+            state: ProcessingState::default(),
+        }
+    }
+
+    /// Create context for embedding request
+    pub fn for_embedding(
+        request: Arc<EmbeddingRequest>,
+        headers: Option<HeaderMap>,
+        model_id: Option<String>,
+        components: Arc<SharedComponents>,
+    ) -> Self {
+        Self {
+            input: RequestInput {
+                request_type: RequestType::Embedding(request),
+                headers,
+                model_id,
+            },
+            components,
+            state: ProcessingState::default(),
+        }
+    }
+
+    /// Create context for classify request
+    pub fn for_classify(
+        request: Arc<ClassifyRequest>,
+        headers: Option<HeaderMap>,
+        model_id: Option<String>,
+        components: Arc<SharedComponents>,
+    ) -> Self {
+        Self {
+            input: RequestInput {
+                request_type: RequestType::Classify(request),
+                headers,
+                model_id,
+            },
+            components,
+            state: ProcessingState::default(),
+        }
+    }
+
+    /// Get chat request (panics if not chat)
+    pub fn chat_request(&self) -> &ChatCompletionRequest {
+        match &self.input.request_type {
+            RequestType::Chat(req) => req.as_ref(),
+            _ => panic!("Expected chat request"),
+        }
+    }
+
+    /// Get Arc clone of chat request (panics if not chat)
+    pub fn chat_request_arc(&self) -> Arc<ChatCompletionRequest> {
+        match &self.input.request_type {
+            RequestType::Chat(req) => Arc::clone(req),
+            _ => panic!("Expected chat request"),
+        }
+    }
+
+    /// Get generate request (panics if not generate)
+    pub fn generate_request(&self) -> &GenerateRequest {
+        match &self.input.request_type {
+            RequestType::Generate(req) => req.as_ref(),
+            _ => panic!("Expected generate request"),
+        }
+    }
+
+    /// Get Arc clone of generate request (panics if not generate)
+    pub fn generate_request_arc(&self) -> Arc<GenerateRequest> {
+        match &self.input.request_type {
+            RequestType::Generate(req) => Arc::clone(req),
+            _ => panic!("Expected generate request"),
+        }
+    }
+
+    /// Get Arc clone of responses request (panics if not responses)
+    pub fn responses_request_arc(&self) -> Arc<ResponsesRequest> {
+        match &self.input.request_type {
+            RequestType::Responses(req) => Arc::clone(req),
+            _ => panic!("Expected responses request"),
+        }
+    }
+
+    /// Check if request is streaming
+    pub fn is_streaming(&self) -> bool {
+        match &self.input.request_type {
+            RequestType::Chat(req) => req.stream,
+            RequestType::Generate(req) => req.stream,
+            RequestType::Responses(req) => req.stream.unwrap_or(false),
+            RequestType::Embedding(_) => false, // Embeddings are never streaming
+            RequestType::Classify(_) => false,  // Classification is never streaming
+        }
+    }
+
+    /// Get the cached tokenizer, cloning the Arc (cheap 8-byte clone)
+    ///
+    /// Returns None if tokenizer hasn't been resolved yet.
+    /// The tokenizer is resolved once in the preparation stage and cached for reuse.
+    pub fn tokenizer_arc(&self) -> Option<Arc<dyn Tokenizer>> {
+        self.state.tokenizer.clone()
+    }
+}
+
+/// Some methods are kept for API completeness even if currently unused.
+#[allow(dead_code)]
+impl WorkerSelection {
+    pub fn is_dual(&self) -> bool {
+        matches!(self, Self::Dual { .. })
+    }
+
+    pub fn single(&self) -> Option<&Arc<dyn Worker>> {
+        match self {
+            Self::Single { worker } => Some(worker),
+            _ => None,
+        }
+    }
+
+    /// Record circuit breaker outcome for all workers
+    pub fn record_outcome(&self, success: bool) {
+        match self {
+            Self::Single { worker } => worker.record_outcome(success),
+            Self::Dual { prefill, decode } => {
+                prefill.record_outcome(success);
+                decode.record_outcome(success);
+            }
+        }
+    }
+
+    /// Record circuit breaker outcomes for dual dispatch (individual tracking)
+    pub fn record_dual_outcomes(&self, prefill_success: bool, decode_success: bool) {
+        if let Self::Dual { prefill, decode } = self {
+            prefill.record_outcome(prefill_success);
+            decode.record_outcome(decode_success);
+        }
+    }
+
+    #[allow(clippy::type_complexity)]
+    pub fn dual(&self) -> Option<(&Arc<dyn Worker>, &Arc<dyn Worker>)> {
+        match self {
+            Self::Dual { prefill, decode } => Some((prefill, decode)),
+            _ => None,
+        }
+    }
+
+    pub fn prefill_worker(&self) -> Option<&Arc<dyn Worker>> {
+        match self {
+            Self::Dual { prefill, .. } => Some(prefill),
+            _ => None,
+        }
+    }
+
+    pub fn decode_worker(&self) -> Option<&Arc<dyn Worker>> {
+        match self {
+            Self::Dual { decode, .. } => Some(decode),
+            _ => None,
+        }
+    }
+}
+
+/// Some methods are kept for API completeness even if currently unused.
+#[allow(dead_code)]
+impl ClientSelection {
+    pub fn single(&self) -> Option<&GrpcClient> {
+        match self {
+            Self::Single { client } => Some(client),
+            _ => None,
+        }
+    }
+
+    pub fn single_mut(&mut self) -> Option<&mut GrpcClient> {
+        match self {
+            Self::Single { client } => Some(client),
+            _ => None,
+        }
+    }
+
+    pub fn dual_mut(&mut self) -> Option<(&mut GrpcClient, &mut GrpcClient)> {
+        match self {
+            Self::Dual { prefill, decode } => Some((prefill, decode)),
+            _ => None,
+        }
+    }
+
+    pub fn prefill_client(&self) -> Option<&GrpcClient> {
+        match self {
+            Self::Dual { prefill, .. } => Some(prefill),
+            _ => None,
+        }
+    }
+
+    pub fn prefill_client_mut(&mut self) -> Option<&mut GrpcClient> {
+        match self {
+            Self::Dual { prefill, .. } => Some(prefill),
+            _ => None,
+        }
+    }
+
+    pub fn decode_client(&self) -> Option<&GrpcClient> {
+        match self {
+            Self::Dual { decode, .. } => Some(decode),
+            _ => None,
+        }
+    }
+
+    pub fn decode_client_mut(&mut self) -> Option<&mut GrpcClient> {
+        match self {
+            Self::Dual { decode, .. } => Some(decode),
+            _ => None,
+        }
+    }
+}
+
+/// Result of request execution (streams from workers)
+/// Uses ProtoStream to automatically abort on cancellation
+pub(crate) enum ExecutionResult {
+    Single {
+        stream: ProtoStream,
+    },
+    Dual {
+        prefill: ProtoStream,
+        decode: Box<ProtoStream>,
+    },
+    /// Embedding requests return a single response, not a stream
+    Embedding {
+        response: ProtoEmbedComplete,
+    },
+}
+
+/// Final processed response
+#[derive(Debug)]
+pub(crate) enum FinalResponse {
+    Chat(ChatCompletionResponse),
+    /// Generate response is a Vec of GenerateResponse (n=1 returns single item, n>1 returns multiple)
+    Generate(Vec<GenerateResponse>),
+    /// Embedding response
+    Embedding(EmbeddingResponse),
+    /// Classification response
+    Classify(ClassifyResponse),
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/builder.rs

@@ -0,0 +1,908 @@
+//! Harmony request builder
+//!
+//! Handles encoding of Chat/Responses requests into Harmony format using openai-harmony library.
+
+use std::sync::OnceLock;
+
+use chrono::Local;
+use openai_harmony::{
+    chat::{
+        Author, ChannelConfig, Content, Conversation, DeveloperContent, Message as HarmonyMessage,
+        ReasoningEffort, Role, SystemContent, TextContent, ToolDescription,
+    },
+    HarmonyEncoding, HarmonyEncodingName,
+};
+use tracing::{debug, trace};
+
+use super::types::HarmonyBuildOutput;
+use crate::protocols::{
+    chat::{ChatCompletionRequest, ChatMessage, MessageContent},
+    common::{ContentPart, Tool},
+    responses::{
+        ReasoningEffort as ResponsesReasoningEffort, ResponseContentPart, ResponseInput,
+        ResponseInputOutputItem, ResponseReasoningContent, ResponseTool, ResponseToolType,
+        ResponsesRequest, StringOrContentParts,
+    },
+};
+
+/// Global Harmony encoding (lazy-initialized)
+static HARMONY_ENCODING: OnceLock<HarmonyEncoding> = OnceLock::new();
+
+/// Get or initialize the Harmony encoding
+///
+/// Uses HarmonyGptOss encoding which supports the gpt-oss model family.
+pub(super) fn get_harmony_encoding() -> &'static HarmonyEncoding {
+    HARMONY_ENCODING.get_or_init(|| {
+        tokio::task::block_in_place(|| {
+            openai_harmony::load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss)
+                .expect("Failed to load Harmony encoding")
+        })
+    })
+}
+
+/// Built-in tools that are added to the system message
+const BUILTIN_TOOLS: &[&str] = &["web_search_preview", "code_interpreter", "container"];
+
+/// Trait for tool-like objects that can be converted to Harmony ToolDescription
+trait ToolLike {
+    /// Check if this is a built-in tool (should be skipped in developer message)
+    #[allow(dead_code)]
+    fn is_builtin(&self) -> bool;
+
+    /// Check if this is a custom tool (function or MCP)
+    fn is_custom(&self) -> bool;
+
+    /// Convert to ToolDescription
+    fn to_tool_description(&self) -> Option<ToolDescription>;
+}
+
+/// Implement ToolLike for Chat Completion Tool
+impl ToolLike for Tool {
+    fn is_builtin(&self) -> bool {
+        matches!(
+            self.tool_type.as_str(),
+            "web_search_preview" | "code_interpreter" | "container"
+        )
+    }
+
+    fn is_custom(&self) -> bool {
+        matches!(self.tool_type.as_str(), "mcp" | "function")
+    }
+
+    fn to_tool_description(&self) -> Option<ToolDescription> {
+        Some(ToolDescription::new(
+            self.function.name.clone(),
+            self.function.description.clone().unwrap_or_default(),
+            Some(self.function.parameters.clone()),
+        ))
+    }
+}
+
+/// Implement ToolLike for Responses API Tool
+impl ToolLike for ResponseTool {
+    fn is_builtin(&self) -> bool {
+        matches!(
+            self.r#type,
+            ResponseToolType::WebSearchPreview | ResponseToolType::CodeInterpreter
+        )
+    }
+
+    fn is_custom(&self) -> bool {
+        matches!(
+            self.r#type,
+            ResponseToolType::Mcp | ResponseToolType::Function
+        )
+    }
+
+    fn to_tool_description(&self) -> Option<ToolDescription> {
+        self.function.as_ref().map(|func| {
+            ToolDescription::new(
+                func.name.clone(),
+                func.description.clone().unwrap_or_default(),
+                Some(func.parameters.clone()),
+            )
+        })
+    }
+}
+
+fn has_custom_tools(tool_types: &[&str]) -> bool {
+    !tool_types.iter().all(|t| BUILTIN_TOOLS.contains(t))
+}
+
+/// Harmony request builder
+///
+/// Converts OpenAI-format requests into Harmony-encoded format with input_ids,
+/// stop tokens, and selection text for worker routing.
+pub(crate) struct HarmonyBuilder {
+    encoding: &'static HarmonyEncoding,
+}
+
+impl HarmonyBuilder {
+    /// Create a new Harmony builder
+    pub fn new() -> Self {
+        Self {
+            encoding: get_harmony_encoding(),
+        }
+    }
+
+    /// Build Harmony request from Chat Completion request
+    ///
+    /// # Arguments
+    ///
+    /// * `request` - The ChatCompletionRequest to encode
+    ///
+    /// # Returns
+    ///
+    /// HarmonyBuildOutput containing input_ids, stop_token_ids, selection_text, and messages
+    pub fn build_from_chat(
+        &self,
+        request: &ChatCompletionRequest,
+    ) -> Result<HarmonyBuildOutput, String> {
+        let mut all_messages = Vec::new();
+
+        let sys_msg = self.build_system_message_from_chat(request);
+        all_messages.push(sys_msg);
+
+        let dev_msg = self.build_developer_message_from_chat(request.tools.as_ref());
+        all_messages.push(dev_msg);
+
+        let mut user_messages = self.convert_chat_messages(&request.messages)?;
+        all_messages.append(&mut user_messages);
+
+        let conversation = Conversation::from_messages(all_messages.clone());
+        let token_ids = self
+            .encoding
+            .render_conversation_for_completion(&conversation, Role::Assistant, None)
+            .map_err(|e| format!("Failed to encode Harmony conversation: {}", e))?;
+
+        let selection_text = self.extract_selection_text(&all_messages);
+
+        // Get stop tokens for Harmony assistant actions (<|return|> and <|call|>)
+        let stop_token_ids: Vec<u32> = self
+            .encoding
+            .stop_tokens_for_assistant_actions()
+            .into_iter()
+            .flat_map(|set| set.into_iter())
+            .collect();
+
+        Ok(HarmonyBuildOutput {
+            input_ids: token_ids,
+            stop_token_ids,
+            selection_text,
+            harmony_messages: all_messages
+                .into_iter()
+                .map(super::types::HarmonyMessage::from_openai_harmony)
+                .collect(),
+        })
+    }
+
+    /// Build Harmony request from Responses request
+    ///
+    /// # Arguments
+    ///
+    /// * `request` - The ResponsesRequest to encode
+    ///
+    /// # Returns
+    ///
+    /// HarmonyBuildOutput containing input_ids, stop_token_ids, selection_text, and messages
+    pub fn build_from_responses(
+        &self,
+        request: &ResponsesRequest,
+    ) -> Result<HarmonyBuildOutput, String> {
+        let all_messages = self.construct_input_messages_with_harmony(request)?;
+
+        let conversation = Conversation::from_messages(all_messages.clone());
+        let token_ids = self
+            .encoding
+            .render_conversation_for_completion(&conversation, Role::Assistant, None)
+            .map_err(|e| format!("Failed to encode Harmony conversation: {}", e))?;
+
+        let selection_text = self.extract_selection_text(&all_messages);
+
+        // Get stop tokens for Harmony assistant actions (<|return|> and <|call|>)
+        let stop_token_ids: Vec<u32> = self
+            .encoding
+            .stop_tokens_for_assistant_actions()
+            .into_iter()
+            .flat_map(|set| set.into_iter())
+            .collect();
+
+        // Decode tokens to see what the model actually receives
+        let decoded_text = self
+            .encoding
+            .tokenizer()
+            .decode_utf8(&token_ids)
+            .unwrap_or_else(|_| "<decode error>".to_string());
+        trace!(
+            token_count = token_ids.len(),
+            token_preview = ?&token_ids[..token_ids.len().min(20)],
+            decoded_length = decoded_text.len(),
+            "Encoded conversation to tokens - decoded text follows:"
+        );
+        trace!("DECODED_TEXT_START\n{}\nDECODED_TEXT_END", decoded_text);
+
+        Ok(HarmonyBuildOutput {
+            input_ids: token_ids,
+            stop_token_ids,
+            selection_text,
+            harmony_messages: all_messages
+                .into_iter()
+                .map(super::types::HarmonyMessage::from_openai_harmony)
+                .collect(),
+        })
+    }
+
+    /// Build system message with common logic
+    ///
+    /// # Arguments
+    /// * `reasoning_effort` - Optional reasoning effort level
+    /// * `has_tools` - Whether custom tools are present
+    fn build_system_message(
+        &self,
+        reasoning_effort: Option<ReasoningEffort>,
+        has_tools: bool,
+    ) -> HarmonyMessage {
+        let mut sys_content = SystemContent::new();
+
+        // Add reasoning_effort if provided
+        if let Some(effort) = reasoning_effort {
+            sys_content = sys_content.with_reasoning_effort(effort);
+        }
+
+        // Set conversation start date (always current date)
+        sys_content =
+            sys_content.with_conversation_start_date(Local::now().format("%Y-%m-%d").to_string());
+
+        // If no tools, remove "commentary" from valid channels
+        if !has_tools {
+            if let Some(channel_config) = &sys_content.channel_config {
+                let valid_channels: Vec<String> = channel_config
+                    .valid_channels
+                    .iter()
+                    .filter(|c| c.as_str() != "commentary")
+                    .cloned()
+                    .collect();
+                sys_content = sys_content
+                    .with_channel_config(ChannelConfig::require_channels(valid_channels));
+            }
+        }
+
+        HarmonyMessage::from_role_and_content(Role::System, sys_content)
+    }
+
+    fn build_system_message_from_chat(&self, request: &ChatCompletionRequest) -> HarmonyMessage {
+        let reasoning_effort = request
+            .reasoning_effort
+            .as_deref()
+            .map(|effort| match effort {
+                "high" => ReasoningEffort::High,
+                "medium" => ReasoningEffort::Medium,
+                "low" => ReasoningEffort::Low,
+                // Harmony does not support minimal reasoning effort
+                "minimal" => ReasoningEffort::Low,
+                _ => ReasoningEffort::Medium,
+            });
+
+        let has_tools = request.tools.is_some();
+        self.build_system_message(reasoning_effort, has_tools)
+    }
+
+    /// Build system message from ResponsesRequest
+    ///
+    /// # Arguments
+    /// * `request` - The ResponsesRequest
+    /// * `with_custom_tools` - Whether custom tools (beyond built-ins) are present
+    fn build_system_message_from_responses(
+        &self,
+        request: &ResponsesRequest,
+        with_custom_tools: bool,
+    ) -> HarmonyMessage {
+        let reasoning_effort = request
+            .reasoning
+            .as_ref()
+            .and_then(|r| r.effort.as_ref())
+            .map(|effort| match effort {
+                ResponsesReasoningEffort::High => ReasoningEffort::High,
+                ResponsesReasoningEffort::Medium => ReasoningEffort::Medium,
+                ResponsesReasoningEffort::Low => ReasoningEffort::Low,
+                ResponsesReasoningEffort::Minimal => ReasoningEffort::Low,
+            });
+
+        self.build_system_message(reasoning_effort, with_custom_tools)
+    }
+
+    /// Build developer message with common logic
+    ///
+    /// Filters out built-in tools and converts custom tools to ToolDescription
+    ///
+    /// # Arguments
+    /// * `tools` - Optional list of tools
+    /// * `instructions` - Optional instructions (Responses API only)
+    fn build_developer_message<T: ToolLike>(
+        &self,
+        tools: Option<&Vec<T>>,
+        instructions: Option<&str>,
+    ) -> HarmonyMessage {
+        let mut dev_content = DeveloperContent::new();
+
+        // Add instructions if provided (Responses API only)
+        if let Some(instructions) = instructions {
+            dev_content = dev_content.with_instructions(instructions.to_string());
+        }
+
+        // Early return if no tools
+        let Some(tools) = tools else {
+            return HarmonyMessage::from_role_and_content(Role::Developer, dev_content);
+        };
+
+        // Filter to custom tools and convert to ToolDescription
+        let tool_descriptions: Vec<ToolDescription> = tools
+            .iter()
+            .filter(|t| t.is_custom())
+            .filter_map(|t| t.to_tool_description())
+            .collect();
+
+        // Add function tools to developer content
+        if !tool_descriptions.is_empty() {
+            dev_content = dev_content.with_function_tools(tool_descriptions);
+        }
+
+        HarmonyMessage::from_role_and_content(Role::Developer, dev_content)
+    }
+
+    fn build_developer_message_from_chat(&self, tools: Option<&Vec<Tool>>) -> HarmonyMessage {
+        self.build_developer_message(tools, None)
+    }
+
+    /// Build developer message from Responses request
+    ///
+    /// # Arguments
+    /// * `instructions` - Optional instructions (Responses API specific)
+    /// * `tools` - Optional list of tools
+    fn build_developer_message_from_responses(
+        &self,
+        instructions: Option<&str>,
+        tools: Option<&Vec<ResponseTool>>,
+    ) -> HarmonyMessage {
+        self.build_developer_message(tools, instructions)
+    }
+
+    /// Construct input messages for Responses API with Harmony
+    ///
+    /// Handles both new conversations and continuations of previous responses.
+    ///
+    /// This handles:
+    /// - New conversation: system message, developer message, and user input
+    /// - Continuing conversation: loads previous messages, cleans up chain-of-thoughts
+    /// - MCP tool allowlisting for special tool types
+    /// - Complex response input parsing with function call tracking
+    ///
+    /// # Arguments
+    /// * `request` - The ResponsesRequest
+    /// * `prev_response` - Optional previous response to continue from
+    fn construct_input_messages_with_harmony(
+        &self,
+        request: &ResponsesRequest,
+    ) -> Result<Vec<HarmonyMessage>, String> {
+        let mut all_messages = Vec::new();
+
+        // Handle new vs continuing conversation
+        if request.previous_response_id.is_none() {
+            // New conversation
+
+            let tool_types: Vec<&str> = request
+                .tools
+                .as_ref()
+                .map(|tools| {
+                    tools
+                        .iter()
+                        .map(|tool| match tool.r#type {
+                            ResponseToolType::Function => "function",
+                            ResponseToolType::WebSearchPreview => "web_search_preview",
+                            ResponseToolType::CodeInterpreter => "code_interpreter",
+                            ResponseToolType::Mcp => "mcp",
+                        })
+                        .collect()
+                })
+                .unwrap_or_default();
+
+            let with_custom_tools = has_custom_tools(&tool_types);
+
+            // Add system message
+            let sys_msg = self.build_system_message_from_responses(request, with_custom_tools);
+            all_messages.push(sys_msg);
+
+            // Add developer message only if we have custom tools
+            if with_custom_tools {
+                let dev_msg = self.build_developer_message_from_responses(
+                    request.instructions.as_deref(),
+                    request.tools.as_ref(),
+                );
+                all_messages.push(dev_msg);
+            }
+        } else {
+            // Continue the previous conversation
+            // NOTE: Previous messages are loaded by serve_harmony_responses() before calling this method.
+            // The request.input will already contain the conversation history when previous_response_id was set.
+            // We just proceed with parsing the input items as normal.
+            debug!("Continuing conversation (history already loaded in request.input)");
+        }
+
+        // Append the new input
+        // Responses API supports simple text inputs without chat format
+        match &request.input {
+            ResponseInput::Text(text) => {
+                let user_msg = HarmonyMessage {
+                    author: Author {
+                        role: Role::User,
+                        name: None,
+                    },
+                    recipient: None,
+                    content: vec![Content::Text(TextContent { text: text.clone() })],
+                    channel: None,
+                    content_type: None,
+                };
+                all_messages.push(user_msg);
+            }
+            ResponseInput::Items(items) => {
+                // Track function calls for looking up call_id → name mapping
+                let mut prev_outputs: Vec<&ResponseInputOutputItem> = Vec::new();
+
+                for item in items {
+                    let msg = self.parse_response_item_to_harmony_message(item, &prev_outputs)?;
+                    all_messages.push(msg);
+
+                    // Track function tool calls so that function_call_output can find the name
+                    if matches!(item, ResponseInputOutputItem::FunctionToolCall { .. }) {
+                        prev_outputs.push(item);
+                    }
+                }
+            }
+        }
+
+        debug!(
+            message_count = all_messages.len(),
+            "Constructed Harmony messages for Responses API"
+        );
+        Ok(all_messages)
+    }
+
+    /// Parse a ResponseInputOutputItem into a HarmonyMessage
+    ///
+    /// Handles conversion of various response item types (messages, function calls, reasoning, etc.)
+    /// to Harmony message format.
+    ///
+    /// # Arguments
+    /// * `item` - The ResponseInputOutputItem to parse
+    /// * `prev_outputs` - Previous items for looking up function call names (for function_call_output)
+    fn parse_response_item_to_harmony_message(
+        &self,
+        item: &ResponseInputOutputItem,
+        prev_outputs: &[&ResponseInputOutputItem],
+    ) -> Result<HarmonyMessage, String> {
+        match item {
+            // Regular message (user or assistant)
+            ResponseInputOutputItem::Message { role, content, .. } => {
+                let harmony_role = match role.as_str() {
+                    "user" => Role::User,
+                    "assistant" => Role::Assistant,
+                    "system" => Role::System,
+                    _ => Role::User, // Default to user for unknown roles
+                };
+
+                // Extract text from content parts
+                let text_parts: Vec<String> = content
+                    .iter()
+                    .filter_map(|part| match part {
+                        ResponseContentPart::OutputText { text, .. } => Some(text.clone()),
+                        ResponseContentPart::InputText { text } => Some(text.clone()),
+                        ResponseContentPart::Unknown => None,
+                    })
+                    .collect();
+
+                let text = text_parts.join("\n");
+
+                Ok(HarmonyMessage {
+                    author: Author {
+                        role: harmony_role,
+                        name: None,
+                    },
+                    recipient: None,
+                    content: vec![Content::Text(TextContent { text })],
+                    channel: None,
+                    content_type: None,
+                })
+            }
+
+            // Reasoning content (chain-of-thought)
+            ResponseInputOutputItem::Reasoning { content, .. } => {
+                // Extract reasoning text
+                let reasoning_texts: Vec<String> = content
+                    .iter()
+                    .map(|rc| match rc {
+                        ResponseReasoningContent::ReasoningText { text } => text.clone(),
+                    })
+                    .collect();
+
+                let text = reasoning_texts.join("\n");
+
+                // Reasoning goes in the "analysis" channel for Harmony
+                Ok(HarmonyMessage {
+                    author: Author {
+                        role: Role::Assistant,
+                        name: None,
+                    },
+                    recipient: None,
+                    content: vec![Content::Text(TextContent { text })],
+                    channel: Some("analysis".to_string()),
+                    content_type: None,
+                })
+            }
+
+            // Function tool call (with optional output)
+            ResponseInputOutputItem::FunctionToolCall {
+                name,
+                arguments,
+                output,
+                ..
+            } => {
+                // If there's an output, this represents the tool result
+                // Otherwise, it's the tool call itself
+                if let Some(output_str) = output {
+                    // Tool result - use Tool role with "functions.{name}" as author name
+                    // IMPORTANT: Must include recipient="assistant" for parser to recognize it.
+                    // We keep channel=None to minimize what the model might copy.
+                    let author_name = format!("functions.{}", name);
+                    debug!(
+                        tool_name = %name,
+                        author_name = %author_name,
+                        output_preview = %output_str.chars().take(100).collect::<String>(),
+                        "Building tool result message with Tool role (recipient=assistant, no channel)"
+                    );
+                    Ok(HarmonyMessage {
+                        author: Author {
+                            role: Role::Tool,
+                            name: Some(author_name),
+                        },
+                        recipient: Some("assistant".to_string()),
+                        content: vec![Content::Text(TextContent {
+                            text: output_str.clone(),
+                        })],
+                        channel: None,
+                        content_type: None,
+                    })
+                } else {
+                    // Tool call - assistant message in commentary channel with recipient
+                    // msg.with_channel("commentary").with_recipient(f"functions.{name}")
+                    let recipient = format!("functions.{}", name);
+                    debug!(
+                        tool_name = %name,
+                        recipient = %recipient,
+                        "Building tool call message with recipient"
+                    );
+                    Ok(HarmonyMessage {
+                        author: Author {
+                            role: Role::Assistant,
+                            name: None,
+                        },
+                        recipient: Some(recipient),
+                        content: vec![Content::Text(TextContent {
+                            text: arguments.clone(),
+                        })],
+                        channel: Some("commentary".to_string()),
+                        content_type: Some("json".to_string()),
+                    })
+                }
+            }
+
+            // Function call output (separate from call) - requires looking up the original call
+            ResponseInputOutputItem::FunctionCallOutput {
+                call_id, output, ..
+            } => {
+                // Search prev_outputs in reverse order to find the matching function call
+                let call = prev_outputs
+                    .iter()
+                    .rev()
+                    .find_map(|item| match item {
+                        ResponseInputOutputItem::FunctionToolCall {
+                            call_id: item_call_id,
+                            name,
+                            ..
+                        } if item_call_id == call_id => Some(name.clone()),
+                        _ => None,
+                    })
+                    .ok_or_else(|| format!("No function call found for call_id: {}", call_id))?;
+
+                // Create Tool message with "functions.{name}" prefix
+                // IMPORTANT: Must include recipient="assistant" for parser to recognize it.
+                // We keep channel=None to minimize what the model might copy.
+                Ok(HarmonyMessage {
+                    author: Author {
+                        role: Role::Tool,
+                        name: Some(format!("functions.{}", call)),
+                    },
+                    recipient: Some("assistant".to_string()),
+                    content: vec![Content::Text(TextContent {
+                        text: output.clone(),
+                    })],
+                    channel: None,
+                    content_type: None,
+                })
+            }
+
+            // Simple input message (usually user message)
+            ResponseInputOutputItem::SimpleInputMessage { content, role, .. } => {
+                let harmony_role = match role.as_str() {
+                    "user" => Role::User,
+                    "assistant" => Role::Assistant,
+                    "system" => Role::System,
+                    _ => Role::User,
+                };
+
+                let text = match content {
+                    StringOrContentParts::String(s) => s.clone(),
+                    StringOrContentParts::Array(parts) => {
+                        // Extract text from content parts
+                        parts
+                            .iter()
+                            .filter_map(|part| match part {
+                                ResponseContentPart::OutputText { text, .. } => Some(text.clone()),
+                                ResponseContentPart::InputText { text } => Some(text.clone()),
+                                ResponseContentPart::Unknown => None,
+                            })
+                            .collect::<Vec<_>>()
+                            .join("\n")
+                    }
+                };
+
+                Ok(HarmonyMessage {
+                    author: Author {
+                        role: harmony_role,
+                        name: None,
+                    },
+                    recipient: None,
+                    content: vec![Content::Text(TextContent { text })],
+                    channel: None,
+                    content_type: None,
+                })
+            }
+        }
+    }
+
+    /// Convert OpenAI ChatMessage format to Harmony messages
+    ///
+    /// - Assistant messages with tool_calls create multiple messages (one per tool call)
+    /// - Tool role messages use Role::Tool with proper author
+    /// - Tool-related messages use channel="commentary"
+    fn convert_chat_messages(
+        &self,
+        messages: &[ChatMessage],
+    ) -> Result<Vec<HarmonyMessage>, String> {
+        let mut harmony_messages = Vec::new();
+
+        // Build a map of tool_call_id -> function_name for tool responses
+        let mut tool_call_map = std::collections::HashMap::new();
+        for msg in messages {
+            if let ChatMessage::Assistant {
+                tool_calls: Some(calls),
+                ..
+            } = msg
+            {
+                for call in calls {
+                    tool_call_map.insert(call.id.clone(), call.function.name.clone());
+                }
+            }
+        }
+
+        for msg in messages {
+            match msg {
+                ChatMessage::System { content, name } => {
+                    // System messages stay as-is
+                    let harmony_msg = HarmonyMessage {
+                        author: Author {
+                            role: Role::System,
+                            name: name.clone(),
+                        },
+                        recipient: None,
+                        content: vec![Content::Text(TextContent {
+                            text: content.to_simple_string(),
+                        })],
+                        channel: None,
+                        content_type: None,
+                    };
+                    harmony_messages.push(harmony_msg);
+                }
+                ChatMessage::Developer {
+                    content,
+                    name,
+                    tools: _,
+                } => {
+                    // Developer messages stay as-is
+                    let harmony_msg = HarmonyMessage {
+                        author: Author {
+                            role: Role::Developer,
+                            name: name.clone(),
+                        },
+                        recipient: None,
+                        content: vec![Content::Text(TextContent {
+                            text: content.to_simple_string(),
+                        })],
+                        channel: None,
+                        content_type: None,
+                    };
+                    harmony_messages.push(harmony_msg);
+                }
+
+                ChatMessage::User { content, name } => {
+                    // Extract text from user content
+                    let text = match content {
+                        MessageContent::Text(text) => text.clone(),
+                        MessageContent::Parts(parts) => {
+                            // For multimodal content, extract text parts
+                            parts
+                                .iter()
+                                .filter_map(|part| {
+                                    if let ContentPart::Text { text } = part {
+                                        Some(text.as_str())
+                                    } else {
+                                        None
+                                    }
+                                })
+                                .collect::<Vec<_>>()
+                                .join("\n")
+                        }
+                    };
+
+                    let harmony_msg = HarmonyMessage {
+                        author: Author {
+                            role: Role::User,
+                            name: name.clone(),
+                        },
+                        recipient: None,
+                        content: vec![Content::Text(TextContent { text })],
+                        channel: None,
+                        content_type: None,
+                    };
+                    harmony_messages.push(harmony_msg);
+                }
+
+                ChatMessage::Assistant {
+                    content,
+                    name,
+                    tool_calls,
+                    reasoning_content,
+                } => {
+                    if let Some(calls) = tool_calls {
+                        // Create one message per tool call with channel="commentary"
+                        for call in calls {
+                            let function_name = &call.function.name;
+                            let arguments = call.function.arguments.clone().unwrap_or_default();
+
+                            let tool_call_msg = HarmonyMessage {
+                                author: Author {
+                                    role: Role::Assistant,
+                                    name: name.clone(),
+                                },
+                                recipient: Some(format!("functions.{}", function_name)),
+                                content: vec![Content::Text(TextContent { text: arguments })],
+                                channel: Some("commentary".to_string()),
+                                content_type: Some("json".to_string()),
+                            };
+                            harmony_messages.push(tool_call_msg);
+                        }
+                    } else {
+                        // Regular assistant message with content
+                        // Combine content with reasoning if present
+                        let mut text = content
+                            .as_ref()
+                            .map(|c| c.to_simple_string())
+                            .unwrap_or_default();
+
+                        if let Some(reasoning) = reasoning_content {
+                            if !text.is_empty() {
+                                text.push('\n');
+                            }
+                            text.push_str(reasoning);
+                        }
+
+                        let harmony_msg = HarmonyMessage {
+                            author: Author {
+                                role: Role::Assistant,
+                                name: name.clone(),
+                            },
+                            recipient: None,
+                            content: vec![Content::Text(TextContent { text })],
+                            channel: Some("final".to_string()),
+                            content_type: None,
+                        };
+                        harmony_messages.push(harmony_msg);
+                    }
+                }
+
+                ChatMessage::Tool {
+                    content,
+                    tool_call_id,
+                } => {
+                    // Look up the function name from the tool_call_id
+                    let function_name = tool_call_map
+                        .get(tool_call_id)
+                        .cloned()
+                        .unwrap_or_else(|| tool_call_id.clone());
+
+                    // Tool result - Must include recipient="assistant" for parser to recognize it.
+                    // We keep channel=None to minimize what the model might copy.
+                    let harmony_msg = HarmonyMessage {
+                        author: Author {
+                            role: Role::Tool,
+                            name: Some(format!("functions.{}", function_name)),
+                        },
+                        recipient: Some("assistant".to_string()),
+                        content: vec![Content::Text(TextContent {
+                            text: content.to_simple_string(),
+                        })],
+                        channel: None,
+                        content_type: None,
+                    };
+                    harmony_messages.push(harmony_msg);
+                }
+
+                ChatMessage::Function { content, name } => {
+                    // Function messages also use Role::Tool
+                    // Tool result - Must include recipient="assistant" for parser to recognize it.
+                    // We keep channel=None to minimize what the model might copy.
+                    let harmony_msg = HarmonyMessage {
+                        author: Author {
+                            role: Role::Tool,
+                            name: Some(format!("functions.{}", name)),
+                        },
+                        recipient: Some("assistant".to_string()),
+                        content: vec![Content::Text(TextContent {
+                            text: content.clone(),
+                        })],
+                        channel: None,
+                        content_type: None,
+                    };
+                    harmony_messages.push(harmony_msg);
+                }
+            }
+        }
+
+        Ok(harmony_messages)
+    }
+
+    /// Extract selection text for worker routing
+    ///
+    /// Uses the last user message for load balancing
+    fn extract_selection_text(&self, messages: &[HarmonyMessage]) -> String {
+        // Find the last user message
+        if let Some(last_user_msg) = messages.iter().rev().find(|m| m.author.role == Role::User) {
+            // Extract full text from content
+            return last_user_msg
+                .content
+                .iter()
+                .filter_map(|c| match c {
+                    Content::Text(tc) => Some(tc.text.as_str()),
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+                .join("");
+        }
+
+        // Fallback: concatenate all text
+        messages
+            .iter()
+            .flat_map(|m| &m.content)
+            .filter_map(|c| match c {
+                Content::Text(tc) => Some(tc.text.as_str()),
+                _ => None,
+            })
+            .collect::<Vec<_>>()
+            .join(" ")
+    }
+}
+
+impl Default for HarmonyBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/detector.rs

@@ -0,0 +1,77 @@
+//! Harmony model detection
+
+use crate::core::{Worker, WorkerRegistry};
+
+/// Harmony model detector
+///
+/// Detects if a model name indicates support for Harmony encoding/parsing.
+pub(crate) struct HarmonyDetector;
+
+impl HarmonyDetector {
+    /// Check if a worker is a Harmony/GPT-OSS model.
+    ///
+    /// Detection priority:
+    /// 1. Check if any model card has architectures containing "GptOssForCausalLM"
+    /// 2. Check if any model card has hf_model_type equal to "gpt_oss"
+    /// 3. Check if model_id contains "gpt-oss" substring (case-insensitive)
+    pub fn is_harmony_worker(worker: &dyn Worker) -> bool {
+        for model_card in worker.models() {
+            // 1. Check architectures for GptOssForCausalLM
+            if model_card
+                .architectures
+                .iter()
+                .any(|arch| arch == "GptOssForCausalLM")
+            {
+                return true;
+            }
+
+            // 2. Check hf_model_type for gpt_oss
+            if let Some(ref model_type) = model_card.hf_model_type {
+                if model_type == "gpt_oss" {
+                    return true;
+                }
+            }
+
+            // 3. Check model id for gpt-oss substring
+            if Self::is_harmony_model(&model_card.id) {
+                return true;
+            }
+        }
+
+        // Fallback: check worker's model_id directly
+        Self::is_harmony_model(worker.model_id())
+    }
+
+    /// Check if a model name contains "gpt-oss" (case-insensitive).
+    pub fn is_harmony_model(model_name: &str) -> bool {
+        // Case-insensitive substring search without heap allocation
+        // More efficient than to_lowercase() which allocates a new String
+        model_name
+            .as_bytes()
+            .windows(7) // "gpt-oss".len()
+            .any(|window| window.eq_ignore_ascii_case(b"gpt-oss"))
+    }
+
+    /// Check if any worker for the given model is a Harmony/GPT-OSS worker.
+    ///
+    /// This method looks up workers from the registry by model name and checks
+    /// if any of them are Harmony workers based on their metadata (architectures,
+    /// hf_model_type).
+    ///
+    /// Falls back to string-based detection if no workers are registered for
+    /// the model (e.g., during startup before workers are discovered).
+    pub fn is_harmony_model_in_registry(registry: &WorkerRegistry, model_name: &str) -> bool {
+        // Get workers for this model
+        let workers = registry.get_by_model(model_name);
+
+        if workers.is_empty() {
+            // No workers found - fall back to string-based detection
+            return Self::is_harmony_model(model_name);
+        }
+
+        // Check if any worker is a Harmony worker
+        workers
+            .iter()
+            .any(|worker| Self::is_harmony_worker(worker.as_ref()))
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/mod.rs

@@ -0,0 +1,48 @@
+//! Harmony pipeline implementation
+//!
+//! This module provides support for GPT-OSS models that use Harmony encoding/parsing.
+//! The Harmony protocol uses a channel-based approach with three channels:
+//! - **analysis**: Reasoning/thinking content (optional)
+//! - **commentary**: Tool calls (optional)
+//! - **final**: Final response text (required)
+//!
+//! ## Architecture
+//!
+//! The Harmony implementation is structured as follows:
+//!
+//! - **detector**: Model detection (is this a Harmony-capable model?)
+//! - **builder**: Request encoding (convert Chat/Responses → input_ids)
+//! - **parser**: Response parsing (output_ids → channels)
+//! - **types**: Shared type definitions
+//!
+//! ## Usage
+//!
+//! ```ignore
+//! use smg::routers::grpc::harmony::{HarmonyDetector, HarmonyBuilder};
+//!
+//! // Detect if model supports Harmony
+//! if HarmonyDetector::is_harmony_model("gpt-4o") {
+//!     // Build Harmony request
+//!     let builder = HarmonyBuilder::new();
+//!     let output = builder.build_from_chat(&request)?;
+//!     // ... use output.input_ids for gRPC request
+//! }
+//! ```
+
+pub(crate) mod builder;
+pub(crate) mod detector;
+pub(crate) mod parser;
+pub(crate) mod processor;
+pub(crate) mod responses;
+pub(crate) mod stages;
+pub(crate) mod streaming;
+pub(crate) mod types;
+
+// Re-export types that are accessed via harmony::TypeName
+pub(crate) use builder::HarmonyBuilder;
+pub(crate) use detector::HarmonyDetector;
+pub(crate) use parser::HarmonyParserAdapter;
+pub(crate) use processor::{HarmonyResponseProcessor, ResponsesIterationResult};
+pub(crate) use responses::{serve_harmony_responses, serve_harmony_responses_stream};
+pub(crate) use streaming::HarmonyStreamingProcessor;
+pub(crate) use types::HarmonyMessage;

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/parser.rs

@@ -0,0 +1,536 @@
+//! Harmony response parser
+//!
+//! Adapter for openai_harmony::StreamableParser that handles channel-based parsing.
+
+use openai_harmony::{chat::Role, HarmonyEncoding, StreamableParser};
+use uuid::Uuid;
+
+use super::types::{HarmonyChannelDelta, HarmonyChannelOutput};
+use crate::protocols::common::{FunctionCallResponse, ToolCall};
+
+/// Get the global Harmony encoding
+///
+/// References the same encoding used by the builder for consistency
+fn get_harmony_encoding() -> &'static HarmonyEncoding {
+    use super::builder::get_harmony_encoding;
+    get_harmony_encoding()
+}
+
+/// Harmony parser adapter
+///
+/// Wraps openai_harmony::StreamableParser and provides methods for parsing
+/// complete responses and streaming chunks.
+pub(crate) struct HarmonyParserAdapter {
+    parser: StreamableParser,
+    prev_recipient: Option<String>,
+    reasoning_token_count: u32,
+}
+
+impl HarmonyParserAdapter {
+    /// Create a new Harmony parser
+    pub fn new() -> Result<Self, String> {
+        let encoding = get_harmony_encoding();
+        let parser = StreamableParser::new(encoding.clone(), Some(Role::Assistant))
+            .map_err(|e| format!("Failed to create StreamableParser: {}", e))?;
+
+        Ok(Self {
+            parser,
+            prev_recipient: None,
+            reasoning_token_count: 0,
+        })
+    }
+
+    /// Extract text from message content (private helper)
+    ///
+    /// Filters text content from a message's content array and joins them into a single string.
+    ///
+    /// # Arguments
+    ///
+    /// * `content` - The content array from a Harmony message
+    ///
+    /// # Returns
+    ///
+    /// Joined text string from all text content items
+    fn extract_text_from_content(content: &[openai_harmony::chat::Content]) -> String {
+        content
+            .iter()
+            .filter_map(|c| match c {
+                openai_harmony::chat::Content::Text(tc) => Some(tc.text.as_str()),
+                _ => None,
+            })
+            .collect::<Vec<_>>()
+            .join("")
+    }
+
+    /// Handle incomplete content from parser state (private helper)
+    ///
+    /// Checks for any remaining incomplete content in the parser and appends it
+    /// to the appropriate channel (analysis or final_text).
+    ///
+    /// # Arguments
+    ///
+    /// * `parser` - Reference to the StreamableParser
+    /// * `analysis` - Mutable reference to analysis content
+    /// * `final_text` - Mutable reference to final text content
+    fn handle_incomplete_content(
+        parser: &StreamableParser,
+        analysis: &mut Option<String>,
+        final_text: &mut String,
+    ) {
+        if let Ok(current_content) = parser.current_content() {
+            if !current_content.is_empty() {
+                let current_channel = parser.current_channel();
+                match current_channel.as_deref() {
+                    Some("analysis") => {
+                        *analysis = Some(current_content);
+                    }
+                    Some("final") | None => {
+                        final_text.push_str(&current_content);
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+
+    /// Parse messages into channel outputs (private helper)
+    ///
+    /// Extracts analysis, commentary (tool calls), and final text from Harmony messages.
+    /// This is the core parsing logic shared by both parse_complete and finalize.
+    ///
+    /// # Arguments
+    ///
+    /// * `messages` - The messages to parse from the Harmony parser
+    ///
+    /// # Returns
+    ///
+    /// Tuple of (analysis, commentary, final_text)
+    pub fn parse_messages(
+        messages: &[openai_harmony::chat::Message],
+    ) -> (Option<String>, Option<Vec<ToolCall>>, String) {
+        let mut analysis: Option<String> = None;
+        let mut commentary: Option<Vec<ToolCall>> = None;
+        let mut final_text = String::new();
+
+        for msg in messages {
+            // Filter: Only process assistant messages
+            if msg.author.role != Role::Assistant {
+                continue;
+            }
+
+            let channel = msg.channel.as_deref().unwrap_or("");
+            let recipient = msg.recipient.as_deref();
+
+            // IMPORTANT: Check recipient FIRST before channel
+            // The model sometimes generates tool calls with channel="analysis" + recipient="functions.*"
+            // instead of channel="commentary" + recipient="functions.*"
+            // We should trust the recipient field to determine if this is a tool call
+            if let Some(recipient_str) = recipient {
+                if recipient_str.starts_with("functions.") {
+                    // This is a tool call, regardless of channel
+                    let function_name = recipient_str.strip_prefix("functions.").unwrap();
+
+                    // Process each content item separately
+                    for content in &msg.content {
+                        if let openai_harmony::chat::Content::Text(tc) = content {
+                            let call_id = format!("call_{}", Uuid::new_v4());
+                            let tool_call = ToolCall {
+                                id: call_id,
+                                tool_type: "function".to_string(),
+                                function: FunctionCallResponse {
+                                    name: function_name.to_string(),
+                                    arguments: Some(tc.text.clone()),
+                                },
+                            };
+
+                            match commentary.as_mut() {
+                                Some(calls) => calls.push(tool_call),
+                                None => commentary = Some(vec![tool_call]),
+                            }
+                        }
+                    }
+                    // Skip further channel processing for this message
+                    continue;
+                } else if recipient_str.starts_with("python")
+                    || recipient_str.starts_with("browser")
+                    || recipient_str.starts_with("container")
+                {
+                    // Built-in tools → treat as reasoning
+                    // For Chat API, we add to analysis content
+                    let text = Self::extract_text_from_content(&msg.content);
+
+                    if !text.is_empty() {
+                        // Append to analysis (built-in tools are reasoning)
+                        match analysis.as_mut() {
+                            Some(existing) => {
+                                existing.push('\n');
+                                existing.push_str(&text);
+                            }
+                            None => analysis = Some(text),
+                        }
+                    }
+                    // Skip further channel processing
+                    continue;
+                }
+            }
+
+            // Now process by channel (only if not already handled by recipient)
+            match channel {
+                "analysis" => {
+                    // Process each content item
+                    // For Chat API, we join them into a single reasoning_content
+                    let text = Self::extract_text_from_content(&msg.content);
+
+                    if !text.is_empty() {
+                        analysis = Some(text);
+                    }
+                }
+                "commentary" => {
+                    // If we reach here, recipient was not "functions.*" or built-in tools
+                    // Commentary channel should always have a recipient
+                    // This is likely a model bug - log warning and treat as reasoning
+                    tracing::warn!(
+                        channel = "commentary",
+                        recipient = ?recipient,
+                        "Commentary message without valid recipient, treating as reasoning"
+                    );
+
+                    let text = Self::extract_text_from_content(&msg.content);
+
+                    if !text.is_empty() {
+                        match analysis.as_mut() {
+                            Some(existing) => {
+                                existing.push('\n');
+                                existing.push_str(&text);
+                            }
+                            None => analysis = Some(text),
+                        }
+                    }
+                }
+                "final" => {
+                    // Process final channel content
+                    let text = Self::extract_text_from_content(&msg.content);
+                    final_text.push_str(&text);
+                }
+                _ => {
+                    // Unknown channel, append to final text as fallback
+                    let text = Self::extract_text_from_content(&msg.content);
+                    final_text.push_str(&text);
+                }
+            }
+        }
+
+        (analysis, commentary, final_text)
+    }
+
+    /// Parse complete response
+    ///
+    /// Parses all output token IDs and returns the complete channel output
+    /// containing analysis, commentary (tool calls), and final text.
+    ///
+    /// # Arguments
+    ///
+    /// * `output_ids` - The complete output token IDs from the model
+    /// * `finish_reason` - The finish reason from GenerateComplete ("stop", "length", etc.)
+    /// * `matched_stop` - Optional matched stop token information from GenerateComplete
+    ///
+    /// # Returns
+    ///
+    /// Complete HarmonyChannelOutput with all three channels parsed
+    pub fn parse_complete(
+        &mut self,
+        output_ids: &[u32],
+        finish_reason: String,
+        matched_stop: Option<serde_json::Value>,
+    ) -> Result<HarmonyChannelOutput, String> {
+        let mut reasoning_token_count = 0u32;
+
+        // Feed all tokens to the parser
+        for &token_id in output_ids {
+            self.parser
+                .process(token_id)
+                .map_err(|e| format!("Failed to process token {}: {}", token_id, e))?;
+
+            // Count reasoning tokens (analysis + commentary channels)
+            if let Some(channel) = self.parser.current_channel() {
+                if channel == "analysis" || channel == "commentary" {
+                    reasoning_token_count += 1;
+                }
+            }
+        }
+
+        // Extract all completed messages from the parser
+        let messages = self.parser.messages();
+
+        // Parse messages into channel outputs using shared helper
+        let (mut analysis, commentary, mut final_text) = Self::parse_messages(messages);
+
+        // Check for incomplete content in parser state
+        Self::handle_incomplete_content(&self.parser, &mut analysis, &mut final_text);
+
+        // Determine finish reason: override to "tool_calls" if commentary has tool calls
+        let final_finish_reason = if commentary.is_some() {
+            "tool_calls".to_string()
+        } else {
+            finish_reason.clone()
+        };
+
+        Ok(HarmonyChannelOutput {
+            analysis,
+            commentary,
+            final_text,
+            finish_reason: final_finish_reason,
+            matched_stop,
+            reasoning_token_count,
+        })
+    }
+
+    /// Get all messages from the parser
+    ///
+    /// Returns the raw messages extracted by the Harmony parser.
+    /// Used for validation checks.
+    pub fn get_messages(&self) -> Vec<openai_harmony::chat::Message> {
+        self.parser.messages().to_vec()
+    }
+
+    /// Extract incomplete commentary content from parser state
+    ///
+    /// When the stream ends, there may be incomplete commentary content in the parser
+    /// that hasn't been finalized into a completed message. This method extracts
+    /// such content and converts it to tool calls.
+    ///
+    /// # Returns
+    ///
+    /// Optional vector of ToolCall if incomplete commentary is found
+    pub fn extract_incomplete_commentary(&self) -> Option<Vec<ToolCall>> {
+        // Check if current channel is commentary
+        let current_channel = self.parser.current_channel();
+        if current_channel.as_deref() != Some("commentary") {
+            return None;
+        }
+
+        // Get current recipient (should be "functions.{name}")
+        let recipient = self.parser.current_recipient()?;
+        if !recipient.starts_with("functions.") {
+            return None;
+        }
+
+        // Get current incomplete content
+        let content = self.parser.current_content().ok()?;
+        if content.is_empty() {
+            return None;
+        }
+
+        // Extract function name from recipient
+        let function_name = recipient.strip_prefix("functions.").unwrap();
+
+        // Create tool call from incomplete content
+        let call_id = format!("call_{}", Uuid::new_v4());
+        let tool_call = ToolCall {
+            id: call_id,
+            tool_type: "function".to_string(),
+            function: FunctionCallResponse {
+                name: function_name.to_string(),
+                arguments: Some(content),
+            },
+        };
+
+        Some(vec![tool_call])
+    }
+
+    /// Parse streaming chunk
+    ///
+    /// Parses incremental token IDs and returns a delta with any new content
+    /// from the analysis, commentary, or final channels.
+    ///
+    /// # Arguments
+    ///
+    /// * `chunk_ids` - New token IDs from the current chunk
+    ///
+    /// # Returns
+    ///
+    /// Optional HarmonyChannelDelta if there's new content to emit
+    pub fn parse_chunk(
+        &mut self,
+        chunk_ids: &[u32],
+    ) -> Result<Option<HarmonyChannelDelta>, String> {
+        let mut has_delta = false;
+        let mut analysis_delta = None;
+        let mut commentary_delta = None;
+        let mut final_delta = None;
+
+        // Track message count before processing
+        let prev_message_count = self.parser.messages().len();
+
+        // Accumulate delta text for commentary channel
+        let mut accumulated_delta = String::new();
+
+        // Process each token
+        for &token_id in chunk_ids {
+            self.parser
+                .process(token_id)
+                .map_err(|e| format!("Failed to process token {}: {}", token_id, e))?;
+
+            // Count reasoning tokens (analysis + commentary channels)
+            if let Some(channel) = self.parser.current_channel() {
+                if channel == "analysis" || channel == "commentary" {
+                    self.reasoning_token_count += 1;
+                }
+            }
+
+            // Check for content delta
+            if let Ok(Some(delta_text)) = self.parser.last_content_delta() {
+                has_delta = true;
+
+                // Determine which channel this delta belongs to
+                let channel = self.parser.current_channel();
+                match channel.as_deref() {
+                    Some("analysis") => {
+                        analysis_delta = Some(delta_text);
+                    }
+                    Some("final") | None => {
+                        final_delta = Some(delta_text);
+                    }
+                    Some("commentary") => {
+                        // Accumulate delta for commentary
+                        accumulated_delta.push_str(&delta_text);
+                    }
+                    _ => {}
+                }
+            }
+        }
+
+        // Handle commentary channel tool call deltas
+        if self.parser.current_channel().as_deref() == Some("commentary") {
+            if let Some(cur_recipient) = self.parser.current_recipient() {
+                if cur_recipient.starts_with("functions.") {
+                    has_delta = true;
+
+                    // Count completed tool calls for index
+                    let base_index = self
+                        .parser
+                        .messages()
+                        .iter()
+                        .filter(|msg| {
+                            msg.channel.as_deref() == Some("commentary")
+                                && msg
+                                    .recipient
+                                    .as_deref()
+                                    .is_some_and(|r| r.starts_with("functions."))
+                        })
+                        .count();
+
+                    // Check if recipient changed (new tool call)
+                    let recipient_changed = self.prev_recipient.as_deref() != Some(&cur_recipient);
+
+                    if recipient_changed {
+                        // NEW tool call: emit name + id
+                        let tool_name = cur_recipient.strip_prefix("functions.").unwrap();
+                        let call_id = format!("call_{}", Uuid::new_v4());
+
+                        commentary_delta = Some(super::types::ToolCallDelta {
+                            index: base_index,
+                            id: Some(call_id),
+                            function: Some(super::types::FunctionDelta {
+                                name: Some(tool_name.to_string()),
+                                arguments: Some(String::new()),
+                            }),
+                        });
+
+                        // Update prev_recipient
+                        self.prev_recipient = Some(cur_recipient);
+                    } else if !accumulated_delta.is_empty() {
+                        // CONTINUING tool call: emit arguments delta
+                        commentary_delta = Some(super::types::ToolCallDelta {
+                            index: base_index,
+                            id: None,
+                            function: Some(super::types::FunctionDelta {
+                                name: None,
+                                arguments: Some(accumulated_delta),
+                            }),
+                        });
+                    }
+                }
+            }
+        }
+
+        // Check if new messages were completed
+        let current_message_count = self.parser.messages().len();
+        let is_final = current_message_count > prev_message_count;
+
+        if has_delta {
+            Ok(Some(HarmonyChannelDelta {
+                analysis_delta,
+                commentary_delta,
+                final_delta,
+                is_final,
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Finalize parsing
+    ///
+    /// Called at the end of streaming to get the final state and any
+    /// remaining content.
+    ///
+    /// # Arguments
+    ///
+    /// * `finish_reason` - The finish reason from GenerateComplete ("stop", "length", etc.)
+    /// * `matched_stop` - Optional matched stop token information from GenerateComplete
+    ///
+    /// # Returns
+    ///
+    /// Final HarmonyChannelOutput with complete parsed content
+    pub fn finalize(
+        &mut self,
+        finish_reason: String,
+        matched_stop: Option<serde_json::Value>,
+    ) -> Result<HarmonyChannelOutput, String> {
+        // Extract all completed messages
+        let messages = self.parser.messages();
+
+        // Parse messages into channel outputs using shared helper
+        let (mut analysis, commentary, mut final_text) = Self::parse_messages(messages);
+
+        // Check for remaining incomplete content
+        Self::handle_incomplete_content(&self.parser, &mut analysis, &mut final_text);
+
+        // Determine finish reason: override to "tool_calls" if commentary has tool calls
+        let final_finish_reason = if commentary.is_some() {
+            "tool_calls".to_string()
+        } else {
+            finish_reason
+        };
+
+        Ok(HarmonyChannelOutput {
+            analysis,
+            commentary,
+            final_text,
+            finish_reason: final_finish_reason,
+            matched_stop,
+            reasoning_token_count: self.reasoning_token_count,
+        })
+    }
+
+    /// Reset parser state
+    ///
+    /// Resets the parser to initial state for reuse
+    #[allow(dead_code)]
+    pub fn reset(&mut self) -> Result<(), String> {
+        // Create a new parser instance (StreamableParser doesn't have a reset method)
+        let encoding = get_harmony_encoding();
+        self.parser = StreamableParser::new(encoding.clone(), Some(Role::Assistant))
+            .map_err(|e| format!("Failed to reset parser: {}", e))?;
+        self.prev_recipient = None;
+        self.reasoning_token_count = 0;
+        Ok(())
+    }
+}
+
+impl Default for HarmonyParserAdapter {
+    fn default() -> Self {
+        Self::new().expect("Failed to create default parser")
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/processor.rs

@@ -0,0 +1,345 @@
+//! Harmony response processor for non-streaming responses
+
+use std::sync::Arc;
+
+use axum::response::Response;
+use smg_grpc_client::sglang_proto::generate_complete::MatchedStop::{
+    MatchedStopStr, MatchedTokenId,
+};
+use tracing::error;
+
+use super::HarmonyParserAdapter;
+use crate::{
+    protocols::{
+        chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse},
+        common::{CompletionTokensDetails, ToolCall, Usage},
+        responses::{
+            OutputTokensDetails, ResponseContentPart, ResponseOutputItem, ResponseReasoningContent,
+            ResponseStatus, ResponseUsage, ResponsesRequest, ResponsesResponse, ResponsesUsage,
+        },
+    },
+    routers::{
+        error,
+        grpc::{
+            common::{response_collection, response_formatting},
+            context::{DispatchMetadata, ExecutionResult},
+        },
+    },
+};
+
+/// Processor for non-streaming Harmony responses
+///
+/// Collects all output tokens from execution and parses them using
+/// HarmonyParserAdapter to extract the complete response.
+pub(crate) struct HarmonyResponseProcessor;
+
+impl HarmonyResponseProcessor {
+    /// Create a new Harmony response processor
+    pub fn new() -> Self {
+        Self
+    }
+
+    /// Process a non-streaming Harmony chat response
+    pub async fn process_non_streaming_chat_response(
+        &self,
+        execution_result: ExecutionResult,
+        chat_request: Arc<ChatCompletionRequest>,
+        dispatch: DispatchMetadata,
+    ) -> Result<ChatCompletionResponse, Response> {
+        // Collect all completed responses (one per choice)
+        let all_responses = response_collection::collect_responses(execution_result, false).await?;
+        if all_responses.is_empty() {
+            return Err(error::internal_error(
+                "no_responses_from_server",
+                "No responses from server",
+            ));
+        }
+
+        // Build choices by parsing output with HarmonyParserAdapter
+        let mut choices: Vec<ChatChoice> = Vec::new();
+        let mut total_reasoning_tokens = 0u32;
+
+        for (index, complete) in all_responses.iter().enumerate() {
+            // Convert matched_stop from proto to JSON
+            let matched_stop = complete.matched_stop().map(|m| match m {
+                MatchedTokenId(id) => {
+                    serde_json::json!(id)
+                }
+                MatchedStopStr(s) => {
+                    serde_json::json!(s)
+                }
+            });
+
+            // Parse Harmony channels with HarmonyParserAdapter
+            let mut parser = HarmonyParserAdapter::new().map_err(|e| {
+                error!(
+                    function = "process_non_streaming_chat_response",
+                    error = %e,
+                    "Failed to create Harmony parser"
+                );
+                error::internal_error(
+                    "create_harmony_parser_failed",
+                    format!("Failed to create Harmony parser: {}", e),
+                )
+            })?;
+
+            // Parse Harmony channels with finish_reason and matched_stop
+            let parsed = parser
+                .parse_complete(
+                    complete.output_ids(),
+                    complete.finish_reason().to_string(),
+                    matched_stop.clone(),
+                )
+                .map_err(|e| {
+                    error!(
+                        function = "process_non_streaming_chat_response",
+                        error = %e,
+                        "Harmony parsing failed on complete response"
+                    );
+                    error::internal_error(
+                        "harmony_parsing_failed",
+                        format!("Harmony parsing failed: {}", e),
+                    )
+                })?;
+
+            // Build response message (assistant)
+            let message = ChatCompletionMessage {
+                role: "assistant".to_string(),
+                content: (!parsed.final_text.is_empty()).then_some(parsed.final_text),
+                tool_calls: parsed.commentary,
+                reasoning_content: parsed.analysis,
+            };
+
+            let finish_reason = parsed.finish_reason;
+
+            // Accumulate reasoning tokens across all responses
+            total_reasoning_tokens += parsed.reasoning_token_count;
+
+            choices.push(ChatChoice {
+                index: index as u32,
+                message,
+                logprobs: None,
+                finish_reason: Some(finish_reason),
+                matched_stop,
+                hidden_states: None,
+            });
+        }
+
+        // Build usage from proto fields
+        let mut usage = response_formatting::build_usage(&all_responses);
+
+        // Add reasoning token count from parsed analysis/commentary channels
+        if total_reasoning_tokens > 0 {
+            usage.completion_tokens_details = Some(CompletionTokensDetails {
+                reasoning_tokens: Some(total_reasoning_tokens),
+            });
+        }
+
+        // Final ChatCompletionResponse
+        Ok(
+            ChatCompletionResponse::builder(&dispatch.request_id, &chat_request.model)
+                .created(dispatch.created)
+                .choices(choices)
+                .usage(usage)
+                .maybe_system_fingerprint(dispatch.weight_version.as_deref())
+                .build(),
+        )
+    }
+}
+
+impl Default for HarmonyResponseProcessor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Result of processing a single Responses API iteration
+///
+/// Used by the MCP tool loop to determine whether to continue
+/// executing tools or return the final response.
+pub(crate) enum ResponsesIterationResult {
+    /// Tool calls found in commentary channel - continue MCP loop
+    ToolCallsFound {
+        tool_calls: Vec<ToolCall>,
+        analysis: Option<String>, // For streaming emission or reasoning output
+        partial_text: String,     // For streaming emission or message output
+        usage: Usage,             // Token usage from this iteration
+        request_id: String,       // Request ID from dispatch
+    },
+    /// No tool calls - return final ResponsesResponse
+    Completed {
+        response: Box<ResponsesResponse>,
+        usage: Usage,
+    },
+}
+
+impl HarmonyResponseProcessor {
+    /// Process a single Responses API iteration
+    ///
+    /// Parses Harmony channels and determines if tool calls are present.
+    /// If tool calls found, returns ToolCallsFound for MCP loop to execute.
+    /// If no tool calls, builds final ResponsesResponse.
+    ///
+    /// # Arguments
+    ///
+    /// * `execution_result` - The execution result from the model
+    /// * `responses_request` - The original Responses API request
+    /// * `dispatch` - Dispatch metadata for request tracking
+    ///
+    /// # Returns
+    ///
+    /// ResponsesIterationResult indicating whether to continue loop or return
+    pub async fn process_responses_iteration(
+        &self,
+        execution_result: ExecutionResult,
+        responses_request: Arc<ResponsesRequest>,
+        dispatch: DispatchMetadata,
+    ) -> Result<ResponsesIterationResult, Response> {
+        // Collect all completed responses
+        let all_responses = response_collection::collect_responses(execution_result, false).await?;
+        if all_responses.is_empty() {
+            return Err(error::internal_error(
+                "no_responses_from_server",
+                "No responses from server",
+            ));
+        }
+
+        // For Responses API, we only process the first response (n=1)
+        let complete = all_responses
+            .first()
+            .ok_or_else(|| error::internal_error("no_complete_response", "No complete response"))?;
+
+        // Parse Harmony channels
+        let mut parser = HarmonyParserAdapter::new().map_err(|e| {
+            error!(
+                function = "process_responses_iteration",
+                error = %e,
+                "Failed to create Harmony parser"
+            );
+            error::internal_error(
+                "create_harmony_parser_failed",
+                format!("Failed to create Harmony parser: {}", e),
+            )
+        })?;
+
+        // Convert matched_stop from proto to JSON
+        let matched_stop = complete.matched_stop().map(|m| match m {
+            MatchedTokenId(id) => {
+                serde_json::json!(id)
+            }
+            MatchedStopStr(s) => {
+                serde_json::json!(s)
+            }
+        });
+
+        let parsed = parser
+            .parse_complete(
+                complete.output_ids(),
+                complete.finish_reason().to_string(),
+                matched_stop,
+            )
+            .map_err(|e| {
+                error!(
+                    function = "process_responses_iteration",
+                    error = %e,
+                    "Harmony parsing failed on complete response"
+                );
+                error::internal_error(
+                    "harmony_parsing_failed",
+                    format!("Harmony parsing failed: {}", e),
+                )
+            })?;
+
+        // VALIDATION: Check if model incorrectly generated Tool role messages
+        // This happens when the model copies the format of tool result messages
+        // instead of continuing as assistant. This is a model hallucination bug.
+        let messages = parser.get_messages();
+        let tool_messages_generated = messages.iter().any(|msg| {
+            msg.author.role == openai_harmony::chat::Role::Tool
+                && msg.recipient.as_deref() == Some("assistant")
+        });
+
+        if tool_messages_generated {
+            tracing::warn!(
+                "Model generated Tool->Assistant message instead of Assistant message. \
+                This is a model hallucination bug where it copies tool result format."
+            );
+        }
+
+        // Build usage (needed for both ToolCallsFound and Completed)
+        let mut usage = response_formatting::build_usage(std::slice::from_ref(complete));
+
+        // Add reasoning token count from parsed analysis/commentary channels
+        if parsed.reasoning_token_count > 0 {
+            usage.completion_tokens_details = Some(CompletionTokensDetails {
+                reasoning_tokens: Some(parsed.reasoning_token_count),
+            });
+        }
+
+        // Check for tool calls in commentary channel
+        if let Some(tool_calls) = parsed.commentary {
+            // Tool calls found - return for MCP loop execution
+            return Ok(ResponsesIterationResult::ToolCallsFound {
+                tool_calls,
+                analysis: parsed.analysis,
+                partial_text: parsed.final_text,
+                usage,
+                request_id: dispatch.request_id.clone(),
+            });
+        }
+
+        // No tool calls - build final ResponsesResponse
+        let mut output: Vec<ResponseOutputItem> = Vec::new();
+
+        // Map analysis channel → ResponseOutputItem::Reasoning
+        if let Some(analysis) = parsed.analysis {
+            let reasoning_item = ResponseOutputItem::Reasoning {
+                id: format!("reasoning_{}", dispatch.request_id),
+                summary: vec![],
+                content: vec![ResponseReasoningContent::ReasoningText { text: analysis }],
+                status: Some("completed".to_string()),
+            };
+            output.push(reasoning_item);
+        }
+
+        // Map final channel → ResponseOutputItem::Message
+        if !parsed.final_text.is_empty() {
+            let message_item = ResponseOutputItem::Message {
+                id: format!("msg_{}", dispatch.request_id),
+                role: "assistant".to_string(),
+                content: vec![ResponseContentPart::OutputText {
+                    text: parsed.final_text,
+                    annotations: vec![],
+                    logprobs: None,
+                }],
+                status: "completed".to_string(),
+            };
+            output.push(message_item);
+        }
+
+        // Build ResponsesResponse with all required fields
+        let response = ResponsesResponse::builder(&dispatch.request_id, &responses_request.model)
+            .copy_from_request(&responses_request)
+            .created_at(dispatch.created as i64)
+            .status(ResponseStatus::Completed)
+            .output(output)
+            .maybe_text(responses_request.text.clone())
+            .usage(ResponsesUsage::Modern(ResponseUsage {
+                input_tokens: usage.prompt_tokens,
+                output_tokens: usage.completion_tokens,
+                total_tokens: usage.total_tokens,
+                input_tokens_details: None,
+                output_tokens_details: usage.completion_tokens_details.as_ref().and_then(|d| {
+                    d.reasoning_tokens.map(|tokens| OutputTokensDetails {
+                        reasoning_tokens: tokens,
+                    })
+                }),
+            }))
+            .build();
+
+        Ok(ResponsesIterationResult::Completed {
+            response: Box::new(response),
+            usage,
+        })
+    }
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs

@@ -0,0 +1,364 @@
+//! Shared helpers and state tracking for Harmony Responses
+
+use axum::response::Response;
+use data_connector::ResponseId;
+use serde_json::{from_value, json, to_string, Value};
+use smg_mcp as mcp;
+use tracing::{debug, error, warn};
+use uuid::Uuid;
+
+use super::execution::ToolResult;
+use crate::{
+    protocols::{
+        common::{ToolCall, ToolChoice, ToolChoiceValue},
+        responses::{
+            McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem,
+            ResponseOutputItem, ResponseReasoningContent, ResponseTool, ResponseToolType,
+            ResponsesRequest, ResponsesResponse, StringOrContentParts,
+        },
+    },
+    routers::{error, grpc::common::responses::ResponsesContext},
+};
+
+/// Record of a single MCP tool call execution
+///
+/// Stores metadata needed to build mcp_call output items for Responses API format
+#[derive(Debug, Clone)]
+pub(super) struct McpCallRecord {
+    /// Tool call ID (stored for potential future use, currently generate new IDs)
+    #[allow(dead_code)]
+    pub call_id: String,
+    /// Tool name
+    pub tool_name: String,
+    /// JSON-encoded arguments
+    pub arguments: String,
+    /// JSON-encoded output/result
+    pub output: String,
+    /// Whether execution succeeded
+    pub success: bool,
+    /// Error message if execution failed
+    pub error: Option<String>,
+}
+
+/// Tracking structure for MCP tool calls across iterations
+///
+/// Accumulates all MCP tool call metadata during multi-turn conversation
+/// so we can build proper mcp_list_tools and mcp_call output items.
+#[derive(Debug, Clone)]
+pub(super) struct McpCallTracking {
+    /// MCP server label (e.g., "sglang-mcp")
+    pub server_label: String,
+    /// All tool call records across all iterations
+    pub tool_calls: Vec<McpCallRecord>,
+}
+
+impl McpCallTracking {
+    pub fn new(server_label: String) -> Self {
+        Self {
+            server_label,
+            tool_calls: Vec::new(),
+        }
+    }
+
+    pub fn record_call(
+        &mut self,
+        call_id: String,
+        tool_name: String,
+        arguments: String,
+        output: String,
+        success: bool,
+        error: Option<String>,
+    ) {
+        self.tool_calls.push(McpCallRecord {
+            call_id,
+            tool_name,
+            arguments,
+            output,
+            success,
+            error,
+        });
+    }
+
+    pub fn total_calls(&self) -> usize {
+        self.tool_calls.len()
+    }
+}
+
+/// Build a HashSet of MCP tool names for O(1) lookup
+///
+/// Creates a HashSet containing the names of all MCP tools in the request,
+/// allowing for efficient O(1) lookups when partitioning tool calls.
+pub(super) fn build_mcp_tool_names_set(
+    request_tools: &[ResponseTool],
+) -> std::collections::HashSet<&str> {
+    request_tools
+        .iter()
+        .filter(|t| t.r#type == ResponseToolType::Mcp)
+        .filter_map(|t| t.function.as_ref().map(|f| f.name.as_str()))
+        .collect()
+}
+
+/// Build next request with tool results appended to history
+///
+/// Constructs a new ResponsesRequest with:
+/// 1. Original input items (preserved)
+/// 2. Assistant message with analysis (reasoning) + partial_text + tool_calls
+/// 3. Tool result messages for each tool execution
+pub(super) fn build_next_request_with_tools(
+    mut request: ResponsesRequest,
+    tool_calls: Vec<ToolCall>,
+    tool_results: Vec<ToolResult>,
+    analysis: Option<String>, // Analysis channel content (becomes reasoning content)
+    partial_text: String,     // Final channel content (becomes message content)
+) -> Result<ResponsesRequest, Box<Response>> {
+    // Get current input items (or empty vec if Text variant)
+    let mut items = match request.input {
+        ResponseInput::Items(items) => items,
+        ResponseInput::Text(text) => {
+            // Convert text to items format
+            vec![ResponseInputOutputItem::SimpleInputMessage {
+                content: StringOrContentParts::String(text),
+                role: "user".to_string(),
+                r#type: None,
+            }]
+        }
+    };
+
+    // Build assistant response item with reasoning + content + tool calls
+    // This represents what the model generated in this iteration
+    let assistant_id = format!("msg_{}", Uuid::new_v4());
+
+    // Add reasoning if present (from analysis channel)
+    if let Some(analysis_text) = analysis {
+        items.push(ResponseInputOutputItem::Reasoning {
+            id: format!("reasoning_{}", assistant_id),
+            summary: vec![],
+            content: vec![ResponseReasoningContent::ReasoningText {
+                text: analysis_text,
+            }],
+            status: Some("completed".to_string()),
+        });
+    }
+
+    // Add message content if present (from final channel)
+    if !partial_text.is_empty() {
+        items.push(ResponseInputOutputItem::Message {
+            id: assistant_id.clone(),
+            role: "assistant".to_string(),
+            content: vec![ResponseContentPart::OutputText {
+                text: partial_text,
+                annotations: vec![],
+                logprobs: None,
+            }],
+            status: Some("completed".to_string()),
+        });
+    }
+
+    // Add function tool calls (from commentary channel)
+    for tool_call in tool_calls {
+        items.push(ResponseInputOutputItem::FunctionToolCall {
+            id: tool_call.id.clone(),
+            call_id: tool_call.id.clone(),
+            name: tool_call.function.name.clone(),
+            arguments: tool_call
+                .function
+                .arguments
+                .unwrap_or_else(|| "{}".to_string()),
+            output: None, // Output will be added next
+            status: Some("in_progress".to_string()),
+        });
+    }
+
+    // Add tool results
+    for tool_result in tool_results {
+        // Serialize tool output to string
+        let output_str = to_string(&tool_result.output).unwrap_or_else(|e| {
+            format!("{{\"error\": \"Failed to serialize tool output: {}\"}}", e)
+        });
+
+        // Update the corresponding tool call with output and completed status
+        // Find and update the matching FunctionToolCall
+        if let Some(ResponseInputOutputItem::FunctionToolCall {
+            output,
+            status,
+            ..
+        }) = items
+            .iter_mut()
+            .find(|item| matches!(item, ResponseInputOutputItem::FunctionToolCall { call_id, .. } if call_id == &tool_result.call_id))
+        {
+            *output = Some(output_str);
+            *status = if tool_result.is_error {
+                Some("failed".to_string())
+            } else {
+                Some("completed".to_string())
+            };
+        }
+    }
+
+    // Update request with new items
+    request.input = ResponseInput::Items(items);
+
+    // Switch tool_choice to "auto" for subsequent iterations
+    // This prevents infinite loops when original tool_choice was "required" or specific function
+    // After receiving tool results, the model should be free to decide whether to call more tools or finish
+    request.tool_choice = Some(ToolChoice::Value(ToolChoiceValue::Auto));
+
+    Ok(request)
+}
+
+/// Inject MCP metadata into final response
+///
+/// Adds mcp_list_tools and mcp_call output items to the response output array.
+/// Following non-Harmony pipeline pattern:
+/// 1. Prepend mcp_list_tools at the beginning
+/// 2. Append all mcp_call items at the end
+pub(super) fn inject_mcp_metadata(
+    response: &mut ResponsesResponse,
+    tracking: &McpCallTracking,
+    mcp_tools: &[mcp::Tool],
+) {
+    // Build mcp_list_tools item
+    let tools = mcp_tools;
+    let tools_info: Vec<McpToolInfo> = tools
+        .iter()
+        .map(|t| McpToolInfo {
+            name: t.name.to_string(),
+            description: t.description.as_ref().map(|d| d.to_string()),
+            input_schema: Value::Object((*t.input_schema).clone()),
+            annotations: Some(json!({
+                "read_only": false
+            })),
+        })
+        .collect();
+
+    let mcp_list_tools = ResponseOutputItem::McpListTools {
+        id: format!("mcpl_{}", Uuid::new_v4()),
+        server_label: tracking.server_label.clone(),
+        tools: tools_info,
+    };
+
+    // Build mcp_call items for each tracked call
+    let mcp_call_items: Vec<ResponseOutputItem> = tracking
+        .tool_calls
+        .iter()
+        .map(|record| ResponseOutputItem::McpCall {
+            id: format!("mcp_{}", Uuid::new_v4()),
+            status: if record.success {
+                "completed"
+            } else {
+                "failed"
+            }
+            .to_string(),
+            approval_request_id: None,
+            arguments: record.arguments.clone(),
+            error: record.error.clone(),
+            name: record.tool_name.clone(),
+            output: record.output.clone(),
+            server_label: tracking.server_label.clone(),
+        })
+        .collect();
+
+    // Inject into response output:
+    // 1. Prepend mcp_list_tools at the beginning
+    response.output.insert(0, mcp_list_tools);
+
+    // 2. Append all mcp_call items at the end
+    response.output.extend(mcp_call_items);
+}
+
+/// Load previous conversation messages from storage
+///
+/// If the request has `previous_response_id`, loads the response chain from storage
+/// and prepends the conversation history to the request input items.
+pub(super) async fn load_previous_messages(
+    ctx: &ResponsesContext,
+    request: ResponsesRequest,
+) -> Result<ResponsesRequest, Response> {
+    let Some(ref prev_id_str) = request.previous_response_id else {
+        // No previous_response_id, return request as-is
+        return Ok(request);
+    };
+
+    let prev_id = ResponseId::from(prev_id_str.as_str());
+
+    // Load response chain from storage
+    let chain = ctx
+        .response_storage
+        .get_response_chain(&prev_id, None)
+        .await
+        .map_err(|e| {
+            error!(
+                function = "load_previous_messages",
+                prev_id = %prev_id_str,
+                error = %e,
+                "Failed to load previous response chain from storage"
+            );
+            error::internal_error(
+                "load_previous_response_chain_failed",
+                format!(
+                    "Failed to load previous response chain for {}: {}",
+                    prev_id_str, e
+                ),
+            )
+        })?;
+
+    // Build conversation history from stored responses
+    let mut history_items = Vec::new();
+
+    // Helper to deserialize and collect items from a JSON array
+    let deserialize_items = |arr: &Value, item_type: &str| -> Vec<ResponseInputOutputItem> {
+        arr.as_array()
+            .into_iter()
+            .flat_map(|items| items.iter())
+            .filter_map(|item| {
+                from_value::<ResponseInputOutputItem>(item.clone())
+                    .map_err(|e| {
+                        warn!(
+                            "Failed to deserialize stored {} item: {}. Item: {}",
+                            item_type, e, item
+                        );
+                    })
+                    .ok()
+            })
+            .collect()
+    };
+
+    for stored in chain.responses.iter() {
+        history_items.extend(deserialize_items(&stored.input, "input"));
+        history_items.extend(deserialize_items(&stored.output, "output"));
+    }
+
+    debug!(
+        previous_response_id = %prev_id_str,
+        history_items_count = history_items.len(),
+        "Loaded conversation history from previous response"
+    );
+
+    // Build modified request with history prepended
+    let mut modified_request = request;
+
+    // Convert current input to items format
+    let all_items = match modified_request.input {
+        ResponseInput::Items(items) => {
+            // Prepend history to existing items
+            let mut combined = history_items;
+            combined.extend(items);
+            combined
+        }
+        ResponseInput::Text(text) => {
+            // Convert text to item and prepend history
+            history_items.push(ResponseInputOutputItem::SimpleInputMessage {
+                content: StringOrContentParts::String(text),
+                role: "user".to_string(),
+                r#type: None,
+            });
+            history_items
+        }
+    };
+
+    // Update request with combined items and clear previous_response_id
+    modified_request.input = ResponseInput::Items(all_items);
+    modified_request.previous_response_id = None;
+
+    Ok(modified_request)
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/responses/execution.rs

@@ -0,0 +1,224 @@
+//! MCP tool execution logic for Harmony Responses
+
+use std::{sync::Arc, time::Instant};
+
+use axum::response::Response;
+use serde_json::{from_str, json, to_string, to_value, Value};
+use smg_mcp::{self as mcp, McpManager};
+use tracing::{debug, error, warn};
+
+use super::common::McpCallTracking;
+use crate::{
+    observability::metrics::{metrics_labels, Metrics},
+    protocols::{
+        common::{Function, ToolCall},
+        responses::{ResponseTool, ResponseToolType},
+    },
+    routers::error,
+};
+
+/// Tool execution result
+///
+/// Contains the result of executing a single MCP tool.
+pub(crate) struct ToolResult {
+    /// Tool call ID (for matching with request)
+    pub call_id: String,
+
+    /// Tool name
+    #[allow(dead_code)] // Kept for documentation and future use
+    pub tool_name: String,
+
+    /// Tool output (JSON value)
+    pub output: Value,
+
+    /// Whether this is an error result
+    pub is_error: bool,
+}
+
+/// Execute MCP tools and collect results
+///
+/// Executes each tool call sequentially via the MCP manager.
+/// Tool execution errors are returned as error results to the model
+/// (allows model to handle gracefully).
+///
+/// Vector of tool results (one per tool call)
+pub(super) async fn execute_mcp_tools(
+    mcp_manager: &Arc<McpManager>,
+    tool_calls: &[ToolCall],
+    tracking: &mut McpCallTracking,
+    model_id: &str,
+) -> Result<Vec<ToolResult>, Response> {
+    let mut results = Vec::new();
+
+    for tool_call in tool_calls {
+        debug!(
+            tool_name = %tool_call.function.name,
+            call_id = %tool_call.id,
+            "Executing MCP tool"
+        );
+
+        // Parse tool arguments from JSON string
+        let args_str = tool_call.function.arguments.as_deref().unwrap_or("{}");
+        let args: Value = from_str(args_str).map_err(|e| {
+            error!(
+                function = "execute_mcp_tools",
+                tool_name = %tool_call.function.name,
+                call_id = %tool_call.id,
+                error = %e,
+                "Failed to parse tool arguments JSON"
+            );
+            error::internal_error(
+                "invalid_tool_args",
+                format!(
+                    "Invalid tool arguments JSON for tool '{}': {}",
+                    tool_call.function.name, e
+                ),
+            )
+        })?;
+
+        // Execute tool via MCP manager
+        let args_map = if let Value::Object(map) = args {
+            Some(map)
+        } else {
+            None
+        };
+
+        let tool_start = Instant::now();
+        let tool_result = mcp_manager
+            .call_tool(&tool_call.function.name, args_map)
+            .await;
+        let tool_duration = tool_start.elapsed();
+
+        match tool_result {
+            Ok(mcp_result) => {
+                debug!(
+                    tool_name = %tool_call.function.name,
+                    call_id = %tool_call.id,
+                    "Tool execution succeeded"
+                );
+
+                // Extract content from MCP result
+                let output = if let Some(content) = mcp_result.content.first() {
+                    // Serialize the entire content item
+                    to_value(content)
+                        .unwrap_or_else(|_| json!({"error": "Failed to serialize tool result"}))
+                } else {
+                    json!({"result": "success"})
+                };
+
+                let is_error = mcp_result.is_error.unwrap_or(false);
+                let output_str = to_string(&output)
+                    .unwrap_or_else(|_| r#"{"error": "Failed to serialize output"}"#.to_string());
+
+                // Record this call in tracking
+                tracking.record_call(
+                    tool_call.id.clone(),
+                    tool_call.function.name.clone(),
+                    args_str.to_string(),
+                    output_str.clone(),
+                    !is_error,
+                    if is_error {
+                        Some(output_str.clone())
+                    } else {
+                        None
+                    },
+                );
+
+                // Record MCP tool metrics
+                Metrics::record_mcp_tool_duration(
+                    model_id,
+                    &tool_call.function.name,
+                    tool_duration,
+                );
+                Metrics::record_mcp_tool_call(
+                    model_id,
+                    &tool_call.function.name,
+                    if is_error {
+                        metrics_labels::RESULT_ERROR
+                    } else {
+                        metrics_labels::RESULT_SUCCESS
+                    },
+                );
+
+                results.push(ToolResult {
+                    call_id: tool_call.id.clone(),
+                    tool_name: tool_call.function.name.clone(),
+                    output,
+                    is_error,
+                });
+            }
+            Err(e) => {
+                warn!(
+                    tool_name = %tool_call.function.name,
+                    call_id = %tool_call.id,
+                    error = %e,
+                    "Tool execution failed"
+                );
+
+                let error_msg = format!("Tool execution failed: {}", e);
+                let error_output = json!({
+                    "error": error_msg.clone()
+                });
+                let error_output_str = to_string(&error_output)
+                    .unwrap_or_else(|_| format!(r#"{{"error": "{}"}}"#, error_msg));
+
+                // Record failed call in tracking
+                tracking.record_call(
+                    tool_call.id.clone(),
+                    tool_call.function.name.clone(),
+                    args_str.to_string(),
+                    error_output_str.clone(),
+                    false,
+                    Some(error_msg),
+                );
+
+                // Record MCP tool metrics
+                Metrics::record_mcp_tool_duration(
+                    model_id,
+                    &tool_call.function.name,
+                    tool_duration,
+                );
+                Metrics::record_mcp_tool_call(
+                    model_id,
+                    &tool_call.function.name,
+                    metrics_labels::RESULT_ERROR,
+                );
+
+                // Return error result to model (let it handle gracefully)
+                results.push(ToolResult {
+                    call_id: tool_call.id.clone(),
+                    tool_name: tool_call.function.name.clone(),
+                    output: error_output,
+                    is_error: true,
+                });
+            }
+        }
+    }
+
+    Ok(results)
+}
+
+/// Convert MCP tools to Responses API tool format
+///
+/// Converts MCP Tool entries (from rmcp SDK) to ResponseTool format so the model
+/// knows about available MCP tools when making tool calls.
+pub(crate) fn convert_mcp_tools_to_response_tools(mcp_tools: &[mcp::Tool]) -> Vec<ResponseTool> {
+    mcp_tools
+        .iter()
+        .map(|tool_info| ResponseTool {
+            r#type: ResponseToolType::Mcp,
+            function: Some(Function {
+                name: tool_info.name.to_string(),
+                description: tool_info.description.as_ref().map(|d| d.to_string()),
+                parameters: Value::Object((*tool_info.input_schema).clone()),
+                strict: None,
+            }),
+            server_url: None, // MCP tools from inventory don't have individual server URLs
+            authorization: None,
+            server_label: None,
+            server_description: tool_info.description.as_ref().map(|d| d.to_string()),
+            require_approval: None,
+            allowed_tools: None,
+        })
+        .collect()
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/responses/mod.rs

@@ -0,0 +1,29 @@
+//! Harmony Responses API implementation with multi-turn MCP tool support
+//!
+//! This module implements the Harmony Responses API orchestration logic,
+//! coordinating full pipeline execution with MCP tool support for multi-turn conversations.
+//!
+//! ## Architecture
+//!
+//! Multi-turn pipeline orchestration (NOT just a tool loop):
+//! - Serves Harmony Responses API requests end-to-end
+//! - Each iteration executes FULL pipeline (worker selection + client acquisition + execution + parsing)
+//! - Handles MCP tool execution and history building between iterations
+//! - Clean separation: serving orchestration vs. pipeline stages (stages/)
+//!
+//! ## Module Structure
+//!
+//! - `non_streaming` - Non-streaming entry point and tool loop
+//! - `streaming` - Streaming entry point and tool loop
+//! - `execution` - MCP tool execution logic
+//! - `common` - Shared helpers and state tracking
+
+pub(crate) mod common;
+pub(crate) mod execution;
+pub(crate) mod non_streaming;
+pub(crate) mod streaming;
+
+// Re-export types accessed via harmony::responses::TypeName
+pub(crate) use execution::ToolResult;
+pub(crate) use non_streaming::serve_harmony_responses;
+pub(crate) use streaming::serve_harmony_responses_stream;

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs

@@ -0,0 +1,463 @@
+//! Non-streaming Harmony Responses API implementation
+
+use std::{
+    sync::Arc,
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+use axum::response::Response;
+use serde_json::{json, to_string};
+use tracing::{debug, error, warn};
+
+use super::{
+    common::{
+        build_mcp_tool_names_set, build_next_request_with_tools, inject_mcp_metadata,
+        load_previous_messages, McpCallTracking,
+    },
+    execution::{convert_mcp_tools_to_response_tools, execute_mcp_tools, ToolResult},
+};
+use crate::{
+    observability::metrics::Metrics,
+    protocols::{
+        common::{ToolCall, Usage},
+        responses::{
+            OutputTokensDetails, ResponseContentPart, ResponseOutputItem, ResponseReasoningContent,
+            ResponseStatus, ResponseUsage, ResponsesRequest, ResponsesResponse, ResponsesUsage,
+        },
+    },
+    routers::{
+        error,
+        grpc::{
+            common::responses::{
+                ensure_mcp_connection, persist_response_if_needed, ResponsesContext,
+            },
+            harmony::processor::ResponsesIterationResult,
+        },
+        mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS},
+    },
+};
+
+/// Execute Harmony Responses API request with multi-turn MCP tool support
+///
+/// This function orchestrates the multi-turn conversation flow:
+/// 1. Execute request through full pipeline
+/// 2. Check for tool calls in commentary channel
+/// 3. If tool calls found:
+///    - Execute MCP tools
+///    - Build next request with tool results
+///    - Repeat from step 1 (full pipeline re-execution)
+/// 4. If no tool calls, return final response
+pub(crate) async fn serve_harmony_responses(
+    ctx: &ResponsesContext,
+    request: ResponsesRequest,
+) -> Result<ResponsesResponse, Response> {
+    // Clone request for persistence
+    let original_request = request.clone();
+
+    // Load previous conversation history if previous_response_id is set
+    let current_request = load_previous_messages(ctx, request).await?;
+
+    // Check MCP connection and get whether MCP tools are present
+    let (has_mcp_tools, server_keys) =
+        ensure_mcp_connection(&ctx.mcp_manager, current_request.tools.as_deref()).await?;
+
+    // Set the server keys in the context
+    {
+        let mut servers = ctx.requested_servers.write().unwrap();
+        *servers = server_keys;
+    }
+
+    let response = if has_mcp_tools {
+        execute_with_mcp_loop(ctx, current_request).await?
+    } else {
+        // No MCP tools - execute pipeline once (may have function tools or no tools)
+        execute_without_mcp_loop(ctx, current_request).await?
+    };
+
+    // Persist response to storage if store=true
+    persist_response_if_needed(
+        ctx.conversation_storage.clone(),
+        ctx.conversation_item_storage.clone(),
+        ctx.response_storage.clone(),
+        &response,
+        &original_request,
+    )
+    .await;
+
+    Ok(response)
+}
+
+/// Execute Harmony Responses with MCP tool loop
+///
+/// Automatically executes MCP tools in a loop until no more tool calls or max iterations
+async fn execute_with_mcp_loop(
+    ctx: &ResponsesContext,
+    mut current_request: ResponsesRequest,
+) -> Result<ResponsesResponse, Response> {
+    let mut iteration_count = 0;
+
+    // Extract server_label from request tools
+    let server_label = extract_server_label(current_request.tools.as_deref(), "sglang-mcp");
+    let mut mcp_tracking = McpCallTracking::new(server_label.clone());
+
+    // Extract user's max_tool_calls limit (if set)
+    let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
+
+    // Add filtered MCP tools (static + requested dynamic) to the request
+    let mcp_tools = {
+        let servers = ctx.requested_servers.read().unwrap();
+        ctx.mcp_manager.list_tools_for_servers(&servers)
+    };
+    if !mcp_tools.is_empty() {
+        let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools);
+
+        let mut all_tools = current_request.tools.clone().unwrap_or_default();
+        all_tools.extend(mcp_response_tools);
+        current_request.tools = Some(all_tools);
+
+        debug!(
+            mcp_tool_count = mcp_tools.len(),
+            total_tool_count = current_request.tools.as_ref().map(|t| t.len()).unwrap_or(0),
+            "MCP client available - added static MCP tools to Harmony Responses request"
+        );
+    }
+
+    loop {
+        iteration_count += 1;
+
+        // Record tool loop iteration metric
+        Metrics::record_mcp_tool_iteration(&current_request.model);
+
+        // Safety check: prevent infinite loops
+        if iteration_count > DEFAULT_MAX_ITERATIONS {
+            error!(
+                function = "execute_with_mcp_loop",
+                iteration_count = iteration_count,
+                max_iterations = DEFAULT_MAX_ITERATIONS,
+                "Maximum tool iterations exceeded"
+            );
+            return Err(error::internal_error(
+                "tool_iterations_exceeded",
+                format!(
+                    "Maximum tool iterations ({}) exceeded",
+                    DEFAULT_MAX_ITERATIONS
+                ),
+            ));
+        }
+
+        debug!(
+            iteration = iteration_count,
+            "Harmony Responses serving iteration"
+        );
+
+        // Execute through full pipeline
+        let iteration_result = ctx
+            .pipeline
+            .execute_harmony_responses(&current_request, ctx)
+            .await?;
+
+        match iteration_result {
+            ResponsesIterationResult::ToolCallsFound {
+                tool_calls,
+                analysis,
+                partial_text,
+                usage,
+                request_id,
+            } => {
+                debug!(
+                    tool_call_count = tool_calls.len(),
+                    has_analysis = analysis.is_some(),
+                    partial_text_len = partial_text.len(),
+                    "Tool calls found - separating MCP and function tools"
+                );
+
+                // Separate MCP and function tool calls based on tool type
+                let request_tools = current_request.tools.as_deref().unwrap_or(&[]);
+                let mcp_tool_names = build_mcp_tool_names_set(request_tools);
+                let (mcp_tool_calls, function_tool_calls): (Vec<_>, Vec<_>) = tool_calls
+                    .into_iter()
+                    .partition(|tc| mcp_tool_names.contains(tc.function.name.as_str()));
+
+                debug!(
+                    mcp_calls = mcp_tool_calls.len(),
+                    function_calls = function_tool_calls.len(),
+                    "Tool calls separated by type"
+                );
+
+                // Check combined limit (user's max_tool_calls vs safety limit)
+                let effective_limit = match max_tool_calls {
+                    Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS),
+                    None => DEFAULT_MAX_ITERATIONS,
+                };
+
+                // Check if we would exceed the limit with these new MCP tool calls
+                let total_calls_after = mcp_tracking.total_calls() + mcp_tool_calls.len();
+                if total_calls_after > effective_limit {
+                    warn!(
+                        current_calls = mcp_tracking.total_calls(),
+                        new_calls = mcp_tool_calls.len() + function_tool_calls.len(),
+                        total_after = total_calls_after,
+                        effective_limit = effective_limit,
+                        user_max = ?max_tool_calls,
+                        "Reached tool call limit - returning incomplete response"
+                    );
+
+                    // Combine back for response
+                    let all_tool_calls: Vec<_> = mcp_tool_calls
+                        .into_iter()
+                        .chain(function_tool_calls)
+                        .collect();
+
+                    // Build response with incomplete status - no tools executed due to limit
+                    let mut response = build_tool_response(
+                        vec![],         // No MCP tools executed
+                        vec![],         // No MCP results
+                        all_tool_calls, // All tools returned as function calls (not executed)
+                        analysis,
+                        partial_text,
+                        usage,
+                        request_id,
+                        Arc::new(current_request),
+                    );
+
+                    // Mark as completed with incomplete_details
+                    response.status = ResponseStatus::Completed;
+                    response.incomplete_details = Some(json!({ "reason": "max_tool_calls" }));
+
+                    // Inject MCP metadata if any calls were executed
+                    if mcp_tracking.total_calls() > 0 {
+                        inject_mcp_metadata(&mut response, &mcp_tracking, &mcp_tools);
+                    }
+
+                    return Ok(response);
+                }
+
+                // Execute MCP tools (if any)
+                let mcp_results = if !mcp_tool_calls.is_empty() {
+                    execute_mcp_tools(
+                        &ctx.mcp_manager,
+                        &mcp_tool_calls,
+                        &mut mcp_tracking,
+                        &current_request.model,
+                    )
+                    .await?
+                } else {
+                    Vec::new()
+                };
+
+                // If there are function tools, exit MCP loop and return response
+                if !function_tool_calls.is_empty() {
+                    debug!(
+                        "Function tool calls present - exiting MCP loop and returning to caller"
+                    );
+
+                    // Build response that includes:
+                    // 1. Reasoning/message from this iteration
+                    // 2. MCP tools as completed (with output) - these were executed
+                    // 3. Function tools as completed (without output) - need caller execution
+                    let mut response = build_tool_response(
+                        mcp_tool_calls,
+                        mcp_results,
+                        function_tool_calls,
+                        analysis,
+                        partial_text,
+                        usage,
+                        request_id,
+                        Arc::new(current_request),
+                    );
+
+                    // Inject MCP metadata for all executed calls
+                    if mcp_tracking.total_calls() > 0 {
+                        inject_mcp_metadata(&mut response, &mcp_tracking, &mcp_tools);
+                    }
+
+                    return Ok(response);
+                }
+
+                // Only MCP tools - continue loop with their results
+                debug!("Only MCP tools - continuing loop with results");
+
+                // Build next request with appended history
+                current_request = build_next_request_with_tools(
+                    current_request,
+                    mcp_tool_calls,
+                    mcp_results,
+                    analysis,
+                    partial_text,
+                )
+                .map_err(|e| *e)?;
+
+                // Continue loop - next iteration will select workers and execute
+            }
+            ResponsesIterationResult::Completed {
+                mut response,
+                usage,
+            } => {
+                debug!(
+                    output_items = response.output.len(),
+                    input_tokens = usage.prompt_tokens,
+                    output_tokens = usage.completion_tokens,
+                    "MCP loop completed - no more tool calls"
+                );
+
+                // Inject MCP metadata into final response
+                inject_mcp_metadata(&mut response, &mcp_tracking, &mcp_tools);
+
+                debug!(
+                    mcp_calls = mcp_tracking.total_calls(),
+                    output_items_after = response.output.len(),
+                    "Injected MCP metadata into final response"
+                );
+
+                // No tool calls - this is the final response
+                return Ok(*response);
+            }
+        }
+    }
+}
+
+/// Execute Harmony Responses without MCP loop (single execution)
+///
+/// For function tools or no tools - executes pipeline once and returns
+async fn execute_without_mcp_loop(
+    ctx: &ResponsesContext,
+    current_request: ResponsesRequest,
+) -> Result<ResponsesResponse, Response> {
+    debug!("Executing Harmony Responses without MCP loop");
+
+    // Execute pipeline once
+    let iteration_result = ctx
+        .pipeline
+        .execute_harmony_responses(&current_request, ctx)
+        .await?;
+
+    match iteration_result {
+        ResponsesIterationResult::ToolCallsFound {
+            tool_calls,
+            analysis,
+            partial_text,
+            usage,
+            request_id,
+        } => {
+            // Function tool calls found - return to caller for execution
+            debug!(
+                tool_call_count = tool_calls.len(),
+                "Function tool calls found - returning to caller"
+            );
+
+            Ok(build_tool_response(
+                vec![],
+                vec![],
+                tool_calls,
+                analysis,
+                partial_text,
+                usage,
+                request_id,
+                Arc::new(current_request),
+            ))
+        }
+        ResponsesIterationResult::Completed { response, usage: _ } => {
+            // No tool calls - return completed response
+            debug!("No tool calls - returning completed response");
+            Ok(*response)
+        }
+    }
+}
+
+/// Build ResponsesResponse with tool calls (MCP and/or function tools)
+#[allow(clippy::too_many_arguments)]
+fn build_tool_response(
+    mcp_tool_calls: Vec<ToolCall>,
+    mcp_results: Vec<ToolResult>,
+    function_tool_calls: Vec<ToolCall>,
+    analysis: Option<String>, // Analysis channel content (reasoning)
+    partial_text: String,     // Final channel content (message)
+    usage: Usage,
+    request_id: String,
+    responses_request: Arc<ResponsesRequest>,
+) -> ResponsesResponse {
+    let mut output: Vec<ResponseOutputItem> = Vec::new();
+
+    // Add reasoning output item if analysis exists
+    if let Some(analysis_text) = analysis {
+        output.push(ResponseOutputItem::Reasoning {
+            id: format!("reasoning_{}", request_id),
+            summary: vec![],
+            content: vec![ResponseReasoningContent::ReasoningText {
+                text: analysis_text,
+            }],
+            status: Some("completed".to_string()),
+        });
+    }
+
+    // Add message output item if partial text exists
+    if !partial_text.is_empty() {
+        output.push(ResponseOutputItem::Message {
+            id: format!("msg_{}", request_id),
+            role: "assistant".to_string(),
+            content: vec![ResponseContentPart::OutputText {
+                text: partial_text,
+                annotations: vec![],
+                logprobs: None,
+            }],
+            status: "completed".to_string(),
+        });
+    }
+
+    // Add MCP tool calls WITH output (these were executed)
+    for (tool_call, result) in mcp_tool_calls.iter().zip(mcp_results.iter()) {
+        let output_str = to_string(&result.output).unwrap_or_else(|e| {
+            format!("{{\"error\": \"Failed to serialize tool output: {}\"}}", e)
+        });
+
+        output.push(ResponseOutputItem::FunctionToolCall {
+            id: tool_call.id.clone(),
+            call_id: tool_call.id.clone(),
+            name: tool_call.function.name.clone(),
+            arguments: tool_call.function.arguments.clone().unwrap_or_default(),
+            output: Some(output_str),
+            status: if result.is_error {
+                "failed"
+            } else {
+                "completed"
+            }
+            .to_string(),
+        });
+    }
+
+    // Add function tool calls WITHOUT output (need caller execution)
+    for tool_call in function_tool_calls {
+        output.push(ResponseOutputItem::FunctionToolCall {
+            id: tool_call.id.clone(),
+            call_id: tool_call.id.clone(),
+            name: tool_call.function.name.clone(),
+            arguments: tool_call.function.arguments.clone().unwrap_or_default(),
+            output: None, // No output = needs execution
+            status: "completed".to_string(),
+        });
+    }
+
+    // Build ResponsesResponse with Completed status
+    let created_at = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs() as i64;
+
+    ResponsesResponse::builder(&request_id, &responses_request.model)
+        .copy_from_request(&responses_request)
+        .created_at(created_at)
+        .status(ResponseStatus::Completed)
+        .output(output)
+        .usage(ResponsesUsage::Modern(ResponseUsage {
+            input_tokens: usage.prompt_tokens,
+            output_tokens: usage.completion_tokens,
+            total_tokens: usage.total_tokens,
+            input_tokens_details: None,
+            output_tokens_details: usage.completion_tokens_details.as_ref().and_then(|d| {
+                d.reasoning_tokens.map(|tokens| OutputTokensDetails {
+                    reasoning_tokens: tokens,
+                })
+            }),
+        }))
+        .build()
+}

third_party/sglang/sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs

@@ -0,0 +1,534 @@
+//! Streaming Harmony Responses API implementation
+
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use axum::response::Response;
+use bytes::Bytes;
+use serde_json::{json, Value};
+use tokio::sync::mpsc;
+use tracing::{debug, warn};
+use uuid::Uuid;
+
+use super::{
+    common::{
+        build_mcp_tool_names_set, build_next_request_with_tools, load_previous_messages,
+        McpCallTracking,
+    },
+    execution::{convert_mcp_tools_to_response_tools, execute_mcp_tools},
+};
+use crate::{
+    observability::metrics::Metrics,
+    protocols::responses::{ResponseToolType, ResponsesRequest},
+    routers::{
+        grpc::{
+            common::responses::{
+                build_sse_response, ensure_mcp_connection, persist_response_if_needed,
+                streaming::{OutputItemType, ResponseStreamEventEmitter},
+                ResponsesContext,
+            },
+            harmony::{processor::ResponsesIterationResult, streaming::HarmonyStreamingProcessor},
+        },
+        mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS},
+    },
+};
+
+/// Serve Harmony Responses API with streaming (SSE)
+///
+/// This is the streaming equivalent of `serve_harmony_responses()`.
+/// Emits SSE events for lifecycle, MCP list_tools, and per-iteration streaming.
+pub(crate) async fn serve_harmony_responses_stream(
+    ctx: &ResponsesContext,
+    request: ResponsesRequest,
+) -> Response {
+    // Load previous conversation history if previous_response_id is set
+    let current_request = match load_previous_messages(ctx, request.clone()).await {
+        Ok(req) => req,
+        Err(err_response) => return err_response,
+    };
+
+    // Check MCP connection BEFORE starting stream and get whether MCP tools are present
+    let (has_mcp_tools, server_keys) =
+        match ensure_mcp_connection(&ctx.mcp_manager, current_request.tools.as_deref()).await {
+            Ok(result) => result,
+            Err(response) => return response,
+        };
+
+    // Set the server keys in the context
+    {
+        let mut servers = ctx.requested_servers.write().unwrap();
+        *servers = server_keys;
+    }
+
+    // Create SSE channel
+    let (tx, rx) = mpsc::unbounded_channel();
+
+    // Create response event emitter
+    let response_id = format!("resp_{}", Uuid::new_v4());
+    let model = current_request.model.clone();
+    let created_at = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+    let mut emitter = ResponseStreamEventEmitter::new(response_id.clone(), model, created_at);
+
+    // Set original request for complete response fields
+    emitter.set_original_request(current_request.clone());
+
+    // Clone context for spawned task
+    let ctx_clone = ctx.clone();
+
+    // Spawn async task to handle streaming
+    tokio::spawn(async move {
+        let ctx = &ctx_clone;
+
+        // Emit initial response.created and response.in_progress events
+        let event = emitter.emit_created();
+        if emitter.send_event(&event, &tx).is_err() {
+            return;
+        }
+        let event = emitter.emit_in_progress();
+        if emitter.send_event(&event, &tx).is_err() {
+            return;
+        }
+
+        if has_mcp_tools {
+            execute_mcp_tool_loop_streaming(ctx, current_request, &request, &mut emitter, &tx)
+                .await;
+        } else {
+            execute_without_mcp_streaming(ctx, &current_request, &request, &mut emitter, &tx).await;
+        }
+    });
+
+    // Return SSE stream response
+    build_sse_response(rx)
+}
+
+/// Execute MCP tool loop with streaming
+///
+/// Handles the full MCP workflow:
+/// - Adds static MCP tools to request
+/// - Emits mcp_list_tools events
+/// - Loops through tool execution iterations
+/// - Emits final response.completed event
+/// - Persists response internally
+async fn execute_mcp_tool_loop_streaming(
+    ctx: &ResponsesContext,
+    mut current_request: ResponsesRequest,
+    original_request: &ResponsesRequest,
+    emitter: &mut ResponseStreamEventEmitter,
+    tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+) {
+    // Extract server_label from request tools
+    let server_label = extract_server_label(current_request.tools.as_deref(), "sglang-mcp");
+
+    // Set server label in emitter for MCP call items
+    emitter.set_mcp_server_label(server_label.clone());
+
+    // Initialize MCP call tracking
+    let mut mcp_tracking = McpCallTracking::new(server_label.clone());
+
+    // Extract user's max_tool_calls limit (if set)
+    let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
+
+    // Add filtered MCP tools (static + requested dynamic) to the request
+    let mcp_tools = {
+        let servers = ctx.requested_servers.read().unwrap();
+        ctx.mcp_manager.list_tools_for_servers(&servers)
+    };
+    if !mcp_tools.is_empty() {
+        let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools);
+        let mut all_tools = current_request.tools.clone().unwrap_or_default();
+        all_tools.extend(mcp_response_tools);
+        current_request.tools = Some(all_tools);
+
+        debug!(
+            mcp_tool_count = mcp_tools.len(),
+            total_tool_count = current_request.tools.as_ref().map(|t| t.len()).unwrap_or(0),
+            "MCP client available - added static MCP tools to Harmony Responses streaming request"
+        );
+    }
+
+    // Build HashSet of MCP tool names for O(1) lookup during streaming
+    // Clone tool names to owned strings to avoid borrowing current_request
+    let mcp_tool_names: std::collections::HashSet<String> = current_request
+        .tools
+        .as_ref()
+        .map(|tools| {
+            tools
+                .iter()
+                .filter(|t| t.r#type == ResponseToolType::Mcp)
+                .filter_map(|t| t.function.as_ref().map(|f| f.name.clone()))
+                .collect()
+        })
+        .unwrap_or_default();
+
+    // Emit mcp_list_tools on first iteration
+    let (output_index, item_id) = emitter.allocate_output_index(OutputItemType::McpListTools);
+
+    // Build tools list for item structure
+    let tool_items: Vec<_> = mcp_tools
+        .iter()
+        .map(|t| {
+            json!({
+                "name": t.name,
+                "description": t.description,
+                "input_schema": Value::Object((*t.input_schema).clone())
+            })
+        })
+        .collect();
+
+    // Build final item with completed status and tools
+    let item_done = json!({
+        "id": item_id,
+        "type": "mcp_list_tools",
+        "server_label": server_label,
+        "status": "completed",
+        "tools": tool_items
+    });
+
+    // Store the completed item data and mark as completed FIRST
+    // This ensures it appears in final response even if event sending fails
+    emitter.emit_output_item_done(output_index, &item_done);
+    emitter.complete_output_item(output_index);
+
+    // Now emit all the events (failures won't affect the stored data)
+    // Emit output_item.added
+    let item = json!({
+        "id": item_id,
+        "type": "mcp_list_tools",
+        "server_label": server_label,
+        "status": "in_progress",
+        "tools": []
+    });
+    let event = emitter.emit_output_item_added(output_index, &item);
+    if emitter.send_event(&event, tx).is_err() {
+        return;
+    }
+
+    // Emit mcp_list_tools.in_progress
+    let event = emitter.emit_mcp_list_tools_in_progress(output_index);
+    if emitter.send_event(&event, tx).is_err() {
+        return;
+    }
+
+    // Emit mcp_list_tools.completed
+    let event = emitter.emit_mcp_list_tools_completed(output_index, &mcp_tools);
+    if emitter.send_event(&event, tx).is_err() {
+        return;
+    }
+
+    // Emit output_item.done
+    let event = emitter.emit_output_item_done(output_index, &item_done);
+    if emitter.send_event(&event, tx).is_err() {
+        return;
+    }
+
+    debug!(
+        tool_count = mcp_tools.len(),
+        "Emitted mcp_list_tools on first iteration"
+    );
+
+    // MCP tool loop (max 10 iterations)
+    let mut iteration_count = 0;
+    loop {
+        iteration_count += 1;
+
+        // Record tool loop iteration metric
+        Metrics::record_mcp_tool_iteration(&current_request.model);
+
+        // Safety check: prevent infinite loops
+        if iteration_count > DEFAULT_MAX_ITERATIONS {
+            emitter.emit_error(
+                &format!(
+                    "Maximum tool iterations ({}) exceeded",
+                    DEFAULT_MAX_ITERATIONS
+                ),
+                Some("max_iterations_exceeded"),
+                tx,
+            );
+            return;
+        }
+
+        debug!(
+            iteration = iteration_count,
+            "Harmony Responses streaming iteration"
+        );
+
+        // Execute pipeline and get stream + load guards
+        let (execution_result, _load_guards) = match ctx
+            .pipeline
+            .execute_harmony_responses_streaming(&current_request, ctx)
+            .await
+        {
+            Ok(result) => result,
+            Err(err_response) => {
+                emitter.emit_error(
+                    &format!("Pipeline execution failed: {:?}", err_response),
+                    Some("pipeline_error"),
+                    tx,
+                );
+                return;
+            }
+        };
+
+        // Process stream with token-level streaming (mixed tools - emits correct events per tool type)
+        // Load guards are held during processing and dropped when iteration completes
+        let iteration_result = match HarmonyStreamingProcessor::process_responses_iteration_stream(
+            execution_result,
+            emitter,
+            tx,
+            &mcp_tool_names,
+        )
+        .await
+        {
+            Ok(result) => result,
+            Err(err_msg) => {
+                emitter.emit_error(&err_msg, Some("processing_error"), tx);
+                return;
+            }
+        };
+
+        // Handle iteration result (tool calls or completion)
+        match iteration_result {
+            ResponsesIterationResult::ToolCallsFound {
+                tool_calls,
+                analysis,
+                partial_text,
+                usage,
+                request_id: _,
+            } => {
+                debug!(
+                    tool_call_count = tool_calls.len(),
+                    has_analysis = analysis.is_some(),
+                    partial_text_len = partial_text.len(),
+                    "Tool calls found - separating MCP and function tools"
+                );
+
+                // Separate MCP and function tool calls based on tool type
+                let request_tools = current_request.tools.as_deref().unwrap_or(&[]);
+                let mcp_tool_names = build_mcp_tool_names_set(request_tools);
+                let (mcp_tool_calls, function_tool_calls): (Vec<_>, Vec<_>) = tool_calls
+                    .into_iter()
+                    .partition(|tc| mcp_tool_names.contains(tc.function.name.as_str()));
+
+                debug!(
+                    mcp_calls = mcp_tool_calls.len(),
+                    function_calls = function_tool_calls.len(),
+                    "Tool calls separated by type in streaming"
+                );
+
+                // Check combined limit (user's max_tool_calls vs safety limit)
+                let effective_limit = match max_tool_calls {
+                    Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS),
+                    None => DEFAULT_MAX_ITERATIONS,
+                };
+
+                // Check if we would exceed the limit with these new MCP tool calls
+                let total_calls_after = mcp_tracking.total_calls() + mcp_tool_calls.len();
+                if total_calls_after > effective_limit {
+                    warn!(
+                        current_calls = mcp_tracking.total_calls(),
+                        new_calls = mcp_tool_calls.len() + function_tool_calls.len(),
+                        total_after = total_calls_after,
+                        effective_limit = effective_limit,
+                        user_max = ?max_tool_calls,
+                        "Reached tool call limit in streaming - emitting completion with incomplete_details"
+                    );
+
+                    // Emit response.completed with incomplete_details and usage
+                    let incomplete_details = json!({ "reason": "max_tool_calls" });
+                    let usage_json = json!({
+                        "input_tokens": usage.prompt_tokens,
+                        "output_tokens": usage.completion_tokens,
+                        "total_tokens": usage.total_tokens,
+                        "incomplete_details": incomplete_details,
+                    });
+                    let event = emitter.emit_completed(Some(&usage_json));
+                    emitter.send_event_best_effort(&event, tx);
+                    return;
+                }
+
+                // Execute MCP tools (if any)
+                let mcp_results = if !mcp_tool_calls.is_empty() {
+                    match execute_mcp_tools(
+                        &ctx.mcp_manager,
+                        &mcp_tool_calls,
+                        &mut mcp_tracking,
+                        &current_request.model,
+                    )
+                    .await
+                    {
+                        Ok(results) => results,
+                        Err(err_response) => {
+                            emitter.emit_error(
+                                &format!("MCP tool execution failed: {:?}", err_response),
+                                Some("mcp_tool_error"),
+                                tx,
+                            );
+                            return;
+                        }
+                    }
+                } else {
+                    Vec::new()
+                };
+
+                // Update mcp_call output items with execution results (if any MCP tools were executed)
+                if !mcp_results.is_empty() {
+                    emitter.update_mcp_call_outputs(&mcp_results);
+                }
+
+                // If there are function tools, exit MCP loop and emit completion
+                if !function_tool_calls.is_empty() {
+                    debug!(
+                        "Function tool calls present - exiting MCP loop and emitting completion"
+                    );
+
+                    // Function tool calls were already emitted during streaming processing
+                    // Just emit response.completed with usage
+                    let usage_json = json!({
+                        "input_tokens": usage.prompt_tokens,
+                        "output_tokens": usage.completion_tokens,
+                        "total_tokens": usage.total_tokens,
+                    });
+                    let event = emitter.emit_completed(Some(&usage_json));
+                    emitter.send_event_best_effort(&event, tx);
+                    return;
+                }
+
+                // Only MCP tools - continue loop with their results
+                debug!("Only MCP tools - continuing loop with results");
+
+                // Build next request with appended history
+                current_request = match build_next_request_with_tools(
+                    current_request,
+                    mcp_tool_calls,
+                    mcp_results,
+                    analysis,
+                    partial_text,
+                ) {
+                    Ok(req) => req,
+                    Err(e) => {
+                        emitter.emit_error(
+                            &format!("Failed to build next request: {:?}", e),
+                            Some("request_building_error"),
+                            tx,
+                        );
+                        return;
+                    }
+                };
+
+                // Continue loop
+            }
+            ResponsesIterationResult::Completed { response, usage } => {
+                debug!(
+                    output_items = response.output.len(),
+                    input_tokens = usage.prompt_tokens,
+                    output_tokens = usage.completion_tokens,
+                    "Harmony Responses streaming completed - no more tool calls"
+                );
+
+                // Finalize response from emitter's accumulated data
+                let final_response = emitter.finalize(Some(usage.clone()));
+
+                // Persist response to storage if store=true
+                persist_response_if_needed(
+                    ctx.conversation_storage.clone(),
+                    ctx.conversation_item_storage.clone(),
+                    ctx.response_storage.clone(),
+                    &final_response,
+                    original_request,
+                )
+                .await;
+
+                // Emit response.completed with usage
+                let usage_json = json!({
+                    "input_tokens": usage.prompt_tokens,
+                    "output_tokens": usage.completion_tokens,
+                    "total_tokens": usage.total_tokens,
+                });
+                let event = emitter.emit_completed(Some(&usage_json));
+                emitter.send_event_best_effort(&event, tx);
+                return;
+            }
+        }
+    }
+}
+
+/// Execute without MCP tool loop (single execution with streaming)
+///
+/// For function tools or no tools - executes pipeline once and emits completion.
+/// The streaming processor handles all output items (reasoning, message, function tool calls).
+async fn execute_without_mcp_streaming(
+    ctx: &ResponsesContext,
+    current_request: &ResponsesRequest,
+    original_request: &ResponsesRequest,
+    emitter: &mut ResponseStreamEventEmitter,
+    tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
+) {
+    debug!("No MCP tools - executing single iteration");
+
+    // Execute pipeline and get stream + load guards
+    let (execution_result, _load_guards) = match ctx
+        .pipeline
+        .execute_harmony_responses_streaming(current_request, ctx)
+        .await
+    {
+        Ok(result) => result,
+        Err(err_response) => {
+            emitter.emit_error(
+                &format!("Pipeline execution failed: {:?}", err_response),
+                Some("pipeline_error"),
+                tx,
+            );
+            return;
+        }
+    };
+
+    // Process stream (emits all output items during streaming - function tool path emits function_call_arguments.* events)
+    // Pass empty HashSet so all tools are treated as function tools (per-tool detection)
+    // Load guards are held during processing and dropped when iteration completes
+    let empty_mcp_tools = std::collections::HashSet::new();
+    let iteration_result = match HarmonyStreamingProcessor::process_responses_iteration_stream(
+        execution_result,
+        emitter,
+        tx,
+        &empty_mcp_tools,
+    )
+    .await
+    {
+        Ok(result) => result,
+        Err(err_msg) => {
+            emitter.emit_error(&err_msg, Some("processing_error"), tx);
+            return;
+        }
+    };
+    // _load_guards dropped here after iteration completes
+
+    // Extract usage from iteration result
+    let usage = match iteration_result {
+        ResponsesIterationResult::ToolCallsFound { usage, .. } => usage,
+        ResponsesIterationResult::Completed { usage, .. } => usage,
+    };
+
+    // Finalize response from emitter's accumulated data
+    let final_response = emitter.finalize(Some(usage.clone()));
+
+    // Persist response to storage if store=true
+    persist_response_if_needed(
+        ctx.conversation_storage.clone(),
+        ctx.conversation_item_storage.clone(),
+        ctx.response_storage.clone(),
+        &final_response,
+        original_request,
+    )
+    .await;
+
+    // Emit response.completed with usage
+    let usage_json = json!({
+        "input_tokens": usage.prompt_tokens,
+        "output_tokens": usage.completion_tokens,
+        "total_tokens": usage.total_tokens,
+    });
+    let event = emitter.emit_completed(Some(&usage_json));
+    emitter.send_event_best_effort(&event, tx);
+}