[model-gateway] refactor: unify worker management into modular workflow structure (#15010)

2025-12-12 15:25:35 -08:00
parent b737a125c5
commit 9a56273ad2
18 changed files with 1796 additions and 1622 deletions
@@ -654,7 +654,7 @@ impl JobQueue {
        // Select workflow based on runtime field
        let workflow_id = match config.runtime.as_deref() {
            Some("external") => WorkflowId::new("external_worker_registration"),
-            _ => WorkflowId::new("worker_registration"),
+            _ => WorkflowId::new("local_worker_registration"),
        };
        engine
@@ -1,591 +0,0 @@
 use std::{collections::HashMap, sync::Arc, 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 crate::{
    app_context::AppContext,
    core::{
        model_card::{ModelCard, ProviderType},
        model_type::ModelType,
        BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, HealthConfig, RuntimeType,
        Worker, WorkerType,
    },
    protocols::worker_spec::WorkerConfigRequest,
    workflow::*,
 };
 // 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"));
 // ============================================================================
 // Model Discovery Types
 // ============================================================================
 /// 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>,
 }
 // ============================================================================
 // Model Discovery Functions
 // ============================================================================
 /// 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.
 ///
 /// This function analyzes the model ID to determine what type of model it is
 /// (LLM, embedding, image generation, audio, etc.).
 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 (DALL-E, Sora, gpt-image)
    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 (TTS, Whisper, realtime, audio)
    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 (models with vision capability)
    // gpt-4o, gpt-4-vision, etc.
    if id_lower.contains("vision") || id_lower.contains("4o") {
        return ModelType::VISION_LLM;
    }
    // Reasoning models (o1, o3, etc.)
    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.
 ///
 /// Returns `None` for models that don't clearly indicate a provider.
 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
 }
 // ============================================================================
 // Workflow Steps
 // ============================================================================
 /// Step 1: Discover models from external /v1/models endpoint
 pub struct DiscoverModelsStep;
 #[async_trait]
 impl StepExecutor for DiscoverModelsStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_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
            );
            context.set::<Vec<ModelCard>>("model_cards", vec![]);
            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.set("model_cards", model_cards);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true // Network issues are retryable
    }
 }
 /// 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
    );
    // Group models into cards (e.g., gpt-4o-2024-08-06 → gpt-4o with aliases)
    let model_cards = group_models_into_cards(models_response.data);
    debug!(
        "Grouped into {} model cards with aliases",
        model_cards.len()
    );
    Ok(model_cards)
 }
 /// Step 2: Create worker objects for each discovered model
 pub struct CreateExternalWorkersStep;
 #[async_trait]
 impl StepExecutor for CreateExternalWorkersStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let model_cards: Arc<Vec<ModelCard>> = context.get_or_err("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,
            }
        };
        // Build labels
        let mut labels = 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>;
            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>;
                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
            );
        }
        context.set("workers", workers);
        context.set("labels", labels);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false // Creation failures are config issues
    }
 }
 /// 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 3: Register workers in registry
 pub struct RegisterExternalWorkersStep;
 #[async_trait]
 impl StepExecutor for RegisterExternalWorkersStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
        let mut worker_ids = Vec::new();
        for worker in workers.iter() {
            let worker_id = app_context.worker_registry.register(Arc::clone(worker));
            worker_ids.push(worker_id.clone());
            debug!(
                "Registered external worker {} (model: {}) with ID {:?}",
                config.url,
                worker.model_id(),
                worker_id
            );
        }
        context.set("worker_ids", worker_ids);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 /// Step 4: Update policy registry
 pub struct UpdateExternalPoliciesStep;
 #[async_trait]
 impl StepExecutor for UpdateExternalPoliciesStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
        let policy_hint = labels.get("policy").map(|s| s.as_str());
        // Each external worker has a different model_id
        for worker in workers.iter() {
            let model_id = worker.model_id().to_string();
            app_context
                .policy_registry
                .on_worker_added(&model_id, policy_hint);
        }
        debug!(
            "Updated policies for {} external workers from {}",
            workers.len(),
            config.url
        );
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 /// Step 5: Activate workers by marking them healthy
 pub struct ActivateExternalWorkersStep;
 #[async_trait]
 impl StepExecutor for ActivateExternalWorkersStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
        for worker in workers.iter() {
            worker.set_healthy(true);
        }
        info!(
            "Activated {} external workers from {} (marked as healthy)",
            workers.len(),
            config.url
        );
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 // ============================================================================
 // Workflow Definition
 // ============================================================================
 /// Create external worker registration workflow definition
 pub fn create_external_worker_registration_workflow() -> WorkflowDefinition {
    WorkflowDefinition::new(
        "external_worker_registration",
        "External Worker Registration",
    )
    .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),
    )
    .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"]),
    )
    .add_step(
        StepDefinition::new(
            "register_workers",
            "Register Workers",
            Arc::new(RegisterExternalWorkersStep),
        )
        .with_timeout(Duration::from_secs(5))
        .with_failure_action(FailureAction::FailWorkflow)
        .depends_on(&["create_workers"]),
    )
    .add_step(
        StepDefinition::new(
            "update_policies",
            "Update Policies",
            Arc::new(UpdateExternalPoliciesStep),
        )
        .with_timeout(Duration::from_secs(5))
        .with_failure_action(FailureAction::ContinueNextStep)
        .depends_on(&["register_workers"]),
    )
    .add_step(
        StepDefinition::new(
            "activate_workers",
            "Activate Workers",
            Arc::new(ActivateExternalWorkersStep),
        )
        .with_timeout(Duration::from_secs(5))
        .with_failure_action(FailureAction::FailWorkflow)
        .depends_on(&["update_policies"]),
    )
 }
@@ -1,26 +1,62 @@
 //! Workflow step implementations
 //!
 //! This module contains concrete step implementations for various workflows:
-//! - Local worker registration (SGLang, vLLM - gRPC or HTTP modes)
+//! - Worker management (registration, removal, future: updates)
 //! - External worker registration (OpenAI, xAI, Anthropic, etc. - HTTPS only)
 //! - Worker removal
 //! - MCP server registration
 //! - WASM module registration and removal
 //! - Future: Tokenizer fetching, LoRA updates, etc.
 pub mod external_worker_registration;
 pub mod mcp_registration;
 pub mod wasm_module_registration;
 pub mod wasm_module_removal;
-pub mod worker_registration;
+pub mod worker;
 pub mod worker_removal;
-pub use external_worker_registration::{
+// Worker management (registration, removal)
-    create_external_worker_registration_workflow, group_models_into_cards,
+#[allow(deprecated)]
-    infer_model_type_from_id, ActivateExternalWorkersStep, CreateExternalWorkersStep,
+pub use worker::create_external_worker_registration_workflow;
-    DiscoverModelsStep, ModelInfo, ModelsResponse, RegisterExternalWorkersStep,
+// Backward compatibility aliases
-    UpdateExternalPoliciesStep,
+#[allow(deprecated)]
 pub use worker::create_worker_registration_workflow;
 pub use worker::{
    // Workflow builders
    create_external_worker_workflow,
    create_local_worker_workflow,
    create_worker_removal_workflow,
    // 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,
    // Removal steps
    FindWorkersToRemoveStep,
    ModelInfo,
    ModelsResponse,
    RegisterWorkersStep,
    RemoveFromPolicyRegistryStep,
    RemoveFromWorkerRegistryStep,
    UpdatePoliciesStep,
    UpdateRemainingPoliciesStep,
    WorkerList,
    WorkerRemovalRequest,
 };
 // Legacy type aliases for backward compatibility
 pub type ActivateWorkerStep = ActivateWorkersStep;
 pub type RegisterWorkerStep = RegisterWorkersStep;
 pub type CreateWorkerStep = CreateLocalWorkerStep;
 pub type ActivateExternalWorkersStep = ActivateWorkersStep;
 pub type RegisterExternalWorkersStep = RegisterWorkersStep;
 pub type UpdateExternalPoliciesStep = UpdatePoliciesStep;
 pub use mcp_registration::{
    create_mcp_registration_workflow, ConnectMcpServerStep, DiscoverMcpInventoryStep,
    McpServerConfigRequest, RegisterMcpServerStep, ValidateRegistrationStep,
@@ -34,12 +70,3 @@ pub use wasm_module_removal::{
    create_wasm_module_removal_workflow, FindModuleToRemoveStep, RemoveModuleStep,
    WasmModuleRemovalRequest,
 };
 pub use worker_registration::{
    create_worker_registration_workflow, ActivateWorkerStep, CreateWorkerStep,
    DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, RegisterWorkerStep,
    UpdatePoliciesStep,
 };
 pub use worker_removal::{
    create_worker_removal_workflow, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
    RemoveFromWorkerRegistryStep, UpdateRemainingPoliciesStep, WorkerRemovalRequest,
 };
@@ -0,0 +1,153 @@
 //! External worker creation step.
 use std::{collections::HashMap, sync::Arc, time::Duration};
 use async_trait::async_trait;
 use tracing::{debug, info};
 use crate::{
    app_context::AppContext,
    core::{
        model_card::ModelCard, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode,
        HealthConfig, RuntimeType, Worker, WorkerType,
    },
    protocols::worker_spec::WorkerConfigRequest,
    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 /// 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 for CreateExternalWorkersStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let model_cards: Arc<Vec<ModelCard>> = context.get_or_err("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,
            }
        };
        // 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>;
            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>;
                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
            );
        }
        context.set("workers", workers);
        context.set("labels", labels);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
@@ -0,0 +1,273 @@
 //! Model discovery step for external API endpoints.
 use std::{collections::HashMap, sync::Arc, 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 crate::{
    core::{
        model_card::{ModelCard, ProviderType},
        model_type::ModelType,
    },
    protocols::worker_spec::WorkerConfigRequest,
    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 // 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 for DiscoverModelsStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_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
            );
            context.set::<Vec<ModelCard>>("model_cards", vec![]);
            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.set("model_cards", model_cards);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
@@ -0,0 +1,104 @@
 //! 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 super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
 use crate::workflow::{
    BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition,
 };
 /// 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 {
    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"]),
    )
 }
@@ -0,0 +1,340 @@
 //! Local worker creation step.
 use std::{collections::HashMap, sync::Arc, time::Duration};
 use async_trait::async_trait;
 use tracing::debug;
 use super::discover_dp::DpInfo;
 use crate::{
    app_context::AppContext,
    core::{
        BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder,
        HealthConfig, ModelCard, RuntimeType, Worker, WorkerType,
    },
    protocols::worker_spec::WorkerConfigRequest,
    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 /// 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 for CreateLocalWorkerStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
        let discovered_labels: Arc<HashMap<String, String>> =
            context.get_or_err("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.as_ref().clone();
        for (key, value) in &config_labels {
            final_labels.insert(key.clone(), value.clone());
        }
        // Determine model_id: config > served_model_name > model_path > "unknown"
        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".to_string());
        if model_id != "unknown" {
            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, &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);
        // 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.as_ref()
            );
        }
        // Create workers - always output as Vec for unified downstream handling
        let workers = if config.dp_aware {
            create_dp_aware_workers(
                context,
                &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,
            )
        };
        context.set("workers", workers);
        context.set("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);
        }
    }
    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,
    context: &WorkflowContext,
    config: &WorkerConfigRequest,
 ) -> RuntimeType {
    if !matches!(connection_mode, ConnectionMode::Grpc { .. }) {
        return RuntimeType::Sglang;
    }
    if let Some(detected_runtime) = context.get::<String>("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) -> 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,
    }
 }
 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(
    context: &WorkflowContext,
    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: Arc<DpInfo> = context.get_or_err("dp_info")?;
    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>;
        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 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>;
    worker.set_healthy(false);
    debug!(
        "Created worker object for {} ({:?}) with {} labels",
        normalized_url,
        connection_mode,
        final_labels.len()
    );
    vec![worker]
 }
@@ -0,0 +1,146 @@
 //! Connection mode detection step.
 use std::{sync::Arc, time::Duration};
 use async_trait::async_trait;
 use reqwest::Client;
 use tracing::debug;
 use crate::{
    app_context::AppContext,
    core::ConnectionMode,
    protocols::worker_spec::WorkerConfigRequest,
    routers::grpc::client::GrpcClient,
    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 /// Strip protocol prefix from URL.
 fn strip_protocol(url: &str) -> String {
    url.trim_start_matches("http://")
        .trim_start_matches("https://")
        .trim_start_matches("grpc://")
        .to_string()
 }
 /// 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 for DetectConnectionModeStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        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.set("connection_mode", connection_mode);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
@@ -0,0 +1,79 @@
 //! Data Parallel (DP) information discovery step.
 use std::sync::Arc;
 use async_trait::async_trait;
 use tracing::debug;
 use super::discover_metadata::get_server_info;
 use crate::{
    protocols::worker_spec::WorkerConfigRequest,
    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 /// DP (Data Parallel) information for a worker.
 #[derive(Debug, Clone)]
 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".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 for DiscoverDPInfoStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_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.set("dp_info", dp_info);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
@@ -0,0 +1,244 @@
 //! Metadata discovery step for local workers.
 use std::{collections::HashMap, sync::Arc, 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 crate::{
    core::ConnectionMode,
    protocols::worker_spec::WorkerConfigRequest,
    routers::grpc::client::GrpcClient,
    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 // 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 dp_size: Option<usize>,
    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>>,
 }
 /// Strip protocol prefix from URL.
 fn strip_protocol(url: &str) -> String {
    url.trim_start_matches("http://")
        .trim_start_matches("https://")
        .trim_start_matches("grpc://")
        .to_string()
 }
 /// 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 !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 !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 for DiscoverMetadataStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
        debug!(
            "Discovering metadata for {} ({:?})",
            config.url, *connection_mode
        );
        let (discovered_labels, detected_runtime) = match connection_mode.as_ref() {
            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);
                    }
                }
                // 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(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)
        });
        debug!(
            "Discovered {} metadata labels for {}",
            discovered_labels.len(),
            config.url
        );
        context.set("discovered_labels", discovered_labels);
        if let Some(runtime) = detected_runtime {
            debug!("Detected runtime type: {}", runtime);
            context.set("detected_runtime_type", runtime);
        }
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
@@ -0,0 +1,195 @@
 mod create_worker;
 mod detect_connection;
 mod discover_dp;
 mod discover_metadata;
 mod removal;
 use std::{sync::Arc, time::Duration};
 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 removal::{
    FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, RemoveFromWorkerRegistryStep,
    UpdateRemainingPoliciesStep, WorkerRemovalRequest,
 };
 use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
 use crate::{
    config::RouterConfig,
    workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
 };
 pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDefinition {
    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"]),
        )
        // 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 {
    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"]),
        )
 }
@@ -1,18 +1,24 @@
-use std::{collections::HashSet, sync::Arc, time::Duration};
+//! Worker removal step implementations.
 use std::{collections::HashSet, sync::Arc};
 use async_trait::async_trait;
 use tracing::{debug, info};
-use crate::{app_context::AppContext, core::Worker, workflow::*};
+use crate::{
    app_context::AppContext,
    core::Worker,
    workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
-/// Request structure for worker removal
+/// Request structure for worker removal.
 #[derive(Debug, Clone)]
 pub struct WorkerRemovalRequest {
    pub url: String,
    pub dp_aware: bool,
 }
-/// Step 1: Find workers to remove based on URL
+/// Step 1: Find workers to remove based on URL.
 pub struct FindWorkersToRemoveStep;
 #[async_trait]
@@ -81,11 +87,11 @@ impl StepExecutor for FindWorkersToRemoveStep {
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false // Worker not found is not retryable
+        false
    }
 }
-/// Step 2: Remove workers from policy registry
+/// Step 2: Remove workers from policy registry.
 pub struct RemoveFromPolicyRegistryStep;
 #[async_trait]
@@ -122,11 +128,11 @@ impl StepExecutor for RemoveFromPolicyRegistryStep {
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false // Policy removal is not retryable
+        false
    }
 }
-/// Step 3: Remove workers from worker registry
+/// Step 3: Remove workers from worker registry.
 pub struct RemoveFromWorkerRegistryStep;
 #[async_trait]
@@ -167,11 +173,11 @@ impl StepExecutor for RemoveFromWorkerRegistryStep {
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false // Worker removal is not retryable
+        false
    }
 }
-/// Step 4: Update cache-aware policies for remaining workers
+/// Step 4: Update cache-aware policies for remaining workers.
 pub struct UpdateRemainingPoliciesStep;
 #[async_trait]
@@ -219,62 +225,6 @@ impl StepExecutor for UpdateRemainingPoliciesStep {
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false // Policy update is not retryable
+        false
    }
 }
 /// Create a worker removal workflow definition
 pub fn create_worker_removal_workflow() -> WorkflowDefinition {
    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"]),
        )
 }
@@ -0,0 +1,17 @@
 pub mod external;
 pub mod local;
 pub mod shared;
 pub use external::{
    create_external_worker_workflow as create_external_worker_registration_workflow,
    create_external_worker_workflow, group_models_into_cards, infer_model_type_from_id,
    CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo, ModelsResponse,
 };
 pub use local::{
    create_local_worker_workflow as create_worker_registration_workflow,
    create_local_worker_workflow, create_worker_removal_workflow, CreateLocalWorkerStep,
    DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, DpInfo,
    FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, RemoveFromWorkerRegistryStep,
    UpdateRemainingPoliciesStep, WorkerRemovalRequest,
 };
 pub use shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep, WorkerList};
@@ -0,0 +1,35 @@
 //! Unified worker activation step.
 use std::sync::Arc;
 use async_trait::async_trait;
 use tracing::info;
 use crate::{
    core::Worker,
    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
 };
 /// Unified step to activate workers by marking them as healthy.
 ///
 /// This is the final step in any worker registration workflow.
 pub struct ActivateWorkersStep;
 #[async_trait]
 impl StepExecutor for ActivateWorkersStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
        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: &crate::workflow::WorkflowError) -> bool {
        false
    }
 }
@@ -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>>;
@@ -0,0 +1,46 @@
 //! Unified worker registration step.
 use std::sync::Arc;
 use async_trait::async_trait;
 use tracing::debug;
 use crate::{
    app_context::AppContext,
    core::Worker,
    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
 };
 /// 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>>`.
 pub struct RegisterWorkersStep;
 #[async_trait]
 impl StepExecutor for RegisterWorkersStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
        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);
        }
        context.set("worker_ids", worker_ids);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
        false
    }
 }
@@ -0,0 +1,78 @@
 //! Unified policy update step.
 use std::{collections::HashMap, sync::Arc};
 use async_trait::async_trait;
 use tracing::debug;
 use crate::{
    app_context::AppContext,
    core::Worker,
    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
 };
 /// 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;
 #[async_trait]
 impl StepExecutor for UpdatePoliciesStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
        let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
        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_fast(&model_id);
            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: &crate::workflow::WorkflowError) -> bool {
        false
    }
 }
@@ -1,943 +0,0 @@
 use std::{collections::HashMap, sync::Arc, 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, info, warn};
 use crate::{
    app_context::AppContext,
    core::{
        BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder,
        HealthConfig, ModelCard, RuntimeType, Worker, WorkerType,
    },
    protocols::worker_spec::WorkerConfigRequest,
    routers::grpc::client::GrpcClient,
    workflow::*,
 };
 // 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)]
 struct ServerInfo {
    #[serde(alias = "model")]
    model_id: Option<String>,
    model_path: Option<String>,
    served_model_name: Option<String>,
    dp_size: Option<usize>,
    version: Option<String>,
    max_batch_size: Option<usize>,
    max_total_tokens: Option<usize>,
    max_prefill_tokens: Option<usize>,
    max_running_requests: Option<usize>,
    max_num_reqs: Option<usize>,
 }
 /// Model information returned from /model_info endpoint
 #[derive(Debug, Clone, Deserialize, Serialize)]
 struct ModelInfo {
    model_path: Option<String>,
    tokenizer_path: Option<String>,
    is_generation: Option<bool>,
    /// HuggingFace model type string (e.g., "llama", "qwen2", "gpt_oss")
    model_type: Option<String>,
    /// Model architectures from HuggingFace config (e.g., ["LlamaForCausalLM"])
    architectures: Option<Vec<String>>,
 }
 #[derive(Debug, Clone)]
 pub struct DpInfo {
    pub dp_size: usize,
    pub model_id: String,
 }
 /// Parse server info from JSON response
 fn parse_server_info(json: Value) -> Result<ServerInfo, String> {
    serde_json::from_value(json).map_err(|e| format!("Failed to parse server info: {}", e))
 }
 /// Get server info from /server_info endpoint
 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 !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))?;
    parse_server_info(json)
 }
 /// Get model info from /model_info endpoint
 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 !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))
 }
 /// Get DP info for a worker URL
 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".to_string());
    Ok(DpInfo { dp_size, model_id })
 }
 /// Strip protocol prefix from URL
 fn strip_protocol(url: &str) -> String {
    url.trim_start_matches("http://")
        .trim_start_matches("https://")
        .trim_start_matches("grpc://")
        .to_string()
 }
 /// 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))
        }
    }
 }
 /// Fetch metadata from gRPC server
 async fn do_fetch_grpc_metadata(
    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())
 }
 /// 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))
    };
    match runtime_type {
        Some(runtime) => {
            let labels = do_fetch_grpc_metadata(&grpc_url, runtime).await?;
            Ok((labels, runtime.to_string()))
        }
        None => {
            // Try SGLang first, then vLLM as fallback
            if let Ok(labels) = do_fetch_grpc_metadata(&grpc_url, "sglang").await {
                return Ok((labels, "sglang".to_string()));
            }
            let labels = do_fetch_grpc_metadata(&grpc_url, "vllm")
                .await
                .map_err(|e| format!("gRPC metadata failed (tried SGLang and vLLM): {}", e))?;
            Ok((labels, "vllm".to_string()))
        }
    }
 }
 /// Step 1: Detect connection mode by probing HTTP and gRPC
 pub struct DetectConnectionModeStep;
 #[async_trait]
 impl StepExecutor for DetectConnectionModeStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        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.set("connection_mode", connection_mode);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
 /// Step 2: Discover metadata from worker
 pub struct DiscoverMetadataStep;
 #[async_trait]
 impl StepExecutor for DiscoverMetadataStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
        debug!(
            "Discovering metadata for {} ({:?})",
            config.url, *connection_mode
        );
        let (discovered_labels, detected_runtime) = match connection_mode.as_ref() {
            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);
                    }
                }
                // Fetch from /model_info for model-related metadata (model_type, architectures)
                if let Ok(model_info) = get_model_info(&config.url, config.api_key.as_deref()).await
                {
                    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)
        });
        debug!(
            "Discovered {} metadata labels for {}",
            discovered_labels.len(),
            config.url
        );
        context.set("discovered_labels", discovered_labels);
        if let Some(runtime) = detected_runtime {
            debug!("Detected runtime type: {}", runtime);
            context.set("detected_runtime_type", runtime);
        }
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
 /// Step 2.5: Discover DP (Data Parallel) information (only for DP-aware workers)
 pub struct DiscoverDPInfoStep;
 #[async_trait]
 impl StepExecutor for DiscoverDPInfoStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_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.set("dp_info", dp_info);
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        true
    }
 }
 /// Step 3: Create worker object with merged configuration + metadata
 pub struct CreateWorkerStep;
 #[async_trait]
 impl StepExecutor for CreateWorkerStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
        let discovered_labels: Arc<HashMap<String, String>> =
            context.get_or_err("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.as_ref().clone();
        for (key, value) in &config_labels {
            final_labels.insert(key.clone(), value.clone());
        }
        // Determine model_id: config > served_model_name > model_path > "unknown"
        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".to_string());
        if model_id != "unknown" {
            debug!("Using model_id: {}", model_id);
        }
        // Create ModelCard
        let model_card = {
            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());
            }
            // Set HuggingFace model type from discovered labels
            if let Some(model_type_str) = final_labels.get("model_type") {
                card = card.with_hf_model_type(model_type_str.clone());
            }
            // Set architectures from discovered labels (JSON array string)
            if let Some(architectures_json) = final_labels.get("architectures") {
                if let Ok(architectures) = serde_json::from_str::<Vec<String>>(architectures_json) {
                    card = card.with_architectures(architectures);
                }
            }
            card
        };
        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 = 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);
        // Get runtime type (for gRPC workers)
        let runtime_type = if matches!(connection_mode.as_ref(), ConnectionMode::Grpc { .. }) {
            if let Some(detected_runtime) = context.get::<String>("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
            }
        } else {
            RuntimeType::Sglang
        };
        // Build circuit breaker config
        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),
            }
        };
        // Build health config
        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,
            }
        };
        // Normalize URL: add protocol prefix only if missing
        let normalized_url = if config.url.starts_with("http://")
            || config.url.starts_with("https://")
            || config.url.starts_with("grpc://")
        {
            config.url.clone()
        } else {
            match connection_mode.as_ref() {
                ConnectionMode::Http => format!("http://{}", config.url),
                ConnectionMode::Grpc { .. } => format!("grpc://{}", config.url),
            }
        };
        if normalized_url != config.url {
            debug!(
                "Normalized worker URL: {} -> {} ({:?})",
                config.url,
                normalized_url,
                connection_mode.as_ref()
            );
        }
        // Handle DP-aware vs non-DP-aware workers
        if config.dp_aware {
            let dp_info: Arc<DpInfo> = context.get_or_err("dp_info")?;
            debug!(
                "Creating {} DP-aware workers for {} (dp_size: {})",
                dp_info.dp_size, config.url, dp_info.dp_size
            );
            let mut workers = Vec::new();
            for rank in 0..dp_info.dp_size {
                let mut builder =
                    DPAwareWorkerBuilder::new(normalized_url.clone(), rank, dp_info.dp_size)
                        .model(model_card.clone())
                        .worker_type(worker_type.clone())
                        .connection_mode(connection_mode.as_ref().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>;
                worker.set_healthy(false);
                workers.push(worker);
                debug!(
                    "Created DP-aware worker {}@{}/{} ({:?})",
                    config.url,
                    rank,
                    dp_info.dp_size,
                    connection_mode.as_ref()
                );
            }
            context.set("workers", workers);
            context.set("labels", final_labels);
        } else {
            let mut builder = BasicWorkerBuilder::new(normalized_url.clone())
                .model(model_card)
                .worker_type(worker_type)
                .connection_mode(connection_mode.as_ref().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>;
            worker.set_healthy(false);
            debug!(
                "Created worker object for {} ({:?}) with {} labels",
                config.url,
                connection_mode.as_ref(),
                final_labels.len()
            );
            context.set("worker", worker);
            context.set("labels", final_labels);
        }
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 /// Step 4: Register worker(s) in registry
 pub struct RegisterWorkerStep;
 #[async_trait]
 impl StepExecutor for RegisterWorkerStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        if config.dp_aware {
            let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
            let mut worker_ids = Vec::new();
            for worker in workers.iter() {
                let worker_id = app_context.worker_registry.register(Arc::clone(worker));
                worker_ids.push(worker_id.clone());
                debug!(
                    "Registered DP-aware worker {} with ID {:?}",
                    config.url, worker_id
                );
            }
            context.set("worker_ids", worker_ids);
        } else {
            let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
            let worker_id = app_context
                .worker_registry
                .register(Arc::clone(worker.as_ref()));
            debug!("Registered worker {} with ID {:?}", config.url, worker_id);
            context.set("worker_id", worker_id);
        }
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 /// Step 5: Update policy registry with worker information
 pub struct UpdatePoliciesStep;
 #[async_trait]
 impl StepExecutor for UpdatePoliciesStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
        let policy_hint = labels.get("policy").map(|s| s.as_str());
        if config.dp_aware {
            let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
            let model_id = workers[0].model_id().to_string();
            for _ in 0..workers.len() {
                app_context
                    .policy_registry
                    .on_worker_added(&model_id, policy_hint);
            }
            // Initialize cache-aware policy if needed
            let all_workers = app_context.worker_registry.get_by_model_fast(&model_id);
            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);
                }
            }
            debug!(
                "Updated policies for {} DP-aware workers {} (model: {})",
                workers.len(),
                config.url,
                model_id
            );
        } else {
            let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
            let model_id = worker.model_id().to_string();
            app_context
                .policy_registry
                .on_worker_added(&model_id, policy_hint);
            // Initialize cache-aware policy if needed
            let all_workers = app_context.worker_registry.get_by_model_fast(&model_id);
            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);
                }
            }
            // Initialize bucket policies for prefill workers
            let prefill_workers = app_context.worker_registry.get_prefill_workers();
            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 worker {} (model: {})",
                config.url, model_id
            );
        }
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 /// Step 6: Activate worker(s) by marking them as healthy
 pub struct ActivateWorkerStep;
 #[async_trait]
 impl StepExecutor for ActivateWorkerStep {
    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
        if config.dp_aware {
            let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
            for worker in workers.iter() {
                worker.set_healthy(true);
            }
            info!(
                "Activated {} DP-aware workers from {} (marked as healthy)",
                workers.len(),
                config.url
            );
        } else {
            let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
            worker.set_healthy(true);
            info!("Activated worker {} (marked as healthy)", config.url);
        }
        Ok(StepResult::Success)
    }
    fn is_retryable(&self, _error: &WorkflowError) -> bool {
        false
    }
 }
 /// Create local worker registration workflow definition
 pub fn create_worker_registration_workflow(
    router_config: &crate::config::RouterConfig,
 ) -> WorkflowDefinition {
    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("worker_registration", "Worker Registration")
        .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),
        )
        .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"]),
        )
        .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"]),
        )
        .add_step(
            StepDefinition::new("create_worker", "Create Worker", Arc::new(CreateWorkerStep))
                .with_timeout(Duration::from_secs(5))
                .with_failure_action(FailureAction::FailWorkflow)
                .depends_on(&["discover_dp_info"]),
        )
        .add_step(
            StepDefinition::new(
                "register_worker",
                "Register Worker",
                Arc::new(RegisterWorkerStep),
            )
            .with_timeout(Duration::from_secs(5))
            .with_failure_action(FailureAction::FailWorkflow)
            .depends_on(&["create_worker"]),
        )
        .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_worker"]),
        )
        .add_step(
            StepDefinition::new(
                "activate_worker",
                "Activate Worker",
                Arc::new(ActivateWorkerStep),
            )
            .with_timeout(Duration::from_secs(5))
            .with_failure_action(FailureAction::FailWorkflow)
            .depends_on(&["update_policies"]),
        )
 }